Repository navigation
Conversation
30822d5 to
2698b37
Compare
Muse Glimmer's vision encoder kept rotate_half as a private helper. Move it to primitives/macos/rope.py, unchanged, so other macOS models can share it without importing another model's private function. Muse Glimmer's exported graph is unchanged.
The macOS (GPU) text and image encoders, plus the Neural Engine float16 variants that use scale-invariant RMSNorm. Includes a parity test for the iOS encoders against the macOS reference, and a shared fixture for the local snapshot.
Adds the `uv run coreai.embedding.export embeddinggemma-2 [--platform iOS] [--skip-vision]` console script, which writes to exports/ by default. Also adds the registry entry, AIModel metadata, and the minimal metadata contract with 8 canonical task prompts. Includes the contract fixtures the Swift runtime reads. uv.lock is unchanged.
Model card covering export (macOS and Neural Engine), the bundle layout, task prompts and usage. Also the catalog entry in models/README.md.
2698b37 to
85db172
Compare
eps * r^2 is now computed as (sqrt(eps) * r)^2. r * r overflowed float16 once r > 256, and the row came out all zeros. Adds a float16-vs-float32 sweep over input magnitudes.
| peak = table[ids.long()].abs().max().item() * math.sqrt(cfg.hidden_size) | ||
| table = table * (3000.0 / peak) | ||
| ios32, _ = text_encoder_from_state_dict(cfg, state, dtype=torch.float32) | ||
| ios16, _ = text_encoder_from_state_dict(cfg, state, dtype=torch.float16) |
There was a problem hiding this comment.
Can we add some real weight fp16 parity? The fp16 text encoder only checked on a tiny random config doesn't catch any possible fp16 bugs, thanks!
| "metric": family.metric, | ||
| "precision": { | ||
| "compute": precision, | ||
| "forbidden": [] if ios else list(family.forbidden_precisions), |
There was a problem hiding this comment.
The contract here somewhat contradicts itself, the family says precision_floor="bfloat16" and forbidden_precisions=("float16",), but the iOS bundle computes in fp16 and writes "forbidden": [] here. Is ScaledRMSNorm what makes fp16 ok on static shape variant?
There was a problem hiding this comment.
Yes, I thought this was readable from comments. Family describes the default GPU path, while ANE needs to have hardware compatible format.
| max_length: int, | ||
| include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, | ||
| ) -> dict[str, AIProgram]: | ||
| """GPU image tower: ``{image_encoder, token_table, encoder}`` programs. |
There was a problem hiding this comment.
Transformer weights appear twice (text encoder and encoder), thats a huge amount (~1GB) of duplicate fp32 this is bc weight sharing only happens inside one program
There was a problem hiding this comment.
Let me see what we can do. Might need to update the signatures.
| from coreai_models.models.ios.embedding_gemma2 import load_embedding_gemma2_text_encoder | ||
|
|
||
| # The token table runs on the GPU in float32; the encoder and image tower use float16. | ||
| gpu_encoder = _load_text_encoder(family, snapshot, torch.float32) |
There was a problem hiding this comment.
Why store the iOS token table in fp32?
There was a problem hiding this comment.
Yes, we no longer need fp32 here
An all-zero row now normalizes to zero instead of NaN, as on the GPU path.
The table is scaled by sqrt(hidden) in float32 and then cast, so the encoder gets the same float16 inputs at half the size (512 MB to 256 MB). The iOS export takes the table from the Neural Engine loader instead of loading the float32 GPU encoder.
Compares it with the float32 GPU encoder per row (cosine >= 0.9999), with full and padded rows. A plain RMSNorm in float16 stays finite but scores 0.91.
The token table and the inputs_embeds encoder are now functions of text_encoder.aimodel instead of separate assets, so the table and the transformer are stored once (about 1 GB less in float32). The metadata points token_table and encoder at text_encoder.aimodel; the encoder functions are named embeds_len* so they stay distinct from the text ones.
|
Updated the bundle to split the repeated weights. |
| The exporter checks them against the snapshot ``config.json`` / ``processor_config.json`` | ||
| and fails closed on a mismatch. |
There was a problem hiding this comment.
so these are model-specific?
| """Component name -> asset file for the bundle. | ||
|
|
||
| In the macOS image bundle, ``token_table`` and ``encoder`` are functions of | ||
| ``text_encoder.aimodel``, so the table and the transformer are stored once. |
There was a problem hiding this comment.
Just to make sure I understand:
iOS: does not need the TEXT_COMPONENT
macOS: does not need TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT unless with_image is True
| precision = check_compute_precision( | ||
| family, default_precision(family, platform), platform=platform | ||
| ) |
There was a problem hiding this comment.
I think this is always giving you the default precision of the platform. I ran:
uv run coreai.embedding.export embeddinggemma-2 --compute-precision bfloat16and got in my metadata:
"precision": {
"compute": "float32",
"forbidden": [
"float16"
],
"compute_units": [
"gpu"
]
},We should pass the optional compute_precision to embedding_block and pass that to check_compute_precision.
| def _converter(include_debug_info: bool) -> Any: | ||
| import coreai_torch | ||
|
|
||
| mode = coreai_torch.TorchConverter.Mode | ||
| return coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) |
There was a problem hiding this comment.
can probably inline this instead of having a thin wrapper
| ) | ||
|
|
||
|
|
||
| def apply_axial_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
double check we don't already have existing definitions of this. Possibly in python/src/coreai_models/primitives/iOS/rope.py
| return None | ||
|
|
||
| @classmethod | ||
| def from_pretrained( |
There was a problem hiding this comment.
can safe tensor loading reuse code from python/src/coreai_models/models/base.py
| return embeds, bias, pool | ||
|
|
||
|
|
||
| def ane_image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, torch.Tensor]: |
There was a problem hiding this comment.
double check what can be generalized between iOS and macOS helper functions
| per_layer = text_config.get("per_layer_config") | ||
| if per_layer is None: | ||
| # HF default: full layers get head_dim 512 / 1 KV head. | ||
| types = kwargs.get("layer_types") or [] |
There was a problem hiding this comment.
I'm a bit confused about what happens if layer_types does not exist. Can you take a look at the different fallbacks and defaults that we might silently fall back to?
| model_dir = _resolve_model_dir(model_dir_or_id) | ||
| config = EmbeddingGemma2TextConfig.from_pretrained(model_dir) | ||
| sd: dict[str, torch.Tensor] = {} | ||
| for file in sorted(model_dir.glob("*.safetensors")): |
There was a problem hiding this comment.
what if there are no safe tensors in there?
embedding_block now takes compute_precision, and the exporter passes the precision it validated, so --compute-precision bfloat16 writes "bfloat16" instead of the platform default.
from_dict needs layer_types, per_layer_config and both rope_parameters thetas. Without layer_types, the per-layer head overrides were built from an empty list before __post_init__ synthesized the types, so full layers got the wrong head_dim; now the config fails with the missing keys.
primitives/ios/rope.apply_rope(ndim=2) is the same axial rotate-half RoPE (bit-identical in float32 and bfloat16), so the local apply_axial_rope goes. The static-shape test now matches data-dependent ops by name, since split adds Python getitem nodes.
ImageTowerFacts has no defaults; EmbeddingGemma 2's values are written out in its family entry. The forbidden-precision message names the floor only when it differs from float32.
- A failed export removes the bundle folder. - The revision error says a Hugging Face cache snapshot directory is needed. - --platform defaults to macOS in the parser; --output-dir help names <repo_root>/exports. - The converter setup is inlined; the dry-run test checks embeds_len128_b1.
The GPU modules run in float32 or bfloat16; the transformers note says the config parser avoids the >=5.19 requirement; the image bias comment has the batch dimension. Every loader raises FileNotFoundError when the snapshot has no .safetensors.
diffusion.gpu gains new_converter and add_function, which export_multifunction now uses; the embedding exporter's GPU shared program and Neural Engine programs build on them instead of their own copies. Converted programs are identical before and after (15 functions hashed across the diffusion, GPU and Neural Engine paths).
…aces image_reference_inputs takes the pixel and position dtypes; the Neural Engine trace passes float16 for both. Converted programs are unchanged.
export_utils holds the snapshot checks, the token table, trace inputs and program saving; export_macos the GPU programs; export_ios the Neural Engine programs. export.py keeps export_embedding_bundle and the CLI. Helpers used across modules are public (save_program, load_text_encoder, load_image_encoder, export_gpu_programs, export_ane_programs). Converted programs, --help and the dry-run plans are identical.
|
@kevchengcodes I added some of the refactors, but completing everything now expands the scope, as there are parallel changes to other runners that need to finish. |
Restores primitives/macos/rope.py and muse_glimmer_vision.py to the base branch.
rope_inv_freq replaces _inv_freq and axial_inv_freq; _conv becomes conv1x1.
read_checkpoint_tensors serves both macOS loaders; the iOS loaders build on them, and load_conv_weights replaces the hand-written key checks.
Drops TraceSpec, add_traces, export_text_encoder_program and the loader dispatchers; the token-table trace is shared, and export_stateless uses new_converter.
rope_head_dim serves the macOS and iOS text encoders; also drops an empty TYPE_CHECKING block.
MaxAbsRMSNorm and MaxAbsRMSNormNoScale take a reduction axis and replace ScaledRMSNorm in the EmbeddingGemma 2 Neural Engine encoders; the magnitude sweep moves with them.
Norms and RoPE run on (B, n, hd, L); drops rotate_half_channels, apply_rope_bc1s and _per_head.
VisionAttention overrides only the RoPE; state_dict names are unchanged.
…figs The macOS text, envelope and asset-sharing asserts move into the matrix and macOS image tests.
HF parity runs on the exported SDPA path; the iOS tests go through the token table and EncoderBandWrapper, which replaces the separate band-wrapper test.
…lines make_converter and validate_compute_precision replace new_converter and check_compute_precision; bundle_assets and the tower pipelines share text_pipeline and IMAGE_PIPELINE.
Rename _resolve_safetensors_files and _load_tensors_for_keys to resolve_safetensors_files and load_tensors_for_keys and switch every caller.
…pers read_checkpoint_tensors resolves model.safetensors or the index.json shards with resolve_safetensors_files and loads with load_tensors_for_keys.
Brings back the sliding-window negative test, the mask unit test, the real-config layout, the sentence-transformers cross-check, the removed-keys guard and the HF SDPA / USE_HF_IMPL variants.
Converts the banded and unbanded encoder, the float16 token table and the image encoder through export_ios, checks the function names and runs them on the Core AI runtime against PyTorch.
assets lists every .aimodel file once and each component names its asset; components gain output_dtype from contract.output_dtype, which the exporters also cast to.
State what each helper does, returns and raises; drop rationale and stale references. Rename the two 'fails closed' tests after what they assert.
Type of Change
Motivation and Context
Retrieval embeddings power on-device semantic search, clustering and RAG, and are not one of the supported model kinds yet. This PR adds the export side of an embedding modality, with EmbeddingGemma 2 as the first model. EmbeddingGemma 2 maps text or an image into one shared space: a single 768-d, L2-normalized vector whose leading 512/256/128 dimensions are valid Matryoshka embeddings. The weights are released by Google DeepMind under Apache-2.0, and the Gemma Prohibited Use Policy applies.
The Swift runtime and a search tool that consume these bundles follow in separate PRs. This one is self-contained: it exports and tests the bundle.
What's added
models/macos/embedding_gemma2{,_vision}.py,models/ios/embedding_gemma2{,_vision}.py): the text encoder and image tower for the GPU (float32 or bfloat16), and float16 variants for the Neural Engine. The macOS text and image encoders share one attention block and sandwich layer. The Neural Engine variants use the max-abs RMSNorm from Gemma 4 so float16 does not overflow.coreai.embedding.export: a console script following the other modalities. A positional registry short name, output underexports/by default, and--platform iOSfor the Neural Engine bundle. The exporter is split intoexport(CLI and bundle assembly),export_utils,export_macos(GPU programs) andexport_ios(Neural Engine programs).coreai_models.embedding.contract): thekind: "embedding"metadata (space id and revision, dimension and Matryoshka widths, task prompts, components and towers), written next to the encoder assets and the tokenizer.embeddingutility entry,embeddinggemma-2.models/embedding_gemma2/README.mdand a new Embedding Models section in the model catalog.Changes outside the new modules
Each is a move or rename with unchanged behavior; every existing caller is switched and its exported graph is unchanged.
primitives/ios/rms_norm.py: the max-absRMSNorm/RMSNormNoScalemove here frommodels/ios/gemma4_vision.py, with anaxisargument for the BC1S layout. Gemma 4 imports them from here.models/base.py:resolve_safetensors_filesandload_tensors_for_keysbecome public (they were_-prefixed and already imported by Muse Glimmer). The EmbeddingGemma 2 loaders use them.diffusion/gpu.py:make_converterandadd_functionare shared by the diffusion and embedding exporters.Bundle layout
text_encoder.aimodel,image_encoder.aimodeltoken_table.aimodel,encoder.aimodel,image_encoder.aimodeltext_encoder(len*_b*functions)token_table -> encoderimage_encoder -> token_table -> encoderimage_encoder -> token_table -> encoderassetslists each file once; each component names itsassetandoutput_dtype. On macOS,token_table(main) andencoder(embeds_len*_b1) are functions oftext_encoder.aimodel, so the transformer weights are stored once.embeds_len*_b1on both platforms; the host reads the names fromcomponents.encoder.functions.token_tableoutput at the placeholder rows before the encoder.Export
The GPU bundle refuses
float16, because the model overflows in float16;bfloat16is accepted.Full exports with the image tower: macOS takes 102 s at 4.6 GB peak host memory (1.0 GB + 640 MB assets); iOS takes 52 s at 3.3 GB (256 + 288 + 293 MB assets).
Numerical parity
The reference is the sentence-transformers / HF float32 pipeline on CPU, compared by cosine similarity.
Tests
torch.exportat a static shape, Neural Engine conversion on the runtime, the metadata contract, the CLI and the registry entry.ruff checkandruff format --check.Testing