Skip to content

Add EmbeddingGemma 2 export: coreai.embedding.export and registry entry - #369

Open
stikves wants to merge 40 commits into
apple:mainfrom
stikves:sukru/embedding-gemma2-export
Open

stikves wants to merge 40 commits into
apple:mainfrom
stikves:sukru/embedding-gemma2-export

Conversation

@stikves

@stikves stikves commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor

Type of Change

  • Bug fix
  • Small feature (agreed in an issue)
  • New model (export recipe only)
  • New model (export recipe and new runtime support)
  • Documentation update

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

  • Encoders (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 under exports/ by default, and --platform iOS for the Neural Engine bundle. The exporter is split into export (CLI and bundle assembly), export_utils, export_macos (GPU programs) and export_ios (Neural Engine programs).
  • Bundle contract (coreai_models.embedding.contract): the kind: "embedding" metadata (space id and revision, dimension and Matryoshka widths, task prompts, components and towers), written next to the encoder assets and the tokenizer.
  • Registry: an embedding utility entry, embeddinggemma-2.
  • Docs: models/embedding_gemma2/README.md and 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-abs RMSNorm / RMSNormNoScale move here from models/ios/gemma4_vision.py, with an axis argument for the BC1S layout. Gemma 4 imports them from here.
  • models/base.py: resolve_safetensors_files and load_tensors_for_keys become public (they were _-prefixed and already imported by Muse Glimmer). The EmbeddingGemma 2 loaders use them.
  • diffusion/gpu.py: make_converter and add_function are shared by the diffusion and embedding exporters.

Bundle layout

macOS (GPU) iOS (Neural Engine)
Assets text_encoder.aimodel, image_encoder.aimodel token_table.aimodel, encoder.aimodel, image_encoder.aimodel
Text tower text_encoder (len*_b* functions) token_table -> encoder
Image tower image_encoder -> token_table -> encoder image_encoder -> token_table -> encoder
Outputs float32 float16
  • assets lists each file once; each component names its asset and output_dtype. On macOS, token_table (main) and encoder (embeds_len*_b1) are functions of text_encoder.aimodel, so the transformer weights are stored once.
  • Encoder functions are named embeds_len*_b1 on both platforms; the host reads the names from components.encoder.functions.
  • The host writes the image soft tokens into the token_table output at the placeholder rows before the encoder.

Export

uv run coreai.embedding.export --list-models
uv run coreai.embedding.export embeddinggemma-2                  # text + image towers, GPU (float32)
uv run coreai.embedding.export embeddinggemma-2 --platform iOS   # Neural Engine (float16)
uv run coreai.embedding.export embeddinggemma-2 --skip-vision    # text tower only

The GPU bundle refuses float16, because the model overflows in float16; bfloat16 is 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.

  • Text encoder vs reference: min cosine ≥ 0.99999 on a mixed-length batch, including an input longer than 1,100 tokens that exercises the sliding-window layers. Padding invariance holds at ≥ 0.99999.
  • Image tower vs reference: soft tokens and the pooled embedding at min cosine ≥ 0.99999 at 768/512/256/128 dimensions.
  • Neural Engine (float16) vs GPU (float32): small random configs for text and image, and the real-weight text encoder at 64 and 1,024 tokens (min cosine ≥ 0.9999). The max-abs RMSNorm is checked against float32 for row peaks from 0 to 6e4.
  • Converted programs: tiny configs go through the real export path, run on the Core AI runtime (CPU and GPU), and match the PyTorch modules (exact for the token table, 1e-2 for the float16 encoders).

Tests

  • The PR's test files run 101 tests and all pass, including the real-weight tiers. The reference tiers need extra packages:
    uv run --frozen --with transformers==5.20.1 --with "torchvision==0.24.*" --with "sentence-transformers>=6.1" \
      pytest python/tests/test_model_units/test_export/test_embedding_export.py \
        python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py \
        python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py \
        python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py \
        python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2_vision.py \
        python/tests/test_model_units/test_primitives/test_ios/test_rms_norm.py
    Without those packages 8 reference tests skip and the remaining 93 pass.
  • Self-contained tiers always run: config parsing, attention masks, padding invariance, torch.export at a static shape, Neural Engine conversion on the runtime, the metadata contract, the CLI and the registry entry.
  • Tests of the moved and renamed helpers' existing callers (Gemma 4 vision, Muse Glimmer, Diffusion Gemma, Gemma 3n) pass unchanged, as do ruff check and ruff format --check.

Testing

  • Tested locally
  • Added/updated tests
  • Added/updated docs

@stikves stikves self-assigned this Oct 23, 2026
@stikves
stikves marked this pull request as ready for review October 23, 2026 01:25
@stikves
stikves requested a review from a team as a code owner October 23, 2026 01:25
@stikves
stikves force-pushed the sukru/embedding-gemma2-export branch from 30822d5 to 2698b37 Compare October 23, 2026 01:26
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.
@stikves
stikves force-pushed the sukru/embedding-gemma2-export branch from 2698b37 to 85db172 Compare October 23, 2026 21:13
stikves and others added 2 commits October 27, 2026 11:37
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why store the iOS token table in fp32?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, we no longer need fp32 here

stikves and others added 6 commits October 27, 2026 17:48
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.
@stikves

stikves commented Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Updated the bundle to split the repeated weights.

Comment on lines +46 to +47
The exporter checks them against the snapshot ``config.json`` / ``processor_config.json``
and fails closed on a mismatch.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so these are model-specific?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes

"""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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +312 to +314
precision = check_compute_precision(
family, default_precision(family, platform), platform=platform
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is always giving you the default precision of the platform. I ran:

uv run coreai.embedding.export embeddinggemma-2 --compute-precision bfloat16

and 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.

Comment on lines +330 to +334
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 []

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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")):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@stikves

stikves commented Oct 29, 2026

Copy link
Copy Markdown
Contributor Author

@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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants