From 374db3331804ae8c1c5b0bd152bef27b3a5fcae9 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 21 Oct 2026 21:42:40 -0700 Subject: [PATCH 01/38] RoPE: public rotate_half in primitives/macos/rope 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. --- .../coreai_models/models/macos/muse_glimmer_vision.py | 11 +++-------- python/src/coreai_models/primitives/macos/rope.py | 6 ++++++ 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/python/src/coreai_models/models/macos/muse_glimmer_vision.py b/python/src/coreai_models/models/macos/muse_glimmer_vision.py index 679dc9f5..f6ddd010 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer_vision.py +++ b/python/src/coreai_models/models/macos/muse_glimmer_vision.py @@ -58,18 +58,13 @@ _load_tensors_for_keys, _resolve_safetensors_files, ) +from coreai_models.primitives.macos.rope import rotate_half # --------------------------------------------------------------------------- # RoPE helpers # --------------------------------------------------------------------------- -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - """Rotate second half of the last dimension: ``cat(-x2, x1)``.""" - mid = x.shape[-1] // 2 - return torch.cat((-x[..., mid:], x[..., :mid]), dim=-1) - - def _apply_rotary_pos_emb_vision( q: torch.Tensor, k: torch.Tensor, @@ -86,8 +81,8 @@ def _apply_rotary_pos_emb_vision( q, k = q.float(), k.float() cos = cos.unsqueeze(-2).float() # [1, seq_len, 1, head_dim] sin = sin.unsqueeze(-2).float() - q_embed = (q * cos) + (_rotate_half(q) * sin) - k_embed = (k * cos) + (_rotate_half(k) * sin) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype) diff --git a/python/src/coreai_models/primitives/macos/rope.py b/python/src/coreai_models/primitives/macos/rope.py index 27992551..d888ae2d 100644 --- a/python/src/coreai_models/primitives/macos/rope.py +++ b/python/src/coreai_models/primitives/macos/rope.py @@ -15,6 +15,12 @@ logger = logging.getLogger(__name__) +def rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate second half of the last dimension: ``cat(-x2, x1)``.""" + mid = x.shape[-1] // 2 + return torch.cat((-x[..., mid:], x[..., :mid]), dim=-1) + + class RoPE(coreai_torch.composite_ops.RoPE): """Apply rotary positional embedding to input tensors.""" From 6184ee974735c11554095a41e0fb7509cba11041 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 20 Oct 2026 00:13:59 -0700 Subject: [PATCH 02/38] EmbeddingGemma 2: text and image encoders for GPU and Neural Engine 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. --- .../models/ios/embedding_gemma2.py | 269 ++++++++++ .../models/ios/embedding_gemma2_vision.py | 233 ++++++++ .../models/macos/embedding_gemma2.py | 447 ++++++++++++++++ .../models/macos/embedding_gemma2_vision.py | 504 ++++++++++++++++++ .../test_models/_embedding_gemma2_common.py | 80 +++ .../test_model_units/test_models/conftest.py | 19 + .../test_ios_layers/test_embedding_gemma2.py | 103 ++++ .../test_embedding_gemma2_vision.py | 108 ++++ .../test_embedding_gemma2.py | 393 ++++++++++++++ .../test_embedding_gemma2_vision.py | 399 ++++++++++++++ 10 files changed, 2555 insertions(+) create mode 100644 python/src/coreai_models/models/ios/embedding_gemma2.py create mode 100644 python/src/coreai_models/models/ios/embedding_gemma2_vision.py create mode 100644 python/src/coreai_models/models/macos/embedding_gemma2.py create mode 100644 python/src/coreai_models/models/macos/embedding_gemma2_vision.py create mode 100644 python/tests/test_model_units/test_models/_embedding_gemma2_common.py create mode 100644 python/tests/test_model_units/test_models/conftest.py create mode 100644 python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py create mode 100644 python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2_vision.py create mode 100644 python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py create mode 100644 python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py new file mode 100644 index 00000000..31f1610e --- /dev/null +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -0,0 +1,269 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Neural Engine float16 EmbeddingGemma 2 text encoder, channels-first (BC1S) conv layout. + +What: the macOS encoder math as 1x1 Conv2d projections with static shapes. +Why: ScaledRMSNorm divides by the row max-abs before squaring, keeping float16 in range. +""" + +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from coreai_models.models.macos.embedding_gemma2 import ( + HF_TEXT_PREFIX, + EmbeddingGemma2TextConfig, + _resolve_model_dir, +) +from coreai_models.models.macos.embedding_gemma2 import ( + EmbeddingGemma2TextEncoder as MacEmbeddingGemma2TextEncoder, +) +from coreai_models.primitives.ios.sdpa import SDPA + +# Smallest positive normal float16; floors the row max-abs so all-zero rows stay finite. +_FP16_TINY = 6.103515625e-05 +MASK_NEG = -30000.0 + + +class ScaledRMSNorm(nn.Module): + """RMSNorm over ``dim`` that normalizes by the row max-abs first, so squares stay at most 1.""" + + def __init__( + self, channels: int, eps: float = 1e-6, with_scale: bool = True, dim: int = 1 + ) -> None: + super().__init__() + self.eps = eps + self.channels = channels + self.dim = dim + self.with_scale = with_scale + if with_scale: + self.weight = nn.Parameter(torch.ones(channels)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = self.dim + m = x.abs().amax(d, keepdim=True).clamp(min=_FP16_TINY) + r = torch.reciprocal(m) + xs = x * r + y = xs * torch.rsqrt((xs * xs).mean(d, keepdim=True) + self.eps * (r * r)) + if self.with_scale: + shape = [1] * x.dim() + shape[d] = self.channels + y = y * self.weight.to(y.dtype).reshape(shape) + return y + + +def rotate_half_channels(x: torch.Tensor, head_dim: int) -> torch.Tensor: + """NeoX rotate-half over the per-head channel axis of ``(B, H*hd, 1, L)``.""" + b, c, one, length = x.shape + x = x.reshape(b, c // head_dim, head_dim, length) + x1, x2 = x[:, :, : head_dim // 2], x[:, :, head_dim // 2 :] + return torch.cat((-x2, x1), dim=2).reshape(b, c, one, length) + + +def apply_rope_bc1s(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, head_dim: int): + """Apply rotary embedding; ``cos``/``sin``: ``(1, 1, hd, L)``; ``x``: ``(B, H*hd, 1, L)``.""" + b, c, one, length = x.shape + xr = x.reshape(b, c // head_dim, head_dim, length) + out = ( + xr * cos + + rotate_half_channels(x, head_dim).reshape(b, c // head_dim, head_dim, length) * sin + ) + return out.reshape(b, c, one, length) + + +def _conv(cin: int, cout: int) -> nn.Conv2d: + return nn.Conv2d(cin, cout, kernel_size=1, bias=False) + + +class EmbeddingGemma2Attention(nn.Module): + """Bidirectional multi-head attention over ``(B, C, 1, L)``.""" + + def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: + super().__init__() + self.n_heads = config.num_attention_heads + self.n_kv = config.layer_num_key_value_heads(layer_idx) + self.head_dim = hd = config.layer_head_dim(layer_idx) + dim, eps = config.hidden_size, config.rms_norm_eps + self.q_proj = _conv(dim, self.n_heads * hd) + self.k_proj = _conv(dim, self.n_kv * hd) + self.v_proj = _conv(dim, self.n_kv * hd) + self.o_proj = _conv(self.n_heads * hd, dim) + self.q_norm = ScaledRMSNorm(hd, eps, dim=2) + self.k_norm = ScaledRMSNorm(hd, eps, dim=2) + self.v_norm = ScaledRMSNorm(hd, eps, with_scale=False, dim=2) + self.sdpa = SDPA(head_dim=hd, scale=1.0) + + def _per_head(self, x: torch.Tensor, n: int, norm: nn.Module) -> torch.Tensor: + b, c, one, length = x.shape + return norm(x.reshape(b, n, self.head_dim, length)).reshape(b, c, one, length) + + def forward(self, x, cos, sin, bias): + hd = self.head_dim + q = self._per_head(self.q_proj(x), self.n_heads, self.q_norm) + k = self._per_head(self.k_proj(x), self.n_kv, self.k_norm) + v = self._per_head(self.v_proj(x), self.n_kv, self.v_norm) + q = apply_rope_bc1s(q, cos, sin, hd) + k = apply_rope_bc1s(k, cos, sin, hd) + return self.o_proj(self.sdpa(q, k, v, bias)) + + +class EmbeddingGemma2EncoderLayer(nn.Module): + """One pre-norm encoder block with the per-layer-input mixing.""" + + def __init__(self, config: EmbeddingGemma2TextConfig, i: int) -> None: + super().__init__() + dim, eps, ple = config.hidden_size, config.rms_norm_eps, config.hidden_size_per_layer_input + self.is_sliding = config.is_sliding(i) + self.self_attn = EmbeddingGemma2Attention(config, i) + self.gate_proj = _conv(dim, config.intermediate_size) + self.up_proj = _conv(dim, config.intermediate_size) + self.down_proj = _conv(config.intermediate_size, dim) + self.per_layer_input_gate = _conv(dim, ple) + self.per_layer_projection = _conv(ple, dim) + self.input_layernorm = ScaledRMSNorm(dim, eps) + self.post_attention_layernorm = ScaledRMSNorm(dim, eps) + self.pre_feedforward_layernorm = ScaledRMSNorm(dim, eps) + self.post_feedforward_layernorm = ScaledRMSNorm(dim, eps) + self.post_per_layer_input_norm = ScaledRMSNorm(dim, eps) + self.register_buffer("layer_scalar", torch.ones(1)) + + def forward(self, x, ple_i, cos, sin, bias): + x = x + self.post_attention_layernorm( + self.self_attn(self.input_layernorm(x), cos, sin, bias) + ) + h = self.pre_feedforward_layernorm(x) + h = self.down_proj(F.gelu(self.gate_proj(h), approximate="tanh") * self.up_proj(h)) + x = x + self.post_feedforward_layernorm(h) + g = F.gelu(self.per_layer_input_gate(x), approximate="tanh") * ple_i + x = x + self.post_per_layer_input_norm(self.per_layer_projection(g)) + return x * self.layer_scalar.to(x.dtype) + + +class EmbeddingGemma2TextEncoder(nn.Module): + """``(inputs_embeds [B,L,H], bias [1,L,1,L], pool_weights [B,L]) -> embedding [B, E]``. + + ``sliding_bias`` optionally replaces ``bias`` on sliding layers. ``pool_weights`` is the + attention mask divided by its row sum, so pooling is a masked mean. + """ + + def __init__(self, config: EmbeddingGemma2TextConfig) -> None: + super().__init__() + self.config = config + dim, eps, L = config.hidden_size, config.rms_norm_eps, config.num_hidden_layers + self.ple_dim = config.hidden_size_per_layer_input + self.per_layer_model_projection = _conv(dim, L * self.ple_dim) + self.per_layer_projection_norm = ScaledRMSNorm(self.ple_dim, eps, dim=2) + self.layers = nn.ModuleList([EmbeddingGemma2EncoderLayer(config, i) for i in range(L)]) + self.norm = ScaledRMSNorm(dim, eps) + self.embedding_projection = _conv(dim, config.embedding_dim) + self._rope_hd = { + "sliding": config.layer_head_dim(next(i for i in range(L) if config.is_sliding(i))), + "full": config.layer_head_dim(next(i for i in range(L) if not config.is_sliding(i))), + } + + def rope( + self, length: int, dtype: torch.dtype, device + ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: + out = {} + for kind, theta in ( + ("sliding", self.config.rope_theta_sliding), + ("full", self.config.rope_theta_full), + ): + hd = self._rope_hd[kind] + inv = MacEmbeddingGemma2TextEncoder._inv_freq(hd, theta, device) + fr = inv[:, None] * torch.arange(length, dtype=torch.float32, device=device)[None, :] + e = torch.cat([fr, fr], 0).reshape(1, 1, hd, length) + out[kind] = (e.cos().to(dtype), e.sin().to(dtype)) + return out + + def forward( + self, + inputs_embeds: torch.Tensor, + bias: torch.Tensor, + pool_weights: torch.Tensor, + sliding_bias: torch.Tensor | None = None, + ) -> torch.Tensor: + b, length, dim = inputs_embeds.shape + x = inputs_embeds.transpose(1, 2).unsqueeze(2) # (B, H, 1, L) + nl = self.config.num_hidden_layers + ple = self.per_layer_model_projection(x) * (dim**-0.5) + ple = self.per_layer_projection_norm(ple.reshape(b, nl, self.ple_dim, length)) + ple = ple.reshape(b, nl * self.ple_dim, 1, length) + ropes = self.rope(length, x.dtype, x.device) + sb = bias if sliding_bias is None else sliding_bias + for i, layer in enumerate(self.layers): + cos, sin = ropes["sliding" if layer.is_sliding else "full"] + ple_i = ple[:, i * self.ple_dim : (i + 1) * self.ple_dim] + x = layer(x, ple_i, cos, sin, sb if layer.is_sliding else bias) + x = self.norm(x) # (B, H, 1, L) + pooled = torch.matmul(x.squeeze(2), pool_weights.to(x.dtype).unsqueeze(-1)) # (B, H, 1) + y = self.embedding_projection(pooled.unsqueeze(-1)).reshape(b, -1) # (B, E) + m = y.abs().amax(-1, keepdim=True).clamp(min=_FP16_TINY) + ys = y * torch.reciprocal(m) + return ys * torch.rsqrt((ys * ys).sum(-1, keepdim=True)) + + +def text_encoder_from_state_dict( + config: EmbeddingGemma2TextConfig, + state_dict: dict[str, torch.Tensor], + dtype: torch.dtype = torch.float16, +) -> tuple[EmbeddingGemma2TextEncoder, torch.Tensor]: + """Build the encoder from HF text weights (``language_model.`` prefix stripped). + + Returns ``(encoder, embed_table_fp32)``; the embedding table stays on the host. + """ + sd = {k: v.float() for k, v in state_dict.items()} + table = sd.pop("embed_tokens.weight") + mapped: dict[str, torch.Tensor] = {} + for k, v in sd.items(): + k2 = k.replace("ple.per_layer_model_projection", "per_layer_model_projection") + k2 = k2.replace("ple.per_layer_projection_norm", "per_layer_projection_norm") + k2 = k2.replace("mlp.", "").replace("ple_block.", "") + mapped[k2] = v + model = EmbeddingGemma2TextEncoder(config) + own = model.state_dict() + for k, v in mapped.items(): + if k not in own: + raise RuntimeError(f"unexpected key {k}") + if own[k].dim() == 4 and v.dim() == 2: + v = v[:, :, None, None] + mapped[k] = v.reshape(own[k].shape) + missing = set(own) - set(mapped) + if missing: + raise RuntimeError(f"missing keys {sorted(missing)[:10]}") + model.load_state_dict(mapped, strict=True) + return model.to(dtype).eval(), table + + +def load_embedding_gemma2_text_encoder( + model_dir_or_id: str | Path = "google/embeddinggemma-2", + dtype: torch.dtype = torch.float16, +) -> tuple[EmbeddingGemma2TextEncoder, torch.Tensor]: + """Load from a local snapshot dir or HF id; returns ``(encoder, embed_table_fp32)``.""" + from safetensors import safe_open + + 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")): + with safe_open(str(file), framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 + if key.startswith(HF_TEXT_PREFIX): + sd[key[len(HF_TEXT_PREFIX) :]] = handle.get_tensor(key) + return text_encoder_from_state_dict(config, sd, dtype) + + +__all__ = [ + "MASK_NEG", + "ScaledRMSNorm", + "EmbeddingGemma2TextEncoder", + "load_embedding_gemma2_text_encoder", + "text_encoder_from_state_dict", +] diff --git a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py new file mode 100644 index 00000000..54c5ea5c --- /dev/null +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -0,0 +1,233 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Neural Engine float16 EmbeddingGemma 2 image encoder, channels-first (BC1S) conv layout. + +What: the vision tower as 1x1 Conv2d projections with a static patch budget and in-graph lookups. +Why: ScaledRMSNorm divides by the row max-abs before squaring, keeping float16 in range. +""" + +from __future__ import annotations + +import json +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open + +from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, ScaledRMSNorm, _conv +from coreai_models.models.macos.embedding_gemma2_vision import ( + EmbeddingGemma2VisionConfig, + axial_inv_freq, +) + + +def one_hot_f16(coord: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Exact one-hot of integer ``coord [..., 1]`` over ``grid [G]``; padding (-1) gives zeros.""" + return torch.relu(1.0 - (coord - grid).abs()) + + +def rotate_half_axial(x: torch.Tensor) -> torch.Tensor: + """Rotate-half inside each spatial half of the head dim; ``x``: ``(B, H, hd, P)``.""" + q = x.shape[2] // 4 + x0, x1, x2, x3 = x[:, :, :q], x[:, :, q : 2 * q], x[:, :, 2 * q : 3 * q], x[:, :, 3 * q :] + return torch.cat((-x1, x0, -x3, x2), dim=2) + + +class VisionAttention(nn.Module): + """Head-batched bidirectional attention, computed over query chunks of ``q_chunk`` rows.""" + + def __init__(self, cfg: EmbeddingGemma2VisionConfig, q_chunk: int = 630) -> None: + super().__init__() + self.q_chunk = q_chunk + self.n_heads, self.head_dim = cfg.num_attention_heads, cfg.head_dim + dim, hd, eps = cfg.hidden_size, cfg.head_dim, cfg.rms_norm_eps + self.q_proj = _conv(dim, self.n_heads * hd) + self.k_proj = _conv(dim, self.n_heads * hd) + self.v_proj = _conv(dim, self.n_heads * hd) + self.o_proj = _conv(self.n_heads * hd, dim) + self.q_norm = ScaledRMSNorm(hd, eps, dim=2) + self.k_norm = ScaledRMSNorm(hd, eps, dim=2) + self.v_norm = ScaledRMSNorm(hd, eps, with_scale=False, dim=2) + + def _heads(self, x: torch.Tensor) -> torch.Tensor: + b, _, _, p = x.shape + return x.reshape(b, self.n_heads, self.head_dim, p) + + def forward(self, x, cos, sin, bias): + b, _, _, p = x.shape + q = self.q_norm(self._heads(self.q_proj(x))) + k = self.k_norm(self._heads(self.k_proj(x))) + v = self.v_norm(self._heads(self.v_proj(x))) + q = q * cos + rotate_half_axial(q) * sin + k = k * cos + rotate_half_axial(k) * sin + c = self.n_heads * self.head_dim + key_bias = bias.reshape(b, 1, 1, p) + qt, vt = q.transpose(2, 3), v.transpose(2, 3) # (B, H, P, hd) + chunks = [] + for s in range(0, p, self.q_chunk): + scores = qt[:, :, s : s + self.q_chunk] @ k + key_bias # (B, H, c, P) + chunks.append(scores.softmax(-1) @ vt) # (B, H, c, hd) + out = torch.cat(chunks, 2) if len(chunks) > 1 else chunks[0] + return self.o_proj(out.transpose(2, 3).reshape(b, c, 1, p)) + + +class VisionEncoderLayer(nn.Module): + """Pre-norm block: attention and a GeGLU-tanh MLP, each with a post-norm residual.""" + + def __init__(self, cfg: EmbeddingGemma2VisionConfig) -> None: + super().__init__() + dim, eps = cfg.hidden_size, cfg.rms_norm_eps + self.self_attn = VisionAttention(cfg) + self.gate_proj = _conv(dim, cfg.intermediate_size) + self.up_proj = _conv(dim, cfg.intermediate_size) + self.down_proj = _conv(cfg.intermediate_size, dim) + self.input_layernorm = ScaledRMSNorm(dim, eps) + self.post_attention_layernorm = ScaledRMSNorm(dim, eps) + self.pre_feedforward_layernorm = ScaledRMSNorm(dim, eps) + self.post_feedforward_layernorm = ScaledRMSNorm(dim, eps) + + def forward(self, x, cos, sin, bias): + x = x + self.post_attention_layernorm( + self.self_attn(self.input_layernorm(x), cos, sin, bias) + ) + h = self.pre_feedforward_layernorm(x) + h = self.down_proj(F.gelu(self.gate_proj(h), approximate="tanh") * self.up_proj(h)) + return x + self.post_feedforward_layernorm(h) + + +class EmbeddingGemma2ImageEncoder(nn.Module): + """``pixel_values [1, P, 768]`` + ``position_ids [1, P, 2]`` (float16) -> ``[1, L, 512]``.""" + + def __init__(self, config: EmbeddingGemma2VisionConfig | None = None): + super().__init__() + self.config = cfg = config or EmbeddingGemma2VisionConfig() + k = cfg.pooling_kernel_size + self.grid = g = cfg.max_soft_tokens * k # largest reachable patch coordinate + 1 + dim = cfg.hidden_size + self.input_proj = _conv(cfg.patch_dim, dim) + self.layers = nn.ModuleList([VisionEncoderLayer(cfg) for _ in range(cfg.num_hidden_layers)]) + self.embedding_pre_projection_norm = ScaledRMSNorm(dim, cfg.rms_norm_eps, with_scale=False) + self.embedding_projection = _conv(dim, cfg.text_hidden_size) + # Per-axis lookup tables, filled by ``bake_tables`` (fp32 math, stored in model dtype): + # [G, hidden] position embedding | [G, hd/4] cos | [G, hd/4] sin. + self.register_buffer("pos_x", torch.zeros(g, dim), persistent=False) + self.register_buffer("pos_y", torch.zeros(g, dim), persistent=False) + q = cfg.head_dim // 4 + self.register_buffer("cos_tab", torch.zeros(g, q), persistent=False) + self.register_buffer("sin_tab", torch.zeros(g, q), persistent=False) + c = torch.arange(g, dtype=torch.float32) + self.register_buffer("grid_ar", c, persistent=False) + self.register_buffer("block_tab", torch.floor(c / k)[:, None], persistent=False) + self.register_buffer("width_tab", torch.floor((c + 1) / k)[:, None], persistent=False) + self.register_buffer( + "slot_ar", torch.arange(cfg.max_soft_tokens, dtype=torch.float32), persistent=False + ) + + @torch.no_grad() + def bake_tables(self, position_embedding_table: torch.Tensor) -> None: + cfg = self.config + g = self.grid + table = position_embedding_table.float() + if table.shape[1] < g: + raise ValueError(f"position table has {table.shape[1]} rows; budget needs {g}") + self.pos_x.copy_(table[0, :g]) + self.pos_y.copy_(table[1, :g]) + inv = axial_inv_freq(cfg) + ang = torch.arange(g, dtype=torch.float32)[:, None] * inv[None, :] + self.cos_tab.copy_(ang.cos()) + self.sin_tab.copy_(ang.sin()) + + def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + dt = self.input_proj.weight.dtype + b, p, _ = pixel_values.shape + pos = position_ids.to(dt) + ox = one_hot_f16(pos[..., 0:1], self.grid_ar.to(dt)) # [B, P, G] + oy = one_hot_f16(pos[..., 1:2], self.grid_ar.to(dt)) + valid = ox.sum(-1, keepdim=True) # [B, P, 1]: 1 real, 0 padding + + # Patch embedding + learned 2-D position embedding (zero for padding rows, as in HF). + x = 2.0 * (pixel_values.to(dt) - 0.5) + x = self.input_proj(x.transpose(1, 2).unsqueeze(2)) # (B, C, 1, P) + pe = ox @ self.pos_x.to(dt) + oy @ self.pos_y.to(dt) # [B, P, C] + x = x + pe.transpose(1, 2).unsqueeze(2) + + # Axial RoPE tables [fx, fx, fy, fy] -> (B, 1, hd, P); padding rows get cos = sin = 0. + cx, sx = ox @ self.cos_tab.to(dt), ox @ self.sin_tab.to(dt) + cy, sy = oy @ self.cos_tab.to(dt), oy @ self.sin_tab.to(dt) + cos = torch.cat((cx, cx, cy, cy), -1).transpose(1, 2).unsqueeze(1) + sin = torch.cat((sx, sx, sy, sy), -1).transpose(1, 2).unsqueeze(1) + + # Additive key-padding bias (1, Pk, 1, 1), broadcast over heads and queries. + bias = ((valid - 1.0) * -MASK_NEG).reshape(b, p, 1, 1) + for layer in self.layers: + x = layer(x, cos, sin, bias) + + # 3x3 pooling as one matmul with exact one-hot slot weights (sum; the scale cancels in + # the scale-free norm below, like the dropped x sqrt(hidden)). + bx = ox @ self.block_tab.to(dt) # [B, P, 1] + by = oy @ self.block_tab.to(dt) + max_x = pos[..., 0:1].amax(1, keepdim=True) # [B, 1, 1] + blocks_per_row = one_hot_f16(max_x, self.grid_ar.to(dt)) @ self.width_tab.to(dt) + slot = bx + blocks_per_row * by + w = one_hot_f16(slot, self.slot_ar.to(dt)) * valid # [B, P, L] + pooled = torch.matmul(x.squeeze(2), w) # (B, C, L) + y = self.embedding_pre_projection_norm(pooled.unsqueeze(2)) # (B, C, 1, L) + y = self.embedding_projection(y) # (B, 512, 1, L) + return y.squeeze(2).transpose(1, 2).contiguous() # [B, L, 512] + + +def load_image_encoder( + model_dir: str, + dtype: torch.dtype = torch.float16, + max_soft_tokens: int | None = None, +) -> EmbeddingGemma2ImageEncoder: + """Load from a local EG2 snapshot directory (``config.json`` + ``*.safetensors``).""" + with open(os.path.join(model_dir, "config.json")) as f: + cfg = EmbeddingGemma2VisionConfig.from_hf_config(json.load(f), max_soft_tokens) + model = EmbeddingGemma2ImageEncoder(cfg) + sd: dict[str, torch.Tensor] = {} + for fn in sorted(x for x in os.listdir(model_dir) if x.endswith(".safetensors")): + with safe_open(os.path.join(model_dir, fn), framework="pt") as h: + for key in h.keys(): # noqa: SIM118 + if key.startswith(("vision_tower.", "embed_vision.")): + sd[key] = h.get_tensor(key).float() + load_hf_state_dict(model, sd) + return model.to(dtype).eval() + + +def load_hf_state_dict( + model: EmbeddingGemma2ImageEncoder, sd: dict[str, torch.Tensor] +) -> EmbeddingGemma2ImageEncoder: + """Load checkpoint-named ``vision_tower.*`` / ``embed_vision.*`` tensors (fail closed).""" + mapped: dict[str, torch.Tensor] = {} + for key, v in sd.items(): + k2 = key.replace(".linear.weight", ".weight") + k2 = k2.removeprefix("vision_tower.").removeprefix("embed_vision.") + k2 = k2.replace("patch_embedder.input_proj", "input_proj") + k2 = k2.replace("encoder.layers.", "layers.").replace("mlp.", "") + mapped[k2] = v.float() + pos_table = mapped.pop("patch_embedder.position_embedding_table") + own = model.state_dict() + for key, v in list(mapped.items()): + if key not in own: + raise RuntimeError(f"unexpected key {key}") + mapped[key] = v.reshape(own[key].shape) + missing = set(own) - set(mapped) + if missing: + raise RuntimeError(f"missing keys {sorted(missing)[:10]}") + model.load_state_dict(mapped, strict=True) + model.bake_tables(pos_table) + return model + + +__all__ = [ + "EmbeddingGemma2ImageEncoder", + "load_hf_state_dict", + "load_image_encoder", + "one_hot_f16", +] diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py new file mode 100644 index 00000000..d7c5e8f1 --- /dev/null +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -0,0 +1,447 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""EmbeddingGemma 2 (EG2) text encoder for Core AI export. + +Bidirectional encoder with per-layer inputs and a sentence-embedding head: masked mean over the +valid tokens (prompt included), bias-free projection to 768, then L2 normalization. Masks derive +from the static sequence length and ``attention_mask``, so the graph has no data-dependent control +flow. The GPU path runs in float32; the float16 variant lives in the iOS modules. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from coreai_models.models.macos.diffusion_gemma import ( + DiffusionGemmaDenseMLP as EmbeddingGemma2MLP, +) +from coreai_models.models.macos.diffusion_gemma import ( + DiffusionGemmaRMSNorm as EmbeddingGemma2RMSNorm, +) +from coreai_models.models.macos.diffusion_gemma import ( + ScaledEmbedding, + apply_rotary_pos_emb, +) +from coreai_models.primitives.macos.sdpa import SDPA + +HF_TEXT_PREFIX = "language_model." + + +@dataclass +class EmbeddingGemma2TextConfig: + """Text-tower config parsed from the HF ``text_config`` block (needs no transformers>=5.19).""" + + vocab_size: int = 262_144 + hidden_size: int = 512 + intermediate_size: int = 2048 + num_hidden_layers: int = 24 + num_attention_heads: int = 4 + num_key_value_heads: int = 2 + head_dim: int = 256 + hidden_activation: str = "gelu_pytorch_tanh" + rms_norm_eps: float = 1e-6 + pad_token_id: int = 0 + attention_bias: bool = False + sliding_window: int = 512 + layer_types: list[str] = field(default_factory=list) + hidden_size_per_layer_input: int = 512 + embedding_dim: int = 768 + rope_theta_sliding: float = 10_000.0 + rope_theta_full: float = 1_000_000.0 + # Per-layer overrides {layer_idx: {"head_dim": int, "num_key_value_heads": int}}. + per_layer_config: dict[int, dict[str, int]] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.layer_types: + self.layer_types = [ + "sliding_attention" if (i + 1) % 6 else "full_attention" + for i in range(self.num_hidden_layers) + ] + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError( + f"layer_types has {len(self.layer_types)} entries, " + f"expected {self.num_hidden_layers}" + ) + if self.hidden_activation != "gelu_pytorch_tanh": + raise ValueError(f"Unsupported hidden_activation {self.hidden_activation!r}") + self.per_layer_config = {int(k): dict(v) for k, v in self.per_layer_config.items()} + + def layer_head_dim(self, layer_idx: int) -> int: + return int(self.per_layer_config.get(layer_idx, {}).get("head_dim", self.head_dim)) + + def layer_num_key_value_heads(self, layer_idx: int) -> int: + override = self.per_layer_config.get(layer_idx, {}) + return int(override.get("num_key_value_heads", self.num_key_value_heads)) + + def is_sliding(self, layer_idx: int) -> bool: + return self.layer_types[layer_idx] == "sliding_attention" + + @classmethod + def from_dict(cls, text_config: dict[str, Any]) -> EmbeddingGemma2TextConfig: + """Build from an HF ``text_config`` dict (``config.json["text_config"]``).""" + rope = text_config.get("rope_parameters") or {} + sliding_theta = (rope.get("sliding_attention") or {}).get("rope_theta", 10_000.0) + full_theta = (rope.get("full_attention") or {}).get("rope_theta", 1_000_000.0) + for kind in ("sliding_attention", "full_attention"): + rope_type = (rope.get(kind) or {}).get("rope_type", "default") + if rope_type != "default": + raise ValueError(f"Unsupported rope_type {rope_type!r} for {kind}") + names = {f.name for f in cls.__dataclass_fields__.values()} + kwargs = {k: v for k, v in text_config.items() if k in names and v is not None} + kwargs["rope_theta_sliding"] = float(sliding_theta) + kwargs["rope_theta_full"] = float(full_theta) + 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 [] + per_layer = { + i: { + "head_dim": text_config.get("global_head_dim", 512), + "num_key_value_heads": text_config.get("num_global_key_value_heads", 1), + } + for i, t in enumerate(types) + if t == "full_attention" + } + kwargs["per_layer_config"] = { + int(k): {kk: int(vv) for kk, vv in dict(v).items() if vv is not None} + for k, v in per_layer.items() + } + return cls(**kwargs) + + @classmethod + def from_pretrained(cls, model_dir: str | Path) -> EmbeddingGemma2TextConfig: + cfg = json.loads((Path(model_dir) / "config.json").read_text()) + return cls.from_dict(cfg.get("text_config", cfg)) + + +def build_attention_masks( + attention_mask: torch.Tensor, sliding_window: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Return boolean ``(full, sliding)`` masks of shape ``[B, 1, L, L]`` (True = attend). + + ``full`` = key_padding + ``sliding`` = ``|i - j| <= sliding_window`` AND key_padding + + Matches ``create_bidirectional_mask`` / ``create_bidirectional_sliding_window_mask`` in + transformers (``masking_utils.sliding_window_bidirectional_overlay`` + + ``padding_mask_function``). Rows of padded queries may be fully masked; the SDPA composite + converts the boolean mask into a finite additive bias, so those rows stay finite and they are + excluded from pooling anyway. + """ + batch, seq_len = attention_mask.shape + key_valid = (attention_mask != 0).reshape(batch, 1, 1, seq_len) + positions = torch.arange(seq_len, device=attention_mask.device, dtype=torch.int32) + distance = (positions.unsqueeze(1) - positions.unsqueeze(0)).abs() + band = (distance <= sliding_window).reshape(1, 1, seq_len, seq_len) + full = key_valid.expand(batch, 1, seq_len, seq_len) + sliding = torch.logical_and(band, key_valid) + return full, sliding + + +class EmbeddingGemma2Attention(nn.Module): + """Bidirectional attention; q/k RMSNorm + RoPE, scale-free v_norm, scaling 1.0.""" + + def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: + super().__init__() + self.n_heads = config.num_attention_heads + self.n_kv_heads = config.layer_num_key_value_heads(layer_idx) + self.head_dim = config.layer_head_dim(layer_idx) + dim, hd = config.hidden_size, self.head_dim + bias = config.attention_bias + + self.q_proj = nn.Linear(dim, self.n_heads * hd, bias=bias) + self.k_proj = nn.Linear(dim, self.n_kv_heads * hd, bias=bias) + self.v_proj = nn.Linear(dim, self.n_kv_heads * hd, bias=bias) + self.o_proj = nn.Linear(self.n_heads * hd, dim, bias=bias) + self.q_norm = EmbeddingGemma2RMSNorm(hd, eps=config.rms_norm_eps) + self.k_norm = EmbeddingGemma2RMSNorm(hd, eps=config.rms_norm_eps) + self.v_norm = EmbeddingGemma2RMSNorm(hd, eps=config.rms_norm_eps, with_scale=False) + # The mask carries both the window and the key padding; the composite itself stays + # windowless and non-causal. + self.sdpa = SDPA(scale=1.0, window_size=0, is_causal=False) + + def forward( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + attn_mask: torch.Tensor, + ) -> torch.Tensor: + batch, seq_len, _ = x.shape + n_h, n_kv, hd = self.n_heads, self.n_kv_heads, self.head_dim + + query = self.q_norm(self.q_proj(x).reshape(batch, seq_len, n_h, hd)) + query = apply_rotary_pos_emb(query, cos, sin, unsqueeze_dim=2).transpose(1, 2) + key = self.k_norm(self.k_proj(x).reshape(batch, seq_len, n_kv, hd)) + key = apply_rotary_pos_emb(key, cos, sin, unsqueeze_dim=2).transpose(1, 2) + value = self.v_norm(self.v_proj(x).reshape(batch, seq_len, n_kv, hd)).transpose(1, 2) + + out = self.sdpa(query=query, key=key, value=value, attn_mask=attn_mask) + out = out.transpose(1, 2).reshape(batch, seq_len, n_h * hd) + return self.o_proj(out) + + +class EmbeddingGemma2PLE(nn.Module): + """Projection-only per-layer embeddings: inputs_embeds -> [B, L, num_layers, ple_dim].""" + + def __init__(self, config: EmbeddingGemma2TextConfig) -> None: + super().__init__() + self.num_layers = config.num_hidden_layers + self.ple_dim = config.hidden_size_per_layer_input + self.per_layer_model_projection = nn.Linear( + config.hidden_size, self.num_layers * self.ple_dim, bias=False + ) + self.scale = config.hidden_size**-0.5 + self.per_layer_projection_norm = EmbeddingGemma2RMSNorm( + self.ple_dim, eps=config.rms_norm_eps + ) + + def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: + batch, seq_len, _ = inputs_embeds.shape + proj = self.per_layer_model_projection(inputs_embeds) * self.scale + proj = proj.reshape(batch, seq_len, self.num_layers, self.ple_dim) + return self.per_layer_projection_norm(proj) + + +class EmbeddingGemma2PLEBlock(nn.Module): + """Third residual sub-block: gate -> GELU -> * per_layer_input -> project -> norm -> add.""" + + def __init__(self, config: EmbeddingGemma2TextConfig) -> None: + super().__init__() + ple_dim = config.hidden_size_per_layer_input + self.per_layer_input_gate = nn.Linear(config.hidden_size, ple_dim, bias=False) + self.per_layer_projection = nn.Linear(ple_dim, config.hidden_size, bias=False) + self.post_per_layer_input_norm = EmbeddingGemma2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward(self, x: torch.Tensor, per_layer_input: torch.Tensor) -> torch.Tensor: + h = F.gelu(self.per_layer_input_gate(x), approximate="tanh") * per_layer_input + return x + self.post_per_layer_input_norm(self.per_layer_projection(h)) + + +class EmbeddingGemma2EncoderLayer(nn.Module): + def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: + super().__init__() + eps = config.rms_norm_eps + self.is_sliding = config.is_sliding(layer_idx) + self.self_attn = EmbeddingGemma2Attention(config, layer_idx) + self.mlp = EmbeddingGemma2MLP(config) + self.input_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) + self.post_attention_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) + self.pre_feedforward_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) + self.post_feedforward_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) + self.ple_block = EmbeddingGemma2PLEBlock(config) + self.register_buffer("layer_scalar", torch.ones(1)) + + def forward( + self, + x: torch.Tensor, + per_layer_input: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + attn_mask: torch.Tensor, + ) -> torch.Tensor: + h = self.self_attn(self.input_layernorm(x), cos, sin, attn_mask) + x = x + self.post_attention_layernorm(h) + h = self.mlp(self.pre_feedforward_layernorm(x)) + x = x + self.post_feedforward_layernorm(h) + x = self.ple_block(x, per_layer_input) + return x * self.layer_scalar + + +class EmbeddingGemma2TextEncoder(nn.Module): + """EG2 text tower + sentence-embedding head (masked mean -> projection -> L2). + + Parameter names mirror the HF ``language_model.*`` namespace (minus the prefix), so a + checkpoint loads with a prefix strip and no renames. + """ + + def __init__(self, config: EmbeddingGemma2TextConfig) -> None: + super().__init__() + self.config = config + self.embed_tokens = ScaledEmbedding( + config.vocab_size, config.hidden_size, embed_scale=config.hidden_size**0.5 + ) + self.ple = EmbeddingGemma2PLE(config) + self.layers = nn.ModuleList( + [EmbeddingGemma2EncoderLayer(config, i) for i in range(config.num_hidden_layers)] + ) + self.norm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.embedding_projection = nn.Linear(config.hidden_size, config.embedding_dim, bias=False) + + sliding_dims = { + config.layer_head_dim(i) for i in range(len(self.layers)) if config.is_sliding(i) + } + full_dims = { + config.layer_head_dim(i) for i in range(len(self.layers)) if not config.is_sliding(i) + } + if len(sliding_dims) > 1 or len(full_dims) > 1: + raise ValueError("RoPE tables assume one head_dim per layer type") + sliding_hd = next(iter(sliding_dims), config.head_dim) + full_hd = next(iter(full_dims), config.head_dim) + self.register_buffer( + "sliding_inv_freq", + self._inv_freq(sliding_hd, config.rope_theta_sliding), + persistent=False, + ) + self.register_buffer( + "full_inv_freq", self._inv_freq(full_hd, config.rope_theta_full), persistent=False + ) + + @staticmethod + def _inv_freq(head_dim: int, base: float, device: torch.device | None = None) -> torch.Tensor: + return 1.0 / ( + base ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim) + ) + + @staticmethod + def _cos_sin( + inv_freq: torch.Tensor, seq_len: int, dtype: torch.dtype, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: + """Half-dim cos/sin ``[1, L, head_dim/2]`` for positions ``0..L-1`` (fp32 math).""" + positions = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = positions.unsqueeze(1) * inv_freq.float().to(device).unsqueeze(0) + freqs = freqs.unsqueeze(0) + return freqs.cos().to(dtype), freqs.sin().to(dtype) + + def encode_tokens( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the tower; returns final-normed hidden states ``[B, L, hidden_size]``.""" + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Pass exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + batch, seq_len, _ = inputs_embeds.shape + if attention_mask is None: + attention_mask = torch.ones( + batch, seq_len, dtype=torch.int32, device=inputs_embeds.device + ) + + per_layer_inputs = self.ple(inputs_embeds) + full_mask, sliding_mask = build_attention_masks(attention_mask, self.config.sliding_window) + dtype, device = inputs_embeds.dtype, inputs_embeds.device + sliding_cs = self._cos_sin(self.sliding_inv_freq, seq_len, dtype, device) + full_cs = self._cos_sin(self.full_inv_freq, seq_len, dtype, device) + + h = inputs_embeds + for i, layer in enumerate(self.layers): + cos, sin = sliding_cs if layer.is_sliding else full_cs + mask = sliding_mask if layer.is_sliding else full_mask + h = layer(h, per_layer_inputs[:, :, i, :], cos, sin, mask) + return self.norm(h) + + def pool(self, hidden: torch.Tensor, pooling_mask: torch.Tensor) -> torch.Tensor: + """Masked mean (fp32) -> projection -> L2 normalize. Returns ``[B, embedding_dim]``.""" + mask = pooling_mask.to(torch.float32).unsqueeze(-1) + summed = (hidden.to(torch.float32) * mask).sum(dim=1) + count = mask.sum(dim=1).clamp(min=1e-9) + mean = summed / count + projected = F.linear(mean, self.embedding_projection.weight.to(torch.float32)) + sq_norm = (projected * projected).sum(dim=-1, keepdim=True) + return projected * torch.rsqrt(sq_norm.clamp(min=1e-24)) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + pooling_mask: torch.Tensor | None = None, + return_token_states: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Return ``sentence_embedding`` ``[B, embedding_dim]`` (unit L2 norm, fp32). + + ``pooling_mask`` defaults to ``attention_mask`` (include_prompt=True); pass a narrower + mask to exclude e.g. prompt tokens from the mean. With ``return_token_states`` the + per-token projected states ``[B, L, embedding_dim]`` (HF ``last_hidden_state`` + equivalent) are returned as a second output for parity checks. + """ + hidden = self.encode_tokens(input_ids, attention_mask, inputs_embeds) + if attention_mask is None: + attention_mask = torch.ones(hidden.shape[:2], dtype=torch.int32, device=hidden.device) + if pooling_mask is None: + pooling_mask = attention_mask + sentence_embedding = self.pool(hidden, pooling_mask) + if return_token_states: + return sentence_embedding, self.embedding_projection(hidden) + return sentence_embedding + + +def extract_text_state_dict( + state_dict: dict[str, torch.Tensor], prefix: str = HF_TEXT_PREFIX +) -> dict[str, torch.Tensor]: + """Keep only ``prefix*`` keys (vision/audio dropped) and strip the prefix.""" + return {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} + + +def load_text_state_dict( + model: EmbeddingGemma2TextEncoder, state_dict: dict[str, torch.Tensor] +) -> None: + """Load an (already prefix-stripped) text state dict; fail closed on any missing/extra key.""" + info = model.load_state_dict(state_dict, strict=False) + if info.missing_keys or info.unexpected_keys: + raise RuntimeError( + f"EG2 text load mismatch: missing={info.missing_keys[:10]} " + f"unexpected={info.unexpected_keys[:10]}" + ) + + +def _resolve_model_dir(model_dir_or_id: str | Path) -> Path: + path = Path(model_dir_or_id).expanduser() + if path.is_dir(): + return path + from huggingface_hub import snapshot_download + + return Path( + snapshot_download( + str(model_dir_or_id), allow_patterns=["config.json", "*.safetensors", "*.json"] + ) + ) + + +def load_embedding_gemma2_text_encoder( + model_dir_or_id: str | Path = "google/embeddinggemma-2", + dtype: torch.dtype = torch.float32, +) -> EmbeddingGemma2TextEncoder: + """Build the text encoder from a local snapshot dir or HF id, loading ``language_model.*``.""" + from safetensors import safe_open + + model_dir = _resolve_model_dir(model_dir_or_id) + config = EmbeddingGemma2TextConfig.from_pretrained(model_dir) + model = EmbeddingGemma2TextEncoder(config) + + state_dict: dict[str, torch.Tensor] = {} + files = sorted(model_dir.glob("*.safetensors")) + if not files: + raise FileNotFoundError(f"No .safetensors in {model_dir}") + for file in files: + with safe_open(str(file), framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + if key.startswith(HF_TEXT_PREFIX): + state_dict[key[len(HF_TEXT_PREFIX) :]] = handle.get_tensor(key).to(dtype) + load_text_state_dict(model, state_dict) + return model.to(dtype).eval() + + +__all__ = [ + "EmbeddingGemma2TextConfig", + "EmbeddingGemma2TextEncoder", + "build_attention_masks", + "extract_text_state_dict", + "load_embedding_gemma2_text_encoder", + "load_text_state_dict", +] diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py new file mode 100644 index 00000000..c3942583 --- /dev/null +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -0,0 +1,504 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""EmbeddingGemma 2 image encoder for Core AI export (vision tower + ``embed_vision``). + +Host-prepared patches ``[B, P, 768]`` and ``position_ids [B, P, 2]`` (-1 for padding) become soft +tokens ``[B, L, 512]`` with P and L static. Padding and the 3x3 pooling slots are derived in-graph +from ``position_ids``, with no data-dependent control flow. The GPU path runs in float32; the +float16 variant lives in the iOS modules. +""" + +from __future__ import annotations + +import json +import math +import os +from collections.abc import Callable, Sequence +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open + +from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaDenseMLP +from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaRMSNorm as RMSNorm +from coreai_models.primitives.macos.rope import rotate_half +from coreai_models.primitives.macos.sdpa import SDPA + +# Special token ids from the EG2 tokenizer / config.json. +EG2_BOS_TOKEN_ID = 2 +EG2_EOS_TOKEN_ID = 1 +EG2_PAD_TOKEN_ID = 0 +EG2_BOI_TOKEN_ID = 255999 +EG2_EOI_TOKEN_ID = 258882 +EG2_IMAGE_TOKEN_ID = 258880 + + +@dataclass +class EmbeddingGemma2VisionConfig: + """Vision tower + ``embed_vision`` hyper-parameters (defaults = google/embeddinggemma-2).""" + + hidden_size: int = 768 + intermediate_size: int = 3072 + num_hidden_layers: int = 16 + num_attention_heads: int = 12 + num_key_value_heads: int = 12 + head_dim: int = 64 + patch_size: int = 16 + pooling_kernel_size: int = 3 + position_embedding_size: int = 10240 + rms_norm_eps: float = 1e-6 + rope_theta: float = 100.0 + text_hidden_size: int = 512 + max_soft_tokens: int = 280 + + @property + def max_patches(self) -> int: + return self.max_soft_tokens * self.pooling_kernel_size**2 + + @property + def patch_dim(self) -> int: + return 3 * self.patch_size**2 + + @classmethod + def from_hf_config( + cls, config: dict, max_soft_tokens: int | None = None + ) -> EmbeddingGemma2VisionConfig: + """Build from an EG2 ``config.json`` dict (top level, with ``vision_config``).""" + vc = config["vision_config"] + if vc.get("use_clipped_linears", False): + raise NotImplementedError("clipped vision linears are not supported") + if vc.get("standardize", False): + raise NotImplementedError( + "vision standardization (std_bias/std_scale) is not supported" + ) + rope = vc.get("rope_parameters") or {} + if rope.get("rope_type", "axial") != "axial": + raise NotImplementedError(f"unsupported vision rope type {rope.get('rope_type')}") + return cls( + hidden_size=vc["hidden_size"], + intermediate_size=vc["intermediate_size"], + num_hidden_layers=vc["num_hidden_layers"], + num_attention_heads=vc["num_attention_heads"], + num_key_value_heads=vc["num_key_value_heads"], + head_dim=vc["head_dim"], + patch_size=vc["patch_size"], + pooling_kernel_size=vc["pooling_kernel_size"], + position_embedding_size=vc["position_embedding_size"], + rms_norm_eps=vc["rms_norm_eps"], + rope_theta=rope.get("rope_theta", 100.0), + text_hidden_size=config["text_config"]["hidden_size"], + max_soft_tokens=max_soft_tokens or config.get("vision_soft_tokens_per_image", 280), + ) + + +def axial_inv_freq(cfg: EmbeddingGemma2VisionConfig) -> torch.Tensor: + """Per-axis RoPE inverse frequencies (``head_dim // 2`` rotary dims per spatial axis).""" + spatial_dim = cfg.head_dim // 2 + return 1.0 / ( + cfg.rope_theta ** (torch.arange(0, spatial_dim, 2, dtype=torch.float32) / spatial_dim) + ) + + +def apply_axial_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Axial 2-D RoPE: the first half of ``head_dim`` rotates by x, the second half by y. + + Args: + x: ``[B, L, H, D]``. + cos, sin: ``[B, L, D]`` laid out ``[fx, fx, fy, fy]`` (see ``axial_rope_cos_sin``). + """ + half = x.shape[-1] // 2 + cos = cos.unsqueeze(2) + sin = sin.unsqueeze(2) + parts = [] + for lo, hi in ((0, half), (half, 2 * half)): + xp, cp, sp = x[..., lo:hi], cos[..., lo:hi], sin[..., lo:hi] + parts.append(xp * cp + rotate_half(xp) * sp) + return torch.cat(parts, dim=-1) + + +class VisionAttention(nn.Module): + """Bidirectional MHA with q/k RMSNorm, scale-free v RMSNorm and axial RoPE; scale 1.0.""" + + def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: + super().__init__() + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + h = config.hidden_size + self.q_proj = nn.Linear(h, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, h, bias=False) + self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) + self.v_norm = RMSNorm(self.head_dim, config.rms_norm_eps, with_scale=False) + self.sdpa = SDPA(scale=1.0, is_causal=False) + + def forward( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + attn_mask: torch.Tensor, + ) -> torch.Tensor: + b, n, _ = x.shape + q = self.q_norm(self.q_proj(x).view(b, n, self.num_heads, self.head_dim)) + k = self.k_norm(self.k_proj(x).view(b, n, self.num_kv_heads, self.head_dim)) + v = self.v_norm(self.v_proj(x).view(b, n, self.num_kv_heads, self.head_dim)) + q = apply_axial_rope(q, cos, sin).transpose(1, 2) + k = apply_axial_rope(k, cos, sin).transpose(1, 2) + v = v.transpose(1, 2) + if self.num_kv_heads != self.num_heads: + rep = self.num_heads // self.num_kv_heads + k = k.repeat_interleave(rep, dim=1) + v = v.repeat_interleave(rep, dim=1) + out = self.sdpa(q, k, v, attn_mask=attn_mask) + return self.o_proj(out.transpose(1, 2).reshape(b, n, self.num_heads * self.head_dim)) + + +class VisionEncoderLayer(nn.Module): + """Sandwich-norm block: ``x + post_norm(attn(pre_norm(x)))``, same for the MLP.""" + + def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: + super().__init__() + h, eps = config.hidden_size, config.rms_norm_eps + self.self_attn = VisionAttention(config) + self.mlp = DiffusionGemmaDenseMLP(config) + self.input_layernorm = RMSNorm(h, eps) + self.post_attention_layernorm = RMSNorm(h, eps) + self.pre_feedforward_layernorm = RMSNorm(h, eps) + self.post_feedforward_layernorm = RMSNorm(h, eps) + + def forward( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + attn_mask: torch.Tensor, + ) -> torch.Tensor: + h = self.self_attn(self.input_layernorm(x), cos, sin, attn_mask) + x = x + self.post_attention_layernorm(h) + h = self.mlp(self.pre_feedforward_layernorm(x)) + return x + self.post_feedforward_layernorm(h) + + +class VisionPatchEmbedder(nn.Module): + """``Linear(2*(p-0.5))`` plus separable learned x/y position embeddings.""" + + def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: + super().__init__() + self.input_proj = nn.Linear(config.patch_dim, config.hidden_size, bias=False) + self.position_embedding_table = nn.Parameter( + torch.ones(2, config.position_embedding_size, config.hidden_size) + ) + + def forward( + self, pixel_values: torch.Tensor, position_ids: torch.Tensor, padding: torch.Tensor + ) -> torch.Tensor: + x = 2.0 * (pixel_values - 0.5) + h = self.input_proj(x.to(self.input_proj.weight.dtype)) + clamped = position_ids.clamp(min=0) + pos = F.embedding(clamped[..., 0], self.position_embedding_table[0]) + F.embedding( + clamped[..., 1], self.position_embedding_table[1] + ) + pos = torch.where(padding.unsqueeze(-1), torch.zeros_like(pos), pos) + return h + pos + + +class VisionEncoder(nn.Module): + """Holds the layer stack (module path ``encoder.layers.N`` mirrors the checkpoint).""" + + def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: + super().__init__() + self.layers = nn.ModuleList( + [VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + + +class MultimodalEmbedder(nn.Module): + """``embed_vision``: scale-free RMSNorm (fp32) then ``Linear(hidden -> text_hidden)``.""" + + def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: + super().__init__() + self.embedding_pre_projection_norm = RMSNorm( + config.hidden_size, config.rms_norm_eps, with_scale=False + ) + self.embedding_projection = nn.Linear( + config.hidden_size, config.text_hidden_size, bias=False + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + normed = self.embedding_pre_projection_norm(x.float()) + w = self.embedding_projection.weight + return F.linear(normed.to(w.dtype), w) + + +class EmbeddingGemma2ImageEncoder(nn.Module): + """gemma4_vision tower + ``embed_vision``: patches -> soft tokens ``[B, L, text_hidden]``.""" + + def __init__(self, config: EmbeddingGemma2VisionConfig | None = None) -> None: + super().__init__() + self.config = config or EmbeddingGemma2VisionConfig() + cfg = self.config + if cfg.head_dim % 4 != 0: + raise ValueError("axial RoPE needs head_dim divisible by 4") + self.patch_embedder = VisionPatchEmbedder(cfg) + self.encoder = VisionEncoder(cfg) + self.embed_vision = MultimodalEmbedder(cfg) + self.register_buffer("rope_inv_freq", axial_inv_freq(cfg), persistent=False) + self.register_buffer( + "pool_slots", torch.arange(cfg.max_soft_tokens, dtype=torch.int64), persistent=False + ) + + # -- pieces (public so tests / activation probes can hook them) --------------------- + + def axial_rope_cos_sin( + self, position_ids: torch.Tensor, dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + """``[B, P, 2]`` positions -> cos/sin ``[B, P, head_dim]`` laid out ``[fx, fx, fy, fy]``. + + Padding rows keep their raw ``-1`` coordinates, as in HF; they are masked as keys + and dropped by the pooler, so their values never reach a real soft token. + """ + freqs = position_ids.unsqueeze(-1).float() * self.rope_inv_freq # [B, P, 2, D/4] + fx, fy = freqs[..., 0, :], freqs[..., 1, :] + angles = torch.cat([fx, fx, fy, fy], dim=-1) + return angles.cos().to(dtype), angles.sin().to(dtype) + + def pool(self, hidden: torch.Tensor, position_ids: torch.Tensor, padding: torch.Tensor): + """Mean over each ``k x k`` patch block, then ``x sqrt(hidden)``; fp32 output. + + Block index = ``x//k + (max_x+1)//k * (y//k)``. With the host contract (a full + ``(W/k) x (H/k)`` block grid, row-major, padded at the end) the real soft tokens are + exactly the first ``num_soft_tokens`` slots. + """ + k = self.config.pooling_kernel_size + h = hidden.float().masked_fill(padding.unsqueeze(-1), 0.0) + clamped = position_ids.clamp(min=0) + max_x = clamped[..., 0].amax(dim=-1, keepdim=True) + 1 + block = torch.div(clamped, k, rounding_mode="floor") + slot = block[..., 0] + torch.div(max_x, k, rounding_mode="floor") * block[..., 1] + weights = (slot.unsqueeze(-1) == self.pool_slots).float() / float(k * k) # [B, P, L] + pooled = torch.matmul(weights.transpose(1, 2), h) # [B, L, hidden] + return pooled * math.sqrt(self.config.hidden_size) + + def encode_pooled(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + """Vision tower output before ``embed_vision``: ``[B, L, hidden]`` (fp32).""" + cfg = self.config + if pixel_values.shape[-2] != cfg.max_patches: + raise ValueError( + f"expected {cfg.max_patches} patches (budget {cfg.max_soft_tokens}), " + f"got {pixel_values.shape[-2]}" + ) + padding = position_ids[..., 0] < 0 + x = self.patch_embedder(pixel_values, position_ids, padding) + cos, sin = self.axial_rope_cos_sin(position_ids, x.dtype) + # Boolean key mask (True = attend), the convention of the repo SDPA primitive: its + # Core AI lowering turns ``~mask`` into a -1e4 additive bias. Do not pass an additive + # float mask here; that path would read it as a bool and invert it. + attn_mask = torch.logical_not(padding)[:, None, None, :] + for layer in self.encoder.layers: + x = layer(x, cos, sin, attn_mask) + return self.pool(x, position_ids, padding) + + def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + """``[B, P, patch_dim]`` + ``[B, P, 2]`` -> soft tokens ``[B, L, text_hidden]``.""" + return self.embed_vision(self.encode_pooled(pixel_values, position_ids)) + + # -- loading --------------------------------------------------------------------------- + + @staticmethod + def map_hf_key(key: str) -> str | None: + """Checkpoint key -> module key (``None`` for keys outside the image path).""" + if key.startswith("vision_tower."): + local = key.removeprefix("vision_tower.") + return local.replace(".linear.weight", ".weight") + if key.startswith("embed_vision."): + return key + return None + + @classmethod + def from_pretrained( + cls, + model_dir: str, + dtype: torch.dtype = torch.float32, + max_soft_tokens: int | None = None, + ) -> EmbeddingGemma2ImageEncoder: + """Load from a local EG2 snapshot directory (``config.json`` + ``model.safetensors``).""" + with open(os.path.join(model_dir, "config.json")) as f: + config = EmbeddingGemma2VisionConfig.from_hf_config(json.load(f), max_soft_tokens) + model = cls(config) + state: dict[str, torch.Tensor] = {} + files = sorted(fn for fn in os.listdir(model_dir) if fn.endswith(".safetensors")) + for fn in files: + with safe_open(os.path.join(model_dir, fn), framework="pt", device="cpu") as f: + for key in f.keys(): # noqa: SIM118 + local = cls.map_hf_key(key) + if local is not None: + state[local] = f.get_tensor(key).to(dtype) + missing, unexpected = model.load_state_dict(state, strict=False) + if missing or unexpected: + raise RuntimeError(f"weight mismatch: missing={missing} unexpected={unexpected}") + return model.to(dtype).eval() + + +def aspect_ratio_preserving_size( + height: int, + width: int, + patch_size: int = 16, + max_soft_tokens: int = 280, + pooling_kernel_size: int = 3, +) -> tuple[int, int]: + """Target ``(height, width)``: the largest size with sides divisible by ``k*patch`` and at + most ``max_soft_tokens * k^2`` patches, preserving aspect ratio (Gemma4ImageProcessor).""" + max_patches = max_soft_tokens * pooling_kernel_size**2 + target_px = max_patches * patch_size**2 + factor = math.sqrt(target_px / (height * width)) + side = pooling_kernel_size * patch_size + th = int(math.floor(factor * height / side)) * side + tw = int(math.floor(factor * width / side)) * side + if th == 0 and tw == 0: + raise ValueError(f"cannot resize {height}x{width} to a non-empty grid") + max_side = (max_patches // pooling_kernel_size**2) * side + if th == 0: + th = side + tw = min(int(math.floor(width / height)) * side, max_side) + elif tw == 0: + tw = side + th = min(int(math.floor(height / width)) * side, max_side) + if th * tw > target_px: + raise ValueError(f"resize {height}x{width} -> {th}x{tw} exceeds {max_patches} patches") + return th, tw + + +def patchify(image: torch.Tensor, patch_size: int = 16) -> torch.Tensor: + """``[3, H, W]`` -> ``[(H/p)*(W/p), p*p*3]``; patches row-major, pixels ``(py, px, c)``.""" + c, h, w = image.shape + ph, pw = h // patch_size, w // patch_size + x = image.reshape(c, ph, patch_size, pw, patch_size).permute(1, 3, 2, 4, 0) + return x.reshape(ph * pw, patch_size * patch_size * c) + + +def patch_position_ids(patch_height: int, patch_width: int) -> torch.Tensor: + """Row-major ``(x, y)`` coordinates ``[ph*pw, 2]``: ``(0,0), (1,0), ..., (pw-1, ph-1)``.""" + ys, xs = torch.meshgrid(torch.arange(patch_height), torch.arange(patch_width), indexing="ij") + return torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1) + + +def pad_to_budget( + patches: torch.Tensor, positions: torch.Tensor, max_patches: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Zero-pad patches and ``-1``-pad positions along dim 0 to ``max_patches``.""" + n = patches.shape[0] + if n > max_patches: + raise ValueError(f"{n} patches exceed the budget of {max_patches}") + pad = max_patches - n + patches = F.pad(patches, (0, 0, 0, pad), value=0.0) + positions = F.pad(positions, (0, 0, 0, pad), value=-1) + return patches, positions + + +def preprocess_resized_image( + image: torch.Tensor, + patch_size: int = 16, + max_soft_tokens: int = 280, + pooling_kernel_size: int = 3, + rescale_factor: float = 1.0 / 255.0, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Already-resized ``[3, H, W]`` RGB (uint8 or 0..255 float) -> model inputs. + + Returns ``(pixel_values [1, P, patch_dim] fp32, position_ids [1, P, 2] int64, + num_soft_tokens)``. Rescale only (no mean/std normalization: the tower does + ``2*(x-0.5)`` itself). + """ + _, h, w = image.shape + side = patch_size * pooling_kernel_size + if h % side or w % side: + raise ValueError(f"resized image {h}x{w} must have sides divisible by {side}") + x = image.to(torch.float32) * rescale_factor + patches = patchify(x, patch_size) + positions = patch_position_ids(h // patch_size, w // patch_size) + num_soft_tokens = patches.shape[0] // pooling_kernel_size**2 + patches, positions = pad_to_budget(patches, positions, max_soft_tokens * pooling_kernel_size**2) + return patches.unsqueeze(0), positions.unsqueeze(0), num_soft_tokens + + +def build_image_input_ids( + num_soft_tokens: int, + *, + bos_token_id: int | None = EG2_BOS_TOKEN_ID, + eos_token_id: int | None = EG2_EOS_TOKEN_ID, + boi_token_id: int = EG2_BOI_TOKEN_ID, + eoi_token_id: int = EG2_EOI_TOKEN_ID, + image_token_id: int = EG2_IMAGE_TOKEN_ID, +) -> list[int]: + """Ids for an image-only input, as the HF processor emits them: + `` * num_soft_tokens ``.""" + ids = [boi_token_id] + [image_token_id] * num_soft_tokens + [eoi_token_id] + if bos_token_id is not None: + ids = [bos_token_id] + ids + if eos_token_id is not None: + ids = ids + [eos_token_id] + return ids + + +def scatter_soft_tokens( + inputs_embeds: torch.Tensor, placeholder_mask: torch.Tensor, soft_tokens: torch.Tensor +) -> torch.Tensor: + """Replace placeholder rows with soft tokens, in order (``masked_scatter`` semantics). + + Gather-based so it exports at a static length: the ``j``-th placeholder in row ``b`` + takes ``soft_tokens[b, j]``. ``soft_tokens`` may carry extra trailing rows (e.g. the + zero padding slots of a fixed budget); they are never read. + + Args: + inputs_embeds: ``[B, S, D]`` text embeddings. + placeholder_mask: ``[B, S]`` bool. + soft_tokens: ``[B, T, D]`` with ``T >=`` placeholders per row. + """ + idx = torch.cumsum(placeholder_mask.to(torch.int64), dim=1) - 1 + idx = idx.clamp(min=0, max=soft_tokens.shape[1] - 1) + gathered = torch.gather( + soft_tokens, 1, idx.unsqueeze(-1).expand(-1, -1, soft_tokens.shape[-1]) + ).to(inputs_embeds.dtype) + return torch.where(placeholder_mask.unsqueeze(-1), gathered, inputs_embeds) + + +def build_image_inputs_embeds( + embed_tokens: Callable[[torch.Tensor], torch.Tensor], + soft_tokens: torch.Tensor, + num_soft_tokens: int | Sequence[int], + *, + image_token_id: int = EG2_IMAGE_TOKEN_ID, + pad_token_id: int = EG2_PAD_TOKEN_ID, + **id_kwargs, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Full text-encoder input for one image per row. + + ``embed_tokens`` maps ids -> *scaled* text embeddings (EG2 multiplies the table by + ``sqrt(hidden)``; HF ``language_model.embed_tokens`` already does). Placeholders are + looked up as ```` (like HF) and then overwritten by the soft tokens. + + Returns ``(inputs_embeds [B, S, D], attention_mask [B, S] int64, input_ids [B, S])``, + right-padded to the longest row. + """ + counts = [num_soft_tokens] if isinstance(num_soft_tokens, int) else list(num_soft_tokens) + if len(counts) != soft_tokens.shape[0]: + raise ValueError("one num_soft_tokens per soft_tokens row") + rows = [build_image_input_ids(n, image_token_id=image_token_id, **id_kwargs) for n in counts] + s = max(len(r) for r in rows) + input_ids = torch.full((len(rows), s), pad_token_id, dtype=torch.long) + attention_mask = torch.zeros((len(rows), s), dtype=torch.long) + for i, r in enumerate(rows): + input_ids[i, : len(r)] = torch.tensor(r) + attention_mask[i, : len(r)] = 1 + placeholder = input_ids == image_token_id + lookup = torch.where(placeholder, torch.full_like(input_ids, pad_token_id), input_ids) + embeds = embed_tokens(lookup.to(soft_tokens.device)) + embeds = scatter_soft_tokens(embeds, placeholder.to(embeds.device), soft_tokens) + return embeds, attention_mask, input_ids diff --git a/python/tests/test_model_units/test_models/_embedding_gemma2_common.py b/python/tests/test_model_units/test_models/_embedding_gemma2_common.py new file mode 100644 index 00000000..38630a95 --- /dev/null +++ b/python/tests/test_model_units/test_models/_embedding_gemma2_common.py @@ -0,0 +1,80 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tiny EmbeddingGemma 2 configs and helpers shared by the macOS and iOS test modules.""" + +import torch +import torch.nn.functional as F + +from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2VisionConfig + +TINY_TEXT_CONFIG = { + "vocab_size": 128, + "hidden_size": 64, + "intermediate_size": 96, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "hidden_activation": "gelu_pytorch_tanh", + "rms_norm_eps": 1e-6, + "pad_token_id": 0, + "sliding_window": 3, + "layer_types": ["sliding_attention", "full_attention", "sliding_attention", "full_attention"], + "per_layer_config": { + "01": {"head_dim": 32, "num_key_value_heads": 1}, + "03": {"head_dim": 32, "num_key_value_heads": 1}, + }, + "rope_parameters": { + "full_attention": {"rope_theta": 1_000_000.0, "rope_type": "default"}, + "sliding_attention": {"rope_theta": 10_000.0, "rope_type": "default"}, + }, + "hidden_size_per_layer_input": 24, + "embedding_dim": 48, + "max_position_embeddings": 512, +} + +TINY_VISION_CONFIG = EmbeddingGemma2VisionConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + patch_size=4, + pooling_kernel_size=3, + position_embedding_size=64, + rms_norm_eps=1e-6, + rope_theta=100.0, + text_hidden_size=32, + max_soft_tokens=12, +) + + +def _cos(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return F.cosine_similarity(a.double(), b.double(), dim=-1) + + +def _st_mean_pool_normalize( + token_embeddings: torch.Tensor, attention_mask: torch.Tensor +) -> torch.Tensor: + """sentence-transformers Pooling(mean, include_prompt=True) + Normalize, re-implemented. + + Source: sentence_transformers 6.1.0 ``sentence_transformer/modules/pooling.py`` (mean branch of + ``_forward_padded``) and ``base/modules/normalize.py``. + """ + mask = attention_mask.unsqueeze(-1).expand_as(token_embeddings).to(token_embeddings.dtype) + mean = (token_embeddings * mask).sum(dim=1) / torch.clamp(mask.sum(dim=1), min=1e-9) + return F.normalize(mean, p=2, dim=-1) + + +def _ragged_batch(lengths: list[int], seq_len: int, vocab: int, seed: int = 1): + gen = torch.Generator().manual_seed(seed) + ids = torch.zeros(len(lengths), seq_len, dtype=torch.int32) + mask = torch.zeros(len(lengths), seq_len, dtype=torch.int32) + for row, n in enumerate(lengths): + ids[row, :n] = torch.randint(1, vocab, (n,), generator=gen, dtype=torch.int32) + mask[row, :n] = 1 + return ids, mask diff --git a/python/tests/test_model_units/test_models/conftest.py b/python/tests/test_model_units/test_models/conftest.py new file mode 100644 index 00000000..19da7304 --- /dev/null +++ b/python/tests/test_model_units/test_models/conftest.py @@ -0,0 +1,19 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +from pathlib import Path + +import pytest + + +@pytest.fixture(scope="session") +def embeddinggemma2_snapshot() -> Path: + """Local google/embeddinggemma-2 snapshot; skips when it is not cached.""" + try: + from huggingface_hub import snapshot_download + + return Path(snapshot_download("google/embeddinggemma-2", local_files_only=True)) + except Exception as error: # noqa: BLE001 + pytest.skip(f"google/embeddinggemma-2 is not in the local Hugging Face cache ({error})") diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py new file mode 100644 index 00000000..8bf9978b --- /dev/null +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -0,0 +1,103 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Neural Engine EG2 text encoder (``models/ios/embedding_gemma2.py``) vs the macOS reference.""" + +import math + +import torch + +from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, text_encoder_from_state_dict +from coreai_models.models.macos.embedding_gemma2 import ( + EmbeddingGemma2TextConfig, + EmbeddingGemma2TextEncoder, +) +from tests.test_model_units.test_models._embedding_gemma2_common import ( + TINY_TEXT_CONFIG, + _cos, + _ragged_batch, +) + + +def _host_inputs( + ids: torch.Tensor, + attention_mask: torch.Tensor, + embed_table: torch.Tensor, + sliding_window: int, + dtype: torch.dtype = torch.float16, +) -> dict[str, torch.Tensor]: + """Graph inputs for one sequence (``ids [1, L]``): scaled embeddings, bias, pooling weights.""" + b, length = ids.shape + emb = embed_table[ids.long()].float() * math.sqrt(embed_table.shape[1]) + valid = attention_mask.bool() + key_ok = valid[0].reshape(length, 1, 1).expand(length, 1, length) # (Lk, 1, Lq) + bias = torch.where(key_ok, 0.0, MASK_NEG).reshape(1, length, 1, length) + out = { + "inputs_embeds": emb.to(dtype), + "bias": bias.to(dtype), + "pool_weights": (attention_mask.float() / attention_mask.float().sum(-1, keepdim=True)).to( + dtype + ), + } + if length > sliding_window + 1: + pos = torch.arange(length) + band = (pos[:, None] - pos[None, :]).abs() <= sliding_window # (Lk, Lq) symmetric + sb = torch.where(key_ok & band.reshape(length, 1, length), 0.0, MASK_NEG) + out["sliding_bias"] = sb.reshape(1, length, 1, length).to(dtype) + return out + + +def _perturb(model: torch.nn.Module, seed: int) -> None: + """Randomize norm weights and layer scalars so the comparison exercises them.""" + gen = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for name, tensor in sorted(model.state_dict().items()): + if name.endswith("layer_scalar"): + tensor.copy_(0.5 + torch.rand(tensor.shape, generator=gen)) + elif "norm" in name: + tensor.copy_(1.0 + 0.3 * torch.randn(tensor.shape, generator=gen)) + + +def _reference_and_state_dict(): + torch.manual_seed(0) + cfg = EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG) + ref = EmbeddingGemma2TextEncoder(cfg) + _perturb(ref, seed=0) + ref.eval() + state = {k: v.detach().clone() for k, v in ref.state_dict().items()} + return cfg, ref, state + + +def test_float32_matches_macos_reference_per_row() -> None: + cfg, ref, state = _reference_and_state_dict() + ios, table = text_encoder_from_state_dict(cfg, state, dtype=torch.float32) + ids, mask = _ragged_batch([12, 7, 4], seq_len=12, vocab=cfg.vocab_size) + with torch.no_grad(): + want = ref(input_ids=ids, attention_mask=mask) + for row in range(ids.shape[0]): + n = int(mask[row].sum()) + inputs = _host_inputs( + ids[row : row + 1], mask[row : row + 1], table, cfg.sliding_window, torch.float32 + ) + got = ios(**inputs) + assert got.shape == (1, cfg.embedding_dim) + cos = _cos(got, want[row : row + 1]) + assert float(cos) >= 0.9999, (row, n, float(cos)) + + +def test_float16_stays_finite_with_large_activations() -> None: + cfg, _, state = _reference_and_state_dict() + _, table = text_encoder_from_state_dict(cfg, state, dtype=torch.float32) + ids, mask = _ragged_batch([12], seq_len=12, vocab=cfg.vocab_size) + # Scale the token embeddings so the hidden stream reaches about 3e3, past sqrt(fp16 max). + 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) + with torch.no_grad(): + want = ios32(**_host_inputs(ids, mask, table, cfg.sliding_window, torch.float32)) + got = ios16(**_host_inputs(ids, mask, table, cfg.sliding_window, torch.float16)) + assert torch.isfinite(got.float()).all() + assert float(_cos(got.float(), want)) >= 0.99 diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2_vision.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2_vision.py new file mode 100644 index 00000000..2a0468fe --- /dev/null +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2_vision.py @@ -0,0 +1,108 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Neural Engine EG2 image encoder (``models/ios/embedding_gemma2_vision.py``) vs macOS.""" + +from pathlib import Path + +import torch +import torch.nn.functional as F + +import coreai_models.models.ios.embedding_gemma2_vision as ios_vision +from coreai_models.models.macos.embedding_gemma2_vision import ( + EmbeddingGemma2ImageEncoder, + pad_to_budget, + patch_position_ids, +) +from tests.test_model_units.test_models._embedding_gemma2_common import TINY_VISION_CONFIG as TINY + + +def _pair(): + torch.manual_seed(0) + ref = EmbeddingGemma2ImageEncoder(TINY) + with torch.no_grad(): + for name, p in ref.named_parameters(): + if name.endswith("norm.weight") or "layernorm" in name: + p.copy_(1.0 + 0.5 * torch.rand_like(p)) + else: + p.normal_(0.0, 0.3) + sd = { + (k if k.startswith("embed_vision.") else "vision_tower." + k): v + for k, v in ref.state_dict().items() + } + ane = ios_vision.EmbeddingGemma2ImageEncoder(TINY) + for layer in ane.layers: + layer.self_attn.q_chunk = 25 # several ragged query chunks at P = 108 + ios_vision.load_hf_state_dict(ane, sd) + return ref.eval(), ane.eval() + + +def _inputs(grids): + pvs, poss, ns = [], [], [] + for gh, gw in grids: + n = gh * gw + pv = torch.rand(n, TINY.patch_dim) + pv, pos = pad_to_budget(pv, patch_position_ids(gh, gw), TINY.max_patches) + pvs.append(pv[None]) + poss.append(pos[None]) + ns.append(n // TINY.pooling_kernel_size**2) + return pvs, poss, ns + + +def test_one_hot_is_exact_including_padding() -> None: + grid = torch.arange(840, dtype=torch.float16) + coord = torch.tensor([[-1.0], [0.0], [3.0], [839.0]], dtype=torch.float16) + oh = ios_vision.one_hot_f16(coord, grid) + assert oh[0].sum() == 0 + for row, c in zip(oh[1:], (0, 3, 839), strict=True): + assert row.sum() == 1 and row[c] == 1 + + +def test_matches_macos_reference_fp32() -> None: + ref, ane = _pair() + pvs, poss, ns = _inputs([(6, 9), (12, 9), (3, 18), (3, 3)]) + for pv, pos, n in zip(pvs, poss, ns, strict=True): + with torch.no_grad(): + want = ref(pv, pos) + got = ane(pv, pos.float()) + assert got.shape == want.shape == (1, TINY.max_soft_tokens, TINY.text_hidden_size) + cos = F.cosine_similarity(got[0, :n].double(), want[0, :n].double(), -1) + assert cos.min() > 1 - 1e-6, (n, cos.min()) + torch.testing.assert_close(got, want, atol=2e-4, rtol=2e-4) + assert (got[0, n:] == 0).all() # padding slots are exact zeros + + +def test_fp16_stays_finite_with_large_activations() -> None: + _, ane = _pair() + ane = ane.half() + pvs, poss, _ = _inputs([(6, 9)]) + with torch.no_grad(): + # Scale the residual stream far past sqrt(fp16 max) = 255.9 via the input projection. + ane.input_proj.weight.mul_(400.0) + out = ane(pvs[0].half(), poss[0].half()) + assert torch.isfinite(out).all() and out.abs().max() > 0 + + +def test_torch_export_static() -> None: + _, ane = _pair() + pvs, poss, _ = _inputs([(6, 9)]) + ep = torch.export.export(ane, args=(pvs[0], poss[0].float())) + torch.testing.assert_close(ep.module()(pvs[0], poss[0].float()), ane(pvs[0], poss[0].float())) + + +def test_real_snapshot_loads_and_pools_finite(embeddinggemma2_snapshot: Path) -> None: + budget = 12 + ane = ios_vision.load_image_encoder( + str(embeddinggemma2_snapshot), dtype=torch.float16, max_soft_tokens=budget + ) + gh, gw = 6, 9 + n_soft = gh * gw // 9 + pv = torch.rand(gh * gw, ane.config.patch_dim) + pv, pos = pad_to_budget(pv, patch_position_ids(gh, gw), budget * 9) + with torch.no_grad(): + out = ane(pv[None].half(), pos[None].half()) + assert out.shape == (1, budget, 512) + assert torch.isfinite(out).all() + assert (out[0, n_soft:] == 0).all() diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py new file mode 100644 index 00000000..32ab4021 --- /dev/null +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -0,0 +1,393 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for the EmbeddingGemma 2 text encoder (``models/macos/embedding_gemma2.py``). + +Tiers +----- +1. Self-contained (always run): config parsing, mask semantics, padding invariance, pooling + mask negative test, inputs_embeds path, torch.export at a static bucket. +2. HF reference (tiny random config): needs ``transformers>=5.19`` (``models/embedding_gemma2``). + The repo pins an older transformers, so run in an overlay:: + + uv run --no-sync --with "transformers>=5.19" --with "sentence-transformers>=6.1" \\ + python -m pytest -q -p no:rerunfailures -p no:xdist -o addopts="" + +3. Real weights (CPU fp32): additionally needs the ``google/embeddinggemma-2`` snapshot in the + local HF cache (the ``embeddinggemma2_snapshot`` fixture skips otherwise). Compared against + sentence-transformers>=6.1 when present; otherwise against HF ``EmbeddingGemma2Model`` + the + mean-pool + L2-normalize re-implementation in ``_st_mean_pool_normalize``. + +Tiers 2/3 skip cleanly when their dependencies are unavailable. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from coreai_models.models.macos.embedding_gemma2 import ( + EmbeddingGemma2TextConfig, + EmbeddingGemma2TextEncoder, + build_attention_masks, + extract_text_state_dict, + load_embedding_gemma2_text_encoder, + load_text_state_dict, +) +from tests.test_model_units.test_models._embedding_gemma2_common import ( + TINY_TEXT_CONFIG, + _cos, + _ragged_batch, + _st_mean_pool_normalize, +) + + +def _randomize_(module: torch.nn.Module, seed: int = 0) -> None: + """Perturb every parameter/buffer (norm weights, layer_scalar included) so bugs are visible.""" + gen = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for name, tensor in sorted(module.state_dict().items()): + if not tensor.is_floating_point(): + continue + if name.endswith("layer_scalar"): + tensor.copy_(0.5 + torch.rand(tensor.shape, generator=gen)) + elif "norm" in name: + tensor.copy_(1.0 + 0.3 * torch.randn(tensor.shape, generator=gen)) + else: + tensor.copy_(0.2 * torch.randn(tensor.shape, generator=gen)) + + +def _tiny_model(seed: int = 0) -> EmbeddingGemma2TextEncoder: + model = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG)) + _randomize_(model, seed) + return model.eval() + + +# --- Tier 1: self-contained --- + + +def test_config_from_dict_resolves_per_layer_overrides(): + cfg = EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG) + assert [cfg.layer_head_dim(i) for i in range(4)] == [16, 32, 16, 32] + assert [cfg.layer_num_key_value_heads(i) for i in range(4)] == [2, 1, 2, 1] + assert (cfg.rope_theta_sliding, cfg.rope_theta_full) == (10_000.0, 1_000_000.0) + + +def test_config_from_snapshot_matches_published_layout(embeddinggemma2_snapshot): + cfg = EmbeddingGemma2TextConfig.from_pretrained(embeddinggemma2_snapshot) + full = [i for i in range(cfg.num_hidden_layers) if not cfg.is_sliding(i)] + assert full == [5, 11, 17, 23] + assert all(cfg.layer_head_dim(i) == 512 and cfg.layer_num_key_value_heads(i) == 1 for i in full) + assert cfg.layer_head_dim(0) == 256 and cfg.layer_num_key_value_heads(0) == 2 + assert (cfg.hidden_size, cfg.embedding_dim, cfg.sliding_window) == (512, 768, 512) + + +def test_masks_bidirectional_window_and_key_padding(): + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0, 0]], dtype=torch.int32) + full, sliding = build_attention_masks(attention_mask, sliding_window=2) + assert full.shape == sliding.shape == (1, 1, 7, 7) + i, j = torch.meshgrid(torch.arange(7), torch.arange(7), indexing="ij") + key_valid = (j < 5).unsqueeze(0).unsqueeze(0) + assert torch.equal(full, key_valid.expand_as(full)) + assert torch.equal(sliding, ((i - j).abs() <= 2).unsqueeze(0).unsqueeze(0) & key_valid) + # Symmetric (bidirectional) on valid positions. + valid = sliding[0, 0, :5, :5] + assert torch.equal(valid, valid.T) + + +def test_forward_shapes_and_unit_norm(): + model = _tiny_model() + ids, mask = _ragged_batch([9, 5], seq_len=12, vocab=128) + emb, tokens = model(input_ids=ids, attention_mask=mask, return_token_states=True) + assert emb.shape == (2, 48) and emb.dtype == torch.float32 + assert tokens.shape == (2, 12, 48) + torch.testing.assert_close(emb.norm(dim=-1), torch.ones(2), atol=1e-6, rtol=0) + + +def test_padding_invariance_tiny(): + """A sequence embedded alone (no padding) == the same sequence inside a padded batch.""" + model = _tiny_model() + ids, mask = _ragged_batch([11, 4, 7], seq_len=16, vocab=128) + with torch.no_grad(): + batched = model(input_ids=ids, attention_mask=mask) + for row, n in enumerate([11, 4, 7]): + alone = model(input_ids=ids[row : row + 1, :n], attention_mask=mask[row : row + 1, :n]) + assert _cos(alone[0], batched[row]).item() > 1 - 1e-6 + torch.testing.assert_close(alone[0], batched[row], atol=2e-5, rtol=0) + + +def test_pooling_mask_excluding_prompt_changes_result(): + model = _tiny_model() + ids, mask = _ragged_batch([10], seq_len=10, vocab=128) + no_prompt = mask.clone() + no_prompt[:, :4] = 0 # pretend the first 4 tokens are the task prompt + with torch.no_grad(): + with_prompt = model(input_ids=ids, attention_mask=mask) + without_prompt = model(input_ids=ids, attention_mask=mask, pooling_mask=no_prompt) + assert _cos(with_prompt, without_prompt).item() < 0.9999 + + +def test_projection_after_pool_equals_project_then_pool(): + model = _tiny_model() + ids, mask = _ragged_batch([8, 3], seq_len=10, vocab=128) + with torch.no_grad(): + emb, tokens = model(input_ids=ids, attention_mask=mask, return_token_states=True) + m = mask.float().unsqueeze(-1) + reference = F.normalize((tokens * m).sum(1) / m.sum(1).clamp(min=1e-9), dim=-1) + torch.testing.assert_close(emb, reference, atol=1e-6, rtol=0) + + +def test_inputs_embeds_path_matches_input_ids(): + model = _tiny_model() + ids, mask = _ragged_batch([6, 6], seq_len=6, vocab=128) + with torch.no_grad(): + a = model(input_ids=ids, attention_mask=mask) + b = model(inputs_embeds=model.embed_tokens(ids), attention_mask=mask) + torch.testing.assert_close(a, b, atol=0, rtol=0) + with pytest.raises(ValueError): + model(input_ids=ids, inputs_embeds=model.embed_tokens(ids)) + + +def test_load_text_state_dict_fails_closed(): + model = _tiny_model() + prefixed = {f"language_model.{k}": v for k, v in model.state_dict().items()} + prefixed["vision_tower.encoder.foo"] = torch.zeros(1) + stripped = extract_text_state_dict(prefixed) + assert "encoder.foo" not in stripped and len(stripped) == len(model.state_dict()) + fresh = EmbeddingGemma2TextEncoder(model.config) + load_text_state_dict(fresh, stripped) + stripped.pop("layers.0.layer_scalar") + with pytest.raises(RuntimeError, match="missing"): + load_text_state_dict(EmbeddingGemma2TextEncoder(model.config), stripped) + + +def test_torch_export_static_bucket(): + """torch.export at a static (B=2, L=128) bucket with int32 ids/mask; outputs match eager.""" + model = _tiny_model() + ids, mask = _ragged_batch([128, 37], seq_len=128, vocab=128) + with torch.no_grad(): + exported = torch.export.export( + model, args=(), kwargs={"input_ids": ids, "attention_mask": mask} + ) + eager = model(input_ids=ids, attention_mask=mask) + traced = exported.module()(input_ids=ids, attention_mask=mask) + assert traced.shape == (2, 48) + torch.testing.assert_close(traced, eager, atol=1e-6, rtol=0) + placeholders = [ + n + for n in exported.graph.nodes + if n.op == "placeholder" and n.name in ("input_ids", "attention_mask") + ] + assert len(placeholders) == 2 + assert all(n.meta["val"].dtype == torch.int32 for n in placeholders) + + +# --- Tier 2: tiny random config vs HF EmbeddingGemma2TextModel --- + + +def _hf_modeling(): + return pytest.importorskip( + "transformers.models.embedding_gemma2.modeling_embedding_gemma2", + reason="needs transformers>=5.19 (run in a uv --with overlay)", + ) + + +def _tiny_hf(attn_impl: str): + """Tiny random HF ``EmbeddingGemma2TextModel`` with perturbed weights.""" + modeling = _hf_modeling() + from transformers.models.embedding_gemma2.configuration_embedding_gemma2 import ( + EmbeddingGemma2Config, + ) + + torch.manual_seed(0) + hf_cfg = EmbeddingGemma2Config(text_config=dict(TINY_TEXT_CONFIG)).text_config + hf_cfg._attn_implementation = attn_impl + hf = modeling.EmbeddingGemma2TextModel(hf_cfg).eval() + _randomize_(hf, seed=3) + return hf + + +@pytest.mark.parametrize("attn_impl", ["eager", "sdpa"]) +def test_tiny_parity_vs_hf(attn_impl): + hf = _tiny_hf(attn_impl) + + ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG)).eval() + load_text_state_dict(ours, dict(hf.state_dict())) + + # Long enough (L=24 > 2*window+1) that the sliding window actually bites, with padding. + ids, mask = _ragged_batch([24, 13, 5], seq_len=24, vocab=128, seed=7) + with torch.no_grad(): + hf_tokens = hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state + hf_emb = _st_mean_pool_normalize(hf_tokens, mask) + emb, tokens = ours(input_ids=ids, attention_mask=mask, return_token_states=True) + + valid = mask.bool() + torch.testing.assert_close(tokens[valid], hf_tokens[valid], atol=1e-4, rtol=1e-4) + assert _cos(emb, hf_emb).min().item() > 1 - 1e-6 + torch.testing.assert_close(emb, hf_emb, atol=1e-5, rtol=0) + + +def test_tiny_sliding_window_matters_vs_hf(): + """Guard against a silently-ignored window: widening ours must break parity with HF.""" + hf = _tiny_hf("eager") + wide = dict(TINY_TEXT_CONFIG, sliding_window=1000) + ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(wide)).eval() + load_text_state_dict(ours, dict(hf.state_dict())) + ids, mask = _ragged_batch([24], seq_len=24, vocab=128, seed=7) + with torch.no_grad(): + hf_emb = _st_mean_pool_normalize( + hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state, mask + ) + emb = ours(input_ids=ids, attention_mask=mask) + assert (emb - hf_emb).abs().max().item() > 1e-3 + + +# --- Tier 3: real weights (CPU fp32) --- + +REAL_TEXTS = [ + ("query", "What is the capital of France?"), + ("document", "Paris is the capital and most populous city of France."), + ("query", "how do neural networks learn"), + ("document", "Gradient descent iteratively updates the weights of a model to minimize a loss."), + ("STS", "A man is playing a guitar on stage."), + ("document", "El aprendizaje automático es una rama de la inteligencia artificial."), + ("CodeRetrieval", "def add(a, b):\n return a + b"), + ("query", "x"), +] + + +@pytest.fixture(scope="module") +def real_setup(embeddinggemma2_snapshot): + snapshot = embeddinggemma2_snapshot + _hf_modeling() + from transformers import AutoConfig, AutoTokenizer + from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import EmbeddingGemma2Model + + prompts = None + st_model = None + try: + from sentence_transformers import SentenceTransformer + + st_model = SentenceTransformer( + str(snapshot), device="cpu", model_kwargs={"dtype": torch.float32} + ) + prompts = st_model.prompts + except Exception: # noqa: BLE001 - optional reference + import json + + prompts = json.loads((snapshot / "config_sentence_transformers.json").read_text())[ + "prompts" + ] + + cfg = AutoConfig.from_pretrained(str(snapshot)) + cfg.vision_config = None + cfg.audio_config = None + hf = EmbeddingGemma2Model.from_pretrained(str(snapshot), config=cfg, dtype=torch.float32).eval() + tokenizer = AutoTokenizer.from_pretrained(str(snapshot)) + ours = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float32) + texts = [prompts[p] + t for p, t in REAL_TEXTS] + if st_model is not None: + tokenize = lambda batch: st_model.tokenize(batch) # noqa: E731 + else: + tokenize = lambda batch: tokenizer(batch, padding=True, return_tensors="pt") # noqa: E731 + return {"st": st_model, "hf": hf, "ours": ours, "texts": texts, "tokenize": tokenize} + + +def test_real_weight_key_coverage(real_setup, embeddinggemma2_snapshot): + from safetensors import safe_open + + with safe_open(str(embeddinggemma2_snapshot / "model.safetensors"), framework="pt") as handle: + all_keys = list(handle.keys()) # noqa: SIM118 - safe_open is not a Mapping + hf_text_keys = { + k[len("language_model.") :] for k in all_keys if k.startswith("language_model.") + } + ours = set(real_setup["ours"].state_dict().keys()) + assert ours == hf_text_keys, (sorted(ours - hf_text_keys)[:5], sorted(hf_text_keys - ours)[:5]) + + +def test_real_weight_parity_and_padding_invariance(real_setup): + ours, hf, texts = real_setup["ours"], real_setup["hf"], real_setup["texts"] + features = real_setup["tokenize"](texts) + ids = features["input_ids"].to(torch.int32) + mask = features["attention_mask"].to(torch.int32) + assert int(mask.sum(1).min()) < int(mask.sum(1).max()), "batch must be mixed-length" + + with torch.inference_mode(): + hf.language_model.config._attn_implementation = "sdpa" + hf_tokens = hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state + ref = _st_mean_pool_normalize(hf_tokens, mask) + if real_setup["st"] is not None: + st_emb = real_setup["st"].encode(texts, batch_size=len(texts), convert_to_tensor=True) + assert _cos(st_emb, ref).min().item() > 0.99999 # our pooling re-impl == ST + ref = st_emb + emb, tokens = ours(input_ids=ids, attention_mask=mask, return_token_states=True) + alone = torch.cat( + [ + ours( + input_ids=ids[r : r + 1, : int(mask[r].sum())], + attention_mask=mask[r : r + 1, : int(mask[r].sum())], + ) + for r in range(len(texts)) + ] + ) + + cos = _cos(emb, ref) + emb_max_abs = (emb - ref).abs().max().item() + pad_cos = _cos(alone, emb) + print( + f"\nEG2 real fp32 CPU: n={len(texts)} lens={mask.sum(1).tolist()} " + f"cos min={cos.min().item():.9f} median={cos.median().item():.9f} " + f"emb max_abs={emb_max_abs:.3e} " + f"padding-invariance cos min={pad_cos.min().item():.9f}" + ) + assert cos.min().item() >= 0.99999 + assert pad_cos.min().item() >= 0.99999 + + +def test_real_weight_excluding_prompt_changes_embedding(real_setup): + ours, texts = real_setup["ours"], real_setup["texts"] + features = real_setup["tokenize"](texts[:1]) + ids = features["input_ids"].to(torch.int32) + mask = features["attention_mask"].to(torch.int32) + prompt_len = int( + real_setup["tokenize"]([texts[0][: texts[0].index("query: ") + len("query:")]])[ + "attention_mask" + ].sum() + ) + no_prompt = mask.clone() + no_prompt[:, :prompt_len] = 0 + with torch.inference_mode(): + a = ours(input_ids=ids, attention_mask=mask) + b = ours(input_ids=ids, attention_mask=mask, pooling_mask=no_prompt) + cos = _cos(a, b).item() + print(f"\nEG2 include_prompt vs exclude_prompt cos={cos:.6f} (prompt_len={prompt_len})") + assert cos < 0.999 + + +def test_real_weight_long_sequence_window_bites(real_setup): + """> 2 * 512 + 1 tokens so the sliding layers' |i - j| <= 512 window is exercised for real.""" + ours, hf = real_setup["ours"], real_setup["hf"] + paragraph = ( + "On-device semantic search embeds documents once and compares vectors at query time. " + ) + features = real_setup["tokenize"]( + [real_setup["texts"][1] + paragraph * 90, real_setup["texts"][0]] + ) + ids = features["input_ids"].to(torch.int32) + mask = features["attention_mask"].to(torch.int32) + assert ids.shape[1] > 1100 + with torch.inference_mode(): + hf.language_model.config._attn_implementation = "sdpa" + ref = _st_mean_pool_normalize( + hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state, mask + ) + emb = ours(input_ids=ids, attention_mask=mask) + cos = _cos(emb, ref) + print( + f"\nEG2 long L={ids.shape[1]} lens={mask.sum(1).tolist()} cos={cos.tolist()} " + f"max_abs={(emb - ref).abs().max().item():.3e}" + ) + assert cos.min().item() >= 0.99999 diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py new file mode 100644 index 00000000..1472c41c --- /dev/null +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py @@ -0,0 +1,399 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Parity tests for the EmbeddingGemma 2 image path (``embedding_gemma2_vision.py``). + +* Tiny-config parity against HF ``Gemma4VisionModel`` / ``Gemma4MultimodalEmbedder`` + (present in the repo's transformers; no checkpoint needed). +* Host preprocessing helpers against HF ``Gemma4ImageProcessor``. +* Scatter / sequence assembly against HF ``masked_scatter`` semantics. +* ``torch.export`` at the static default budget (280 soft tokens / 2520 patches). +* Real-weight parity (soft tokens and end-to-end image embedding) against HF + ``EmbeddingGemma2Model``. Needs transformers >= 5.19 and the local snapshot; skipped + otherwise. Run with an overlay, e.g.:: + + uv run --no-sync --with "transformers>=5.19" --with "torchvision==0.24.*" \\ + python -m pytest -q -p no:rerunfailures -p no:xdist -o addopts="" +""" + +import math +import os + +import pytest +import torch +import torch.nn.functional as F + +from coreai_models.models.macos.embedding_gemma2_vision import ( + EG2_BOI_TOKEN_ID, + EG2_BOS_TOKEN_ID, + EG2_EOI_TOKEN_ID, + EG2_EOS_TOKEN_ID, + EG2_IMAGE_TOKEN_ID, + EmbeddingGemma2ImageEncoder, + EmbeddingGemma2VisionConfig, + aspect_ratio_preserving_size, + build_image_input_ids, + build_image_inputs_embeds, + pad_to_budget, + patch_position_ids, + patchify, + preprocess_resized_image, + scatter_soft_tokens, +) +from tests.test_model_units.test_models._embedding_gemma2_common import ( + TINY_VISION_CONFIG as TINY, +) +from tests.test_model_units.test_models._embedding_gemma2_common import ( + _cos, + _st_mean_pool_normalize, +) + +gemma4_modeling = pytest.importorskip("transformers.models.gemma4.modeling_gemma4") +gemma4_config = pytest.importorskip("transformers.models.gemma4.configuration_gemma4") + + +def _hf_tiny_vision(cfg: EmbeddingGemma2VisionConfig): + # transformers >= 5.19 names the vision RoPE "axial"; older 5.x builds call the same + # (mathematically identical) per-axis RoPE "default". + rot = gemma4_modeling.Gemma4VisionRotaryEmbedding + rope_type = "axial" if hasattr(rot, "compute_axial_rope_parameters") else "default" + hf_cfg = gemma4_config.Gemma4VisionConfig( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size, + num_hidden_layers=cfg.num_hidden_layers, + num_attention_heads=cfg.num_attention_heads, + num_key_value_heads=cfg.num_key_value_heads, + head_dim=cfg.head_dim, + global_head_dim=cfg.head_dim, + patch_size=cfg.patch_size, + pooling_kernel_size=cfg.pooling_kernel_size, + position_embedding_size=cfg.position_embedding_size, + rms_norm_eps=cfg.rms_norm_eps, + rope_parameters={"rope_type": rope_type, "rope_theta": cfg.rope_theta}, + hidden_activation="gelu_pytorch_tanh", + use_clipped_linears=False, + standardize=False, + ) + hf_cfg._attn_implementation = "eager" + torch.manual_seed(0) + model = gemma4_modeling.Gemma4VisionModel(hf_cfg).eval() + # HF init leaves norms at 1 / zeros; randomize every float tensor so the test bites. + with torch.no_grad(): + for name, p in model.named_parameters(): + if "norm" in name: + p.copy_(1.0 + 0.2 * torch.randn_like(p)) + else: + p.copy_(0.2 * torch.randn_like(p)) + return hf_cfg, model + + +def _load_from_hf_vision(ours: EmbeddingGemma2ImageEncoder, hf_vision, hf_embedder=None): + state = {} + for k, v in hf_vision.state_dict().items(): + local = EmbeddingGemma2ImageEncoder.map_hf_key("vision_tower." + k) + if local is not None and local in dict(ours.named_parameters()): + state[local] = v + if hf_embedder is not None: + for k, v in hf_embedder.state_dict().items(): + state["embed_vision." + k] = v + missing, unexpected = ours.load_state_dict(state, strict=False) + expected_missing = ( + [] if hf_embedder is not None else ["embed_vision.embedding_projection.weight"] + ) + assert sorted(missing) == expected_missing + assert not unexpected + + +def _random_inputs(cfg: EmbeddingGemma2VisionConfig, grids: list[tuple[int, int]], seed: int = 0): + """Random patches for images whose block grid is ``(bh, bw)`` blocks, padded to budget.""" + g = torch.Generator().manual_seed(seed) + k = cfg.pooling_kernel_size + pv, pos, counts = [], [], [] + for bh, bw in grids: + ph, pw = bh * k, bw * k + img = torch.rand(3, ph * cfg.patch_size, pw * cfg.patch_size, generator=g) + patches = patchify(img, cfg.patch_size) + positions = patch_position_ids(ph, pw) + patches, positions = pad_to_budget(patches, positions, cfg.max_patches) + pv.append(patches) + pos.append(positions) + counts.append(bh * bw) + return torch.stack(pv), torch.stack(pos), counts + + +# --- Tiny-config parity vs HF --- + + +@pytest.mark.parametrize("use_hf_sdpa", ["true", "false"]) +@pytest.mark.parametrize("grids", [[(3, 4)], [(2, 3)], [(1, 5), (4, 3)]]) +def test_tiny_vision_tower_matches_hf(grids, use_hf_sdpa, monkeypatch) -> None: + # The directory conftest forces USE_HF_IMPL=true; also cover the Core AI lowering path + # of the SDPA primitive (bool-mask -> -1e4 bias), which is what gets exported. + monkeypatch.setenv("USE_HF_IMPL", use_hf_sdpa) + _, hf = _hf_tiny_vision(TINY) + ours = EmbeddingGemma2ImageEncoder(TINY).eval() + _load_from_hf_vision(ours, hf) + pixel_values, position_ids, counts = _random_inputs(TINY, grids) + with torch.no_grad(): + ref = hf(pixel_values=pixel_values, pixel_position_ids=position_ids).last_hidden_state + got = ours.encode_pooled(pixel_values, position_ids) + # HF strips padding slots and concatenates images; ours keeps the static budget. + got_valid = torch.cat([got[i, :n] for i, n in enumerate(counts)]) + assert got_valid.shape == ref.shape + torch.testing.assert_close(got_valid, ref.float(), rtol=1e-5, atol=1e-5) + # Padding slots are exactly zero. + for i, n in enumerate(counts): + assert torch.count_nonzero(got[i, n:]) == 0 + + +def test_tiny_soft_tokens_match_hf_embedder() -> None: + hf_cfg, hf = _hf_tiny_vision(TINY) + text_cfg = gemma4_config.Gemma4TextConfig(hidden_size=TINY.text_hidden_size) + torch.manual_seed(1) + embedder = gemma4_modeling.Gemma4MultimodalEmbedder(hf_cfg, text_cfg).eval() + ours = EmbeddingGemma2ImageEncoder(TINY).eval() + _load_from_hf_vision(ours, hf, embedder) + pixel_values, position_ids, counts = _random_inputs(TINY, [(4, 3)], seed=3) + with torch.no_grad(): + ref = embedder( + hf(pixel_values=pixel_values, pixel_position_ids=position_ids).last_hidden_state + ) + got = ours(pixel_values, position_ids) + assert got.shape == (1, TINY.max_soft_tokens, TINY.text_hidden_size) + torch.testing.assert_close(got[0, : counts[0]], ref, rtol=1e-5, atol=1e-5) + assert torch.count_nonzero(got[0, counts[0] :]) == 0 + + +def test_rejects_wrong_patch_count() -> None: + ours = EmbeddingGemma2ImageEncoder(TINY) + with pytest.raises(ValueError, match="expected"): + ours(torch.zeros(1, TINY.max_patches - 9, TINY.patch_dim), torch.zeros(1, 99, 2).long()) + + +# --- Host preprocessing vs Gemma4ImageProcessor --- + + +def _hf_aspect_fn(): + from transformers.models.gemma4.image_processing_pil_gemma4 import ( + get_aspect_ratio_preserving_size, + ) + + return get_aspect_ratio_preserving_size + + +@pytest.mark.parametrize("budget", [70, 140, 280, 560, 1120]) +def test_aspect_ratio_size_matches_hf(budget: int) -> None: + hf_fn = _hf_aspect_fn() + sizes = [(300, 451), (451, 300), (512, 512), (48, 48), (20, 20), (1080, 1920), (1920, 1080)] + sizes += [(50, 2000), (2000, 50), (31, 1000), (3000, 40), (17, 33), (4032, 3024)] + for h, w in sizes: + try: + ref = hf_fn( + height=h, width=w, patch_size=16, max_patches=budget * 9, pooling_kernel_size=3 + ) + except ValueError: + with pytest.raises(ValueError): + aspect_ratio_preserving_size(h, w, 16, budget, 3) + continue + assert aspect_ratio_preserving_size(h, w, 16, budget, 3) == ref, (h, w) + + +def test_preprocess_matches_hf_processor_without_resize() -> None: + """At the target size HF skips the resize, so patching/positions/padding must be exact.""" + np = pytest.importorskip("numpy") + try: + from transformers.models.gemma4.image_processing_pil_gemma4 import ( + Gemma4ImageProcessorPil as Proc, + ) + except ImportError: + pytest.skip("Gemma4ImageProcessorPil unavailable") + pil = pytest.importorskip("PIL.Image") + # Fixed points of the resize policy (the policy is not idempotent in general, e.g. + # 480x240 -> 1104x528 -> 1152x528), so HF leaves these untouched. + for h, w in [(624, 960), (960, 624), (768, 768), (576, 1056)]: + assert aspect_ratio_preserving_size(h, w) == (h, w) + arr = (np.random.default_rng(h + w).integers(0, 256, size=(h, w, 3))).astype("uint8") + out = Proc(max_soft_tokens=280)(images=[pil.fromarray(arr)], return_tensors="pt") + pv, pos, n = preprocess_resized_image(torch.from_numpy(arr).permute(2, 0, 1)) + assert n == int(out["num_soft_tokens_per_image"][0]) + torch.testing.assert_close(pv, out["pixel_values"].float(), rtol=0, atol=1e-7) + assert torch.equal(pos, out["image_position_ids"].long()) + + +# --- Sequence assembly --- + + +def test_build_image_input_ids_layout() -> None: + ids = build_image_input_ids(3) + assert ids == [ + EG2_BOS_TOKEN_ID, + EG2_BOI_TOKEN_ID, + EG2_IMAGE_TOKEN_ID, + EG2_IMAGE_TOKEN_ID, + EG2_IMAGE_TOKEN_ID, + EG2_EOI_TOKEN_ID, + EG2_EOS_TOKEN_ID, + ] + + +def test_scatter_matches_masked_scatter() -> None: + torch.manual_seed(0) + embeds = torch.randn(2, 9, 8) + mask = torch.zeros(2, 9, dtype=torch.bool) + mask[0, 2:6] = True + mask[1, 1:4] = True + soft = torch.randn(2, 6, 8) # extra trailing rows must be ignored + got = scatter_soft_tokens(embeds, mask, soft) + for b, n in ((0, 4), (1, 3)): + ref = embeds[b : b + 1].masked_scatter(mask[b : b + 1, :, None], soft[b, :n]) + torch.testing.assert_close(got[b : b + 1], ref, rtol=0, atol=0) + + +def test_build_image_inputs_embeds() -> None: + table = torch.randn(300000, 4) + + def embed(ids): + return F.embedding(ids, table) * 2.0 + + soft = torch.randn(2, 5, 4) + embeds, mask, ids = build_image_inputs_embeds(embed, soft, [5, 3]) + assert ids.shape == (2, 9) and mask[1].tolist() == [1] * 7 + [0, 0] + torch.testing.assert_close(embeds[0, 2:7], soft[0]) + torch.testing.assert_close(embeds[1, 2:5], soft[1, :3]) + torch.testing.assert_close(embeds[1, 0], table[EG2_BOS_TOKEN_ID] * 2.0) + torch.testing.assert_close(embeds[1, 5], table[EG2_EOI_TOKEN_ID] * 2.0) + + +# --- torch.export at the static default budget --- + + +def test_torch_export_default_budget_static_shape(monkeypatch) -> None: + monkeypatch.setenv("USE_HF_IMPL", "false") # the lowering path that gets exported + cfg = EmbeddingGemma2VisionConfig(num_hidden_layers=2) # real widths, 2 layers for speed + torch.manual_seed(0) + model = EmbeddingGemma2ImageEncoder(cfg).eval() + with torch.no_grad(): + for p in model.parameters(): + p.normal_(0.0, 0.02) + k = cfg.pooling_kernel_size + img = torch.randint(0, 256, (3, 15 * k * 16, 18 * k * 16), dtype=torch.uint8) + pixel_values, position_ids, n = preprocess_resized_image(img) + assert pixel_values.shape == (1, 2520, 768) and n == 270 + with torch.no_grad(): + ep = torch.export.export(model, (pixel_values, position_ids)) + eager = model(pixel_values, position_ids) + exported = ep.module()(pixel_values, position_ids) + assert exported.shape == (1, 280, 512) + torch.testing.assert_close(exported, eager, rtol=1e-5, atol=1e-5) + # No dynamic dims and no data-dependent ops (nonzero / masked_select / item) in the graph. + targets = {str(node.target) for node in ep.graph.nodes if node.op == "call_function"} + assert not any(t for t in targets if "nonzero" in t or "masked_select" in t or "item" in t) + + +# --- Real weights (transformers >= 5.19 + local snapshot) --- + + +def _synthetic_images(): + """Deterministic natural-ish RGB images (smooth gradients + shapes + noise).""" + np = pytest.importorskip("numpy") + pil = pytest.importorskip("PIL.Image") + rng = np.random.default_rng(0) + out = [] + for h, w in [(300, 451), (180, 640), (777, 333)]: + yy, xx = np.mgrid[0:h, 0:w].astype("float32") + img = np.stack( + [ + 128 + 100 * np.sin(xx / (17 + h % 7)), + 128 + 100 * np.cos(yy / 23), + 255 * (xx + yy) / (h + w), + ], + -1, + ) + cy, cx, r = h / 2, w / 3, min(h, w) / 4 + img[(yy - cy) ** 2 + (xx - cx) ** 2 < r * r] = [230, 40, 40] + img += rng.normal(0, 8, img.shape) + out.append(pil.fromarray(np.clip(img, 0, 255).astype("uint8"))) + return out + + +@pytest.fixture(scope="module") +def eg2_hf(embeddinggemma2_snapshot): + pytest.importorskip( + "transformers.models.embedding_gemma2.modeling_embedding_gemma2", + reason="needs transformers>=5.19", + ) + from transformers import AutoProcessor + from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import ( + EmbeddingGemma2Model, + ) + + snapshot = str(embeddinggemma2_snapshot) + processor = AutoProcessor.from_pretrained(snapshot) + model = EmbeddingGemma2Model.from_pretrained( + snapshot, dtype=torch.float32, audio_config=None + ).eval() + # Build ours on the SDPA lowering path that gets exported (not the HF-eager shim). + saved = os.environ.get("USE_HF_IMPL") + os.environ["USE_HF_IMPL"] = "false" + try: + ours = EmbeddingGemma2ImageEncoder.from_pretrained(snapshot, dtype=torch.float32) + finally: + if saved is None: + os.environ.pop("USE_HF_IMPL", None) + else: + os.environ["USE_HF_IMPL"] = saved + return processor, model, ours + + +@pytest.fixture(scope="module") +def eg2_runs(eg2_hf): + processor, model, ours = eg2_hf + runs = [] + for img in _synthetic_images(): + inputs = processor(images=[img], return_tensors="pt") + # The processor drops num_soft_tokens_per_image; count the placeholders instead. + n = int((inputs["input_ids"][0] == EG2_IMAGE_TOKEN_ID).sum()) + with torch.no_grad(): + ref_soft = model.get_image_features( + inputs["pixel_values"], inputs["image_position_ids"], return_dict=True + ).pooler_output[0] + got_soft = ours(inputs["pixel_values"], inputs["image_position_ids"]) + out = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + pixel_values=inputs["pixel_values"], + image_position_ids=inputs["image_position_ids"], + ) + runs.append((inputs, n, ref_soft, got_soft, out.last_hidden_state)) + return runs + + +def test_real_soft_tokens_match_hf(eg2_runs) -> None: + for inputs, n, ref_soft, got_soft, _ in eg2_runs: + assert ref_soft.shape == (n, 512) + assert got_soft.shape == (1, 280, 512) + cos = _cos(got_soft[0, :n], ref_soft) + assert float(cos.min()) >= 0.99999, float(cos.min()) + assert torch.count_nonzero(got_soft[0, n:]) == 0 + assert inputs["input_ids"][0].tolist() == build_image_input_ids(n) + + +def test_real_end_to_end_image_embedding_matches_hf(eg2_hf, eg2_runs) -> None: + _, model, _ = eg2_hf + lm = model.language_model + for inputs, n, _, got_soft, ref_hidden in eg2_runs: + with torch.no_grad(): + embeds, mask, ids = build_image_inputs_embeds(lm.embed_tokens, got_soft, n) + hidden = lm(inputs_embeds=embeds, attention_mask=mask).last_hidden_state + assert torch.equal(ids, inputs["input_ids"]) + got = _st_mean_pool_normalize(hidden, mask) + ref = _st_mean_pool_normalize(ref_hidden, inputs["attention_mask"]) + for dim in (768, 512, 256, 128): + c = F.cosine_similarity( + F.normalize(got[:, :dim].double(), dim=-1), + F.normalize(ref[:, :dim].double(), dim=-1), + ) + assert float(c.min()) >= 0.99999, (dim, float(c.min())) + assert math.isclose(float(got.norm()), 1.0, rel_tol=1e-5) From 03bb9748d057f976e4d124f29674c9e3a03a9bfb Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 20 Oct 2026 00:13:59 -0700 Subject: [PATCH 03/38] coreai.embedding.export: embedding bundle exporter 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. --- python/pyproject.toml | 1 + .../src/coreai_models/embedding/__init__.py | 6 + .../src/coreai_models/embedding/contract.py | 459 ++++++++++ python/src/coreai_models/embedding/export.py | 819 ++++++++++++++++++ python/src/coreai_models/export/metadata.py | 11 + python/src/coreai_models/model_registry.py | 23 +- .../test_export/test_embedding_export.py | 667 ++++++++++++++ .../test_export/test_include_debug_info.py | 10 + .../test_model_units/test_model_registry.py | 39 + 9 files changed, 2031 insertions(+), 4 deletions(-) create mode 100644 python/src/coreai_models/embedding/__init__.py create mode 100644 python/src/coreai_models/embedding/contract.py create mode 100644 python/src/coreai_models/embedding/export.py create mode 100644 python/tests/test_model_units/test_export/test_embedding_export.py diff --git a/python/pyproject.toml b/python/pyproject.toml index a7e429b7..60e50074 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -49,6 +49,7 @@ dependencies = [ "coreai.llm.export" = "coreai_models.llm.export:main" "coreai.llm.eval" = "coreai_models.llm.eval:main" "coreai.vlm.export" = "coreai_models.vlm.export:main" +"coreai.embedding.export" = "coreai_models.embedding.export:main" "coreai.diffusion.export" = "coreai_models.diffusion.export:main" "coreai.model.registry" = "coreai_models.model_registry:main" diff --git a/python/src/coreai_models/embedding/__init__.py b/python/src/coreai_models/embedding/__init__.py new file mode 100644 index 00000000..12737658 --- /dev/null +++ b/python/src/coreai_models/embedding/__init__.py @@ -0,0 +1,6 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Retrieval-embedding bundle export; CLI: `uv run coreai.embedding.export`.""" diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py new file mode 100644 index 00000000..08cf60b8 --- /dev/null +++ b/python/src/coreai_models/embedding/contract.py @@ -0,0 +1,459 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Builds the ``embedding`` block and ``metadata.json`` for ``kind: "embedding"`` bundles. + +Also holds the per-family export facts the exporter and the Swift runtime both read. +EmbeddingGemma 2 overflows in float16 on the GPU, so macOS bundles compute in float32 or +bfloat16, while Neural Engine (iOS) bundles use float16 with scale-invariant norms. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any + +from coreai_models._constants import METADATA_VERSION + +BUNDLE_KIND = "embedding" +OUTPUT_NAME = "embedding" +TEXT_INPUT_NAMES = ("input_ids", "attention_mask") +TOKEN_TABLE_INPUT_NAMES = ("input_ids",) +TOKEN_TABLE_OUTPUT_NAME = "inputs_embeds" +ENCODER_INPUT_NAMES = ("inputs_embeds", "attention_mask") +ANE_ENCODER_INPUT_NAMES = ("inputs_embeds", "bias", "pool_weights") +IMAGE_INPUT_NAMES = ("pixel_values", "position_ids") +IMAGE_OUTPUT_NAME = "soft_tokens" +TEXT_COMPONENT = "text_encoder" +TOKEN_TABLE_COMPONENT = "token_table" +ENCODER_COMPONENT = "encoder" +IMAGE_COMPONENT = "image_encoder" +SUPPORTED_COMPUTE_PRECISIONS = ("float32", "bfloat16", "float16") +PLATFORMS = ("macOS", "iOS") + + +class EmbeddingPrecisionError(ValueError): + """Raised when an export asks for a compute precision the platform or family forbids.""" + + +@dataclass(frozen=True) +class ImageTowerFacts: + """Facts the image tower needs on the host (merge ids, preprocessing, soft-token budget). + + The exporter checks them against the snapshot ``config.json`` / ``processor_config.json`` + and fails closed on a mismatch. + """ + + placeholder_id: int + boi_token_id: int + eoi_token_id: int + patch_size: int = 16 + pooling_kernel_size: int = 3 + rescale_factor: float = 1.0 / 255.0 + normalize: bool = False + image_mean: tuple[float, ...] = (0.0, 0.0, 0.0) + image_std: tuple[float, ...] = (1.0, 1.0, 1.0) + budget: int = 280 + soft_token_dim: int = 512 + + @property + def max_patches(self) -> int: + return self.budget * self.pooling_kernel_size**2 + + @property + def patch_dim(self) -> int: + return 3 * self.patch_size**2 + + def sequence_length(self) -> int: + """Longest image-only sequence: `` *budget ``.""" + return self.budget + 4 + + +@dataclass(frozen=True) +class EmbeddingFamily: + """Facts about one retrieval-embedding family that the bundle must carry. + + Values come from the HF snapshot files and model card; the space *revision* is the snapshot + SHA of the weights actually exported. + """ + + key: str + hf_model_id: str + dimension: int + matryoshka_dimensions: tuple[int, ...] + # canonical prompt name -> prompt key in the snapshot's config_sentence_transformers.json + prompt_names: dict[str, str] + metric: str = "cosine" + normalized: bool = True + renormalize_after_truncation: bool = True + pooling_mode: str = "mean" + include_prompt: bool = True + default_compute_precision: str = "float32" + precision_floor: str | None = None + forbidden_precisions: tuple[str, ...] = () + model_max_tokens: int = 512 + special_tokens: dict[str, bool] = field(default_factory=dict) + bos_token_id: int | None = None + eos_token_id: int | None = None + pad_token_id: int | None = None + padding_side: str = "right" + truncation_keep_eos: bool = True + document_template: str | None = None + image: ImageTowerFacts | None = None + length_buckets: tuple[int, ...] = (128, 512, 2048) + batch_sizes: tuple[int, ...] = (1, 8, 32) + max_batch_tokens: int = 16384 + ane_length_buckets: tuple[int, ...] = (128, 512, 1024) + + +EMBEDDINGGEMMA2 = EmbeddingFamily( + key="embeddinggemma-2", + hf_model_id="google/embeddinggemma-2", + dimension=768, + # Model card only (no HF JSON lists them); quality is near-lossless down to 256. + matryoshka_dimensions=(768, 512, 256, 128), + prompt_names={ + "query": "query", + "document": "document", + "sentence_similarity": "STS", + "classification": "Classification", + "clustering": "Clustering", + "question_answering": "QuestionAnswering", + "fact_checking": "FactChecking", + "code_retrieval": "CodeRetrieval", + }, + precision_floor="bfloat16", + forbidden_precisions=("float16",), + model_max_tokens=8192, + # tokenizer.json post-processor is TemplateProcessing " A ". + special_tokens={"bos": True, "eos": True}, + bos_token_id=2, + eos_token_id=1, + pad_token_id=0, + document_template="title: {title} | text: {text}", + # config.json image/boi/eoi token ids; processor_config.json image_processor. + image=ImageTowerFacts(placeholder_id=258880, boi_token_id=255999, eoi_token_id=258882), +) + +EMBEDDING_FAMILIES: dict[str, EmbeddingFamily] = {EMBEDDINGGEMMA2.key: EMBEDDINGGEMMA2} + + +def family_for_model(name: str) -> EmbeddingFamily: + """Return the family whose key or HF id matches ``name`` (case-insensitive).""" + wanted = name.lower() + for family in EMBEDDING_FAMILIES.values(): + if wanted in (family.key.lower(), family.hf_model_id.lower()): + return family + raise KeyError(f"No embedding family {name!r}; known: {sorted(EMBEDDING_FAMILIES)}") + + +def default_precision(family: EmbeddingFamily, platform: str) -> str: + return "float16" if platform == "iOS" else family.default_compute_precision + + +def check_compute_precision( + family: EmbeddingFamily, precision: str, *, platform: str = "macOS" +) -> str: + """Validate ``precision`` for ``family`` on ``platform`` and return it. + + Raises: + EmbeddingPrecisionError: if the precision is unknown, not float16 on iOS, or forbidden + for the family on macOS. + """ + if precision not in SUPPORTED_COMPUTE_PRECISIONS: + raise EmbeddingPrecisionError( + f"Unknown compute precision {precision!r}; expected one of " + f"{list(SUPPORTED_COMPUTE_PRECISIONS)}" + ) + if platform == "iOS": + if precision != "float16": + raise EmbeddingPrecisionError( + f"iOS (Neural Engine) bundles compute in float16; got {precision}." + ) + return precision + if precision in family.forbidden_precisions: + raise EmbeddingPrecisionError( + f"{family.hf_model_id} forbids {precision} compute on macOS (it overflows: NaN or " + f"silently degraded vectors). Use float32 (default) or " + f"{family.precision_floor or 'float32'}." + ) + return precision + + +def consolidate_prompts( + family: EmbeddingFamily, snapshot_prompts: dict[str, str] +) -> dict[str, str]: + """Collapse the snapshot's prompt table to canonical name -> prefix, in family order. + + Fails closed: a mapped snapshot key that is missing, or a snapshot prefix that no canonical + prompt carries (a new upstream task), raises instead of being dropped. + """ + missing = [key for key in family.prompt_names.values() if key not in snapshot_prompts] + if missing: + raise ValueError(f"snapshot prompts are missing {missing}") + canonical = {name: snapshot_prompts[key] for name, key in family.prompt_names.items()} + unknown = sorted( + key for key, prefix in snapshot_prompts.items() if prefix not in canonical.values() + ) + if unknown: + raise ValueError(f"snapshot prompts {unknown} use prefixes outside the canonical set") + return canonical + + +@dataclass(frozen=True) +class FunctionSpec: + """One static function; text inputs are ``int32[batch, length]``.""" + + name: str + batch: int + length: int + + def as_dict(self) -> dict[str, Any]: + return {"function": self.name, "batch": self.batch, "length": self.length} + + +def function_name(length: int, batch: int) -> str: + return f"len{length}_b{batch}" + + +def bucket_plan( + length_buckets: list[int] | tuple[int, ...], + batch_sizes: list[int] | tuple[int, ...], + max_batch_tokens: int, +) -> list[FunctionSpec]: + """Cross length buckets with batch sizes, dropping pairs above ``max_batch_tokens``. + + Batch 1 is always kept for every length (the unbatched path). Attention cost is O(B * L^2), so + the token cap is what keeps the large buckets small. + """ + lengths = sorted(set(int(x) for x in length_buckets)) + batches = sorted(set(int(x) for x in batch_sizes) | {1}) + if not lengths or lengths[0] < 2: + raise ValueError(f"length_buckets must be >= 2 (BOS + EOS), got {length_buckets}") + if batches[0] < 1: + raise ValueError(f"batch_sizes must be >= 1, got {batch_sizes}") + return [ + FunctionSpec(function_name(length, batch), batch, length) + for length in lengths + for batch in batches + if batch == 1 or batch * length <= max_batch_tokens + ] + + +def text_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec]: + """Text-tower functions: batched buckets on macOS, batch 1 per Neural Engine length on iOS.""" + if platform == "iOS": + return bucket_plan(family.ane_length_buckets, (1,), family.max_batch_tokens) + return bucket_plan(family.length_buckets, family.batch_sizes, family.max_batch_tokens) + + +def encoder_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec]: + """Batch-1 ``inputs_embeds`` functions of the ``encoder`` component.""" + if platform == "iOS": + lengths = sorted(set(family.ane_length_buckets)) + else: + assert family.image is not None + first = -(-family.image.sequence_length() // 32) * 32 + lengths = sorted({first} | {n for n in family.length_buckets if n > first}) + return [FunctionSpec(function_name(n, 1), 1, n) for n in lengths] + + +def bundle_name(family: EmbeddingFamily, platform: str, with_image: bool) -> str: + return family.key + ("-ios" if platform == "iOS" else "") + ("" if with_image else "-text") + + +def bundle_assets(platform: str, with_image: bool) -> dict[str, str]: + """Component name -> asset file for the bundle.""" + if platform == "iOS": + components = [TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT] + if with_image: + components.append(IMAGE_COMPONENT) + elif with_image: + components = [TEXT_COMPONENT, IMAGE_COMPONENT, TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT] + else: + components = [TEXT_COMPONENT] + return {name: f"{name}.aimodel" for name in components} + + +def embedding_block( + family: EmbeddingFamily, + *, + platform: str, + with_image: bool, + revision: str, + prompts: dict[str, str], + default_prompt_name: str | None = None, +) -> dict[str, Any]: + """Build the ``embedding`` block of ``metadata.json``. + + macOS: a ``text_encoder`` graph (ids in, vector out), plus the image tower + ``image_encoder -> token_table -> encoder`` when ``with_image``. iOS: the text tower runs + ``token_table -> encoder`` on the Neural Engine. The host merges soft tokens into the table + output at the placeholder rows before the encoder. + """ + ios = platform == "iOS" + if platform not in PLATFORMS: + raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") + precision = check_compute_precision( + family, default_precision(family, platform), platform=platform + ) + if not revision: + raise ValueError("revision (snapshot SHA) is required to stamp the embedding space") + if not prompts: + raise ValueError("prompts must not be empty") + if default_prompt_name is not None and default_prompt_name not in prompts: + raise ValueError(f"default_prompt_name {default_prompt_name!r} not in prompts") + if family.matryoshka_dimensions and family.matryoshka_dimensions[0] != family.dimension: + raise ValueError("matryoshka_dimensions must start with the native dimension") + image = family.image + if with_image and image is None: + raise ValueError(f"{family.hf_model_id} has no image tower") + + text_fns = text_functions(family, platform) + text_lengths = sorted({f.length for f in text_fns}) + if text_lengths[-1] > family.model_max_tokens: + raise ValueError( + f"length bucket {text_lengths[-1]} exceeds model_max_tokens {family.model_max_tokens}" + ) + enc_fns: list[FunctionSpec] = [] + if with_image or ios: + enc_fns = encoder_functions(family, platform) + functions_json = [f.as_dict() for f in enc_fns] + + specs: dict[str, Any] = { + TEXT_COMPONENT: { + "functions": [f.as_dict() for f in text_fns], + "inputs": list(TEXT_INPUT_NAMES), + "output": OUTPUT_NAME, + "compute_units": ["gpu"], + }, + IMAGE_COMPONENT: { + "function": "main", + "inputs": list(IMAGE_INPUT_NAMES), + "output": IMAGE_OUTPUT_NAME, + "compute_units": ["neural_engine"] if ios else ["gpu"], + }, + TOKEN_TABLE_COMPONENT: { + "function": "main", + "inputs": list(TOKEN_TABLE_INPUT_NAMES), + "output": TOKEN_TABLE_OUTPUT_NAME, + "compute_units": ["gpu"], + }, + ENCODER_COMPONENT: { + "functions": functions_json, + "inputs": list(ANE_ENCODER_INPUT_NAMES if ios else ENCODER_INPUT_NAMES), + "output": OUTPUT_NAME, + "mask_encoding": "additive_bias" if ios else "attention_mask", + "compute_units": ["neural_engine"] if ios else ["gpu"], + }, + } + components = {name: specs[name] for name in bundle_assets(platform, with_image)} + + text_tower: dict[str, Any] = { + "modality": "text", + "pipeline": [TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT] if ios else [TEXT_COMPONENT], + "inputs": {"input_ids": "int32[B,L]", "attention_mask": "int32[B,L]"}, + "output": OUTPUT_NAME, + "batch_sizes": sorted({f.batch for f in text_fns}), + "length_buckets": text_lengths, + } + if not ios: + text_tower["max_batch_tokens"] = family.max_batch_tokens + towers = [text_tower] + if with_image: + assert image is not None + image_lengths = sorted(f.length for f in enc_fns if f.length >= image.sequence_length()) + if not image_lengths: + raise ValueError( + f"no encoder function fits the longest image sequence ({image.sequence_length()})" + ) + towers.append( + { + "modality": "image", + "pipeline": [IMAGE_COMPONENT, TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT], + "output": OUTPUT_NAME, + "batch_sizes": [1], + "length_buckets": image_lengths, + } + ) + + special_tokens: dict[str, Any] = dict(family.special_tokens) + for key, value in ( + ("bos_token_id", family.bos_token_id), + ("eos_token_id", family.eos_token_id), + ("pad_token_id", family.pad_token_id), + ): + if value is not None: + special_tokens[key] = value + + block: dict[str, Any] = { + "space": {"id": family.hf_model_id, "revision": revision}, + "dimension": family.dimension, + "matryoshka_dimensions": list(family.matryoshka_dimensions), + "renormalize_after_truncation": family.renormalize_after_truncation, + "normalized": family.normalized, + "metric": family.metric, + "precision": { + "compute": precision, + "forbidden": [] if ios else list(family.forbidden_precisions), + "compute_units": ["neural_engine", "gpu"] if ios else ["gpu"], + }, + "components": components, + "towers": towers, + "text": { + "tokenizer": "tokenizer", + "padding_side": family.padding_side, + "special_tokens": special_tokens, + "truncation": {"keep_eos": family.truncation_keep_eos, "max_tokens": text_lengths[-1]}, + "model_max_tokens": family.model_max_tokens, + "prompts": dict(prompts), + "default_prompt_name": default_prompt_name, + "document_template": family.document_template, + }, + } + if with_image: + assert image is not None + block["merge"] = { + "placeholder_ids": {"image": image.placeholder_id}, + "token_ids": {"boi": image.boi_token_id, "eoi": image.eoi_token_id}, + } + block["preprocessing"] = { + "image": { + "patch_size": image.patch_size, + "rescale_factor": image.rescale_factor, + "normalize": image.normalize, + "image_mean": list(image.image_mean), + "image_std": list(image.image_std), + "max_soft_tokens": image.budget, + "pooling_kernel_size": image.pooling_kernel_size, + } + } + return block + + +def build_bundle_metadata( + family: EmbeddingFamily, + block: dict[str, Any], + *, + name: str, + assets: dict[str, str], + hf_revision: str | None, +) -> dict[str, Any]: + """Wrap an :func:`embedding_block` in the 0.2 ``metadata.json`` envelope.""" + return { + "metadata_version": METADATA_VERSION, + "kind": BUNDLE_KIND, + "name": name, + "assets": assets, + "embedding": block, + "source": { + "model_definition": "torch", + "hf_model_id": family.hf_model_id, + "hf_revision": hf_revision, + }, + "compression": None, + "compilation": {"date": datetime.now().astimezone().isoformat(), "targets": []}, + } diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py new file mode 100644 index 00000000..d41c7fbb --- /dev/null +++ b/python/src/coreai_models/embedding/export.py @@ -0,0 +1,819 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Export a retrieval-embedding model to a ``kind: "embedding"`` bundle. + + uv run coreai.embedding.export embeddinggemma-2 # macOS GPU bundle: text + image + uv run coreai.embedding.export embeddinggemma-2 --platform iOS # Neural Engine bundle + uv run coreai.embedding.export embeddinggemma-2 --skip-vision # text only + +The bundle contract (``metadata.json`` and the component assets) is defined in +``coreai_models.embedding.contract``. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import shutil +import sys +import time +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn + +from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO +from coreai_models.embedding.contract import ( + ANE_ENCODER_INPUT_NAMES, + EMBEDDING_FAMILIES, + ENCODER_COMPONENT, + ENCODER_INPUT_NAMES, + IMAGE_COMPONENT, + IMAGE_INPUT_NAMES, + IMAGE_OUTPUT_NAME, + OUTPUT_NAME, + PLATFORMS, + TEXT_COMPONENT, + TEXT_INPUT_NAMES, + TOKEN_TABLE_COMPONENT, + TOKEN_TABLE_INPUT_NAMES, + TOKEN_TABLE_OUTPUT_NAME, + EmbeddingFamily, + EmbeddingPrecisionError, + FunctionSpec, + ImageTowerFacts, + build_bundle_metadata, + bundle_assets, + bundle_name, + check_compute_precision, + consolidate_prompts, + default_precision, + embedding_block, + encoder_functions, + family_for_model, + text_functions, +) + +if TYPE_CHECKING: + from coreai.authoring import AIProgram + +logger = logging.getLogger(__name__) + +TOKENIZER_FILES = ( + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "special_tokens_map.json", + "added_tokens.json", +) +SNAPSHOT_PATTERNS = [ + "*.json", + "*.safetensors", + "tokenizer.model", + "1_Pooling/*", + "2_Normalize/*", +] +_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16} + + +def _find_repo_root() -> Path | None: + d = Path(__file__).resolve().parent + while d != d.parent: + if (d / "pyproject.toml").exists() and (d / "python").exists(): + return d + d = d.parent + return None + + +def _default_output_dir() -> str: + root = _find_repo_root() + return str(root / "exports") if root is not None else "exports" + + +def resolve_snapshot(hf_model_id: str, model_dir: str | None = None) -> tuple[Path, str]: + """Return ``(snapshot_dir, revision_sha)`` for a local dir or a hub download. + + The SHA is the HF cache snapshot directory name (``.../snapshots/``); a plain directory + outside the cache yields ``""``. + """ + if model_dir is not None: + path = Path(model_dir).expanduser().resolve() + if not path.is_dir(): + raise FileNotFoundError(f"model dir {path} is not a directory") + else: + from huggingface_hub import snapshot_download + + path = Path(snapshot_download(hf_model_id, allow_patterns=SNAPSHOT_PATTERNS)).resolve() + sha = path.name if path.parent.name == "snapshots" else "" + return path, sha + + +def read_prompts(snapshot: Path) -> tuple[dict[str, str], str | None]: + """Read ``prompts`` / ``default_prompt_name`` from ``config_sentence_transformers.json``.""" + cfg_path = snapshot / "config_sentence_transformers.json" + if not cfg_path.exists(): + return {}, None + cfg = json.loads(cfg_path.read_text()) + return dict(cfg.get("prompts") or {}), cfg.get("default_prompt_name") + + +def verify_pooling_config(snapshot: Path, family: EmbeddingFamily) -> None: + """Fail closed if the snapshot's sentence-transformers pooling disagrees with the family.""" + pooling = snapshot / "1_Pooling" / "config.json" + if not pooling.exists(): + logger.warning("No 1_Pooling/config.json in %s; trusting the family defaults", snapshot) + return + cfg = json.loads(pooling.read_text()) + mode = cfg.get("pooling_mode") + if mode is None: + mode = "mean" if cfg.get("pooling_mode_mean_tokens") else None + problems = [] + if mode != family.pooling_mode: + problems.append(f"pooling_mode {mode!r} != {family.pooling_mode!r}") + if bool(cfg.get("include_prompt", True)) != family.include_prompt: + problems.append(f"include_prompt {cfg.get('include_prompt')} != {family.include_prompt}") + if int(cfg.get("embedding_dimension", family.dimension)) != family.dimension: + problems.append( + f"embedding_dimension {cfg.get('embedding_dimension')} != {family.dimension}" + ) + if family.normalized and not (snapshot / "2_Normalize").exists(): + problems.append("family is normalized but the snapshot has no 2_Normalize module") + if problems: + raise ValueError(f"Snapshot pooling config disagrees with the family: {problems}") + + +def copy_tokenizer(snapshot: Path, dest: Path) -> list[str]: + """Copy the tokenizer files verbatim (no re-serialization through a tokenizer class).""" + dest.mkdir(parents=True, exist_ok=True) + copied = [] + for name in TOKENIZER_FILES: + src = snapshot / name + if src.exists(): + shutil.copy2(src, dest / name) + copied.append(name) + if "tokenizer.json" not in copied: + raise FileNotFoundError(f"{snapshot} has no tokenizer.json") + return copied + + +class TextEmbeddingExportWrapper(nn.Module): + """``(input_ids, attention_mask) -> embedding`` with float32 output. + + The wrapped encoder owns pooling (masked mean incl. prompt), projection and L2 normalization. + """ + + def __init__(self, encoder: nn.Module) -> None: + super().__init__() + self.encoder = encoder + + def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + return self.encoder(input_ids=input_ids, attention_mask=attention_mask).to(torch.float32) + + +class TokenTableExportWrapper(nn.Module): + """``input_ids [B, L] int32 -> inputs_embeds [B, L, H]`` float32 (scaled token embeddings).""" + + def __init__(self, embed_tokens: nn.Module) -> None: + super().__init__() + self.embed_tokens = embed_tokens + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids).to(torch.float32) + + +class InputsEmbedsEncoderWrapper(nn.Module): + """``(inputs_embeds, attention_mask) -> embedding`` for merged (e.g. image) inputs.""" + + def __init__(self, encoder: nn.Module) -> None: + super().__init__() + self.encoder = encoder + + def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + dtype = next(self.encoder.parameters()).dtype + return self.encoder( + inputs_embeds=inputs_embeds.to(dtype), attention_mask=attention_mask + ).to(torch.float32) + + +class EncoderBandWrapper(nn.Module): + """``(inputs_embeds, bias, pool_weights) -> embedding`` at one static length (Neural Engine). + + When ``length > sliding_window + 1`` the sliding layers also need the ``|i - j| <= window`` + band; it is a per-length constant folded in here (``min(bias, band)``), so every length keeps + the same three inputs. + """ + + def __init__(self, encoder: nn.Module, length: int, sliding_window: int) -> None: + super().__init__() + from coreai_models.models.ios.embedding_gemma2 import MASK_NEG + + self.encoder = encoder + self.banded = length > sliding_window + 1 + if self.banded: + pos = torch.arange(length) + band = torch.where((pos[:, None] - pos[None, :]).abs() <= sliding_window, 0.0, MASK_NEG) + dtype = next(encoder.parameters()).dtype + self.register_buffer("band", band.reshape(1, length, 1, length).to(dtype)) + + def forward( + self, inputs_embeds: torch.Tensor, bias: torch.Tensor, pool_weights: torch.Tensor + ) -> torch.Tensor: + sliding = torch.minimum(bias, self.band) if self.banded else None + return self.encoder(inputs_embeds, bias, pool_weights, sliding) + + +class ImageEncoderExportWrapper(nn.Module): + """``(pixel_values f32, position_ids int32) -> soft_tokens f32`` at the module's dtype.""" + + def __init__(self, image_encoder: nn.Module) -> None: + super().__init__() + self.image_encoder = image_encoder + + def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + dtype = next(self.image_encoder.parameters()).dtype + return self.image_encoder(pixel_values.to(dtype), position_ids).to(torch.float32) + + +@contextmanager +def without_token_table(encoder: nn.Module) -> Iterator[nn.Module]: + """Temporarily drop ``encoder.embed_tokens`` so an ``inputs_embeds`` trace carries no table.""" + table = encoder.embed_tokens + encoder.embed_tokens = nn.Identity() + try: + yield encoder + finally: + encoder.embed_tokens = table + + +def reference_inputs(spec: FunctionSpec, vocab_size: int) -> tuple[torch.Tensor, torch.Tensor]: + """Trace inputs for one function (values do not matter; shapes and dtypes do).""" + generator = torch.Generator().manual_seed(spec.batch * 100_003 + spec.length) + ids = torch.randint(3, vocab_size, (spec.batch, spec.length), generator=generator) + mask = torch.ones(spec.batch, spec.length, dtype=torch.int32) + return ids.to(torch.int32), mask + + +def image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, torch.Tensor]: + """Image-encoder trace inputs: a full block grid padded to the budget (int32 positions).""" + k = image.pooling_kernel_size + n = image.budget + # Square-ish grid of n soft tokens (blocks) -> (bw*k) x (bh*k) patches, row-major. + bw = max(1, int(n**0.5)) + bh = max(1, n // bw) + pw, ph = bw * k, bh * k + ys, xs = torch.meshgrid(torch.arange(ph), torch.arange(pw), indexing="ij") + real = torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1).to(torch.int32) + positions = torch.full((1, image.max_patches, 2), -1, dtype=torch.int32) + positions[0, : real.shape[0]] = real + generator = torch.Generator().manual_seed(image.budget) + pixels = torch.rand((1, image.max_patches, image.patch_dim), generator=generator) + pixels[0, real.shape[0] :] = 0.0 + return pixels, positions + + +def ane_text_reference_inputs( + length: int, hidden: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Neural Engine encoder trace inputs (all valid tokens): embeds, zero bias, uniform pooling.""" + generator = torch.Generator().manual_seed(length) + embeds = torch.randn((1, length, hidden), generator=generator).to(torch.float16) + bias = torch.zeros((1, length, 1, length), dtype=torch.float16) + pool = torch.full((1, length), 1.0 / length, dtype=torch.float16) + return embeds, bias, pool + + +def ane_image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, torch.Tensor]: + """Neural Engine image trace inputs (float16): a block grid padded to the budget (-1 pad).""" + k = image.pooling_kernel_size + a = int(image.budget**0.5) + gh, gw = k * a, k * (image.budget // a) + generator = torch.Generator().manual_seed(image.budget) + pixels = torch.rand((1, image.max_patches, image.patch_dim), generator=generator) + ys, xs = torch.meshgrid(torch.arange(gh), torch.arange(gw), indexing="ij") + positions = torch.full((1, image.max_patches, 2), -1.0) + positions[0, : gh * gw] = torch.stack([xs.reshape(-1), ys.reshape(-1)], -1).float() + return pixels.to(torch.float16), positions.to(torch.float16) + + +@dataclass +class TraceSpec: + """One function of a Neural Engine program: example args + IO names.""" + + name: str + args: tuple[torch.Tensor, ...] + input_names: tuple[str, ...] + output_names: tuple[str, ...] + + +def convert_ane_functions( + traces: list[tuple[nn.Module, TraceSpec]], + *, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """``torch.export`` every ``(module, trace)``; convert them into one multi-function program. + + Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. + """ + import coreai_torch + + from coreai_models.export.mlir_ops import register_custom_torch_lowering + + converter = coreai_torch.TorchConverter( + mode=( + coreai_torch.TorchConverter.Mode.DEBUG + if include_debug_info + else coreai_torch.TorchConverter.Mode.RELEASE + ) + ) + register_custom_torch_lowering(converter) + decomp = torch.export.default_decompositions() + for module, trace in traces: + module.eval() + start = time.time() + with torch.no_grad(): + exported = torch.export.export(module, args=trace.args) + converter.add_exported_program( + exported.run_decompositions(decomp), + input_names=list(trace.input_names), + output_names=list(trace.output_names), + entrypoint_name=trace.name, + ) + logger.info( + "traced %s %s in %.1fs", + trace.name, + [tuple(a.shape) for a in trace.args], + time.time() - start, + ) + return converter.to_coreai() + + +def export_text_encoder_program( + wrapper: nn.Module, + functions: list[FunctionSpec], + vocab_size: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """GPU ``text_encoder``: one static (batch, length) function per spec, weights shared.""" + from coreai_models.diffusion.gpu import export_multifunction + + # Static (batch, length) traces; the host pads to the next listed batch. + return export_multifunction( + [(spec.name, wrapper, reference_inputs(spec, vocab_size)) for spec in functions], + TEXT_INPUT_NAMES, + (OUTPUT_NAME,), + include_debug_info, + ) + + +def export_token_table_program( + embed_tokens: nn.Module, + max_length: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """``token_table``: function ``main``, ids -> scaled embeddings, dynamic length (batch 1).""" + from coreai_models.diffusion.gpu import export_stateless + + # Placeholder ids look up like any id; the host overwrites those rows with soft tokens. + length = torch.export.Dim("length", min=2, max=max_length) + ids = torch.zeros((1, 64), dtype=torch.int32) + return export_stateless( + TokenTableExportWrapper(embed_tokens), + (ids,), + TOKEN_TABLE_INPUT_NAMES, + (TOKEN_TABLE_OUTPUT_NAME,), + dynamic_shapes=({1: length},), + include_debug_info=include_debug_info, + ) + + +def export_image_programs( + text_encoder: nn.Module, + image_encoder: nn.Module, + image: ImageTowerFacts, + functions: list[FunctionSpec], + max_length: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> dict[str, AIProgram]: + """GPU image tower: ``{image_encoder, token_table, encoder}`` programs. + + ``max_length`` bounds the token table's dynamic sequence dimension. + """ + from coreai_models.diffusion.gpu import export_multifunction, export_stateless + + programs: dict[str, AIProgram] = { + TOKEN_TABLE_COMPONENT: export_token_table_program( + text_encoder.embed_tokens, max_length, include_debug_info + ) + } + with without_token_table(text_encoder) as encoder: + wrapper = InputsEmbedsEncoderWrapper(encoder) + traces = [] + for spec in functions: + generator = torch.Generator().manual_seed(spec.length) + embeds = torch.randn( + (spec.batch, spec.length, image.soft_token_dim), generator=generator + ) + mask = torch.ones((spec.batch, spec.length), dtype=torch.int32) + traces.append((spec.name, wrapper, (embeds, mask))) + programs[ENCODER_COMPONENT] = export_multifunction( + traces, ENCODER_INPUT_NAMES, (OUTPUT_NAME,), include_debug_info + ) + programs[IMAGE_COMPONENT] = export_stateless( + ImageEncoderExportWrapper(image_encoder), + image_reference_inputs(image), + IMAGE_INPUT_NAMES, + (IMAGE_OUTPUT_NAME,), + include_debug_info=include_debug_info, + ) + return programs + + +def export_ane_encoder_program( + encoder: nn.Module, + functions: list[FunctionSpec], + sliding_window: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """Neural Engine ``encoder``: one static batch-1 float16 function per length.""" + hidden = int(encoder.config.hidden_size) + traces = [ + ( + EncoderBandWrapper(encoder, spec.length, sliding_window), + TraceSpec( + spec.name, + ane_text_reference_inputs(spec.length, hidden), + ANE_ENCODER_INPUT_NAMES, + (OUTPUT_NAME,), + ), + ) + for spec in functions + ] + return convert_ane_functions(traces, include_debug_info=include_debug_info) + + +def export_ane_image_encoder_program( + image_encoder: nn.Module, + image: ImageTowerFacts, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """Neural Engine ``image_encoder``: function ``main`` at the static budget.""" + trace = TraceSpec( + "main", ane_image_reference_inputs(image), IMAGE_INPUT_NAMES, (IMAGE_OUTPUT_NAME,) + ) + return convert_ane_functions([(image_encoder, trace)], include_debug_info=include_debug_info) + + +def verify_image_config(snapshot: Path, family: EmbeddingFamily) -> None: + """Fail closed if the snapshot's merge ids / image processor disagree with the family facts.""" + image = family.image + if image is None: + raise ValueError(f"{family.hf_model_id} has no image tower") + problems = [] + config_path = snapshot / "config.json" + if config_path.exists(): + cfg = json.loads(config_path.read_text()) + for key, want in ( + ("image_token_id", image.placeholder_id), + ("boi_token_id", image.boi_token_id), + ("eoi_token_id", image.eoi_token_id), + ): + if key in cfg and int(cfg[key]) != want: + problems.append(f"config.json {key} {cfg[key]} != {want}") + text_hidden = (cfg.get("text_config") or {}).get("hidden_size") + if text_hidden is not None and int(text_hidden) != image.soft_token_dim: + problems.append(f"text hidden_size {text_hidden} != {image.soft_token_dim}") + processor_path = snapshot / "processor_config.json" + if processor_path.exists(): + proc = json.loads(processor_path.read_text()).get("image_processor") or {} + checks = ( + ("patch_size", image.patch_size), + ("pooling_kernel_size", image.pooling_kernel_size), + ("max_soft_tokens", image.budget), + ("do_normalize", image.normalize), + ("resample", 3), # PIL/torchvision BICUBIC + ) + for key, want in checks: + if key in proc and proc[key] != want: + problems.append(f"processor {key} {proc[key]!r} != {want!r}") + if "rescale_factor" in proc and abs(proc["rescale_factor"] - image.rescale_factor) > 1e-12: + problems.append(f"processor rescale_factor {proc['rescale_factor']} != 1/255") + if problems: + raise ValueError(f"Snapshot image config disagrees with the family: {problems}") + + +def _load_text_encoder(family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype) -> nn.Module: + if family.key == "embeddinggemma-2": + from coreai_models.models.macos.embedding_gemma2 import load_embedding_gemma2_text_encoder + + return load_embedding_gemma2_text_encoder(snapshot, dtype=dtype) + raise NotImplementedError(f"No text-encoder loader for family {family.key!r}") + + +def _load_image_encoder( + family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype, max_soft_tokens: int +) -> nn.Module: + if family.key == "embeddinggemma-2": + from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder + + return EmbeddingGemma2ImageEncoder.from_pretrained( + str(snapshot), dtype=dtype, max_soft_tokens=max_soft_tokens + ) + raise NotImplementedError(f"No image-encoder loader for family {family.key!r}") + + +def _save(program: AIProgram, out: Path, family: EmbeddingFamily, component: str) -> None: + from coreai_models.export import metadata as export_metadata + + program.save_asset( + out / f"{component}.aimodel", + export_metadata.build_aimodel_metadata(family.hf_model_id, component=component), + ) + + +def _export_gpu_programs( + family: EmbeddingFamily, + snapshot: Path, + out: Path, + *, + with_image: bool, + compute_precision: str, + include_debug_info: bool, +) -> None: + dtype = _DTYPES[compute_precision] + encoder = _load_text_encoder(family, snapshot, dtype) + program = export_text_encoder_program( + TextEmbeddingExportWrapper(encoder), + text_functions(family, "macOS"), + int(encoder.config.vocab_size), + include_debug_info, + ) + _save(program, out, family, TEXT_COMPONENT) + del program + if not with_image: + return + assert family.image is not None + image_encoder = _load_image_encoder(family, snapshot, dtype, family.image.budget) + programs = export_image_programs( + encoder, + image_encoder, + family.image, + encoder_functions(family, "macOS"), + max_length=family.model_max_tokens, + include_debug_info=include_debug_info, + ) + for component, component_program in programs.items(): + _save(component_program, out, family, component) + + +def _export_ane_programs( + family: EmbeddingFamily, + snapshot: Path, + out: Path, + *, + with_image: bool, + include_debug_info: bool, +) -> None: + if family.key != "embeddinggemma-2": + raise NotImplementedError(f"No Neural Engine export for family {family.key!r}") + 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) + _save( + export_token_table_program( + gpu_encoder.embed_tokens, family.model_max_tokens, include_debug_info + ), + out, + family, + TOKEN_TABLE_COMPONENT, + ) + del gpu_encoder + + encoder, _table = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float16) + del _table + _save( + export_ane_encoder_program( + encoder, + encoder_functions(family, "iOS"), + int(encoder.config.sliding_window), + include_debug_info, + ), + out, + family, + ENCODER_COMPONENT, + ) + del encoder + if with_image: + assert family.image is not None + from coreai_models.models.ios.embedding_gemma2_vision import load_image_encoder + + image_encoder = load_image_encoder( + str(snapshot), torch.float16, max_soft_tokens=family.image.budget + ) + _save( + export_ane_image_encoder_program(image_encoder, family.image, include_debug_info), + out, + family, + IMAGE_COMPONENT, + ) + + +def export_embedding_bundle( + family: EmbeddingFamily, + output_dir: Path, + *, + platform: str = "macOS", + with_image: bool = True, + compute_precision: str | None = None, + overwrite: bool = False, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, + model_dir: str | None = None, +) -> dict[str, Any]: + """Export one bundle to ``output_dir`` and return its ``metadata.json`` content. + + macOS writes a GPU bundle (``text_encoder``, plus ``image_encoder`` / ``token_table`` / + ``encoder`` with images); iOS writes the Neural Engine bundle (``token_table``, ``encoder``, + plus ``image_encoder``). Both carry the tokenizer and a ``metadata.json``. + """ + if platform not in PLATFORMS: + raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") + precision = check_compute_precision( + family, compute_precision or default_precision(family, platform), platform=platform + ) # before touching disk + if with_image and family.image is None: + raise ValueError(f"{family.hf_model_id} has no image tower") + out = Path(output_dir) + if out.exists(): + if not overwrite: + raise FileExistsError(f"{out} exists (pass --overwrite)") + shutil.rmtree(out) + out.mkdir(parents=True) + + snapshot, sha = resolve_snapshot(family.hf_model_id, model_dir) + if not sha: + raise ValueError(f"Could not infer the snapshot SHA from {snapshot}") + verify_pooling_config(snapshot, family) + if with_image: + verify_image_config(snapshot, family) + snapshot_prompts, default_prompt = read_prompts(snapshot) + prompts = consolidate_prompts(family, snapshot_prompts) + copy_tokenizer(snapshot, out / "tokenizer") + + if platform == "iOS": + _export_ane_programs( + family, + snapshot, + out, + with_image=with_image, + include_debug_info=include_debug_info, + ) + else: + _export_gpu_programs( + family, + snapshot, + out, + with_image=with_image, + compute_precision=precision, + include_debug_info=include_debug_info, + ) + + block = embedding_block( + family, + platform=platform, + with_image=with_image, + revision=sha[:16], + prompts=prompts, + default_prompt_name=default_prompt, + ) + metadata = build_bundle_metadata( + family, + block, + name=bundle_name(family, platform, with_image), + assets=bundle_assets(platform, with_image), + hf_revision=sha, + ) + (out / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + logger.info("wrote %s", out) + return metadata + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="coreai.embedding.export", + description="Export a retrieval-embedding model to a kind=embedding bundle.", + ) + parser.add_argument( + "model", + nargs="?", + help="Registry short name (e.g. embeddinggemma-2) or HuggingFace model ID", + ) + parser.add_argument( + "--platform", + choices=list(PLATFORMS), + default=None, + help="Target platform. macOS: GPU bundle (float32). iOS: Neural Engine bundle " + "(float16); also runs on a Mac.", + ) + parser.add_argument( + "--skip-vision", + action="store_true", + help="Export the text tower only (skip the image encoder).", + ) + parser.add_argument( + "--compute-precision", + choices=["float32", "bfloat16", "float16"], + default=None, + help="Default: float32 on macOS, float16 on iOS.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=None, + help="Parent directory for the bundle (default: /exports).", + ) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--dry-run", action="store_true", help="Print the export plan and exit.") + parser.add_argument("--list-models", action="store_true") + parser.add_argument( + "--include-debug-info", + action="store_true", + help=( + "Embed debug information in the exported .aimodel for debugging a conversion. " + "Default: off, which embeds minimum debug information and makes the " + "exported asset smaller." + ), + ) + parser.add_argument("-v", "--verbose", action="store_true") + return parser + + +def _print_models() -> None: + for family in EMBEDDING_FAMILIES.values(): + towers = "text,image" if family.image else "text" + print(f"{family.key:24s} {family.hf_model_id:32s} {towers:12s} {','.join(PLATFORMS)}") + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + if args.list_models: + _print_models() + return 0 + if args.model is None: + parser.error("a model is required (e.g. embeddinggemma-2); see --list-models") + try: + family = family_for_model(args.model) + except KeyError: + parser.error( + f"unknown model {args.model!r}. Run `uv run coreai.embedding.export --list-models`." + ) + platform = args.platform or "macOS" + with_image = not args.skip_vision + try: + precision = check_compute_precision( + family, + args.compute_precision or default_precision(family, platform), + platform=platform, + ) + except EmbeddingPrecisionError as error: + parser.error(str(error)) + out = Path(args.output_dir or _default_output_dir()) / bundle_name(family, platform, with_image) + if args.dry_run: + assets = bundle_assets(platform, with_image) + towers = "text,image" if with_image else "text" + print(f"platform: {platform}") + print(f"precision: {precision}") + print(f"towers: {towers}") + print(f"assets: {', '.join(assets.values())}") + if platform == "macOS": + print(f"text functions: {[f.name for f in text_functions(family, platform)]}") + if with_image or platform == "iOS": + print(f"encoder functions: {[f.name for f in encoder_functions(family, platform)]}") + print(f"output: {out}") + return 0 + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + ) + export_embedding_bundle( + family, + out, + platform=platform, + with_image=with_image, + compute_precision=precision, + overwrite=args.overwrite, + include_debug_info=args.include_debug_info, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/src/coreai_models/export/metadata.py b/python/src/coreai_models/export/metadata.py index 61cab40d..84381c24 100644 --- a/python/src/coreai_models/export/metadata.py +++ b/python/src/coreai_models/export/metadata.py @@ -362,6 +362,17 @@ class AIModelMetadataFields: "Source: https://huggingface.co/Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers" ), ), + # ---- Embedding ---- + "google/embeddinggemma-2": AIModelMetadataFields( + author="Google DeepMind", + license="Apache-2.0", + model_description=( + "EmbeddingGemma 2 is a multimodal retrieval embedding model from Google DeepMind " + "producing 768-d L2-normalized vectors (Matryoshka prefixes 512/256/128) for text " + "and images in one space. " + "Source: https://huggingface.co/google/embeddinggemma-2" + ), + ), # ---- Segmentation ---- "facebook/sam3": AIModelMetadataFields( author="N. Carion et al.", diff --git a/python/src/coreai_models/model_registry.py b/python/src/coreai_models/model_registry.py index 018f9c9c..4b7c8600 100644 --- a/python/src/coreai_models/model_registry.py +++ b/python/src/coreai_models/model_registry.py @@ -600,6 +600,13 @@ class UtilityModel: UTILITY_PRESETS: list[UtilityModel] = [ # --- Embedding --- + UtilityModel( + "embeddinggemma-2", + "google/embeddinggemma-2", + "embeddinggemma", + "embedding", + "coreai.embedding.export", + ), UtilityModel( "clip-vit-b32", "openai/clip-vit-base-patch32", @@ -1011,9 +1018,17 @@ def _format_diffusion_header() -> str: return f"{'SHORT_NAME':32s} {'COMPRESSION':18s} HF_ID" -def _utility_to_export_args(model: UtilityModel) -> list[str]: - """Build the shell command to export a utility model.""" - return ["uv", "run", model.export_script, "--model", model.hf_id] +def _utility_to_export_args(model: UtilityModel, platform: str | None = None) -> list[str]: + """Build the shell command to export a utility model. + + Scripts (``*.py``) take ``--model ``; console scripts take the short name and + ``--platform iOS`` when targeting iOS. + """ + if model.export_script.endswith(".py"): + return ["uv", "run", model.export_script, "--model", model.hf_id] + return ["uv", "run", model.export_script, model.short_name] + ( + ["--platform", "iOS"] if platform == "iOS" else [] + ) def _format_utility_row(u: UtilityModel) -> str: @@ -1318,7 +1333,7 @@ def _action_utility_model_info(args: argparse.Namespace) -> None: sys.stderr.write(f"Error: model {args.model_info!r} has no --platform {args.platform!r}\n") sys.exit(1) if args.format == "export-args": - print(" ".join(_utility_to_export_args(model))) + print(" ".join(_utility_to_export_args(model, args.platform))) elif args.format == "json": print(json.dumps(asdict(model), indent=2)) elif args.format == "tsv": diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py new file mode 100644 index 00000000..c433d08f --- /dev/null +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -0,0 +1,667 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for the ``kind: "embedding"`` bundle contract and exporter (``coreai_models.embedding``). + +Metadata, precision and CLI tests need no hardware or weights. The conversion tests use tiny random +EmbeddingGemma 2 configs and skip when ``coreai_torch`` is unavailable. +""" + +from __future__ import annotations + +import dataclasses +import json +from pathlib import Path + +import pytest +import tomllib +import torch + +from coreai_models.embedding import export as embedding_export +from coreai_models.embedding.contract import ( + EMBEDDING_FAMILIES, + EmbeddingPrecisionError, + FunctionSpec, + bucket_plan, + build_bundle_metadata, + bundle_assets, + bundle_name, + check_compute_precision, + consolidate_prompts, + embedding_block, + encoder_functions, + family_for_model, + text_functions, +) + +EG2 = family_for_model("embeddinggemma-2") +IMAGE = EG2.image +REVISION = "914f7f89142e33e7" +PROMPTS = { + "query": "task: search result | query: ", + "document": "title: none | text: ", + "sentence_similarity": "task: sentence similarity | query: ", + "classification": "task: classification | query: ", + "clustering": "task: clustering | query: ", + "question_answering": "task: question answering | query: ", + "fact_checking": "task: fact checking | query: ", + "code_retrieval": "task: code retrieval | query: ", +} +# The snapshot's 20-entry prompt table (several keys share one prefix). +SNAPSHOT_PROMPTS = { + **dict.fromkeys( + ["query", "Retrieval-query", "SearchQuery", "Retrieval", "Reranking", "BitextMining"], + PROMPTS["query"], + ), + **dict.fromkeys(["document", "Document", "Retrieval-document"], PROMPTS["document"]), + **dict.fromkeys( + ["STS", "SentenceSimilarity", "PairClassification", "Summarization"], + PROMPTS["sentence_similarity"], + ), + **dict.fromkeys(["Classification", "MultilabelClassification"], PROMPTS["classification"]), + "Clustering": PROMPTS["clustering"], + "QuestionAnswering": PROMPTS["question_answering"], + "FactChecking": PROMPTS["fact_checking"], + **dict.fromkeys(["CodeRetrieval", "InstructionRetrieval"], PROMPTS["code_retrieval"]), +} +REMOVED_KEYS = { + "pooling", + "attention", + "granularity", + "budgets", + "max_sequence_length", + "prompt_applies_to", + "side", + "strategy", + "wrap", + "soft_tokens_scaled", + "floor", + "norm", + "dynamic_batch", + "dynamic_length", + "hidden_size", + "note", + "max_patches", + "soft_token_dim", + "convert_rgb", + "exif_transpose", + "resize", + "resize_policy", + "patch_pixel_order", + "position_order", + "padding_position", + "prompt", + "truncation_side", +} + + +def _block(platform="macOS", with_image=False, **overrides): + kwargs = dict(revision=REVISION, prompts=PROMPTS) + kwargs.update(overrides) + return embedding_block(EG2, platform=platform, with_image=with_image, **kwargs) + + +def _keys(node): + if isinstance(node, dict): + for key, value in node.items(): + yield key + yield from _keys(value) + elif isinstance(node, list): + for value in node: + yield from _keys(value) + + +# --- Precision guard --- + + +def test_ios_requires_float16(): + assert check_compute_precision(EG2, "float16", platform="iOS") == "float16" + for precision in ("float32", "bfloat16"): + with pytest.raises(EmbeddingPrecisionError, match="float16"): + check_compute_precision(EG2, precision, platform="iOS") + + +def test_macos_refuses_float16_and_allows_bfloat16(): + with pytest.raises(EmbeddingPrecisionError, match="forbids float16"): + check_compute_precision(EG2, "float16") + assert check_compute_precision(EG2, "bfloat16") == "bfloat16" + assert check_compute_precision(EG2, "float32") == "float32" + + +def test_unknown_precision_is_refused(): + with pytest.raises(EmbeddingPrecisionError, match="Unknown compute precision"): + check_compute_precision(EG2, "int8") + + +def test_export_bundle_refuses_float16_before_touching_out(tmp_path): + out = tmp_path / "bundle" + with pytest.raises(EmbeddingPrecisionError): + embedding_export.export_embedding_bundle(EG2, out, compute_precision="float16") + assert not out.exists() + + +# --- Prompts --- + + +def test_consolidate_prompts_maps_the_snapshot_table_to_eight_names(): + assert len(SNAPSHOT_PROMPTS) == 20 + prompts = consolidate_prompts(EG2, SNAPSHOT_PROMPTS) + assert prompts == PROMPTS + assert list(prompts) == list(EG2.prompt_names) + + +def test_consolidate_prompts_fails_closed(): + with pytest.raises(ValueError, match="canonical"): + consolidate_prompts(EG2, {**SNAPSHOT_PROMPTS, "Translation": "task: translate | query: "}) + missing = {k: v for k, v in SNAPSHOT_PROMPTS.items() if k != "FactChecking"} + with pytest.raises(ValueError, match="FactChecking"): + consolidate_prompts(EG2, missing) + + +# --- Bucket plan --- + + +def test_bucket_plan_respects_token_cap_and_keeps_batch_one(): + plan = bucket_plan([2048, 128, 512], [32, 8], 16384) + shapes = [(f.length, f.batch) for f in plan] + assert shapes == [ + (128, 1), + (128, 8), + (128, 32), + (512, 1), + (512, 8), + (512, 32), + (2048, 1), + (2048, 8), + ] + assert all(f.batch == 1 or f.batch * f.length <= 16384 for f in plan) + assert plan[0].name == "len128_b1" + assert len({f.name for f in plan}) == len(plan) + + +def test_bucket_plan_rejects_bad_lengths(): + with pytest.raises(ValueError): + bucket_plan([1], [1], 1024) + + +def test_function_plans_per_platform(): + assert [f.length for f in encoder_functions(EG2, "macOS")] == [288, 512, 2048] + assert [f.length for f in encoder_functions(EG2, "iOS")] == [128, 512, 1024] + assert [(f.length, f.batch) for f in text_functions(EG2, "iOS")] == [ + (128, 1), + (512, 1), + (1024, 1), + ] + assert IMAGE.sequence_length() == 284 # 280 x + + +# --- embedding_block schema --- + + +def test_embedding_block_macos_text(): + block = _block() + assert list(block) == [ + "space", + "dimension", + "matryoshka_dimensions", + "renormalize_after_truncation", + "normalized", + "metric", + "precision", + "components", + "towers", + "text", + ] + assert block["space"] == {"id": "google/embeddinggemma-2", "revision": REVISION} + assert block["matryoshka_dimensions"] == [768, 512, 256, 128] + assert block["precision"] == { + "compute": "float32", + "forbidden": ["float16"], + "compute_units": ["gpu"], + } + assert set(block["components"]) == {"text_encoder"} + component = block["components"]["text_encoder"] + assert component["output"] == "embedding" and component["compute_units"] == ["gpu"] + assert {"function": "len2048_b8", "batch": 8, "length": 2048} in component["functions"] + (tower,) = block["towers"] + assert tower["pipeline"] == ["text_encoder"] + assert tower["batch_sizes"] == [1, 8, 32] + assert tower["length_buckets"] == [128, 512, 2048] + assert tower["max_batch_tokens"] == 16384 + text = block["text"] + assert text["truncation"] == {"keep_eos": True, "max_tokens": 2048} + assert text["special_tokens"] == { + "bos": True, + "eos": True, + "bos_token_id": 2, + "eos_token_id": 1, + "pad_token_id": 0, + } + assert text["prompts"] == PROMPTS and text["default_prompt_name"] is None + json.dumps(block) + + +@pytest.mark.parametrize( + ("platform", "with_image", "components", "pipelines", "units"), + [ + ("macOS", False, ["text_encoder"], [["text_encoder"]], {"text_encoder": ["gpu"]}), + ( + "macOS", + True, + ["text_encoder", "image_encoder", "token_table", "encoder"], + [["text_encoder"], ["image_encoder", "token_table", "encoder"]], + { + "text_encoder": ["gpu"], + "image_encoder": ["gpu"], + "token_table": ["gpu"], + "encoder": ["gpu"], + }, + ), + ( + "iOS", + False, + ["token_table", "encoder"], + [["token_table", "encoder"]], + {"token_table": ["gpu"], "encoder": ["neural_engine"]}, + ), + ( + "iOS", + True, + ["token_table", "encoder", "image_encoder"], + [["token_table", "encoder"], ["image_encoder", "token_table", "encoder"]], + { + "token_table": ["gpu"], + "encoder": ["neural_engine"], + "image_encoder": ["neural_engine"], + }, + ), + ], +) +def test_embedding_block_platform_matrix(platform, with_image, components, pipelines, units): + block = _block(platform, with_image) + assert list(block["components"]) == components + assert [t["pipeline"] for t in block["towers"]] == pipelines + assert {k: v["compute_units"] for k, v in block["components"].items()} == units + assert ("merge" in block) == with_image and ("preprocessing" in block) == with_image + assert set(bundle_assets(platform, with_image)) == set(components) + assert not REMOVED_KEYS & set(_keys(block)), REMOVED_KEYS & set(_keys(block)) + assert all("batch" not in tower and "output" in tower for tower in block["towers"]) + assert "output" not in block["precision"] + json.dumps(block) + + +def test_embedding_block_ios_details(): + block = _block("iOS", True) + enc = block["components"]["encoder"] + assert enc["inputs"] == ["inputs_embeds", "bias", "pool_weights"] + assert enc["mask_encoding"] == "additive_bias" + assert [f["length"] for f in enc["functions"]] == [128, 512, 1024] + assert block["precision"] == { + "compute": "float16", + "forbidden": [], + "compute_units": ["neural_engine", "gpu"], + } + text_tower, image_tower = block["towers"] + assert "max_batch_tokens" not in text_tower + assert text_tower["length_buckets"] == [128, 512, 1024] + assert image_tower["length_buckets"] == [512, 1024] + assert image_tower["batch_sizes"] == [1] + assert block["text"]["truncation"]["max_tokens"] == 1024 + assert block["space"] == _block("macOS")["space"] + + +def test_embedding_block_macos_image_details(): + block = _block("macOS", True) + enc = block["components"]["encoder"] + assert enc["inputs"] == ["inputs_embeds", "attention_mask"] + assert enc["mask_encoding"] == "attention_mask" + assert block["towers"][1]["length_buckets"] == [288, 512, 2048] + assert block["merge"] == { + "placeholder_ids": {"image": 258880}, + "token_ids": {"boi": 255999, "eoi": 258882}, + } + pre = block["preprocessing"]["image"] + assert pre["patch_size"] == 16 and pre["pooling_kernel_size"] == 3 + assert pre["max_soft_tokens"] == 280 + assert pre["rescale_factor"] == pytest.approx(1 / 255) + assert pre["normalize"] is False + + +def test_embedding_block_validation(): + with pytest.raises(ValueError, match="revision"): + _block(revision="") + with pytest.raises(ValueError, match="prompts"): + _block(prompts={}) + with pytest.raises(ValueError, match="default_prompt_name"): + _block(default_prompt_name="missing") + with pytest.raises(ValueError, match="platform"): + _block("watchOS") + no_image = dataclasses.replace(EG2, image=None) + with pytest.raises(ValueError, match="no image tower"): + embedding_block(no_image, platform="macOS", with_image=True, revision="x", prompts=PROMPTS) + + +def test_bundle_metadata_envelope(): + meta = build_bundle_metadata( + EG2, + _block("iOS", True), + name=bundle_name(EG2, "iOS", True), + assets=bundle_assets("iOS", True), + hf_revision="914f7f89142e33e77833254d9c9b90c3cef7303b", + ) + assert meta["metadata_version"] == "0.2" and meta["kind"] == "embedding" + assert meta["name"] == "embeddinggemma-2-ios" + pipelines = {c for t in meta["embedding"]["towers"] for c in t["pipeline"]} + assert pipelines == set(meta["assets"]) + assert meta["source"]["hf_model_id"] == "google/embeddinggemma-2" + + +def test_family_registry(): + assert "embeddinggemma-2" in EMBEDDING_FAMILIES + assert family_for_model("Google/EmbeddingGemma-2") is EG2 + with pytest.raises(KeyError, match="embeddinggemma-2"): + family_for_model("google/embeddinggemma-300m") + + +# --- Snapshot helpers --- + + +def _fake_snapshot(root: Path, pooling: dict) -> Path: + snap = root / "models--google--embeddinggemma-2" / "snapshots" / ("a" * 40) + (snap / "1_Pooling").mkdir(parents=True) + (snap / "2_Normalize").mkdir() + (snap / "1_Pooling" / "config.json").write_text(json.dumps(pooling)) + (snap / "config_sentence_transformers.json").write_text( + json.dumps({"prompts": SNAPSHOT_PROMPTS, "default_prompt_name": None}) + ) + (snap / "tokenizer.json").write_text("{}") + (snap / "tokenizer_config.json").write_text("{}") + return snap + + +def test_snapshot_helpers(tmp_path): + snap = _fake_snapshot( + tmp_path, {"embedding_dimension": 768, "pooling_mode": "mean", "include_prompt": True} + ) + path, sha = embedding_export.resolve_snapshot("google/embeddinggemma-2", str(snap)) + assert sha == "a" * 40 and path == snap.resolve() + assert embedding_export.read_prompts(snap) == (SNAPSHOT_PROMPTS, None) + embedding_export.verify_pooling_config(snap, EG2) + copied = embedding_export.copy_tokenizer(snap, tmp_path / "bundle" / "tokenizer") + assert copied == ["tokenizer.json", "tokenizer_config.json"] + + +def test_verify_pooling_config_fails_closed(tmp_path): + snap = _fake_snapshot( + tmp_path, {"embedding_dimension": 768, "pooling_mode": "mean", "include_prompt": False} + ) + with pytest.raises(ValueError, match="include_prompt"): + embedding_export.verify_pooling_config(snap, EG2) + + +def _image_snapshot(root: Path, **config_overrides) -> Path: + snap = root / "snap" + snap.mkdir(parents=True) + cfg = { + "image_token_id": 258880, + "boi_token_id": 255999, + "eoi_token_id": 258882, + "text_config": {"hidden_size": 512}, + } + cfg.update(config_overrides) + (snap / "config.json").write_text(json.dumps(cfg)) + proc = { + "image_processor": { + "patch_size": 16, + "pooling_kernel_size": 3, + "max_soft_tokens": 280, + "do_normalize": False, + "resample": 3, + "rescale_factor": 0.00392156862745098, + } + } + (snap / "processor_config.json").write_text(json.dumps(proc)) + return snap + + +def test_verify_image_config(tmp_path): + embedding_export.verify_image_config(_image_snapshot(tmp_path / "ok"), EG2) + with pytest.raises(ValueError, match="boi_token_id"): + embedding_export.verify_image_config(_image_snapshot(tmp_path / "bad", boi_token_id=1), EG2) + + +# --- Conversion (tiny random config) --- + +TINY_CONFIG = { + "vocab_size": 64, + "hidden_size": 32, + "intermediate_size": 48, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 16, + "sliding_window": 3, + "layer_types": ["sliding_attention", "full_attention"], + "hidden_size_per_layer_input": 8, + "embedding_dim": 24, +} + + +def test_tiny_multifunction_conversion(): + pytest.importorskip("coreai_torch") + wrapper = embedding_export.TextEmbeddingExportWrapper(_tiny_text_encoder()) + plan = bucket_plan([16], [1, 4], 1024) + ids, mask = embedding_export.reference_inputs(plan[1], 64) + assert ids.dtype == torch.int32 and mask.dtype == torch.int32 and ids.shape == (4, 16) + out = wrapper(ids, mask) + assert out.dtype == torch.float32 and out.shape == (4, 24) + torch.testing.assert_close(out.norm(dim=-1), torch.ones(4)) + assert embedding_export.export_text_encoder_program(wrapper, plan, 64) is not None + + +def _tiny_text_encoder(): + from coreai_models.models.macos.embedding_gemma2 import ( + EmbeddingGemma2TextConfig, + EmbeddingGemma2TextEncoder, + ) + + config = EmbeddingGemma2TextConfig.from_dict(TINY_CONFIG) + torch.manual_seed(0) + return EmbeddingGemma2TextEncoder(config).eval() + + +def _tiny_image_encoder(): + from coreai_models.models.macos.embedding_gemma2_vision import ( + EmbeddingGemma2ImageEncoder, + EmbeddingGemma2VisionConfig, + ) + + config = EmbeddingGemma2VisionConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + patch_size=4, + pooling_kernel_size=3, + position_embedding_size=64, + text_hidden_size=32, + max_soft_tokens=12, + ) + torch.manual_seed(1) + return EmbeddingGemma2ImageEncoder(config).eval() + + +TINY_IMAGE = dataclasses.replace(IMAGE, patch_size=4, budget=12, soft_token_dim=32) + + +def test_inputs_embeds_path_equals_ids_path(): + """token_table(ids) -> encoder(inputs_embeds) == text_encoder(ids) (the merge is lossless).""" + text = _tiny_text_encoder() + ids = torch.randint(3, 64, (1, 16), dtype=torch.int32) + mask = torch.ones(1, 16, dtype=torch.int32) + mask[0, 11:] = 0 + with torch.no_grad(): + direct = embedding_export.TextEmbeddingExportWrapper(text)(ids, mask) + embeds = embedding_export.TokenTableExportWrapper(text.embed_tokens)(ids) + with embedding_export.without_token_table(text) as encoder: + merged = embedding_export.InputsEmbedsEncoderWrapper(encoder)(embeds, mask) + assert embeds.dtype == torch.float32 and embeds.shape == (1, 16, 32) + torch.testing.assert_close(merged, direct, atol=1e-6, rtol=0) + assert isinstance(text.embed_tokens, torch.nn.Embedding) # restored + + +def test_image_reference_inputs_are_a_valid_padded_grid(): + pixels, positions = embedding_export.image_reference_inputs(TINY_IMAGE) + assert pixels.shape == (1, 108, 48) and pixels.dtype == torch.float32 + assert positions.shape == (1, 108, 2) and positions.dtype == torch.int32 + real = positions[0, :, 0] >= 0 + assert int(real.sum()) % 9 == 0 + assert torch.all(pixels[0, ~real] == 0) + soft = embedding_export.ImageEncoderExportWrapper(_tiny_image_encoder())(pixels, positions) + n = int(real.sum()) // 9 + assert soft.shape == (1, 12, 32) and soft.dtype == torch.float32 + assert torch.all(soft[0, n:] == 0) and torch.all(soft[0, :n].abs().sum(-1) > 0) + + +def test_tiny_image_components_conversion(): + pytest.importorskip("coreai_torch") + programs = embedding_export.export_image_programs( + _tiny_text_encoder(), + _tiny_image_encoder(), + TINY_IMAGE, + [FunctionSpec("len32_b1", 1, 32)], + max_length=64, + ) + assert set(programs) == {"image_encoder", "token_table", "encoder"} + assert all(p is not None for p in programs.values()) + + +def test_ane_reference_inputs_shapes(): + embeds, bias, pool = embedding_export.ane_text_reference_inputs(128, 512) + assert embeds.shape == (1, 128, 512) and bias.shape == (1, 128, 1, 128) + assert pool.shape == (1, 128) and {t.dtype for t in (embeds, bias, pool)} == {torch.float16} + pixels, positions = embedding_export.ane_image_reference_inputs(IMAGE) + assert pixels.shape == (1, 2520, 768) and positions.shape == (1, 2520, 2) + assert pixels.dtype == positions.dtype == torch.float16 + real = positions[0, :, 0] >= 0 + assert int(real.sum()) % 9 == 0 and torch.all(positions[0, ~real] == -1) + + +def test_ane_band_wrapper_matches_explicit_sliding_bias(): + """The in-graph band equals an explicit sliding-window bias built on the host.""" + from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, EmbeddingGemma2TextEncoder + from coreai_models.models.macos.embedding_gemma2 import EmbeddingGemma2TextConfig + + config = EmbeddingGemma2TextConfig.from_dict(TINY_CONFIG) + torch.manual_seed(0) + model = EmbeddingGemma2TextEncoder(config).eval() + length, valid = 16, 12 + pos = torch.arange(length) + key_ok = (pos < valid)[:, None] # (key, query) + bias = torch.where(key_ok, 0.0, MASK_NEG).expand(length, length) + band = (pos[:, None] - pos[None, :]).abs() <= 3 + sliding = torch.where(key_ok & band, 0.0, MASK_NEG) + bias, sliding = (t.reshape(1, length, 1, length) for t in (bias, sliding)) + embeds = torch.randn(1, length, 32) + pool = (pos < valid).float().reshape(1, length) / valid + wrapper = embedding_export.EncoderBandWrapper(model, length, 3) + with torch.no_grad(): + want = model(embeds, bias, pool, sliding) + got = wrapper(embeds, bias, pool) + torch.testing.assert_close(got, want) + assert not embedding_export.EncoderBandWrapper(model, 4, 3).banded + + +# --- CLI --- + + +@pytest.fixture +def no_export(monkeypatch): + def fail(*args, **kwargs): # pragma: no cover - must not be reached + raise AssertionError("export must not start") + + monkeypatch.setattr(embedding_export, "export_embedding_bundle", fail) + + +def test_cli_precision_conflicts_exit_before_loading(tmp_path, no_export, capsys): + out = str(tmp_path) + with pytest.raises(SystemExit) as macos: + embedding_export.main( + ["embeddinggemma-2", "--compute-precision", "float16", "--output-dir", out] + ) + assert macos.value.code == 2 and "forbids float16" in capsys.readouterr().err + with pytest.raises(SystemExit) as ios: + embedding_export.main( + ["embeddinggemma-2", "--platform", "iOS", "--compute-precision", "float32"] + ) + assert ios.value.code == 2 and "float16" in capsys.readouterr().err + assert not list(tmp_path.iterdir()) + + +def test_cli_unknown_model_exits_2(no_export, capsys): + with pytest.raises(SystemExit) as info: + embedding_export.main(["acme/none"]) + assert info.value.code == 2 + assert "--list-models" in capsys.readouterr().err + + +def test_cli_missing_model_says_model_required(no_export, capsys): + for argv in ([], ["--dry-run"]): + with pytest.raises(SystemExit) as missing: + embedding_export.main(argv) + assert missing.value.code == 2 + assert "a model is required" in capsys.readouterr().err + + +@pytest.mark.parametrize( + ("argv", "name"), + [ + ([], "embeddinggemma-2"), + (["--skip-vision"], "embeddinggemma-2-text"), + (["--platform", "iOS"], "embeddinggemma-2-ios"), + (["--platform", "iOS", "--skip-vision"], "embeddinggemma-2-ios-text"), + ], +) +def test_cli_default_output_names(argv, name, tmp_path, monkeypatch): + seen = {} + + def fake(family, out, **kwargs): + seen["out"] = Path(out) + seen.update(kwargs) + return {} + + monkeypatch.setattr(embedding_export, "export_embedding_bundle", fake) + rc = embedding_export.main(["embeddinggemma-2", *argv, "--output-dir", str(tmp_path)]) + assert rc == 0 + assert seen["out"] == tmp_path / name + + +def test_cli_default_output_is_under_exports(monkeypatch): + seen = {} + monkeypatch.setattr( + embedding_export, "export_embedding_bundle", lambda f, out, **kw: seen.update(out=out) + ) + assert embedding_export.main(["embeddinggemma-2"]) == 0 + assert Path(seen["out"]).parent.name == "exports" + + +def test_cli_dry_run_never_exports(no_export, capsys): + assert embedding_export.main(["embeddinggemma-2", "--platform", "iOS", "--dry-run"]) == 0 + out = capsys.readouterr().out + assert "iOS" in out and "float16" in out and "encoder.aimodel" in out + assert "len128_b1" in out and "embeddinggemma-2-ios" in out + + +def test_cli_list_models(capsys): + assert embedding_export.main(["--list-models"]) == 0 + assert "embeddinggemma-2" in capsys.readouterr().out + + +def test_console_script_is_registered(): + root = Path(__file__).resolve().parents[4] + scripts = tomllib.loads((root / "python" / "pyproject.toml").read_text())["project"]["scripts"] + assert scripts["coreai.embedding.export"] == "coreai_models.embedding.export:main" diff --git a/python/tests/test_model_units/test_export/test_include_debug_info.py b/python/tests/test_model_units/test_export/test_include_debug_info.py index 2fa6f7d7..fab89e00 100644 --- a/python/tests/test_model_units/test_export/test_include_debug_info.py +++ b/python/tests/test_model_units/test_export/test_include_debug_info.py @@ -20,6 +20,7 @@ from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO from coreai_models.diffusion.pipeline import DiffusionExportConfig +from coreai_models.embedding.export import build_parser as build_embedding_parser from coreai_models.export.pipeline import ExportConfig from coreai_models.llm.export import build_parser from coreai_models.segmentation.export import build_parser as build_segmentation_parser @@ -78,6 +79,15 @@ def test_llm_cli_resolves_the_flag_onto_the_export_config() -> None: assert config.include_debug_info is expected +# --- coreai.embedding.export ------------------------------------------ + + +def test_embedding_cli_defaults_to_off_and_accepts_the_flag() -> None: + assert build_embedding_parser().parse_args(["embeddinggemma-2"]).include_debug_info is False + args = build_embedding_parser().parse_args(["embeddinggemma-2", "--include-debug-info"]) + assert args.include_debug_info is True + + # --- standalone models/*/export.py recipes ---------------------------- # ``models/sam3/export.py`` and ``models/sam3_video/export.py`` delegate to the diff --git a/python/tests/test_model_units/test_model_registry.py b/python/tests/test_model_units/test_model_registry.py index c0421244..6c55cd3b 100644 --- a/python/tests/test_model_units/test_model_registry.py +++ b/python/tests/test_model_units/test_model_registry.py @@ -15,9 +15,11 @@ ModelPreset, _preset_to_export_args, _preset_to_output_name, + _utility_to_export_args, filter_presets, filter_utility_models, lookup_preset, + lookup_utility_model, presets_for_type, try_lookup_preset, ) @@ -210,3 +212,40 @@ def test_utility_platform_filter_excludes_macos_only() -> None: names = [u.short_name for u in ios_models] assert "depth-anything-3-small" not in names assert "clip-vit-b32" in names + + +def test_embedding_utility_models_have_a_matching_family() -> None: + from coreai_models.embedding.contract import EMBEDDING_FAMILIES + from coreai_models.model_registry import UTILITY_PRESETS + + rows = [u for u in UTILITY_PRESETS if u.export_script == "coreai.embedding.export"] + assert rows + for row in rows: + assert row.short_name in EMBEDDING_FAMILIES + assert EMBEDDING_FAMILIES[row.short_name].hf_model_id == row.hf_id + + +def test_embedding_export_args_use_the_console_script() -> None: + model = lookup_utility_model("embeddinggemma-2") + assert model is not None + assert _utility_to_export_args(model, "iOS") == [ + "uv", + "run", + "coreai.embedding.export", + "embeddinggemma-2", + "--platform", + "iOS", + ] + assert _utility_to_export_args(model, "macOS")[-1] == "embeddinggemma-2" + + +def test_script_utility_export_args_are_unchanged() -> None: + model = lookup_utility_model("clip-vit-b32") + assert model is not None + assert _utility_to_export_args(model) == [ + "uv", + "run", + "models/clip/export.py", + "--model", + "openai/clip-vit-base-patch32", + ] From 85db172aa990978e76a57676a564f69a420cfa79 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 20 Oct 2026 00:13:59 -0700 Subject: [PATCH 04/38] EmbeddingGemma 2: model README and catalog entry Model card covering export (macOS and Neural Engine), the bundle layout, task prompts and usage. Also the catalog entry in models/README.md. --- models/README.md | 18 +++++++++++-- models/embedding_gemma2/README.md | 42 +++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 2 deletions(-) create mode 100644 models/embedding_gemma2/README.md diff --git a/models/README.md b/models/README.md index 51e975eb..27d913df 100644 --- a/models/README.md +++ b/models/README.md @@ -133,7 +133,7 @@ Pass `--include-debug-info` to switch the converter to `DEBUG` mode, which embed uv run coreai.llm.export Qwen/Qwen3-0.6B --include-debug-info ``` -The flag is available on `coreai.llm.export`, `coreai.vlm.export`, `coreai.diffusion.export`, `coreai.segmentation.export`, and every standalone `models//export.py` recipe, so all export paths produce assets carrying the same debug information by default. +The flag is available on `coreai.llm.export`, `coreai.vlm.export`, `coreai.diffusion.export`, `coreai.segmentation.export`, `coreai.embedding.export`, and every standalone `models//export.py` recipe, so all export paths produce assets carrying the same debug information by default. **Note:** `--include-debug-info` is independent of `--verbose`/`-v`. `--verbose` only raises the console log level; it does not change what goes into the asset. @@ -189,6 +189,16 @@ decoder (`main`), token-embedding lookup (`embedding`), vision encoder (`vision`), tokenizer, and `metadata.json`. Pass `--skip-vision` to export the text portion only. +### Embedding Models + +```bash +uv run coreai.embedding.export --list-models +uv run coreai.embedding.export embeddinggemma-2 # text + image (macOS) +uv run coreai.embedding.export embeddinggemma-2 --platform iOS # Neural Engine +``` + +See the [EmbeddingGemma 2](embedding_gemma2) model card. + ### Standalone Export Scripts Models with a standalone `export.py` are run directly: @@ -245,9 +255,13 @@ uv run models//export.py --include-debug-info # embed debug information - [RoBERTa](roberta) - [T5](t5) +### Embedding Models + +- [EmbeddingGemma 2](embedding_gemma2) + ## Adding a Model -To make a new model exportable via short-name, add a `ModelPreset(...)` entry to `LLM_PRESETS` or `DIFFUSION_PRESETS` in [`python/src/coreai_models/model_registry.py`](../python/src/coreai_models/model_registry.py). Set the short name, HuggingFace ID, family, variant, and the export defaults (compression, compute precision, max context length). +To make a new model exportable via short-name, add a `ModelPreset(...)` entry to `LLM_PRESETS` or `DIFFUSION_PRESETS` in [`python/src/coreai_models/model_registry.py`](../python/src/coreai_models/model_registry.py). Set the short name, HuggingFace ID, family, variant, and the export defaults (compression, compute precision, max context length). Embedding models are registered as `embedding` utility models in the registry and as a family in `coreai_models/embedding/contract.py`. For models with bespoke export logic that doesn't fit the standard `coreai.llm.export` / `coreai.diffusion.export` flow, write a standalone recipe under `models//export.py` — see existing recipes for the [PEP 723](https://peps.python.org/pep-0723/) pattern and `models/README.md` for the contribution checklist. diff --git a/models/embedding_gemma2/README.md b/models/embedding_gemma2/README.md new file mode 100644 index 00000000..3a4b2e70 --- /dev/null +++ b/models/embedding_gemma2/README.md @@ -0,0 +1,42 @@ +# EmbeddingGemma 2 + +Multimodal retrieval embedding model that maps text and images into one shared space: each input becomes a single 768-dimensional, L2-normalized vector, and its leading 512, 256 or 128 dimensions are valid smaller embeddings.[^1] + +## Setup + +If you haven't installed `uv`, install it by + +```bash +brew install uv +``` + +## Export + +```sh +uv run coreai.embedding.export embeddinggemma-2 +``` + +Saves to `/exports/embeddinggemma-2/` (text and image towers, GPU). Pass `--output-dir ` to override the destination. + +```sh +uv run coreai.embedding.export --help +``` + +**Options:** + +| Flag | Description | Default | +| --------------------- | ---------------------------------------------------- | ------------------------------------ | +| `--platform` | `macOS` (GPU bundle) or `iOS` (Neural Engine bundle) | `macOS` | +| `--skip-vision` | Export the text tower only | — | +| `--compute-precision` | `float32`, `bfloat16` on macOS; `float16` on iOS | `float32` on macOS, `float16` on iOS | +| `--output-dir` | Parent directory for the bundle | `/exports/` | +| `--overwrite` | Overwrite an existing bundle | — | +| `--dry-run` | Print the export plan without writing | — | + +**Supported models:** + +| Model | HuggingFace ID | Parameters | +| ---------------- | ----------------------- | ---------------------------- | +| embeddinggemma-2 | google/embeddinggemma-2 | 439M (text 271M, image 168M) | + +[^1]: [Documentation](https://ai.google.dev/gemma/docs/embeddinggemma) · [HuggingFace](https://huggingface.co/google/embeddinggemma-2) From 3b1517387f4888695db3c69623134f4bf69363c8 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 27 Oct 2026 11:42:00 -0700 Subject: [PATCH 05/38] Keep ScaledRMSNorm finite in float16 for rows with max-abs below 1/256 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. --- .../models/ios/embedding_gemma2.py | 5 ++++- .../test_ios_layers/test_embedding_gemma2.py | 20 ++++++++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index 31f1610e..24243ffb 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -51,7 +51,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: m = x.abs().amax(d, keepdim=True).clamp(min=_FP16_TINY) r = torch.reciprocal(m) xs = x * r - y = xs * torch.rsqrt((xs * xs).mean(d, keepdim=True) + self.eps * (r * r)) + # eps * r^2 as (sqrt(eps) * r)^2. r goes up to 1 / _FP16_TINY, and r * r overflows + # float16 once r > 256, which zeroes the row. + e = self.eps**0.5 * r + y = xs * torch.rsqrt((xs * xs).mean(d, keepdim=True) + e * e) if self.with_scale: shape = [1] * x.dim() shape[d] = self.channels diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py index 8bf9978b..09048b09 100644 --- a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -7,9 +7,14 @@ import math +import pytest import torch -from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, text_encoder_from_state_dict +from coreai_models.models.ios.embedding_gemma2 import ( + MASK_NEG, + ScaledRMSNorm, + text_encoder_from_state_dict, +) from coreai_models.models.macos.embedding_gemma2 import ( EmbeddingGemma2TextConfig, EmbeddingGemma2TextEncoder, @@ -101,3 +106,16 @@ def test_float16_stays_finite_with_large_activations() -> None: got = ios16(**_host_inputs(ids, mask, table, cfg.sliding_window, torch.float16)) assert torch.isfinite(got.float()).all() assert float(_cos(got.float(), want)) >= 0.99 + + +@pytest.mark.parametrize("peak", [0.0, 1e-7, 1e-5, 1e-3, 3.9e-3, 4e-3, 1.0, 1e2, 1e4, 6e4]) +def test_float16_scaled_rms_norm_matches_float32_across_magnitudes(peak: float) -> None: + eps = 1e-6 + x = torch.randn(8, 256, generator=torch.Generator().manual_seed(0)) + x = x / x.abs().amax(1, keepdim=True) * peak + x16 = x.half() + want = x16.float() * torch.rsqrt((x16.float() ** 2).mean(1, keepdim=True) + eps) + with torch.no_grad(): + got = ScaledRMSNorm(256, eps, with_scale=False).half()(x16).float() + assert torch.isfinite(got).all() + torch.testing.assert_close(got, want, rtol=2e-3, atol=2e-3) From 5b5e7f1192360081cc74441e56d56c11502a0776 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 27 Oct 2026 17:48:19 -0700 Subject: [PATCH 06/38] Clamp the squared norm in the Neural Engine L2 normalize An all-zero row now normalizes to zero instead of NaN, as on the GPU path. --- .../src/coreai_models/models/ios/embedding_gemma2.py | 5 +++-- .../test_ios_layers/test_embedding_gemma2.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index 24243ffb..eb499135 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -27,7 +27,8 @@ ) from coreai_models.primitives.ios.sdpa import SDPA -# Smallest positive normal float16; floors the row max-abs so all-zero rows stay finite. +# Smallest positive normal float16; floors the row max-abs and the squared norm so all-zero rows +# stay finite. _FP16_TINY = 6.103515625e-05 MASK_NEG = -30000.0 @@ -210,7 +211,7 @@ def forward( y = self.embedding_projection(pooled.unsqueeze(-1)).reshape(b, -1) # (B, E) m = y.abs().amax(-1, keepdim=True).clamp(min=_FP16_TINY) ys = y * torch.reciprocal(m) - return ys * torch.rsqrt((ys * ys).sum(-1, keepdim=True)) + return ys * torch.rsqrt((ys * ys).sum(-1, keepdim=True).clamp(min=_FP16_TINY)) def text_encoder_from_state_dict( diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py index 09048b09..7afccad2 100644 --- a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -119,3 +119,15 @@ def test_float16_scaled_rms_norm_matches_float32_across_magnitudes(peak: float) got = ScaledRMSNorm(256, eps, with_scale=False).half()(x16).float() assert torch.isfinite(got).all() torch.testing.assert_close(got, want, rtol=2e-3, atol=2e-3) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_all_zero_pool_weights_give_a_zero_vector(dtype: torch.dtype) -> None: + cfg, _, state = _reference_and_state_dict() + ios, table = text_encoder_from_state_dict(cfg, state, dtype=dtype) + ids, mask = _ragged_batch([6], seq_len=6, vocab=cfg.vocab_size) + inputs = _host_inputs(ids, mask, table, cfg.sliding_window, dtype) + inputs["pool_weights"] = torch.zeros_like(inputs["pool_weights"]) + with torch.no_grad(): + got = ios(**inputs) + assert torch.equal(got, torch.zeros_like(got)) From 0ea9996bd52d95b75d9f46c14f39d6486c7e7b2a Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 27 Oct 2026 17:49:28 -0700 Subject: [PATCH 07/38] Store the Neural Engine token table in float16 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. --- python/src/coreai_models/embedding/export.py | 29 ++++++++++++------- .../test_export/test_embedding_export.py | 24 +++++++++++++++ 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index d41c7fbb..00998ff8 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -179,14 +179,20 @@ def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torc class TokenTableExportWrapper(nn.Module): - """``input_ids [B, L] int32 -> inputs_embeds [B, L, H]`` float32 (scaled token embeddings).""" + """``input_ids [B, L] int32 -> inputs_embeds [B, L, H]`` as ``dtype`` (scaled embeddings).""" - def __init__(self, embed_tokens: nn.Module) -> None: + def __init__(self, embed_tokens: nn.Module, dtype: torch.dtype = torch.float32) -> None: super().__init__() self.embed_tokens = embed_tokens + self.dtype = dtype def forward(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids).to(torch.float32) + return self.embed_tokens(input_ids).to(self.dtype) + + +def scaled_token_table(table: torch.Tensor, scale: float, dtype: torch.dtype) -> nn.Embedding: + """``table * scale`` stored as ``dtype``. Scaling in float32 first keeps the cast exact.""" + return nn.Embedding.from_pretrained((table.float() * scale).to(dtype)) class InputsEmbedsEncoderWrapper(nn.Module): @@ -377,6 +383,7 @@ def export_token_table_program( embed_tokens: nn.Module, max_length: int, include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, + dtype: torch.dtype = torch.float32, ) -> AIProgram: """``token_table``: function ``main``, ids -> scaled embeddings, dynamic length (batch 1).""" from coreai_models.diffusion.gpu import export_stateless @@ -385,7 +392,7 @@ def export_token_table_program( length = torch.export.Dim("length", min=2, max=max_length) ids = torch.zeros((1, 64), dtype=torch.int32) return export_stateless( - TokenTableExportWrapper(embed_tokens), + TokenTableExportWrapper(embed_tokens, dtype), (ids,), TOKEN_TABLE_INPUT_NAMES, (TOKEN_TABLE_OUTPUT_NAME,), @@ -585,20 +592,20 @@ def _export_ane_programs( raise NotImplementedError(f"No Neural Engine export for family {family.key!r}") 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) + # The token table runs on the GPU and stores the scaled table in float16, the encoder's input + # precision. + encoder, table = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float16) + token_table = scaled_token_table(table, encoder.config.hidden_size**0.5, torch.float16) + del table _save( export_token_table_program( - gpu_encoder.embed_tokens, family.model_max_tokens, include_debug_info + token_table, family.model_max_tokens, include_debug_info, dtype=torch.float16 ), out, family, TOKEN_TABLE_COMPONENT, ) - del gpu_encoder - - encoder, _table = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float16) - del _table + del token_table _save( export_ane_encoder_program( encoder, diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index c433d08f..60c0ea26 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -514,6 +514,30 @@ def test_inputs_embeds_path_equals_ids_path(): assert isinstance(text.embed_tokens, torch.nn.Embedding) # restored +def test_float16_scaled_token_table_equals_the_float32_table_cast(): + """The float16 token table gives the encoder the same inputs as the float32 one.""" + text = _tiny_text_encoder() + ids = torch.randint(3, 64, (1, 16), dtype=torch.int32) + table = embedding_export.scaled_token_table( + text.embed_tokens.weight, text.config.hidden_size**0.5, torch.float16 + ) + with torch.no_grad(): + want = embedding_export.TokenTableExportWrapper(text.embed_tokens)(ids).half() + got = embedding_export.TokenTableExportWrapper(table, torch.float16)(ids) + assert got.dtype == torch.float16 + assert torch.equal(got, want) + + +def test_tiny_float16_token_table_conversion(): + pytest.importorskip("coreai_torch") + text = _tiny_text_encoder() + table = embedding_export.scaled_token_table( + text.embed_tokens.weight, text.config.hidden_size**0.5, torch.float16 + ) + program = embedding_export.export_token_table_program(table, 64, dtype=torch.float16) + assert program is not None + + def test_image_reference_inputs_are_a_valid_padded_grid(): pixels, positions = embedding_export.image_reference_inputs(TINY_IMAGE) assert pixels.shape == (1, 108, 48) and pixels.dtype == torch.float32 From 18b2ce69e497e9bedc2e916e947547ffb272c91d Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 27 Oct 2026 17:49:43 -0700 Subject: [PATCH 08/38] Say which encoder the family precision fields describe --- python/src/coreai_models/embedding/contract.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py index 08cf60b8..ddc4669c 100644 --- a/python/src/coreai_models/embedding/contract.py +++ b/python/src/coreai_models/embedding/contract.py @@ -91,6 +91,8 @@ class EmbeddingFamily: renormalize_after_truncation: bool = True pooling_mode: str = "mean" include_prompt: bool = True + # Precision of the GPU encoder, which runs the reference math. The Neural Engine encoder + # rescales its norms to compute in float16. default_compute_precision: str = "float32" precision_floor: str | None = None forbidden_precisions: tuple[str, ...] = () @@ -398,6 +400,7 @@ def embedding_block( "metric": family.metric, "precision": { "compute": precision, + # The Neural Engine graphs compute in float16. "forbidden": [] if ios else list(family.forbidden_precisions), "compute_units": ["neural_engine", "gpu"] if ios else ["gpu"], }, From 6b00b9c8237fcab24ba6554ddebffeaacf297c43 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 27 Oct 2026 17:50:45 -0700 Subject: [PATCH 09/38] Check the float16 Neural Engine text encoder against real weights 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. --- .../test_ios_layers/test_embedding_gemma2.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py index 7afccad2..7451957a 100644 --- a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -13,12 +13,16 @@ from coreai_models.models.ios.embedding_gemma2 import ( MASK_NEG, ScaledRMSNorm, + load_embedding_gemma2_text_encoder, text_encoder_from_state_dict, ) from coreai_models.models.macos.embedding_gemma2 import ( EmbeddingGemma2TextConfig, EmbeddingGemma2TextEncoder, ) +from coreai_models.models.macos.embedding_gemma2 import ( + load_embedding_gemma2_text_encoder as load_macos_text_encoder, +) from tests.test_model_units.test_models._embedding_gemma2_common import ( TINY_TEXT_CONFIG, _cos, @@ -131,3 +135,18 @@ def test_all_zero_pool_weights_give_a_zero_vector(dtype: torch.dtype) -> None: with torch.no_grad(): got = ios(**inputs) assert torch.equal(got, torch.zeros_like(got)) + + +@pytest.mark.parametrize(("length", "valid"), [(8, 8), (64, 40), (128, 128)]) +def test_real_weight_float16_matches_float32(embeddinggemma2_snapshot, length, valid) -> None: + reference = load_macos_text_encoder(embeddinggemma2_snapshot, dtype=torch.float32).eval() + ios16, table = load_embedding_gemma2_text_encoder(embeddinggemma2_snapshot, torch.float16) + gen = torch.Generator().manual_seed(length) + ids = torch.randint(3, table.shape[0], (1, length), generator=gen, dtype=torch.int32) + ids[0, 0], ids[0, valid - 1] = 2, 1 # BOS ... EOS + mask = (torch.arange(length) < valid).to(torch.int32)[None] + with torch.no_grad(): + want = reference(input_ids=ids, attention_mask=mask) + got = ios16(**_host_inputs(ids, mask, table, ios16.config.sliding_window, torch.float16)) + assert torch.isfinite(got).all() + assert float(_cos(got.float(), want)) >= 0.9999 From ce3b01866e3bdfcedac23d4107b4d6e18c7bacaa Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 27 Oct 2026 17:55:44 -0700 Subject: [PATCH 10/38] Store the macOS text weights once when the bundle has an image tower 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. --- .../src/coreai_models/embedding/contract.py | 32 ++- python/src/coreai_models/embedding/export.py | 189 ++++++++++++------ .../test_export/test_embedding_export.py | 22 +- 3 files changed, 160 insertions(+), 83 deletions(-) diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py index ddc4669c..a61aa9a6 100644 --- a/python/src/coreai_models/embedding/contract.py +++ b/python/src/coreai_models/embedding/contract.py @@ -217,8 +217,8 @@ def as_dict(self) -> dict[str, Any]: return {"function": self.name, "batch": self.batch, "length": self.length} -def function_name(length: int, batch: int) -> str: - return f"len{length}_b{batch}" +def function_name(length: int, batch: int, prefix: str = "") -> str: + return f"{prefix}len{length}_b{batch}" def bucket_plan( @@ -253,14 +253,14 @@ def text_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec] def encoder_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec]: - """Batch-1 ``inputs_embeds`` functions of the ``encoder`` component.""" + """Batch-1 ``inputs_embeds`` functions of the ``encoder`` component, named ``embeds_len*``.""" if platform == "iOS": lengths = sorted(set(family.ane_length_buckets)) else: assert family.image is not None first = -(-family.image.sequence_length() // 32) * 32 lengths = sorted({first} | {n for n in family.length_buckets if n > first}) - return [FunctionSpec(function_name(n, 1), 1, n) for n in lengths] + return [FunctionSpec(function_name(n, 1, "embeds_"), 1, n) for n in lengths] def bundle_name(family: EmbeddingFamily, platform: str, with_image: bool) -> str: @@ -268,16 +268,25 @@ def bundle_name(family: EmbeddingFamily, platform: str, with_image: bool) -> str def bundle_assets(platform: str, with_image: bool) -> dict[str, str]: - """Component name -> asset file for the bundle.""" + """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. + """ if platform == "iOS": components = [TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT] if with_image: components.append(IMAGE_COMPONENT) - elif with_image: - components = [TEXT_COMPONENT, IMAGE_COMPONENT, TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT] - else: - components = [TEXT_COMPONENT] - return {name: f"{name}.aimodel" for name in components} + return {name: f"{name}.aimodel" for name in components} + text = f"{TEXT_COMPONENT}.aimodel" + if not with_image: + return {TEXT_COMPONENT: text} + return { + TEXT_COMPONENT: text, + IMAGE_COMPONENT: f"{IMAGE_COMPONENT}.aimodel", + TOKEN_TABLE_COMPONENT: text, + ENCODER_COMPONENT: text, + } def embedding_block( @@ -292,7 +301,8 @@ def embedding_block( """Build the ``embedding`` block of ``metadata.json``. macOS: a ``text_encoder`` graph (ids in, vector out), plus the image tower - ``image_encoder -> token_table -> encoder`` when ``with_image``. iOS: the text tower runs + ``image_encoder -> token_table -> encoder`` when ``with_image``; ``token_table`` and + ``encoder`` are functions of ``text_encoder.aimodel``. iOS: the text tower runs ``token_table -> encoder`` on the Neural Engine. The host merges soft tokens into the table output at the placeholder rows before the encoder. """ diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index 00998ff8..adf24954 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -267,6 +267,13 @@ def reference_inputs(spec: FunctionSpec, vocab_size: int) -> tuple[torch.Tensor, return ids.to(torch.int32), mask +def embeds_reference_inputs(spec: FunctionSpec, hidden: int) -> tuple[torch.Tensor, torch.Tensor]: + """Trace inputs for one ``inputs_embeds`` function (values do not matter).""" + generator = torch.Generator().manual_seed(spec.length) + embeds = torch.randn((spec.batch, spec.length, hidden), generator=generator) + return embeds, torch.ones((spec.batch, spec.length), dtype=torch.int32) + + def image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, torch.Tensor]: """Image-encoder trace inputs: a full block grid padded to the budget (int32 positions).""" k = image.pooling_kernel_size @@ -311,41 +318,33 @@ def ane_image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, to @dataclass class TraceSpec: - """One function of a Neural Engine program: example args + IO names.""" + """One function of a multi-function program: example args, IO names, dynamic dims.""" name: str args: tuple[torch.Tensor, ...] input_names: tuple[str, ...] output_names: tuple[str, ...] + dynamic_shapes: tuple[dict[int, Any] | None, ...] | None = None -def convert_ane_functions( - traces: list[tuple[nn.Module, TraceSpec]], - *, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """``torch.export`` every ``(module, trace)``; convert them into one multi-function program. - - Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. - """ +def _converter(include_debug_info: bool) -> Any: import coreai_torch - from coreai_models.export.mlir_ops import register_custom_torch_lowering + mode = coreai_torch.TorchConverter.Mode + return coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) - converter = coreai_torch.TorchConverter( - mode=( - coreai_torch.TorchConverter.Mode.DEBUG - if include_debug_info - else coreai_torch.TorchConverter.Mode.RELEASE - ) - ) - register_custom_torch_lowering(converter) - decomp = torch.export.default_decompositions() + +def _add_functions( + converter: Any, traces: list[tuple[nn.Module, TraceSpec]], decomp: dict[Any, Any] +) -> None: + """``torch.export`` every ``(module, trace)`` and add it to ``converter`` as one function.""" for module, trace in traces: module.eval() start = time.time() with torch.no_grad(): - exported = torch.export.export(module, args=trace.args) + exported = torch.export.export( + module, args=trace.args, dynamic_shapes=trace.dynamic_shapes + ) converter.add_exported_program( exported.run_decompositions(decomp), input_names=list(trace.input_names), @@ -358,6 +357,22 @@ def convert_ane_functions( [tuple(a.shape) for a in trace.args], time.time() - start, ) + + +def convert_ane_functions( + traces: list[tuple[nn.Module, TraceSpec]], + *, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """Convert every ``(module, trace)`` into one multi-function Neural Engine program. + + Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. + """ + from coreai_models.export.mlir_ops import register_custom_torch_lowering + + converter = _converter(include_debug_info) + register_custom_torch_lowering(converter) + _add_functions(converter, traces, torch.export.default_decompositions()) return converter.to_coreai() @@ -401,46 +416,85 @@ def export_token_table_program( ) -def export_image_programs( - text_encoder: nn.Module, - image_encoder: nn.Module, - image: ImageTowerFacts, - functions: list[FunctionSpec], +def export_shared_text_program( + encoder: nn.Module, + text_fns: list[FunctionSpec], + encoder_fns: list[FunctionSpec], max_length: int, include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> dict[str, AIProgram]: - """GPU image tower: ``{image_encoder, token_table, encoder}`` programs. +) -> AIProgram: + """GPU ``text_encoder`` that also holds the ``token_table`` and ``encoder`` functions. - ``max_length`` bounds the token table's dynamic sequence dimension. + Text functions take ids; ``main`` maps ids to ``inputs_embeds`` (dynamic length up to + ``max_length``); ``embeds_*`` functions take ``inputs_embeds``. The functions share one copy + of the table and the transformer. """ - from coreai_models.diffusion.gpu import export_multifunction, export_stateless + import coreai_torch - programs: dict[str, AIProgram] = { - TOKEN_TABLE_COMPONENT: export_token_table_program( - text_encoder.embed_tokens, max_length, include_debug_info + vocab_size = int(encoder.config.vocab_size) + hidden = int(encoder.config.hidden_size) + length = torch.export.Dim("length", min=2, max=max_length) + text = TextEmbeddingExportWrapper(encoder) + traces = [ + ( + text, + TraceSpec( + spec.name, reference_inputs(spec, vocab_size), TEXT_INPUT_NAMES, (OUTPUT_NAME,) + ), ) - } - with without_token_table(text_encoder) as encoder: - wrapper = InputsEmbedsEncoderWrapper(encoder) - traces = [] - for spec in functions: - generator = torch.Generator().manual_seed(spec.length) - embeds = torch.randn( - (spec.batch, spec.length, image.soft_token_dim), generator=generator - ) - mask = torch.ones((spec.batch, spec.length), dtype=torch.int32) - traces.append((spec.name, wrapper, (embeds, mask))) - programs[ENCODER_COMPONENT] = export_multifunction( - traces, ENCODER_INPUT_NAMES, (OUTPUT_NAME,), include_debug_info + for spec in text_fns + ] + traces.append( + ( + TokenTableExportWrapper(encoder.embed_tokens), + TraceSpec( + "main", + (torch.zeros((1, 64), dtype=torch.int32),), + TOKEN_TABLE_INPUT_NAMES, + (TOKEN_TABLE_OUTPUT_NAME,), + ({1: length},), + ), ) - programs[IMAGE_COMPONENT] = export_stateless( + ) + converter = _converter(include_debug_info) + decomp = coreai_torch.get_decomp_table() + _add_functions(converter, traces, decomp) + with without_token_table(encoder) as inner: + embeds = InputsEmbedsEncoderWrapper(inner) + _add_functions( + converter, + [ + ( + embeds, + TraceSpec( + spec.name, + embeds_reference_inputs(spec, hidden), + ENCODER_INPUT_NAMES, + (OUTPUT_NAME,), + ), + ) + for spec in encoder_fns + ], + decomp, + ) + return converter.to_coreai() + + +def export_image_encoder_program( + image_encoder: nn.Module, + image: ImageTowerFacts, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """GPU ``image_encoder``: function ``main``, pixels and positions -> soft tokens.""" + from coreai_models.diffusion.gpu import export_stateless + + return export_stateless( ImageEncoderExportWrapper(image_encoder), image_reference_inputs(image), IMAGE_INPUT_NAMES, (IMAGE_OUTPUT_NAME,), include_debug_info=include_debug_info, ) - return programs def export_ane_encoder_program( @@ -556,28 +610,33 @@ def _export_gpu_programs( ) -> None: dtype = _DTYPES[compute_precision] encoder = _load_text_encoder(family, snapshot, dtype) - program = export_text_encoder_program( - TextEmbeddingExportWrapper(encoder), - text_functions(family, "macOS"), - int(encoder.config.vocab_size), - include_debug_info, - ) + if with_image: + program = export_shared_text_program( + encoder, + text_functions(family, "macOS"), + encoder_functions(family, "macOS"), + family.model_max_tokens, + include_debug_info, + ) + else: + program = export_text_encoder_program( + TextEmbeddingExportWrapper(encoder), + text_functions(family, "macOS"), + int(encoder.config.vocab_size), + include_debug_info, + ) _save(program, out, family, TEXT_COMPONENT) - del program + del program, encoder if not with_image: return assert family.image is not None image_encoder = _load_image_encoder(family, snapshot, dtype, family.image.budget) - programs = export_image_programs( - encoder, - image_encoder, - family.image, - encoder_functions(family, "macOS"), - max_length=family.model_max_tokens, - include_debug_info=include_debug_info, + _save( + export_image_encoder_program(image_encoder, family.image, include_debug_info), + out, + family, + IMAGE_COMPONENT, ) - for component, component_program in programs.items(): - _save(component_program, out, family, component) def _export_ane_programs( @@ -799,7 +858,7 @@ def main(argv: list[str] | None = None) -> int: print(f"platform: {platform}") print(f"precision: {precision}") print(f"towers: {towers}") - print(f"assets: {', '.join(assets.values())}") + print(f"assets: {', '.join(dict.fromkeys(assets.values()))}") if platform == "macOS": print(f"text functions: {[f.name for f in text_functions(family, platform)]}") if with_image or platform == "iOS": diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 60c0ea26..a53136f3 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -551,17 +551,25 @@ def test_image_reference_inputs_are_a_valid_padded_grid(): assert torch.all(soft[0, n:] == 0) and torch.all(soft[0, :n].abs().sum(-1) > 0) -def test_tiny_image_components_conversion(): +def test_tiny_shared_text_program_conversion(): pytest.importorskip("coreai_torch") - programs = embedding_export.export_image_programs( + program = embedding_export.export_shared_text_program( _tiny_text_encoder(), - _tiny_image_encoder(), - TINY_IMAGE, - [FunctionSpec("len32_b1", 1, 32)], + [FunctionSpec("len16_b1", 1, 16), FunctionSpec("len32_b1", 1, 32)], + [FunctionSpec("embeds_len32_b1", 1, 32)], max_length=64, ) - assert set(programs) == {"image_encoder", "token_table", "encoder"} - assert all(p is not None for p in programs.values()) + assert program is not None + image = embedding_export.export_image_encoder_program(_tiny_image_encoder(), TINY_IMAGE) + assert image is not None + + +def test_macos_image_bundle_shares_the_text_encoder_asset(): + assets = bundle_assets("macOS", True) + assert assets["token_table"] == assets["encoder"] == assets["text_encoder"] + assert assets["image_encoder"] == "image_encoder.aimodel" + names = [f.name for f in text_functions(EG2, "macOS") + encoder_functions(EG2, "macOS")] + assert len(names) == len(set(names)) and "main" not in names def test_ane_reference_inputs_shapes(): From 0f054ec91e046ac8d2239e6bfdca37e77d86f828 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 17:23:21 -0700 Subject: [PATCH 11/38] Record the requested compute precision in the bundle metadata 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. --- .../src/coreai_models/embedding/contract.py | 6 +++-- python/src/coreai_models/embedding/export.py | 1 + .../test_export/test_embedding_export.py | 27 +++++++++++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py index a61aa9a6..b8a743e1 100644 --- a/python/src/coreai_models/embedding/contract.py +++ b/python/src/coreai_models/embedding/contract.py @@ -297,6 +297,7 @@ def embedding_block( revision: str, prompts: dict[str, str], default_prompt_name: str | None = None, + compute_precision: str | None = None, ) -> dict[str, Any]: """Build the ``embedding`` block of ``metadata.json``. @@ -304,13 +305,14 @@ def embedding_block( ``image_encoder -> token_table -> encoder`` when ``with_image``; ``token_table`` and ``encoder`` are functions of ``text_encoder.aimodel``. iOS: the text tower runs ``token_table -> encoder`` on the Neural Engine. The host merges soft tokens into the table - output at the placeholder rows before the encoder. + output at the placeholder rows before the encoder. ``compute_precision`` defaults to the + platform's precision. """ ios = platform == "iOS" if platform not in PLATFORMS: raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") precision = check_compute_precision( - family, default_precision(family, platform), platform=platform + family, compute_precision or default_precision(family, platform), platform=platform ) if not revision: raise ValueError("revision (snapshot SHA) is required to stamp the embedding space") diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index adf24954..a39d06e7 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -758,6 +758,7 @@ def export_embedding_bundle( revision=sha[:16], prompts=prompts, default_prompt_name=default_prompt, + compute_precision=precision, ) metadata = build_bundle_metadata( family, diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index a53136f3..4e275d42 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -343,6 +343,33 @@ def test_embedding_block_validation(): embedding_block(no_image, platform="macOS", with_image=True, revision="x", prompts=PROMPTS) +def test_embedding_block_records_the_compute_precision(): + assert _block()["precision"]["compute"] == "float32" + assert _block(compute_precision="bfloat16")["precision"]["compute"] == "bfloat16" + assert _block("iOS")["precision"]["compute"] == "float16" + with pytest.raises(EmbeddingPrecisionError, match="forbids float16"): + _block(compute_precision="float16") + + +def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch): + seen = {} + monkeypatch.setattr( + embedding_export, "resolve_snapshot", lambda *a: (tmp_path, "0123456789abcdef0123") + ) + monkeypatch.setattr(embedding_export, "verify_pooling_config", lambda *a: None) + monkeypatch.setattr(embedding_export, "read_prompts", lambda *a: (PROMPTS, None)) + monkeypatch.setattr(embedding_export, "consolidate_prompts", lambda family, prompts: prompts) + monkeypatch.setattr(embedding_export, "copy_tokenizer", lambda *a: ["tokenizer.json"]) + monkeypatch.setattr(embedding_export, "_export_gpu_programs", lambda *a, **kw: seen.update(kw)) + meta = embedding_export.export_embedding_bundle( + EG2, tmp_path / "out", with_image=False, compute_precision="bfloat16" + ) + assert seen["compute_precision"] == "bfloat16" + assert meta["embedding"]["precision"]["compute"] == "bfloat16" + written = json.loads((tmp_path / "out" / "metadata.json").read_text()) + assert written["embedding"]["precision"]["compute"] == "bfloat16" + + def test_bundle_metadata_envelope(): meta = build_bundle_metadata( EG2, From 98966aa20aff7a844bfb4178fa8efdb884262662 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 19:22:33 -0700 Subject: [PATCH 12/38] Require the layer layout in the EmbeddingGemma 2 text config 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. --- .../models/macos/embedding_gemma2.py | 46 +++++++++---------- .../test_export/test_embedding_export.py | 5 ++ .../test_embedding_gemma2.py | 13 ++++++ 3 files changed, 40 insertions(+), 24 deletions(-) diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index d7c5e8f1..fa9a964f 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -62,11 +62,6 @@ class EmbeddingGemma2TextConfig: per_layer_config: dict[int, dict[str, int]] = field(default_factory=dict) def __post_init__(self) -> None: - if not self.layer_types: - self.layer_types = [ - "sliding_attention" if (i + 1) % 6 else "full_attention" - for i in range(self.num_hidden_layers) - ] if len(self.layer_types) != self.num_hidden_layers: raise ValueError( f"layer_types has {len(self.layer_types)} entries, " @@ -88,30 +83,33 @@ def is_sliding(self, layer_idx: int) -> bool: @classmethod def from_dict(cls, text_config: dict[str, Any]) -> EmbeddingGemma2TextConfig: - """Build from an HF ``text_config`` dict (``config.json["text_config"]``).""" - rope = text_config.get("rope_parameters") or {} - sliding_theta = (rope.get("sliding_attention") or {}).get("rope_theta", 10_000.0) - full_theta = (rope.get("full_attention") or {}).get("rope_theta", 1_000_000.0) + """Build from an HF ``text_config`` dict (``config.json["text_config"]``). + + ``layer_types``, ``per_layer_config`` and both ``rope_parameters`` entries define the layer + layout, so they are required; the other fields default to EmbeddingGemma 2's values. + """ + missing = [ + key + for key in ("layer_types", "per_layer_config", "rope_parameters") + if not text_config.get(key) + ] + if missing: + raise ValueError(f"text_config is missing {missing}") + rope = text_config["rope_parameters"] + thetas = {} for kind in ("sliding_attention", "full_attention"): - rope_type = (rope.get(kind) or {}).get("rope_type", "default") + params = rope.get(kind) or {} + if "rope_theta" not in params: + raise ValueError(f"rope_parameters[{kind!r}] has no rope_theta") + rope_type = params.get("rope_type", "default") if rope_type != "default": raise ValueError(f"Unsupported rope_type {rope_type!r} for {kind}") + thetas[kind] = float(params["rope_theta"]) names = {f.name for f in cls.__dataclass_fields__.values()} kwargs = {k: v for k, v in text_config.items() if k in names and v is not None} - kwargs["rope_theta_sliding"] = float(sliding_theta) - kwargs["rope_theta_full"] = float(full_theta) - 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 [] - per_layer = { - i: { - "head_dim": text_config.get("global_head_dim", 512), - "num_key_value_heads": text_config.get("num_global_key_value_heads", 1), - } - for i, t in enumerate(types) - if t == "full_attention" - } + kwargs["rope_theta_sliding"] = thetas["sliding_attention"] + kwargs["rope_theta_full"] = thetas["full_attention"] + per_layer = text_config["per_layer_config"] kwargs["per_layer_config"] = { int(k): {kk: int(vv) for kk, vv in dict(v).items() if vv is not None} for k, v in per_layer.items() diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 4e275d42..76bd8901 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -471,6 +471,11 @@ def test_verify_image_config(tmp_path): "head_dim": 16, "sliding_window": 3, "layer_types": ["sliding_attention", "full_attention"], + "per_layer_config": {"01": {"head_dim": 32, "num_key_value_heads": 1}}, + "rope_parameters": { + "full_attention": {"rope_theta": 1_000_000.0, "rope_type": "default"}, + "sliding_attention": {"rope_theta": 10_000.0, "rope_type": "default"}, + }, "hidden_size_per_layer_input": 8, "embedding_dim": 24, } diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py index 32ab4021..ed7bcdda 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -76,6 +76,19 @@ def test_config_from_dict_resolves_per_layer_overrides(): assert (cfg.rope_theta_sliding, cfg.rope_theta_full) == (10_000.0, 1_000_000.0) +@pytest.mark.parametrize("key", ["layer_types", "per_layer_config", "rope_parameters"]) +def test_config_requires_the_layer_layout(key): + config = {k: v for k, v in TINY_TEXT_CONFIG.items() if k != key} + with pytest.raises(ValueError, match=key): + EmbeddingGemma2TextConfig.from_dict(config) + + +def test_config_requires_both_rope_thetas(): + rope = {"sliding_attention": {"rope_theta": 10_000.0}, "full_attention": {}} + with pytest.raises(ValueError, match="full_attention"): + EmbeddingGemma2TextConfig.from_dict({**TINY_TEXT_CONFIG, "rope_parameters": rope}) + + def test_config_from_snapshot_matches_published_layout(embeddinggemma2_snapshot): cfg = EmbeddingGemma2TextConfig.from_pretrained(embeddinggemma2_snapshot) full = [i for i in range(cfg.num_hidden_layers) if not cfg.is_sliding(i)] From 7f38ce4380174f92021c69cfa9c6b8787a3f2c13 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 19:22:33 -0700 Subject: [PATCH 13/38] Use the shared 2-D RoPE in the EmbeddingGemma 2 image encoder 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. --- .../models/macos/embedding_gemma2_vision.py | 25 ++++--------------- .../test_embedding_gemma2_vision.py | 3 ++- 2 files changed, 7 insertions(+), 21 deletions(-) diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index c3942583..120d2e89 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -26,7 +26,7 @@ from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaDenseMLP from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaRMSNorm as RMSNorm -from coreai_models.primitives.macos.rope import rotate_half +from coreai_models.primitives.ios.rope import apply_rope from coreai_models.primitives.macos.sdpa import SDPA # Special token ids from the EG2 tokenizer / config.json. @@ -104,23 +104,6 @@ def axial_inv_freq(cfg: EmbeddingGemma2VisionConfig) -> torch.Tensor: ) -def apply_axial_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: - """Axial 2-D RoPE: the first half of ``head_dim`` rotates by x, the second half by y. - - Args: - x: ``[B, L, H, D]``. - cos, sin: ``[B, L, D]`` laid out ``[fx, fx, fy, fy]`` (see ``axial_rope_cos_sin``). - """ - half = x.shape[-1] // 2 - cos = cos.unsqueeze(2) - sin = sin.unsqueeze(2) - parts = [] - for lo, hi in ((0, half), (half, 2 * half)): - xp, cp, sp = x[..., lo:hi], cos[..., lo:hi], sin[..., lo:hi] - parts.append(xp * cp + rotate_half(xp) * sp) - return torch.cat(parts, dim=-1) - - class VisionAttention(nn.Module): """Bidirectional MHA with q/k RMSNorm, scale-free v RMSNorm and axial RoPE; scale 1.0.""" @@ -150,8 +133,10 @@ def forward( q = self.q_norm(self.q_proj(x).view(b, n, self.num_heads, self.head_dim)) k = self.k_norm(self.k_proj(x).view(b, n, self.num_kv_heads, self.head_dim)) v = self.v_norm(self.v_proj(x).view(b, n, self.num_kv_heads, self.head_dim)) - q = apply_axial_rope(q, cos, sin).transpose(1, 2) - k = apply_axial_rope(k, cos, sin).transpose(1, 2) + # Axial 2-D RoPE: half of head_dim rotates by x, half by y (cos/sin from + # axial_rope_cos_sin, laid out [fx, fx, fy, fy]). + q = apply_rope(q, cos, sin, head_axis=2, ndim=2).transpose(1, 2) + k = apply_rope(k, cos, sin, head_axis=2, ndim=2).transpose(1, 2) v = v.transpose(1, 2) if self.num_kv_heads != self.num_heads: rep = self.num_heads // self.num_kv_heads diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py index 1472c41c..076b0112 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py @@ -289,7 +289,8 @@ def test_torch_export_default_budget_static_shape(monkeypatch) -> None: torch.testing.assert_close(exported, eager, rtol=1e-5, atol=1e-5) # No dynamic dims and no data-dependent ops (nonzero / masked_select / item) in the graph. targets = {str(node.target) for node in ep.graph.nodes if node.op == "call_function"} - assert not any(t for t in targets if "nonzero" in t or "masked_select" in t or "item" in t) + data_dependent = ("aten.nonzero", "aten.masked_select", "aten.item", "aten._local_scalar_dense") + assert not any(op in t for t in targets for op in data_dependent), sorted(targets) # --- Real weights (transformers >= 5.19 + local snapshot) --- From 9b45f9f496ed54d817cd91aab3281e79ad16f685 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 19:28:41 -0700 Subject: [PATCH 14/38] Set the image-tower facts per family and fix the precision message 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. --- .../src/coreai_models/embedding/contract.py | 39 +++++++++++++------ .../test_export/test_embedding_export.py | 8 ++++ 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py index b8a743e1..37706b26 100644 --- a/python/src/coreai_models/embedding/contract.py +++ b/python/src/coreai_models/embedding/contract.py @@ -41,7 +41,7 @@ class EmbeddingPrecisionError(ValueError): @dataclass(frozen=True) class ImageTowerFacts: - """Facts the image tower needs on the host (merge ids, preprocessing, soft-token budget). + """One family's image-tower facts for the host (merge ids, preprocessing, soft-token budget). The exporter checks them against the snapshot ``config.json`` / ``processor_config.json`` and fails closed on a mismatch. @@ -50,14 +50,14 @@ class ImageTowerFacts: placeholder_id: int boi_token_id: int eoi_token_id: int - patch_size: int = 16 - pooling_kernel_size: int = 3 - rescale_factor: float = 1.0 / 255.0 - normalize: bool = False - image_mean: tuple[float, ...] = (0.0, 0.0, 0.0) - image_std: tuple[float, ...] = (1.0, 1.0, 1.0) - budget: int = 280 - soft_token_dim: int = 512 + patch_size: int + pooling_kernel_size: int + rescale_factor: float + normalize: bool + image_mean: tuple[float, ...] + image_std: tuple[float, ...] + budget: int + soft_token_dim: int @property def max_patches(self) -> int: @@ -137,7 +137,19 @@ class EmbeddingFamily: pad_token_id=0, document_template="title: {title} | text: {text}", # config.json image/boi/eoi token ids; processor_config.json image_processor. - image=ImageTowerFacts(placeholder_id=258880, boi_token_id=255999, eoi_token_id=258882), + image=ImageTowerFacts( + placeholder_id=258880, + boi_token_id=255999, + eoi_token_id=258882, + patch_size=16, + pooling_kernel_size=3, + rescale_factor=1.0 / 255.0, + normalize=False, + image_mean=(0.0, 0.0, 0.0), + image_std=(1.0, 1.0, 1.0), + budget=280, + soft_token_dim=512, + ), ) EMBEDDING_FAMILIES: dict[str, EmbeddingFamily] = {EMBEDDINGGEMMA2.key: EMBEDDINGGEMMA2} @@ -177,10 +189,13 @@ def check_compute_precision( ) return precision if precision in family.forbidden_precisions: + floor = family.precision_floor + alternatives = "float32 (default)" + ( + f" or {floor}" if floor and floor != "float32" else "" + ) raise EmbeddingPrecisionError( f"{family.hf_model_id} forbids {precision} compute on macOS (it overflows: NaN or " - f"silently degraded vectors). Use float32 (default) or " - f"{family.precision_floor or 'float32'}." + f"silently degraded vectors). Use {alternatives}." ) return precision diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 76bd8901..35233e37 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -370,6 +370,14 @@ def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch assert written["embedding"]["precision"]["compute"] == "bfloat16" +def test_forbidden_precision_message_lists_the_allowed_ones(): + with pytest.raises(EmbeddingPrecisionError, match=r"Use float32 \(default\) or bfloat16\."): + check_compute_precision(EG2, "float16") + no_floor = dataclasses.replace(EG2, precision_floor=None) + with pytest.raises(EmbeddingPrecisionError, match=r"Use float32 \(default\)\.$"): + check_compute_precision(no_floor, "float16") + + def test_bundle_metadata_envelope(): meta = build_bundle_metadata( EG2, From 4fbd256706eed249a4ea073410c4f27aabc18342 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 19:28:42 -0700 Subject: [PATCH 15/38] Tidy the embedding export CLI and clean up after a failed export - 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 /exports. - The converter setup is inlined; the dry-run test checks embeds_len128_b1. --- python/src/coreai_models/embedding/export.py | 112 +++++++++--------- .../test_export/test_embedding_export.py | 19 ++- 2 files changed, 76 insertions(+), 55 deletions(-) diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index a39d06e7..47a3d0fd 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -327,13 +327,6 @@ class TraceSpec: dynamic_shapes: tuple[dict[int, Any] | None, ...] | None = None -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) - - def _add_functions( converter: Any, traces: list[tuple[nn.Module, TraceSpec]], decomp: dict[Any, Any] ) -> None: @@ -368,9 +361,12 @@ def convert_ane_functions( Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. """ + import coreai_torch + from coreai_models.export.mlir_ops import register_custom_torch_lowering - converter = _converter(include_debug_info) + mode = coreai_torch.TorchConverter.Mode + converter = coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) register_custom_torch_lowering(converter) _add_functions(converter, traces, torch.export.default_decompositions()) return converter.to_coreai() @@ -456,7 +452,8 @@ def export_shared_text_program( ), ) ) - converter = _converter(include_debug_info) + mode = coreai_torch.TorchConverter.Mode + converter = coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) decomp = coreai_torch.get_decomp_table() _add_functions(converter, traces, decomp) with without_token_table(encoder) as inner: @@ -722,52 +719,59 @@ def export_embedding_bundle( raise FileExistsError(f"{out} exists (pass --overwrite)") shutil.rmtree(out) out.mkdir(parents=True) + try: + snapshot, sha = resolve_snapshot(family.hf_model_id, model_dir) + if not sha: + raise ValueError( + f"Could not read the snapshot revision from {snapshot}; pass a Hugging Face " + "cache snapshot directory (.../snapshots/), whose name is the revision the " + "bundle records" + ) + verify_pooling_config(snapshot, family) + if with_image: + verify_image_config(snapshot, family) + snapshot_prompts, default_prompt = read_prompts(snapshot) + prompts = consolidate_prompts(family, snapshot_prompts) + copy_tokenizer(snapshot, out / "tokenizer") + + if platform == "iOS": + _export_ane_programs( + family, + snapshot, + out, + with_image=with_image, + include_debug_info=include_debug_info, + ) + else: + _export_gpu_programs( + family, + snapshot, + out, + with_image=with_image, + compute_precision=precision, + include_debug_info=include_debug_info, + ) - snapshot, sha = resolve_snapshot(family.hf_model_id, model_dir) - if not sha: - raise ValueError(f"Could not infer the snapshot SHA from {snapshot}") - verify_pooling_config(snapshot, family) - if with_image: - verify_image_config(snapshot, family) - snapshot_prompts, default_prompt = read_prompts(snapshot) - prompts = consolidate_prompts(family, snapshot_prompts) - copy_tokenizer(snapshot, out / "tokenizer") - - if platform == "iOS": - _export_ane_programs( + block = embedding_block( family, - snapshot, - out, + platform=platform, with_image=with_image, - include_debug_info=include_debug_info, + revision=sha[:16], + prompts=prompts, + default_prompt_name=default_prompt, + compute_precision=precision, ) - else: - _export_gpu_programs( + metadata = build_bundle_metadata( family, - snapshot, - out, - with_image=with_image, - compute_precision=precision, - include_debug_info=include_debug_info, + block, + name=bundle_name(family, platform, with_image), + assets=bundle_assets(platform, with_image), + hf_revision=sha, ) - - block = embedding_block( - family, - platform=platform, - with_image=with_image, - revision=sha[:16], - prompts=prompts, - default_prompt_name=default_prompt, - compute_precision=precision, - ) - metadata = build_bundle_metadata( - family, - block, - name=bundle_name(family, platform, with_image), - assets=bundle_assets(platform, with_image), - hf_revision=sha, - ) - (out / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + (out / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + except BaseException: + shutil.rmtree(out, ignore_errors=True) # a failed export leaves no partial bundle + raise logger.info("wrote %s", out) return metadata @@ -785,9 +789,9 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--platform", choices=list(PLATFORMS), - default=None, - help="Target platform. macOS: GPU bundle (float32). iOS: Neural Engine bundle " - "(float16); also runs on a Mac.", + default="macOS", + help="Target platform. macOS: GPU bundle (float32 or bfloat16). iOS: Neural Engine " + "bundle (float16); also runs on a Mac.", ) parser.add_argument( "--skip-vision", @@ -804,7 +808,7 @@ def build_parser() -> argparse.ArgumentParser: "--output-dir", type=Path, default=None, - help="Parent directory for the bundle (default: /exports).", + help="Parent directory for the bundle (default: /exports).", ) parser.add_argument("--overwrite", action="store_true") parser.add_argument("--dry-run", action="store_true", help="Print the export plan and exit.") @@ -842,7 +846,7 @@ def main(argv: list[str] | None = None) -> int: parser.error( f"unknown model {args.model!r}. Run `uv run coreai.embedding.export --list-models`." ) - platform = args.platform or "macOS" + platform = args.platform with_image = not args.skip_vision try: precision = check_compute_precision( diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 35233e37..98789cc3 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -370,6 +370,23 @@ def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch assert written["embedding"]["precision"]["compute"] == "bfloat16" +def test_failed_export_removes_the_bundle_dir(tmp_path, monkeypatch): + def fail(*args, **kwargs): + raise RuntimeError("conversion failed") + + monkeypatch.setattr( + embedding_export, "resolve_snapshot", lambda *a: (tmp_path, "0123456789abcdef0123") + ) + monkeypatch.setattr(embedding_export, "verify_pooling_config", lambda *a: None) + monkeypatch.setattr(embedding_export, "read_prompts", lambda *a: (PROMPTS, None)) + monkeypatch.setattr(embedding_export, "consolidate_prompts", lambda family, prompts: prompts) + monkeypatch.setattr(embedding_export, "copy_tokenizer", lambda *a: ["tokenizer.json"]) + monkeypatch.setattr(embedding_export, "_export_gpu_programs", fail) + with pytest.raises(RuntimeError, match="conversion failed"): + embedding_export.export_embedding_bundle(EG2, tmp_path / "out", with_image=False) + assert not (tmp_path / "out").exists() + + def test_forbidden_precision_message_lists_the_allowed_ones(): with pytest.raises(EmbeddingPrecisionError, match=r"Use float32 \(default\) or bfloat16\."): check_compute_precision(EG2, "float16") @@ -725,7 +742,7 @@ def test_cli_dry_run_never_exports(no_export, capsys): assert embedding_export.main(["embeddinggemma-2", "--platform", "iOS", "--dry-run"]) == 0 out = capsys.readouterr().out assert "iOS" in out and "float16" in out and "encoder.aimodel" in out - assert "len128_b1" in out and "embeddinggemma-2-ios" in out + assert "embeds_len128_b1" in out and "embeddinggemma-2-ios" in out def test_cli_list_models(capsys): From 59dc4b951212d5fc947c4732e12ac3335f21d890 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 19:28:42 -0700 Subject: [PATCH 16/38] Correct the EmbeddingGemma 2 docstrings and report missing weight files 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. --- python/src/coreai_models/models/ios/embedding_gemma2.py | 5 ++++- .../coreai_models/models/ios/embedding_gemma2_vision.py | 7 +++++-- python/src/coreai_models/models/macos/embedding_gemma2.py | 8 ++++++-- .../coreai_models/models/macos/embedding_gemma2_vision.py | 6 ++++-- 4 files changed, 19 insertions(+), 7 deletions(-) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index eb499135..05fb2767 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -256,7 +256,10 @@ def load_embedding_gemma2_text_encoder( 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")): + files = sorted(model_dir.glob("*.safetensors")) + if not files: + raise FileNotFoundError(f"No .safetensors in {model_dir}") + for file in files: with safe_open(str(file), framework="pt") as handle: for key in handle.keys(): # noqa: SIM118 if key.startswith(HF_TEXT_PREFIX): diff --git a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py index 54c5ea5c..04a1c799 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -162,7 +162,7 @@ def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> tor cos = torch.cat((cx, cx, cy, cy), -1).transpose(1, 2).unsqueeze(1) sin = torch.cat((sx, sx, sy, sy), -1).transpose(1, 2).unsqueeze(1) - # Additive key-padding bias (1, Pk, 1, 1), broadcast over heads and queries. + # Additive key-padding bias (B, Pk, 1, 1), broadcast over heads and queries. bias = ((valid - 1.0) * -MASK_NEG).reshape(b, p, 1, 1) for layer in self.layers: x = layer(x, cos, sin, bias) @@ -191,7 +191,10 @@ def load_image_encoder( cfg = EmbeddingGemma2VisionConfig.from_hf_config(json.load(f), max_soft_tokens) model = EmbeddingGemma2ImageEncoder(cfg) sd: dict[str, torch.Tensor] = {} - for fn in sorted(x for x in os.listdir(model_dir) if x.endswith(".safetensors")): + files = sorted(x for x in os.listdir(model_dir) if x.endswith(".safetensors")) + if not files: + raise FileNotFoundError(f"No .safetensors in {model_dir}") + for fn in files: with safe_open(os.path.join(model_dir, fn), framework="pt") as h: for key in h.keys(): # noqa: SIM118 if key.startswith(("vision_tower.", "embed_vision.")): diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index fa9a964f..35b3deaf 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -8,7 +8,7 @@ Bidirectional encoder with per-layer inputs and a sentence-embedding head: masked mean over the valid tokens (prompt included), bias-free projection to 768, then L2 normalization. Masks derive from the static sequence length and ``attention_mask``, so the graph has no data-dependent control -flow. The GPU path runs in float32; the float16 variant lives in the iOS modules. +flow. The GPU path runs in float32 or bfloat16; the float16 variant lives in the iOS modules. """ from __future__ import annotations @@ -39,7 +39,11 @@ @dataclass class EmbeddingGemma2TextConfig: - """Text-tower config parsed from the HF ``text_config`` block (needs no transformers>=5.19).""" + """Text-tower config parsed from the HF ``text_config`` block. + + Reads ``config.json`` directly, so export runs on the repo's pinned transformers; the HF + EmbeddingGemma 2 classes need transformers>=5.19. + """ vocab_size: int = 262_144 hidden_size: int = 512 diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index 120d2e89..55d4704b 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -7,8 +7,8 @@ Host-prepared patches ``[B, P, 768]`` and ``position_ids [B, P, 2]`` (-1 for padding) become soft tokens ``[B, L, 512]`` with P and L static. Padding and the 3x3 pooling slots are derived in-graph -from ``position_ids``, with no data-dependent control flow. The GPU path runs in float32; the -float16 variant lives in the iOS modules. +from ``position_ids``, with no data-dependent control flow. The GPU path runs in float32 or +bfloat16; the float16 variant lives in the iOS modules. """ from __future__ import annotations @@ -320,6 +320,8 @@ def from_pretrained( model = cls(config) state: dict[str, torch.Tensor] = {} files = sorted(fn for fn in os.listdir(model_dir) if fn.endswith(".safetensors")) + if not files: + raise FileNotFoundError(f"No .safetensors in {model_dir}") for fn in files: with safe_open(os.path.join(model_dir, fn), framework="pt", device="cpu") as f: for key in f.keys(): # noqa: SIM118 From d1c0f033b502226ced0f900a4b81b912ec849e76 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 20:32:42 -0700 Subject: [PATCH 17/38] Share the converter setup and per-function tracing with diffusion.gpu 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). --- python/src/coreai_models/diffusion/gpu.py | 63 ++++++++++++-------- python/src/coreai_models/embedding/export.py | 48 +++++++-------- 2 files changed, 62 insertions(+), 49 deletions(-) diff --git a/python/src/coreai_models/diffusion/gpu.py b/python/src/coreai_models/diffusion/gpu.py index b61751af..61db98d9 100644 --- a/python/src/coreai_models/diffusion/gpu.py +++ b/python/src/coreai_models/diffusion/gpu.py @@ -99,6 +99,42 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram: return program +def new_converter( + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> coreai_torch.TorchConverter: + """A converter in ``DEBUG`` mode (full debug info) or ``RELEASE`` mode (smaller asset).""" + mode = coreai_torch.TorchConverter.Mode + return coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) + + +def add_function( + converter: coreai_torch.TorchConverter, + name: str, + module: torch.nn.Module, + args: tuple[torch.Tensor, ...], + input_names: tuple[str, ...] | list[str], + output_names: tuple[str, ...] | list[str], + *, + dynamic_shapes: tuple[dict[int, torch.export.Dim] | None, ...] | None = None, + decompositions: dict | None = None, +) -> None: + """Trace ``module`` on ``args`` and add it to ``converter`` as function ``name``. + + Functions added to one converter share their weights. ``decompositions`` defaults to Core AI's + table. + """ + module.eval() + with torch.no_grad(): + exported = torch.export.export(module, args=args, dynamic_shapes=dynamic_shapes) + table = coreai_torch.get_decomp_table() if decompositions is None else decompositions + converter.add_exported_program( + exported.run_decompositions(table), + input_names=input_names, + output_names=output_names, + entrypoint_name=name, + ) + + def export_multifunction( functions: list[tuple[str, torch.nn.Module, tuple[torch.Tensor, ...]]], input_names: tuple[str, ...], @@ -122,29 +158,8 @@ def export_multifunction( Returns: An optimized AIProgram with multiple named functions. """ - converter = coreai_torch.TorchConverter( - mode=( - coreai_torch.TorchConverter.Mode.DEBUG - if include_debug_info - else coreai_torch.TorchConverter.Mode.RELEASE - ) - ) - coreai_decomp_table = coreai_torch.get_decomp_table() - + converter = new_converter(include_debug_info) for name, wrapper, dummy_inputs in functions: - wrapper.eval() logger.info(f"Tracing function '{name}' (seq dims: {[t.shape for t in dummy_inputs[:2]]})") - - with torch.no_grad(): - exported = torch.export.export(wrapper, args=dummy_inputs) - decomposed = exported.run_decompositions(coreai_decomp_table) - - converter.add_exported_program( - decomposed, - input_names=input_names, - output_names=output_names, - entrypoint_name=name, - ) - - program = converter.to_coreai() - return program + add_function(converter, name, wrapper, dummy_inputs, input_names, output_names) + return converter.to_coreai() diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index 47a3d0fd..ae88664a 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -327,22 +327,25 @@ class TraceSpec: dynamic_shapes: tuple[dict[int, Any] | None, ...] | None = None -def _add_functions( - converter: Any, traces: list[tuple[nn.Module, TraceSpec]], decomp: dict[Any, Any] +def add_traces( + converter: Any, + traces: list[tuple[nn.Module, TraceSpec]], + decompositions: dict[Any, Any] | None = None, ) -> None: - """``torch.export`` every ``(module, trace)`` and add it to ``converter`` as one function.""" + """Add each ``(module, trace)`` to ``converter`` as a function (Core AI's table by default).""" + from coreai_models.diffusion.gpu import add_function + for module, trace in traces: - module.eval() start = time.time() - with torch.no_grad(): - exported = torch.export.export( - module, args=trace.args, dynamic_shapes=trace.dynamic_shapes - ) - converter.add_exported_program( - exported.run_decompositions(decomp), - input_names=list(trace.input_names), - output_names=list(trace.output_names), - entrypoint_name=trace.name, + add_function( + converter, + trace.name, + module, + trace.args, + list(trace.input_names), + list(trace.output_names), + dynamic_shapes=trace.dynamic_shapes, + decompositions=decompositions, ) logger.info( "traced %s %s in %.1fs", @@ -361,14 +364,12 @@ def convert_ane_functions( Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. """ - import coreai_torch - + from coreai_models.diffusion.gpu import new_converter from coreai_models.export.mlir_ops import register_custom_torch_lowering - mode = coreai_torch.TorchConverter.Mode - converter = coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) + converter = new_converter(include_debug_info) register_custom_torch_lowering(converter) - _add_functions(converter, traces, torch.export.default_decompositions()) + add_traces(converter, traces, torch.export.default_decompositions()) return converter.to_coreai() @@ -425,7 +426,7 @@ def export_shared_text_program( ``max_length``); ``embeds_*`` functions take ``inputs_embeds``. The functions share one copy of the table and the transformer. """ - import coreai_torch + from coreai_models.diffusion.gpu import new_converter vocab_size = int(encoder.config.vocab_size) hidden = int(encoder.config.hidden_size) @@ -452,13 +453,11 @@ def export_shared_text_program( ), ) ) - mode = coreai_torch.TorchConverter.Mode - converter = coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) - decomp = coreai_torch.get_decomp_table() - _add_functions(converter, traces, decomp) + converter = new_converter(include_debug_info) + add_traces(converter, traces) with without_token_table(encoder) as inner: embeds = InputsEmbedsEncoderWrapper(inner) - _add_functions( + add_traces( converter, [ ( @@ -472,7 +471,6 @@ def export_shared_text_program( ) for spec in encoder_fns ], - decomp, ) return converter.to_coreai() From 44dd5a5031428dd1f12eeb292c502c662fdbfd1e Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 20:34:06 -0700 Subject: [PATCH 18/38] Use one image reference-input helper for the GPU and Neural Engine traces image_reference_inputs takes the pixel and position dtypes; the Neural Engine trace passes float16 for both. Converted programs are unchanged. --- python/src/coreai_models/embedding/export.py | 36 +++++++++---------- .../test_export/test_embedding_export.py | 2 +- 2 files changed, 18 insertions(+), 20 deletions(-) diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index ae88664a..3d863be0 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -274,8 +274,16 @@ def embeds_reference_inputs(spec: FunctionSpec, hidden: int) -> tuple[torch.Tens return embeds, torch.ones((spec.batch, spec.length), dtype=torch.int32) -def image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, torch.Tensor]: - """Image-encoder trace inputs: a full block grid padded to the budget (int32 positions).""" +def image_reference_inputs( + image: ImageTowerFacts, + dtype: torch.dtype = torch.float32, + position_dtype: torch.dtype = torch.int32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Image-encoder trace inputs: a block grid padded to the budget (-1 positions, zero pixels). + + The GPU program takes float32 pixels and int32 positions; the Neural Engine takes float16 for + both. + """ k = image.pooling_kernel_size n = image.budget # Square-ish grid of n soft tokens (blocks) -> (bw*k) x (bh*k) patches, row-major. @@ -283,13 +291,13 @@ def image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, torch. bh = max(1, n // bw) pw, ph = bw * k, bh * k ys, xs = torch.meshgrid(torch.arange(ph), torch.arange(pw), indexing="ij") - real = torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1).to(torch.int32) + real = torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1) positions = torch.full((1, image.max_patches, 2), -1, dtype=torch.int32) - positions[0, : real.shape[0]] = real + positions[0, : real.shape[0]] = real.to(torch.int32) generator = torch.Generator().manual_seed(image.budget) pixels = torch.rand((1, image.max_patches, image.patch_dim), generator=generator) pixels[0, real.shape[0] :] = 0.0 - return pixels, positions + return pixels.to(dtype), positions.to(position_dtype) def ane_text_reference_inputs( @@ -303,19 +311,6 @@ def ane_text_reference_inputs( return embeds, bias, pool -def ane_image_reference_inputs(image: ImageTowerFacts) -> tuple[torch.Tensor, torch.Tensor]: - """Neural Engine image trace inputs (float16): a block grid padded to the budget (-1 pad).""" - k = image.pooling_kernel_size - a = int(image.budget**0.5) - gh, gw = k * a, k * (image.budget // a) - generator = torch.Generator().manual_seed(image.budget) - pixels = torch.rand((1, image.max_patches, image.patch_dim), generator=generator) - ys, xs = torch.meshgrid(torch.arange(gh), torch.arange(gw), indexing="ij") - positions = torch.full((1, image.max_patches, 2), -1.0) - positions[0, : gh * gw] = torch.stack([xs.reshape(-1), ys.reshape(-1)], -1).float() - return pixels.to(torch.float16), positions.to(torch.float16) - - @dataclass class TraceSpec: """One function of a multi-function program: example args, IO names, dynamic dims.""" @@ -522,7 +517,10 @@ def export_ane_image_encoder_program( ) -> AIProgram: """Neural Engine ``image_encoder``: function ``main`` at the static budget.""" trace = TraceSpec( - "main", ane_image_reference_inputs(image), IMAGE_INPUT_NAMES, (IMAGE_OUTPUT_NAME,) + "main", + image_reference_inputs(image, torch.float16, torch.float16), + IMAGE_INPUT_NAMES, + (IMAGE_OUTPUT_NAME,), ) return convert_ane_functions([(image_encoder, trace)], include_debug_info=include_debug_info) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 98789cc3..16fbbc5c 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -633,7 +633,7 @@ def test_ane_reference_inputs_shapes(): embeds, bias, pool = embedding_export.ane_text_reference_inputs(128, 512) assert embeds.shape == (1, 128, 512) and bias.shape == (1, 128, 1, 128) assert pool.shape == (1, 128) and {t.dtype for t in (embeds, bias, pool)} == {torch.float16} - pixels, positions = embedding_export.ane_image_reference_inputs(IMAGE) + pixels, positions = embedding_export.image_reference_inputs(IMAGE, torch.float16, torch.float16) assert pixels.shape == (1, 2520, 768) and positions.shape == (1, 2520, 2) assert pixels.dtype == positions.dtype == torch.float16 real = positions[0, :, 0] >= 0 From 3747ec44f400b38c183ddad55550c2646976e8ae Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 20:38:20 -0700 Subject: [PATCH 19/38] Split the embedding exporter into shared, GPU and Neural Engine modules 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. --- python/src/coreai_models/embedding/export.py | 643 +----------------- .../src/coreai_models/embedding/export_ios.py | 190 ++++++ .../coreai_models/embedding/export_macos.py | 273 ++++++++ .../coreai_models/embedding/export_utils.py | 270 ++++++++ .../test_export/test_embedding_export.py | 59 +- 5 files changed, 779 insertions(+), 656 deletions(-) create mode 100644 python/src/coreai_models/embedding/export_ios.py create mode 100644 python/src/coreai_models/embedding/export_macos.py create mode 100644 python/src/coreai_models/embedding/export_utils.py diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index 3d863be0..894f6165 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -20,36 +20,15 @@ import logging import shutil import sys -import time -from collections.abc import Iterator -from contextlib import contextmanager -from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any -import torch -import torch.nn as nn - from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO from coreai_models.embedding.contract import ( - ANE_ENCODER_INPUT_NAMES, EMBEDDING_FAMILIES, - ENCODER_COMPONENT, - ENCODER_INPUT_NAMES, - IMAGE_COMPONENT, - IMAGE_INPUT_NAMES, - IMAGE_OUTPUT_NAME, - OUTPUT_NAME, PLATFORMS, - TEXT_COMPONENT, - TEXT_INPUT_NAMES, - TOKEN_TABLE_COMPONENT, - TOKEN_TABLE_INPUT_NAMES, - TOKEN_TABLE_OUTPUT_NAME, EmbeddingFamily, EmbeddingPrecisionError, - FunctionSpec, - ImageTowerFacts, build_bundle_metadata, bundle_assets, bundle_name, @@ -61,28 +40,25 @@ family_for_model, text_functions, ) +from coreai_models.embedding.export_ios import ( + export_ane_programs, +) +from coreai_models.embedding.export_macos import ( + export_gpu_programs, +) +from coreai_models.embedding.export_utils import ( + copy_tokenizer, + read_prompts, + resolve_snapshot, + verify_image_config, + verify_pooling_config, +) if TYPE_CHECKING: - from coreai.authoring import AIProgram + pass logger = logging.getLogger(__name__) -TOKENIZER_FILES = ( - "tokenizer.json", - "tokenizer_config.json", - "tokenizer.model", - "special_tokens_map.json", - "added_tokens.json", -) -SNAPSHOT_PATTERNS = [ - "*.json", - "*.safetensors", - "tokenizer.model", - "1_Pooling/*", - "2_Normalize/*", -] -_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16} - def _find_repo_root() -> Path | None: d = Path(__file__).resolve().parent @@ -98,593 +74,6 @@ def _default_output_dir() -> str: return str(root / "exports") if root is not None else "exports" -def resolve_snapshot(hf_model_id: str, model_dir: str | None = None) -> tuple[Path, str]: - """Return ``(snapshot_dir, revision_sha)`` for a local dir or a hub download. - - The SHA is the HF cache snapshot directory name (``.../snapshots/``); a plain directory - outside the cache yields ``""``. - """ - if model_dir is not None: - path = Path(model_dir).expanduser().resolve() - if not path.is_dir(): - raise FileNotFoundError(f"model dir {path} is not a directory") - else: - from huggingface_hub import snapshot_download - - path = Path(snapshot_download(hf_model_id, allow_patterns=SNAPSHOT_PATTERNS)).resolve() - sha = path.name if path.parent.name == "snapshots" else "" - return path, sha - - -def read_prompts(snapshot: Path) -> tuple[dict[str, str], str | None]: - """Read ``prompts`` / ``default_prompt_name`` from ``config_sentence_transformers.json``.""" - cfg_path = snapshot / "config_sentence_transformers.json" - if not cfg_path.exists(): - return {}, None - cfg = json.loads(cfg_path.read_text()) - return dict(cfg.get("prompts") or {}), cfg.get("default_prompt_name") - - -def verify_pooling_config(snapshot: Path, family: EmbeddingFamily) -> None: - """Fail closed if the snapshot's sentence-transformers pooling disagrees with the family.""" - pooling = snapshot / "1_Pooling" / "config.json" - if not pooling.exists(): - logger.warning("No 1_Pooling/config.json in %s; trusting the family defaults", snapshot) - return - cfg = json.loads(pooling.read_text()) - mode = cfg.get("pooling_mode") - if mode is None: - mode = "mean" if cfg.get("pooling_mode_mean_tokens") else None - problems = [] - if mode != family.pooling_mode: - problems.append(f"pooling_mode {mode!r} != {family.pooling_mode!r}") - if bool(cfg.get("include_prompt", True)) != family.include_prompt: - problems.append(f"include_prompt {cfg.get('include_prompt')} != {family.include_prompt}") - if int(cfg.get("embedding_dimension", family.dimension)) != family.dimension: - problems.append( - f"embedding_dimension {cfg.get('embedding_dimension')} != {family.dimension}" - ) - if family.normalized and not (snapshot / "2_Normalize").exists(): - problems.append("family is normalized but the snapshot has no 2_Normalize module") - if problems: - raise ValueError(f"Snapshot pooling config disagrees with the family: {problems}") - - -def copy_tokenizer(snapshot: Path, dest: Path) -> list[str]: - """Copy the tokenizer files verbatim (no re-serialization through a tokenizer class).""" - dest.mkdir(parents=True, exist_ok=True) - copied = [] - for name in TOKENIZER_FILES: - src = snapshot / name - if src.exists(): - shutil.copy2(src, dest / name) - copied.append(name) - if "tokenizer.json" not in copied: - raise FileNotFoundError(f"{snapshot} has no tokenizer.json") - return copied - - -class TextEmbeddingExportWrapper(nn.Module): - """``(input_ids, attention_mask) -> embedding`` with float32 output. - - The wrapped encoder owns pooling (masked mean incl. prompt), projection and L2 normalization. - """ - - def __init__(self, encoder: nn.Module) -> None: - super().__init__() - self.encoder = encoder - - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.encoder(input_ids=input_ids, attention_mask=attention_mask).to(torch.float32) - - -class TokenTableExportWrapper(nn.Module): - """``input_ids [B, L] int32 -> inputs_embeds [B, L, H]`` as ``dtype`` (scaled embeddings).""" - - def __init__(self, embed_tokens: nn.Module, dtype: torch.dtype = torch.float32) -> None: - super().__init__() - self.embed_tokens = embed_tokens - self.dtype = dtype - - def forward(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids).to(self.dtype) - - -def scaled_token_table(table: torch.Tensor, scale: float, dtype: torch.dtype) -> nn.Embedding: - """``table * scale`` stored as ``dtype``. Scaling in float32 first keeps the cast exact.""" - return nn.Embedding.from_pretrained((table.float() * scale).to(dtype)) - - -class InputsEmbedsEncoderWrapper(nn.Module): - """``(inputs_embeds, attention_mask) -> embedding`` for merged (e.g. image) inputs.""" - - def __init__(self, encoder: nn.Module) -> None: - super().__init__() - self.encoder = encoder - - def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - dtype = next(self.encoder.parameters()).dtype - return self.encoder( - inputs_embeds=inputs_embeds.to(dtype), attention_mask=attention_mask - ).to(torch.float32) - - -class EncoderBandWrapper(nn.Module): - """``(inputs_embeds, bias, pool_weights) -> embedding`` at one static length (Neural Engine). - - When ``length > sliding_window + 1`` the sliding layers also need the ``|i - j| <= window`` - band; it is a per-length constant folded in here (``min(bias, band)``), so every length keeps - the same three inputs. - """ - - def __init__(self, encoder: nn.Module, length: int, sliding_window: int) -> None: - super().__init__() - from coreai_models.models.ios.embedding_gemma2 import MASK_NEG - - self.encoder = encoder - self.banded = length > sliding_window + 1 - if self.banded: - pos = torch.arange(length) - band = torch.where((pos[:, None] - pos[None, :]).abs() <= sliding_window, 0.0, MASK_NEG) - dtype = next(encoder.parameters()).dtype - self.register_buffer("band", band.reshape(1, length, 1, length).to(dtype)) - - def forward( - self, inputs_embeds: torch.Tensor, bias: torch.Tensor, pool_weights: torch.Tensor - ) -> torch.Tensor: - sliding = torch.minimum(bias, self.band) if self.banded else None - return self.encoder(inputs_embeds, bias, pool_weights, sliding) - - -class ImageEncoderExportWrapper(nn.Module): - """``(pixel_values f32, position_ids int32) -> soft_tokens f32`` at the module's dtype.""" - - def __init__(self, image_encoder: nn.Module) -> None: - super().__init__() - self.image_encoder = image_encoder - - def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - dtype = next(self.image_encoder.parameters()).dtype - return self.image_encoder(pixel_values.to(dtype), position_ids).to(torch.float32) - - -@contextmanager -def without_token_table(encoder: nn.Module) -> Iterator[nn.Module]: - """Temporarily drop ``encoder.embed_tokens`` so an ``inputs_embeds`` trace carries no table.""" - table = encoder.embed_tokens - encoder.embed_tokens = nn.Identity() - try: - yield encoder - finally: - encoder.embed_tokens = table - - -def reference_inputs(spec: FunctionSpec, vocab_size: int) -> tuple[torch.Tensor, torch.Tensor]: - """Trace inputs for one function (values do not matter; shapes and dtypes do).""" - generator = torch.Generator().manual_seed(spec.batch * 100_003 + spec.length) - ids = torch.randint(3, vocab_size, (spec.batch, spec.length), generator=generator) - mask = torch.ones(spec.batch, spec.length, dtype=torch.int32) - return ids.to(torch.int32), mask - - -def embeds_reference_inputs(spec: FunctionSpec, hidden: int) -> tuple[torch.Tensor, torch.Tensor]: - """Trace inputs for one ``inputs_embeds`` function (values do not matter).""" - generator = torch.Generator().manual_seed(spec.length) - embeds = torch.randn((spec.batch, spec.length, hidden), generator=generator) - return embeds, torch.ones((spec.batch, spec.length), dtype=torch.int32) - - -def image_reference_inputs( - image: ImageTowerFacts, - dtype: torch.dtype = torch.float32, - position_dtype: torch.dtype = torch.int32, -) -> tuple[torch.Tensor, torch.Tensor]: - """Image-encoder trace inputs: a block grid padded to the budget (-1 positions, zero pixels). - - The GPU program takes float32 pixels and int32 positions; the Neural Engine takes float16 for - both. - """ - k = image.pooling_kernel_size - n = image.budget - # Square-ish grid of n soft tokens (blocks) -> (bw*k) x (bh*k) patches, row-major. - bw = max(1, int(n**0.5)) - bh = max(1, n // bw) - pw, ph = bw * k, bh * k - ys, xs = torch.meshgrid(torch.arange(ph), torch.arange(pw), indexing="ij") - real = torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1) - positions = torch.full((1, image.max_patches, 2), -1, dtype=torch.int32) - positions[0, : real.shape[0]] = real.to(torch.int32) - generator = torch.Generator().manual_seed(image.budget) - pixels = torch.rand((1, image.max_patches, image.patch_dim), generator=generator) - pixels[0, real.shape[0] :] = 0.0 - return pixels.to(dtype), positions.to(position_dtype) - - -def ane_text_reference_inputs( - length: int, hidden: int -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Neural Engine encoder trace inputs (all valid tokens): embeds, zero bias, uniform pooling.""" - generator = torch.Generator().manual_seed(length) - embeds = torch.randn((1, length, hidden), generator=generator).to(torch.float16) - bias = torch.zeros((1, length, 1, length), dtype=torch.float16) - pool = torch.full((1, length), 1.0 / length, dtype=torch.float16) - return embeds, bias, pool - - -@dataclass -class TraceSpec: - """One function of a multi-function program: example args, IO names, dynamic dims.""" - - name: str - args: tuple[torch.Tensor, ...] - input_names: tuple[str, ...] - output_names: tuple[str, ...] - dynamic_shapes: tuple[dict[int, Any] | None, ...] | None = None - - -def add_traces( - converter: Any, - traces: list[tuple[nn.Module, TraceSpec]], - decompositions: dict[Any, Any] | None = None, -) -> None: - """Add each ``(module, trace)`` to ``converter`` as a function (Core AI's table by default).""" - from coreai_models.diffusion.gpu import add_function - - for module, trace in traces: - start = time.time() - add_function( - converter, - trace.name, - module, - trace.args, - list(trace.input_names), - list(trace.output_names), - dynamic_shapes=trace.dynamic_shapes, - decompositions=decompositions, - ) - logger.info( - "traced %s %s in %.1fs", - trace.name, - [tuple(a.shape) for a in trace.args], - time.time() - start, - ) - - -def convert_ane_functions( - traces: list[tuple[nn.Module, TraceSpec]], - *, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """Convert every ``(module, trace)`` into one multi-function Neural Engine program. - - Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. - """ - from coreai_models.diffusion.gpu import new_converter - from coreai_models.export.mlir_ops import register_custom_torch_lowering - - converter = new_converter(include_debug_info) - register_custom_torch_lowering(converter) - add_traces(converter, traces, torch.export.default_decompositions()) - return converter.to_coreai() - - -def export_text_encoder_program( - wrapper: nn.Module, - functions: list[FunctionSpec], - vocab_size: int, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """GPU ``text_encoder``: one static (batch, length) function per spec, weights shared.""" - from coreai_models.diffusion.gpu import export_multifunction - - # Static (batch, length) traces; the host pads to the next listed batch. - return export_multifunction( - [(spec.name, wrapper, reference_inputs(spec, vocab_size)) for spec in functions], - TEXT_INPUT_NAMES, - (OUTPUT_NAME,), - include_debug_info, - ) - - -def export_token_table_program( - embed_tokens: nn.Module, - max_length: int, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, - dtype: torch.dtype = torch.float32, -) -> AIProgram: - """``token_table``: function ``main``, ids -> scaled embeddings, dynamic length (batch 1).""" - from coreai_models.diffusion.gpu import export_stateless - - # Placeholder ids look up like any id; the host overwrites those rows with soft tokens. - length = torch.export.Dim("length", min=2, max=max_length) - ids = torch.zeros((1, 64), dtype=torch.int32) - return export_stateless( - TokenTableExportWrapper(embed_tokens, dtype), - (ids,), - TOKEN_TABLE_INPUT_NAMES, - (TOKEN_TABLE_OUTPUT_NAME,), - dynamic_shapes=({1: length},), - include_debug_info=include_debug_info, - ) - - -def export_shared_text_program( - encoder: nn.Module, - text_fns: list[FunctionSpec], - encoder_fns: list[FunctionSpec], - max_length: int, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """GPU ``text_encoder`` that also holds the ``token_table`` and ``encoder`` functions. - - Text functions take ids; ``main`` maps ids to ``inputs_embeds`` (dynamic length up to - ``max_length``); ``embeds_*`` functions take ``inputs_embeds``. The functions share one copy - of the table and the transformer. - """ - from coreai_models.diffusion.gpu import new_converter - - vocab_size = int(encoder.config.vocab_size) - hidden = int(encoder.config.hidden_size) - length = torch.export.Dim("length", min=2, max=max_length) - text = TextEmbeddingExportWrapper(encoder) - traces = [ - ( - text, - TraceSpec( - spec.name, reference_inputs(spec, vocab_size), TEXT_INPUT_NAMES, (OUTPUT_NAME,) - ), - ) - for spec in text_fns - ] - traces.append( - ( - TokenTableExportWrapper(encoder.embed_tokens), - TraceSpec( - "main", - (torch.zeros((1, 64), dtype=torch.int32),), - TOKEN_TABLE_INPUT_NAMES, - (TOKEN_TABLE_OUTPUT_NAME,), - ({1: length},), - ), - ) - ) - converter = new_converter(include_debug_info) - add_traces(converter, traces) - with without_token_table(encoder) as inner: - embeds = InputsEmbedsEncoderWrapper(inner) - add_traces( - converter, - [ - ( - embeds, - TraceSpec( - spec.name, - embeds_reference_inputs(spec, hidden), - ENCODER_INPUT_NAMES, - (OUTPUT_NAME,), - ), - ) - for spec in encoder_fns - ], - ) - return converter.to_coreai() - - -def export_image_encoder_program( - image_encoder: nn.Module, - image: ImageTowerFacts, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """GPU ``image_encoder``: function ``main``, pixels and positions -> soft tokens.""" - from coreai_models.diffusion.gpu import export_stateless - - return export_stateless( - ImageEncoderExportWrapper(image_encoder), - image_reference_inputs(image), - IMAGE_INPUT_NAMES, - (IMAGE_OUTPUT_NAME,), - include_debug_info=include_debug_info, - ) - - -def export_ane_encoder_program( - encoder: nn.Module, - functions: list[FunctionSpec], - sliding_window: int, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """Neural Engine ``encoder``: one static batch-1 float16 function per length.""" - hidden = int(encoder.config.hidden_size) - traces = [ - ( - EncoderBandWrapper(encoder, spec.length, sliding_window), - TraceSpec( - spec.name, - ane_text_reference_inputs(spec.length, hidden), - ANE_ENCODER_INPUT_NAMES, - (OUTPUT_NAME,), - ), - ) - for spec in functions - ] - return convert_ane_functions(traces, include_debug_info=include_debug_info) - - -def export_ane_image_encoder_program( - image_encoder: nn.Module, - image: ImageTowerFacts, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """Neural Engine ``image_encoder``: function ``main`` at the static budget.""" - trace = TraceSpec( - "main", - image_reference_inputs(image, torch.float16, torch.float16), - IMAGE_INPUT_NAMES, - (IMAGE_OUTPUT_NAME,), - ) - return convert_ane_functions([(image_encoder, trace)], include_debug_info=include_debug_info) - - -def verify_image_config(snapshot: Path, family: EmbeddingFamily) -> None: - """Fail closed if the snapshot's merge ids / image processor disagree with the family facts.""" - image = family.image - if image is None: - raise ValueError(f"{family.hf_model_id} has no image tower") - problems = [] - config_path = snapshot / "config.json" - if config_path.exists(): - cfg = json.loads(config_path.read_text()) - for key, want in ( - ("image_token_id", image.placeholder_id), - ("boi_token_id", image.boi_token_id), - ("eoi_token_id", image.eoi_token_id), - ): - if key in cfg and int(cfg[key]) != want: - problems.append(f"config.json {key} {cfg[key]} != {want}") - text_hidden = (cfg.get("text_config") or {}).get("hidden_size") - if text_hidden is not None and int(text_hidden) != image.soft_token_dim: - problems.append(f"text hidden_size {text_hidden} != {image.soft_token_dim}") - processor_path = snapshot / "processor_config.json" - if processor_path.exists(): - proc = json.loads(processor_path.read_text()).get("image_processor") or {} - checks = ( - ("patch_size", image.patch_size), - ("pooling_kernel_size", image.pooling_kernel_size), - ("max_soft_tokens", image.budget), - ("do_normalize", image.normalize), - ("resample", 3), # PIL/torchvision BICUBIC - ) - for key, want in checks: - if key in proc and proc[key] != want: - problems.append(f"processor {key} {proc[key]!r} != {want!r}") - if "rescale_factor" in proc and abs(proc["rescale_factor"] - image.rescale_factor) > 1e-12: - problems.append(f"processor rescale_factor {proc['rescale_factor']} != 1/255") - if problems: - raise ValueError(f"Snapshot image config disagrees with the family: {problems}") - - -def _load_text_encoder(family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype) -> nn.Module: - if family.key == "embeddinggemma-2": - from coreai_models.models.macos.embedding_gemma2 import load_embedding_gemma2_text_encoder - - return load_embedding_gemma2_text_encoder(snapshot, dtype=dtype) - raise NotImplementedError(f"No text-encoder loader for family {family.key!r}") - - -def _load_image_encoder( - family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype, max_soft_tokens: int -) -> nn.Module: - if family.key == "embeddinggemma-2": - from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder - - return EmbeddingGemma2ImageEncoder.from_pretrained( - str(snapshot), dtype=dtype, max_soft_tokens=max_soft_tokens - ) - raise NotImplementedError(f"No image-encoder loader for family {family.key!r}") - - -def _save(program: AIProgram, out: Path, family: EmbeddingFamily, component: str) -> None: - from coreai_models.export import metadata as export_metadata - - program.save_asset( - out / f"{component}.aimodel", - export_metadata.build_aimodel_metadata(family.hf_model_id, component=component), - ) - - -def _export_gpu_programs( - family: EmbeddingFamily, - snapshot: Path, - out: Path, - *, - with_image: bool, - compute_precision: str, - include_debug_info: bool, -) -> None: - dtype = _DTYPES[compute_precision] - encoder = _load_text_encoder(family, snapshot, dtype) - if with_image: - program = export_shared_text_program( - encoder, - text_functions(family, "macOS"), - encoder_functions(family, "macOS"), - family.model_max_tokens, - include_debug_info, - ) - else: - program = export_text_encoder_program( - TextEmbeddingExportWrapper(encoder), - text_functions(family, "macOS"), - int(encoder.config.vocab_size), - include_debug_info, - ) - _save(program, out, family, TEXT_COMPONENT) - del program, encoder - if not with_image: - return - assert family.image is not None - image_encoder = _load_image_encoder(family, snapshot, dtype, family.image.budget) - _save( - export_image_encoder_program(image_encoder, family.image, include_debug_info), - out, - family, - IMAGE_COMPONENT, - ) - - -def _export_ane_programs( - family: EmbeddingFamily, - snapshot: Path, - out: Path, - *, - with_image: bool, - include_debug_info: bool, -) -> None: - if family.key != "embeddinggemma-2": - raise NotImplementedError(f"No Neural Engine export for family {family.key!r}") - from coreai_models.models.ios.embedding_gemma2 import load_embedding_gemma2_text_encoder - - # The token table runs on the GPU and stores the scaled table in float16, the encoder's input - # precision. - encoder, table = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float16) - token_table = scaled_token_table(table, encoder.config.hidden_size**0.5, torch.float16) - del table - _save( - export_token_table_program( - token_table, family.model_max_tokens, include_debug_info, dtype=torch.float16 - ), - out, - family, - TOKEN_TABLE_COMPONENT, - ) - del token_table - _save( - export_ane_encoder_program( - encoder, - encoder_functions(family, "iOS"), - int(encoder.config.sliding_window), - include_debug_info, - ), - out, - family, - ENCODER_COMPONENT, - ) - del encoder - if with_image: - assert family.image is not None - from coreai_models.models.ios.embedding_gemma2_vision import load_image_encoder - - image_encoder = load_image_encoder( - str(snapshot), torch.float16, max_soft_tokens=family.image.budget - ) - _save( - export_ane_image_encoder_program(image_encoder, family.image, include_debug_info), - out, - family, - IMAGE_COMPONENT, - ) - - def export_embedding_bundle( family: EmbeddingFamily, output_dir: Path, @@ -731,7 +120,7 @@ def export_embedding_bundle( copy_tokenizer(snapshot, out / "tokenizer") if platform == "iOS": - _export_ane_programs( + export_ane_programs( family, snapshot, out, @@ -739,7 +128,7 @@ def export_embedding_bundle( include_debug_info=include_debug_info, ) else: - _export_gpu_programs( + export_gpu_programs( family, snapshot, out, diff --git a/python/src/coreai_models/embedding/export_ios.py b/python/src/coreai_models/embedding/export_ios.py new file mode 100644 index 00000000..c29b885f --- /dev/null +++ b/python/src/coreai_models/embedding/export_ios.py @@ -0,0 +1,190 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Neural Engine (iOS) programs of an embedding bundle. + +The static-length encoder, the float16 token table, and the image encoder. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn + +from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO +from coreai_models.embedding.contract import ( + ANE_ENCODER_INPUT_NAMES, + ENCODER_COMPONENT, + IMAGE_COMPONENT, + IMAGE_INPUT_NAMES, + IMAGE_OUTPUT_NAME, + OUTPUT_NAME, + TOKEN_TABLE_COMPONENT, + EmbeddingFamily, + FunctionSpec, + ImageTowerFacts, + encoder_functions, +) +from coreai_models.embedding.export_utils import ( + TraceSpec, + add_traces, + export_token_table_program, + image_reference_inputs, + save_program, + scaled_token_table, +) + +if TYPE_CHECKING: + from coreai.authoring import AIProgram + + +class EncoderBandWrapper(nn.Module): + """``(inputs_embeds, bias, pool_weights) -> embedding`` at one static length (Neural Engine). + + When ``length > sliding_window + 1`` the sliding layers also need the ``|i - j| <= window`` + band; it is a per-length constant folded in here (``min(bias, band)``), so every length keeps + the same three inputs. + """ + + def __init__(self, encoder: nn.Module, length: int, sliding_window: int) -> None: + super().__init__() + from coreai_models.models.ios.embedding_gemma2 import MASK_NEG + + self.encoder = encoder + self.banded = length > sliding_window + 1 + if self.banded: + pos = torch.arange(length) + band = torch.where((pos[:, None] - pos[None, :]).abs() <= sliding_window, 0.0, MASK_NEG) + dtype = next(encoder.parameters()).dtype + self.register_buffer("band", band.reshape(1, length, 1, length).to(dtype)) + + def forward( + self, inputs_embeds: torch.Tensor, bias: torch.Tensor, pool_weights: torch.Tensor + ) -> torch.Tensor: + sliding = torch.minimum(bias, self.band) if self.banded else None + return self.encoder(inputs_embeds, bias, pool_weights, sliding) + + +def ane_text_reference_inputs( + length: int, hidden: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Neural Engine encoder trace inputs (all valid tokens): embeds, zero bias, uniform pooling.""" + generator = torch.Generator().manual_seed(length) + embeds = torch.randn((1, length, hidden), generator=generator).to(torch.float16) + bias = torch.zeros((1, length, 1, length), dtype=torch.float16) + pool = torch.full((1, length), 1.0 / length, dtype=torch.float16) + return embeds, bias, pool + + +def convert_ane_functions( + traces: list[tuple[nn.Module, TraceSpec]], + *, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """Convert every ``(module, trace)`` into one multi-function Neural Engine program. + + Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. + """ + from coreai_models.diffusion.gpu import new_converter + from coreai_models.export.mlir_ops import register_custom_torch_lowering + + converter = new_converter(include_debug_info) + register_custom_torch_lowering(converter) + add_traces(converter, traces, torch.export.default_decompositions()) + return converter.to_coreai() + + +def export_ane_encoder_program( + encoder: nn.Module, + functions: list[FunctionSpec], + sliding_window: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """Neural Engine ``encoder``: one static batch-1 float16 function per length.""" + hidden = int(encoder.config.hidden_size) + traces = [ + ( + EncoderBandWrapper(encoder, spec.length, sliding_window), + TraceSpec( + spec.name, + ane_text_reference_inputs(spec.length, hidden), + ANE_ENCODER_INPUT_NAMES, + (OUTPUT_NAME,), + ), + ) + for spec in functions + ] + return convert_ane_functions(traces, include_debug_info=include_debug_info) + + +def export_ane_image_encoder_program( + image_encoder: nn.Module, + image: ImageTowerFacts, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """Neural Engine ``image_encoder``: function ``main`` at the static budget.""" + trace = TraceSpec( + "main", + image_reference_inputs(image, torch.float16, torch.float16), + IMAGE_INPUT_NAMES, + (IMAGE_OUTPUT_NAME,), + ) + return convert_ane_functions([(image_encoder, trace)], include_debug_info=include_debug_info) + + +def export_ane_programs( + family: EmbeddingFamily, + snapshot: Path, + out: Path, + *, + with_image: bool, + include_debug_info: bool, +) -> None: + if family.key != "embeddinggemma-2": + raise NotImplementedError(f"No Neural Engine export for family {family.key!r}") + from coreai_models.models.ios.embedding_gemma2 import load_embedding_gemma2_text_encoder + + # The token table runs on the GPU and stores the scaled table in float16, the encoder's input + # precision. + encoder, table = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float16) + token_table = scaled_token_table(table, encoder.config.hidden_size**0.5, torch.float16) + del table + save_program( + export_token_table_program( + token_table, family.model_max_tokens, include_debug_info, dtype=torch.float16 + ), + out, + family, + TOKEN_TABLE_COMPONENT, + ) + del token_table + save_program( + export_ane_encoder_program( + encoder, + encoder_functions(family, "iOS"), + int(encoder.config.sliding_window), + include_debug_info, + ), + out, + family, + ENCODER_COMPONENT, + ) + del encoder + if with_image: + assert family.image is not None + from coreai_models.models.ios.embedding_gemma2_vision import load_image_encoder + + image_encoder = load_image_encoder( + str(snapshot), torch.float16, max_soft_tokens=family.image.budget + ) + save_program( + export_ane_image_encoder_program(image_encoder, family.image, include_debug_info), + out, + family, + IMAGE_COMPONENT, + ) diff --git a/python/src/coreai_models/embedding/export_macos.py b/python/src/coreai_models/embedding/export_macos.py new file mode 100644 index 00000000..ecb82a8f --- /dev/null +++ b/python/src/coreai_models/embedding/export_macos.py @@ -0,0 +1,273 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""GPU (macOS) programs of an embedding bundle. + +The text encoder, the shared text / token-table / encoder program, and the image encoder. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn + +from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO +from coreai_models.embedding.contract import ( + ENCODER_INPUT_NAMES, + IMAGE_COMPONENT, + IMAGE_INPUT_NAMES, + IMAGE_OUTPUT_NAME, + OUTPUT_NAME, + TEXT_COMPONENT, + TEXT_INPUT_NAMES, + TOKEN_TABLE_INPUT_NAMES, + TOKEN_TABLE_OUTPUT_NAME, + EmbeddingFamily, + FunctionSpec, + ImageTowerFacts, + encoder_functions, + text_functions, +) +from coreai_models.embedding.export_utils import ( + TokenTableExportWrapper, + TraceSpec, + add_traces, + image_reference_inputs, + save_program, +) + +if TYPE_CHECKING: + from coreai.authoring import AIProgram + + +_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16} + + +class TextEmbeddingExportWrapper(nn.Module): + """``(input_ids, attention_mask) -> embedding`` with float32 output. + + The wrapped encoder owns pooling (masked mean incl. prompt), projection and L2 normalization. + """ + + def __init__(self, encoder: nn.Module) -> None: + super().__init__() + self.encoder = encoder + + def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + return self.encoder(input_ids=input_ids, attention_mask=attention_mask).to(torch.float32) + + +class InputsEmbedsEncoderWrapper(nn.Module): + """``(inputs_embeds, attention_mask) -> embedding`` for merged (e.g. image) inputs.""" + + def __init__(self, encoder: nn.Module) -> None: + super().__init__() + self.encoder = encoder + + def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + dtype = next(self.encoder.parameters()).dtype + return self.encoder( + inputs_embeds=inputs_embeds.to(dtype), attention_mask=attention_mask + ).to(torch.float32) + + +class ImageEncoderExportWrapper(nn.Module): + """``(pixel_values f32, position_ids int32) -> soft_tokens f32`` at the module's dtype.""" + + def __init__(self, image_encoder: nn.Module) -> None: + super().__init__() + self.image_encoder = image_encoder + + def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + dtype = next(self.image_encoder.parameters()).dtype + return self.image_encoder(pixel_values.to(dtype), position_ids).to(torch.float32) + + +@contextmanager +def without_token_table(encoder: nn.Module) -> Iterator[nn.Module]: + """Temporarily drop ``encoder.embed_tokens`` so an ``inputs_embeds`` trace carries no table.""" + table = encoder.embed_tokens + encoder.embed_tokens = nn.Identity() + try: + yield encoder + finally: + encoder.embed_tokens = table + + +def reference_inputs(spec: FunctionSpec, vocab_size: int) -> tuple[torch.Tensor, torch.Tensor]: + """Trace inputs for one function (values do not matter; shapes and dtypes do).""" + generator = torch.Generator().manual_seed(spec.batch * 100_003 + spec.length) + ids = torch.randint(3, vocab_size, (spec.batch, spec.length), generator=generator) + mask = torch.ones(spec.batch, spec.length, dtype=torch.int32) + return ids.to(torch.int32), mask + + +def embeds_reference_inputs(spec: FunctionSpec, hidden: int) -> tuple[torch.Tensor, torch.Tensor]: + """Trace inputs for one ``inputs_embeds`` function (values do not matter).""" + generator = torch.Generator().manual_seed(spec.length) + embeds = torch.randn((spec.batch, spec.length, hidden), generator=generator) + return embeds, torch.ones((spec.batch, spec.length), dtype=torch.int32) + + +def export_text_encoder_program( + wrapper: nn.Module, + functions: list[FunctionSpec], + vocab_size: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """GPU ``text_encoder``: one static (batch, length) function per spec, weights shared.""" + from coreai_models.diffusion.gpu import export_multifunction + + # Static (batch, length) traces; the host pads to the next listed batch. + return export_multifunction( + [(spec.name, wrapper, reference_inputs(spec, vocab_size)) for spec in functions], + TEXT_INPUT_NAMES, + (OUTPUT_NAME,), + include_debug_info, + ) + + +def export_shared_text_program( + encoder: nn.Module, + text_fns: list[FunctionSpec], + encoder_fns: list[FunctionSpec], + max_length: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """GPU ``text_encoder`` that also holds the ``token_table`` and ``encoder`` functions. + + Text functions take ids; ``main`` maps ids to ``inputs_embeds`` (dynamic length up to + ``max_length``); ``embeds_*`` functions take ``inputs_embeds``. The functions share one copy + of the table and the transformer. + """ + from coreai_models.diffusion.gpu import new_converter + + vocab_size = int(encoder.config.vocab_size) + hidden = int(encoder.config.hidden_size) + length = torch.export.Dim("length", min=2, max=max_length) + text = TextEmbeddingExportWrapper(encoder) + traces = [ + ( + text, + TraceSpec( + spec.name, reference_inputs(spec, vocab_size), TEXT_INPUT_NAMES, (OUTPUT_NAME,) + ), + ) + for spec in text_fns + ] + traces.append( + ( + TokenTableExportWrapper(encoder.embed_tokens), + TraceSpec( + "main", + (torch.zeros((1, 64), dtype=torch.int32),), + TOKEN_TABLE_INPUT_NAMES, + (TOKEN_TABLE_OUTPUT_NAME,), + ({1: length},), + ), + ) + ) + converter = new_converter(include_debug_info) + add_traces(converter, traces) + with without_token_table(encoder) as inner: + embeds = InputsEmbedsEncoderWrapper(inner) + add_traces( + converter, + [ + ( + embeds, + TraceSpec( + spec.name, + embeds_reference_inputs(spec, hidden), + ENCODER_INPUT_NAMES, + (OUTPUT_NAME,), + ), + ) + for spec in encoder_fns + ], + ) + return converter.to_coreai() + + +def export_image_encoder_program( + image_encoder: nn.Module, + image: ImageTowerFacts, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """GPU ``image_encoder``: function ``main``, pixels and positions -> soft tokens.""" + from coreai_models.diffusion.gpu import export_stateless + + return export_stateless( + ImageEncoderExportWrapper(image_encoder), + image_reference_inputs(image), + IMAGE_INPUT_NAMES, + (IMAGE_OUTPUT_NAME,), + include_debug_info=include_debug_info, + ) + + +def load_text_encoder(family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype) -> nn.Module: + if family.key == "embeddinggemma-2": + from coreai_models.models.macos.embedding_gemma2 import load_embedding_gemma2_text_encoder + + return load_embedding_gemma2_text_encoder(snapshot, dtype=dtype) + raise NotImplementedError(f"No text-encoder loader for family {family.key!r}") + + +def load_image_encoder( + family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype, max_soft_tokens: int +) -> nn.Module: + if family.key == "embeddinggemma-2": + from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder + + return EmbeddingGemma2ImageEncoder.from_pretrained( + str(snapshot), dtype=dtype, max_soft_tokens=max_soft_tokens + ) + raise NotImplementedError(f"No image-encoder loader for family {family.key!r}") + + +def export_gpu_programs( + family: EmbeddingFamily, + snapshot: Path, + out: Path, + *, + with_image: bool, + compute_precision: str, + include_debug_info: bool, +) -> None: + dtype = _DTYPES[compute_precision] + encoder = load_text_encoder(family, snapshot, dtype) + if with_image: + program = export_shared_text_program( + encoder, + text_functions(family, "macOS"), + encoder_functions(family, "macOS"), + family.model_max_tokens, + include_debug_info, + ) + else: + program = export_text_encoder_program( + TextEmbeddingExportWrapper(encoder), + text_functions(family, "macOS"), + int(encoder.config.vocab_size), + include_debug_info, + ) + save_program(program, out, family, TEXT_COMPONENT) + del program, encoder + if not with_image: + return + assert family.image is not None + image_encoder = load_image_encoder(family, snapshot, dtype, family.image.budget) + save_program( + export_image_encoder_program(image_encoder, family.image, include_debug_info), + out, + family, + IMAGE_COMPONENT, + ) diff --git a/python/src/coreai_models/embedding/export_utils.py b/python/src/coreai_models/embedding/export_utils.py new file mode 100644 index 00000000..bc5d558a --- /dev/null +++ b/python/src/coreai_models/embedding/export_utils.py @@ -0,0 +1,270 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Helpers shared by the GPU and Neural Engine embedding exports. + +Snapshot checks, the token table, trace inputs, and program assembly and saving. +""" + +from __future__ import annotations + +import json +import logging +import shutil +import time +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn + +from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO +from coreai_models.embedding.contract import ( + TOKEN_TABLE_INPUT_NAMES, + TOKEN_TABLE_OUTPUT_NAME, + EmbeddingFamily, + ImageTowerFacts, +) + +if TYPE_CHECKING: + from coreai.authoring import AIProgram + +logger = logging.getLogger(__name__) + + +TOKENIZER_FILES = ( + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "special_tokens_map.json", + "added_tokens.json", +) + + +SNAPSHOT_PATTERNS = [ + "*.json", + "*.safetensors", + "tokenizer.model", + "1_Pooling/*", + "2_Normalize/*", +] + + +def resolve_snapshot(hf_model_id: str, model_dir: str | None = None) -> tuple[Path, str]: + """Return ``(snapshot_dir, revision_sha)`` for a local dir or a hub download. + + The SHA is the HF cache snapshot directory name (``.../snapshots/``); a plain directory + outside the cache yields ``""``. + """ + if model_dir is not None: + path = Path(model_dir).expanduser().resolve() + if not path.is_dir(): + raise FileNotFoundError(f"model dir {path} is not a directory") + else: + from huggingface_hub import snapshot_download + + path = Path(snapshot_download(hf_model_id, allow_patterns=SNAPSHOT_PATTERNS)).resolve() + sha = path.name if path.parent.name == "snapshots" else "" + return path, sha + + +def read_prompts(snapshot: Path) -> tuple[dict[str, str], str | None]: + """Read ``prompts`` / ``default_prompt_name`` from ``config_sentence_transformers.json``.""" + cfg_path = snapshot / "config_sentence_transformers.json" + if not cfg_path.exists(): + return {}, None + cfg = json.loads(cfg_path.read_text()) + return dict(cfg.get("prompts") or {}), cfg.get("default_prompt_name") + + +def verify_pooling_config(snapshot: Path, family: EmbeddingFamily) -> None: + """Fail closed if the snapshot's sentence-transformers pooling disagrees with the family.""" + pooling = snapshot / "1_Pooling" / "config.json" + if not pooling.exists(): + logger.warning("No 1_Pooling/config.json in %s; trusting the family defaults", snapshot) + return + cfg = json.loads(pooling.read_text()) + mode = cfg.get("pooling_mode") + if mode is None: + mode = "mean" if cfg.get("pooling_mode_mean_tokens") else None + problems = [] + if mode != family.pooling_mode: + problems.append(f"pooling_mode {mode!r} != {family.pooling_mode!r}") + if bool(cfg.get("include_prompt", True)) != family.include_prompt: + problems.append(f"include_prompt {cfg.get('include_prompt')} != {family.include_prompt}") + if int(cfg.get("embedding_dimension", family.dimension)) != family.dimension: + problems.append( + f"embedding_dimension {cfg.get('embedding_dimension')} != {family.dimension}" + ) + if family.normalized and not (snapshot / "2_Normalize").exists(): + problems.append("family is normalized but the snapshot has no 2_Normalize module") + if problems: + raise ValueError(f"Snapshot pooling config disagrees with the family: {problems}") + + +def copy_tokenizer(snapshot: Path, dest: Path) -> list[str]: + """Copy the tokenizer files verbatim (no re-serialization through a tokenizer class).""" + dest.mkdir(parents=True, exist_ok=True) + copied = [] + for name in TOKENIZER_FILES: + src = snapshot / name + if src.exists(): + shutil.copy2(src, dest / name) + copied.append(name) + if "tokenizer.json" not in copied: + raise FileNotFoundError(f"{snapshot} has no tokenizer.json") + return copied + + +class TokenTableExportWrapper(nn.Module): + """``input_ids [B, L] int32 -> inputs_embeds [B, L, H]`` as ``dtype`` (scaled embeddings).""" + + def __init__(self, embed_tokens: nn.Module, dtype: torch.dtype = torch.float32) -> None: + super().__init__() + self.embed_tokens = embed_tokens + self.dtype = dtype + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids).to(self.dtype) + + +def scaled_token_table(table: torch.Tensor, scale: float, dtype: torch.dtype) -> nn.Embedding: + """``table * scale`` stored as ``dtype``. Scaling in float32 first keeps the cast exact.""" + return nn.Embedding.from_pretrained((table.float() * scale).to(dtype)) + + +def image_reference_inputs( + image: ImageTowerFacts, + dtype: torch.dtype = torch.float32, + position_dtype: torch.dtype = torch.int32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Image-encoder trace inputs: a block grid padded to the budget (-1 positions, zero pixels). + + The GPU program takes float32 pixels and int32 positions; the Neural Engine takes float16 for + both. + """ + k = image.pooling_kernel_size + n = image.budget + # Square-ish grid of n soft tokens (blocks) -> (bw*k) x (bh*k) patches, row-major. + bw = max(1, int(n**0.5)) + bh = max(1, n // bw) + pw, ph = bw * k, bh * k + ys, xs = torch.meshgrid(torch.arange(ph), torch.arange(pw), indexing="ij") + real = torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1) + positions = torch.full((1, image.max_patches, 2), -1, dtype=torch.int32) + positions[0, : real.shape[0]] = real.to(torch.int32) + generator = torch.Generator().manual_seed(image.budget) + pixels = torch.rand((1, image.max_patches, image.patch_dim), generator=generator) + pixels[0, real.shape[0] :] = 0.0 + return pixels.to(dtype), positions.to(position_dtype) + + +@dataclass +class TraceSpec: + """One function of a multi-function program: example args, IO names, dynamic dims.""" + + name: str + args: tuple[torch.Tensor, ...] + input_names: tuple[str, ...] + output_names: tuple[str, ...] + dynamic_shapes: tuple[dict[int, Any] | None, ...] | None = None + + +def add_traces( + converter: Any, + traces: list[tuple[nn.Module, TraceSpec]], + decompositions: dict[Any, Any] | None = None, +) -> None: + """Add each ``(module, trace)`` to ``converter`` as a function (Core AI's table by default).""" + from coreai_models.diffusion.gpu import add_function + + for module, trace in traces: + start = time.time() + add_function( + converter, + trace.name, + module, + trace.args, + list(trace.input_names), + list(trace.output_names), + dynamic_shapes=trace.dynamic_shapes, + decompositions=decompositions, + ) + logger.info( + "traced %s %s in %.1fs", + trace.name, + [tuple(a.shape) for a in trace.args], + time.time() - start, + ) + + +def export_token_table_program( + embed_tokens: nn.Module, + max_length: int, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, + dtype: torch.dtype = torch.float32, +) -> AIProgram: + """``token_table``: function ``main``, ids -> scaled embeddings, dynamic length (batch 1).""" + from coreai_models.diffusion.gpu import export_stateless + + # Placeholder ids look up like any id; the host overwrites those rows with soft tokens. + length = torch.export.Dim("length", min=2, max=max_length) + ids = torch.zeros((1, 64), dtype=torch.int32) + return export_stateless( + TokenTableExportWrapper(embed_tokens, dtype), + (ids,), + TOKEN_TABLE_INPUT_NAMES, + (TOKEN_TABLE_OUTPUT_NAME,), + dynamic_shapes=({1: length},), + include_debug_info=include_debug_info, + ) + + +def verify_image_config(snapshot: Path, family: EmbeddingFamily) -> None: + """Fail closed if the snapshot's merge ids / image processor disagree with the family facts.""" + image = family.image + if image is None: + raise ValueError(f"{family.hf_model_id} has no image tower") + problems = [] + config_path = snapshot / "config.json" + if config_path.exists(): + cfg = json.loads(config_path.read_text()) + for key, want in ( + ("image_token_id", image.placeholder_id), + ("boi_token_id", image.boi_token_id), + ("eoi_token_id", image.eoi_token_id), + ): + if key in cfg and int(cfg[key]) != want: + problems.append(f"config.json {key} {cfg[key]} != {want}") + text_hidden = (cfg.get("text_config") or {}).get("hidden_size") + if text_hidden is not None and int(text_hidden) != image.soft_token_dim: + problems.append(f"text hidden_size {text_hidden} != {image.soft_token_dim}") + processor_path = snapshot / "processor_config.json" + if processor_path.exists(): + proc = json.loads(processor_path.read_text()).get("image_processor") or {} + checks = ( + ("patch_size", image.patch_size), + ("pooling_kernel_size", image.pooling_kernel_size), + ("max_soft_tokens", image.budget), + ("do_normalize", image.normalize), + ("resample", 3), # PIL/torchvision BICUBIC + ) + for key, want in checks: + if key in proc and proc[key] != want: + problems.append(f"processor {key} {proc[key]!r} != {want!r}") + if "rescale_factor" in proc and abs(proc["rescale_factor"] - image.rescale_factor) > 1e-12: + problems.append(f"processor rescale_factor {proc['rescale_factor']} != 1/255") + if problems: + raise ValueError(f"Snapshot image config disagrees with the family: {problems}") + + +def save_program(program: AIProgram, out: Path, family: EmbeddingFamily, component: str) -> None: + from coreai_models.export import metadata as export_metadata + + program.save_asset( + out / f"{component}.aimodel", + export_metadata.build_aimodel_metadata(family.hf_model_id, component=component), + ) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 16fbbc5c..f05dc6d6 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -20,6 +20,7 @@ import torch from coreai_models.embedding import export as embedding_export +from coreai_models.embedding import export_ios, export_macos, export_utils from coreai_models.embedding.contract import ( EMBEDDING_FAMILIES, EmbeddingPrecisionError, @@ -360,7 +361,7 @@ def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch monkeypatch.setattr(embedding_export, "read_prompts", lambda *a: (PROMPTS, None)) monkeypatch.setattr(embedding_export, "consolidate_prompts", lambda family, prompts: prompts) monkeypatch.setattr(embedding_export, "copy_tokenizer", lambda *a: ["tokenizer.json"]) - monkeypatch.setattr(embedding_export, "_export_gpu_programs", lambda *a, **kw: seen.update(kw)) + monkeypatch.setattr(embedding_export, "export_gpu_programs", lambda *a, **kw: seen.update(kw)) meta = embedding_export.export_embedding_bundle( EG2, tmp_path / "out", with_image=False, compute_precision="bfloat16" ) @@ -381,7 +382,7 @@ def fail(*args, **kwargs): monkeypatch.setattr(embedding_export, "read_prompts", lambda *a: (PROMPTS, None)) monkeypatch.setattr(embedding_export, "consolidate_prompts", lambda family, prompts: prompts) monkeypatch.setattr(embedding_export, "copy_tokenizer", lambda *a: ["tokenizer.json"]) - monkeypatch.setattr(embedding_export, "_export_gpu_programs", fail) + monkeypatch.setattr(embedding_export, "export_gpu_programs", fail) with pytest.raises(RuntimeError, match="conversion failed"): embedding_export.export_embedding_bundle(EG2, tmp_path / "out", with_image=False) assert not (tmp_path / "out").exists() @@ -437,11 +438,11 @@ def test_snapshot_helpers(tmp_path): snap = _fake_snapshot( tmp_path, {"embedding_dimension": 768, "pooling_mode": "mean", "include_prompt": True} ) - path, sha = embedding_export.resolve_snapshot("google/embeddinggemma-2", str(snap)) + path, sha = export_utils.resolve_snapshot("google/embeddinggemma-2", str(snap)) assert sha == "a" * 40 and path == snap.resolve() - assert embedding_export.read_prompts(snap) == (SNAPSHOT_PROMPTS, None) - embedding_export.verify_pooling_config(snap, EG2) - copied = embedding_export.copy_tokenizer(snap, tmp_path / "bundle" / "tokenizer") + assert export_utils.read_prompts(snap) == (SNAPSHOT_PROMPTS, None) + export_utils.verify_pooling_config(snap, EG2) + copied = export_utils.copy_tokenizer(snap, tmp_path / "bundle" / "tokenizer") assert copied == ["tokenizer.json", "tokenizer_config.json"] @@ -450,7 +451,7 @@ def test_verify_pooling_config_fails_closed(tmp_path): tmp_path, {"embedding_dimension": 768, "pooling_mode": "mean", "include_prompt": False} ) with pytest.raises(ValueError, match="include_prompt"): - embedding_export.verify_pooling_config(snap, EG2) + export_utils.verify_pooling_config(snap, EG2) def _image_snapshot(root: Path, **config_overrides) -> Path: @@ -479,9 +480,9 @@ def _image_snapshot(root: Path, **config_overrides) -> Path: def test_verify_image_config(tmp_path): - embedding_export.verify_image_config(_image_snapshot(tmp_path / "ok"), EG2) + export_utils.verify_image_config(_image_snapshot(tmp_path / "ok"), EG2) with pytest.raises(ValueError, match="boi_token_id"): - embedding_export.verify_image_config(_image_snapshot(tmp_path / "bad", boi_token_id=1), EG2) + export_utils.verify_image_config(_image_snapshot(tmp_path / "bad", boi_token_id=1), EG2) # --- Conversion (tiny random config) --- @@ -508,14 +509,14 @@ def test_verify_image_config(tmp_path): def test_tiny_multifunction_conversion(): pytest.importorskip("coreai_torch") - wrapper = embedding_export.TextEmbeddingExportWrapper(_tiny_text_encoder()) + wrapper = export_macos.TextEmbeddingExportWrapper(_tiny_text_encoder()) plan = bucket_plan([16], [1, 4], 1024) - ids, mask = embedding_export.reference_inputs(plan[1], 64) + ids, mask = export_macos.reference_inputs(plan[1], 64) assert ids.dtype == torch.int32 and mask.dtype == torch.int32 and ids.shape == (4, 16) out = wrapper(ids, mask) assert out.dtype == torch.float32 and out.shape == (4, 24) torch.testing.assert_close(out.norm(dim=-1), torch.ones(4)) - assert embedding_export.export_text_encoder_program(wrapper, plan, 64) is not None + assert export_macos.export_text_encoder_program(wrapper, plan, 64) is not None def _tiny_text_encoder(): @@ -562,10 +563,10 @@ def test_inputs_embeds_path_equals_ids_path(): mask = torch.ones(1, 16, dtype=torch.int32) mask[0, 11:] = 0 with torch.no_grad(): - direct = embedding_export.TextEmbeddingExportWrapper(text)(ids, mask) - embeds = embedding_export.TokenTableExportWrapper(text.embed_tokens)(ids) - with embedding_export.without_token_table(text) as encoder: - merged = embedding_export.InputsEmbedsEncoderWrapper(encoder)(embeds, mask) + direct = export_macos.TextEmbeddingExportWrapper(text)(ids, mask) + embeds = export_utils.TokenTableExportWrapper(text.embed_tokens)(ids) + with export_macos.without_token_table(text) as encoder: + merged = export_macos.InputsEmbedsEncoderWrapper(encoder)(embeds, mask) assert embeds.dtype == torch.float32 and embeds.shape == (1, 16, 32) torch.testing.assert_close(merged, direct, atol=1e-6, rtol=0) assert isinstance(text.embed_tokens, torch.nn.Embedding) # restored @@ -575,12 +576,12 @@ def test_float16_scaled_token_table_equals_the_float32_table_cast(): """The float16 token table gives the encoder the same inputs as the float32 one.""" text = _tiny_text_encoder() ids = torch.randint(3, 64, (1, 16), dtype=torch.int32) - table = embedding_export.scaled_token_table( + table = export_utils.scaled_token_table( text.embed_tokens.weight, text.config.hidden_size**0.5, torch.float16 ) with torch.no_grad(): - want = embedding_export.TokenTableExportWrapper(text.embed_tokens)(ids).half() - got = embedding_export.TokenTableExportWrapper(table, torch.float16)(ids) + want = export_utils.TokenTableExportWrapper(text.embed_tokens)(ids).half() + got = export_utils.TokenTableExportWrapper(table, torch.float16)(ids) assert got.dtype == torch.float16 assert torch.equal(got, want) @@ -588,21 +589,21 @@ def test_float16_scaled_token_table_equals_the_float32_table_cast(): def test_tiny_float16_token_table_conversion(): pytest.importorskip("coreai_torch") text = _tiny_text_encoder() - table = embedding_export.scaled_token_table( + table = export_utils.scaled_token_table( text.embed_tokens.weight, text.config.hidden_size**0.5, torch.float16 ) - program = embedding_export.export_token_table_program(table, 64, dtype=torch.float16) + program = export_utils.export_token_table_program(table, 64, dtype=torch.float16) assert program is not None def test_image_reference_inputs_are_a_valid_padded_grid(): - pixels, positions = embedding_export.image_reference_inputs(TINY_IMAGE) + pixels, positions = export_utils.image_reference_inputs(TINY_IMAGE) assert pixels.shape == (1, 108, 48) and pixels.dtype == torch.float32 assert positions.shape == (1, 108, 2) and positions.dtype == torch.int32 real = positions[0, :, 0] >= 0 assert int(real.sum()) % 9 == 0 assert torch.all(pixels[0, ~real] == 0) - soft = embedding_export.ImageEncoderExportWrapper(_tiny_image_encoder())(pixels, positions) + soft = export_macos.ImageEncoderExportWrapper(_tiny_image_encoder())(pixels, positions) n = int(real.sum()) // 9 assert soft.shape == (1, 12, 32) and soft.dtype == torch.float32 assert torch.all(soft[0, n:] == 0) and torch.all(soft[0, :n].abs().sum(-1) > 0) @@ -610,14 +611,14 @@ def test_image_reference_inputs_are_a_valid_padded_grid(): def test_tiny_shared_text_program_conversion(): pytest.importorskip("coreai_torch") - program = embedding_export.export_shared_text_program( + program = export_macos.export_shared_text_program( _tiny_text_encoder(), [FunctionSpec("len16_b1", 1, 16), FunctionSpec("len32_b1", 1, 32)], [FunctionSpec("embeds_len32_b1", 1, 32)], max_length=64, ) assert program is not None - image = embedding_export.export_image_encoder_program(_tiny_image_encoder(), TINY_IMAGE) + image = export_macos.export_image_encoder_program(_tiny_image_encoder(), TINY_IMAGE) assert image is not None @@ -630,10 +631,10 @@ def test_macos_image_bundle_shares_the_text_encoder_asset(): def test_ane_reference_inputs_shapes(): - embeds, bias, pool = embedding_export.ane_text_reference_inputs(128, 512) + embeds, bias, pool = export_ios.ane_text_reference_inputs(128, 512) assert embeds.shape == (1, 128, 512) and bias.shape == (1, 128, 1, 128) assert pool.shape == (1, 128) and {t.dtype for t in (embeds, bias, pool)} == {torch.float16} - pixels, positions = embedding_export.image_reference_inputs(IMAGE, torch.float16, torch.float16) + pixels, positions = export_utils.image_reference_inputs(IMAGE, torch.float16, torch.float16) assert pixels.shape == (1, 2520, 768) and positions.shape == (1, 2520, 2) assert pixels.dtype == positions.dtype == torch.float16 real = positions[0, :, 0] >= 0 @@ -657,12 +658,12 @@ def test_ane_band_wrapper_matches_explicit_sliding_bias(): bias, sliding = (t.reshape(1, length, 1, length) for t in (bias, sliding)) embeds = torch.randn(1, length, 32) pool = (pos < valid).float().reshape(1, length) / valid - wrapper = embedding_export.EncoderBandWrapper(model, length, 3) + wrapper = export_ios.EncoderBandWrapper(model, length, 3) with torch.no_grad(): want = model(embeds, bias, pool, sliding) got = wrapper(embeds, bias, pool) torch.testing.assert_close(got, want) - assert not embedding_export.EncoderBandWrapper(model, 4, 3).banded + assert not export_ios.EncoderBandWrapper(model, 4, 3).banded # --- CLI --- From d2c748544ea49e3152e28e1e3df806f246cf24ff Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Wed, 28 Oct 2026 20:57:06 -0700 Subject: [PATCH 20/38] Note the host-gathered token table as a follow-up for the Neural Engine bundle --- python/src/coreai_models/embedding/export_ios.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/src/coreai_models/embedding/export_ios.py b/python/src/coreai_models/embedding/export_ios.py index c29b885f..0a090587 100644 --- a/python/src/coreai_models/embedding/export_ios.py +++ b/python/src/coreai_models/embedding/export_ios.py @@ -151,6 +151,9 @@ def export_ane_programs( # The token table runs on the GPU and stores the scaled table in float16, the encoder's input # precision. + # TODO: write the table as a safetensors sidecar that the host gathers from, as the Gemma 4 + # per-layer embeddings do (SafeTensorsReader), instead of a GPU program; lands with the Swift + # embedding runtime. encoder, table = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float16) token_table = scaled_token_table(table, encoder.config.hidden_size**0.5, torch.float16) del table From a8364dc6e6892ee26467eb06bccb09fbc7a037e6 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 21/38] Revert the Muse Glimmer rotate_half move Restores primitives/macos/rope.py and muse_glimmer_vision.py to the base branch. --- .../coreai_models/models/macos/muse_glimmer_vision.py | 11 ++++++++--- python/src/coreai_models/primitives/macos/rope.py | 6 ------ 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/python/src/coreai_models/models/macos/muse_glimmer_vision.py b/python/src/coreai_models/models/macos/muse_glimmer_vision.py index f6ddd010..679dc9f5 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer_vision.py +++ b/python/src/coreai_models/models/macos/muse_glimmer_vision.py @@ -58,13 +58,18 @@ _load_tensors_for_keys, _resolve_safetensors_files, ) -from coreai_models.primitives.macos.rope import rotate_half # --------------------------------------------------------------------------- # RoPE helpers # --------------------------------------------------------------------------- +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate second half of the last dimension: ``cat(-x2, x1)``.""" + mid = x.shape[-1] // 2 + return torch.cat((-x[..., mid:], x[..., :mid]), dim=-1) + + def _apply_rotary_pos_emb_vision( q: torch.Tensor, k: torch.Tensor, @@ -81,8 +86,8 @@ def _apply_rotary_pos_emb_vision( q, k = q.float(), k.float() cos = cos.unsqueeze(-2).float() # [1, seq_len, 1, head_dim] sin = sin.unsqueeze(-2).float() - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) + q_embed = (q * cos) + (_rotate_half(q) * sin) + k_embed = (k * cos) + (_rotate_half(k) * sin) return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype) diff --git a/python/src/coreai_models/primitives/macos/rope.py b/python/src/coreai_models/primitives/macos/rope.py index d888ae2d..27992551 100644 --- a/python/src/coreai_models/primitives/macos/rope.py +++ b/python/src/coreai_models/primitives/macos/rope.py @@ -15,12 +15,6 @@ logger = logging.getLogger(__name__) -def rotate_half(x: torch.Tensor) -> torch.Tensor: - """Rotate second half of the last dimension: ``cat(-x2, x1)``.""" - mid = x.shape[-1] // 2 - return torch.cat((-x[..., mid:], x[..., :mid]), dim=-1) - - class RoPE(coreai_torch.composite_ops.RoPE): """Apply rotary positional embedding to input tensors.""" From 188d4bcff1c780d14bee205503f27196a4e589c3 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 22/38] Share the RoPE frequency and 1x1 conv helpers publicly rope_inv_freq replaces _inv_freq and axial_inv_freq; _conv becomes conv1x1. --- .../models/ios/embedding_gemma2.py | 32 +++++++++---------- .../models/ios/embedding_gemma2_vision.py | 28 ++++++++-------- .../models/macos/embedding_gemma2.py | 16 +++++----- .../models/macos/embedding_gemma2_vision.py | 14 +++----- 4 files changed, 42 insertions(+), 48 deletions(-) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index 05fb2767..45f6fb3f 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -21,9 +21,7 @@ HF_TEXT_PREFIX, EmbeddingGemma2TextConfig, _resolve_model_dir, -) -from coreai_models.models.macos.embedding_gemma2 import ( - EmbeddingGemma2TextEncoder as MacEmbeddingGemma2TextEncoder, + rope_inv_freq, ) from coreai_models.primitives.ios.sdpa import SDPA @@ -82,7 +80,8 @@ def apply_rope_bc1s(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, head_ return out.reshape(b, c, one, length) -def _conv(cin: int, cout: int) -> nn.Conv2d: +def conv1x1(cin: int, cout: int) -> nn.Conv2d: + """Bias-free 1x1 ``Conv2d``, the BC1S form of a bias-free ``Linear``.""" return nn.Conv2d(cin, cout, kernel_size=1, bias=False) @@ -95,10 +94,10 @@ def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: self.n_kv = config.layer_num_key_value_heads(layer_idx) self.head_dim = hd = config.layer_head_dim(layer_idx) dim, eps = config.hidden_size, config.rms_norm_eps - self.q_proj = _conv(dim, self.n_heads * hd) - self.k_proj = _conv(dim, self.n_kv * hd) - self.v_proj = _conv(dim, self.n_kv * hd) - self.o_proj = _conv(self.n_heads * hd, dim) + self.q_proj = conv1x1(dim, self.n_heads * hd) + self.k_proj = conv1x1(dim, self.n_kv * hd) + self.v_proj = conv1x1(dim, self.n_kv * hd) + self.o_proj = conv1x1(self.n_heads * hd, dim) self.q_norm = ScaledRMSNorm(hd, eps, dim=2) self.k_norm = ScaledRMSNorm(hd, eps, dim=2) self.v_norm = ScaledRMSNorm(hd, eps, with_scale=False, dim=2) @@ -126,11 +125,11 @@ def __init__(self, config: EmbeddingGemma2TextConfig, i: int) -> None: dim, eps, ple = config.hidden_size, config.rms_norm_eps, config.hidden_size_per_layer_input self.is_sliding = config.is_sliding(i) self.self_attn = EmbeddingGemma2Attention(config, i) - self.gate_proj = _conv(dim, config.intermediate_size) - self.up_proj = _conv(dim, config.intermediate_size) - self.down_proj = _conv(config.intermediate_size, dim) - self.per_layer_input_gate = _conv(dim, ple) - self.per_layer_projection = _conv(ple, dim) + self.gate_proj = conv1x1(dim, config.intermediate_size) + self.up_proj = conv1x1(dim, config.intermediate_size) + self.down_proj = conv1x1(config.intermediate_size, dim) + self.per_layer_input_gate = conv1x1(dim, ple) + self.per_layer_projection = conv1x1(ple, dim) self.input_layernorm = ScaledRMSNorm(dim, eps) self.post_attention_layernorm = ScaledRMSNorm(dim, eps) self.pre_feedforward_layernorm = ScaledRMSNorm(dim, eps) @@ -162,11 +161,11 @@ def __init__(self, config: EmbeddingGemma2TextConfig) -> None: self.config = config dim, eps, L = config.hidden_size, config.rms_norm_eps, config.num_hidden_layers self.ple_dim = config.hidden_size_per_layer_input - self.per_layer_model_projection = _conv(dim, L * self.ple_dim) + self.per_layer_model_projection = conv1x1(dim, L * self.ple_dim) self.per_layer_projection_norm = ScaledRMSNorm(self.ple_dim, eps, dim=2) self.layers = nn.ModuleList([EmbeddingGemma2EncoderLayer(config, i) for i in range(L)]) self.norm = ScaledRMSNorm(dim, eps) - self.embedding_projection = _conv(dim, config.embedding_dim) + self.embedding_projection = conv1x1(dim, config.embedding_dim) self._rope_hd = { "sliding": config.layer_head_dim(next(i for i in range(L) if config.is_sliding(i))), "full": config.layer_head_dim(next(i for i in range(L) if not config.is_sliding(i))), @@ -181,7 +180,7 @@ def rope( ("full", self.config.rope_theta_full), ): hd = self._rope_hd[kind] - inv = MacEmbeddingGemma2TextEncoder._inv_freq(hd, theta, device) + inv = rope_inv_freq(hd, theta, device) fr = inv[:, None] * torch.arange(length, dtype=torch.float32, device=device)[None, :] e = torch.cat([fr, fr], 0).reshape(1, 1, hd, length) out[kind] = (e.cos().to(dtype), e.sin().to(dtype)) @@ -271,6 +270,7 @@ def load_embedding_gemma2_text_encoder( "MASK_NEG", "ScaledRMSNorm", "EmbeddingGemma2TextEncoder", + "conv1x1", "load_embedding_gemma2_text_encoder", "text_encoder_from_state_dict", ] diff --git a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py index 04a1c799..44658df5 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -19,11 +19,9 @@ import torch.nn.functional as F from safetensors import safe_open -from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, ScaledRMSNorm, _conv -from coreai_models.models.macos.embedding_gemma2_vision import ( - EmbeddingGemma2VisionConfig, - axial_inv_freq, -) +from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, ScaledRMSNorm, conv1x1 +from coreai_models.models.macos.embedding_gemma2 import rope_inv_freq +from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2VisionConfig def one_hot_f16(coord: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: @@ -46,10 +44,10 @@ def __init__(self, cfg: EmbeddingGemma2VisionConfig, q_chunk: int = 630) -> None self.q_chunk = q_chunk self.n_heads, self.head_dim = cfg.num_attention_heads, cfg.head_dim dim, hd, eps = cfg.hidden_size, cfg.head_dim, cfg.rms_norm_eps - self.q_proj = _conv(dim, self.n_heads * hd) - self.k_proj = _conv(dim, self.n_heads * hd) - self.v_proj = _conv(dim, self.n_heads * hd) - self.o_proj = _conv(self.n_heads * hd, dim) + self.q_proj = conv1x1(dim, self.n_heads * hd) + self.k_proj = conv1x1(dim, self.n_heads * hd) + self.v_proj = conv1x1(dim, self.n_heads * hd) + self.o_proj = conv1x1(self.n_heads * hd, dim) self.q_norm = ScaledRMSNorm(hd, eps, dim=2) self.k_norm = ScaledRMSNorm(hd, eps, dim=2) self.v_norm = ScaledRMSNorm(hd, eps, with_scale=False, dim=2) @@ -83,9 +81,9 @@ def __init__(self, cfg: EmbeddingGemma2VisionConfig) -> None: super().__init__() dim, eps = cfg.hidden_size, cfg.rms_norm_eps self.self_attn = VisionAttention(cfg) - self.gate_proj = _conv(dim, cfg.intermediate_size) - self.up_proj = _conv(dim, cfg.intermediate_size) - self.down_proj = _conv(cfg.intermediate_size, dim) + self.gate_proj = conv1x1(dim, cfg.intermediate_size) + self.up_proj = conv1x1(dim, cfg.intermediate_size) + self.down_proj = conv1x1(cfg.intermediate_size, dim) self.input_layernorm = ScaledRMSNorm(dim, eps) self.post_attention_layernorm = ScaledRMSNorm(dim, eps) self.pre_feedforward_layernorm = ScaledRMSNorm(dim, eps) @@ -109,10 +107,10 @@ def __init__(self, config: EmbeddingGemma2VisionConfig | None = None): k = cfg.pooling_kernel_size self.grid = g = cfg.max_soft_tokens * k # largest reachable patch coordinate + 1 dim = cfg.hidden_size - self.input_proj = _conv(cfg.patch_dim, dim) + self.input_proj = conv1x1(cfg.patch_dim, dim) self.layers = nn.ModuleList([VisionEncoderLayer(cfg) for _ in range(cfg.num_hidden_layers)]) self.embedding_pre_projection_norm = ScaledRMSNorm(dim, cfg.rms_norm_eps, with_scale=False) - self.embedding_projection = _conv(dim, cfg.text_hidden_size) + self.embedding_projection = conv1x1(dim, cfg.text_hidden_size) # Per-axis lookup tables, filled by ``bake_tables`` (fp32 math, stored in model dtype): # [G, hidden] position embedding | [G, hd/4] cos | [G, hd/4] sin. self.register_buffer("pos_x", torch.zeros(g, dim), persistent=False) @@ -137,7 +135,7 @@ def bake_tables(self, position_embedding_table: torch.Tensor) -> None: raise ValueError(f"position table has {table.shape[1]} rows; budget needs {g}") self.pos_x.copy_(table[0, :g]) self.pos_y.copy_(table[1, :g]) - inv = axial_inv_freq(cfg) + inv = rope_inv_freq(cfg.head_dim // 2, cfg.rope_theta) ang = torch.arange(g, dtype=torch.float32)[:, None] * inv[None, :] self.cos_tab.copy_(ang.cos()) self.sin_tab.copy_(ang.sin()) diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index 35b3deaf..93211b88 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -37,6 +37,11 @@ HF_TEXT_PREFIX = "language_model." +def rope_inv_freq(dim: int, base: float, device: torch.device | None = None) -> torch.Tensor: + """RoPE inverse frequencies ``base ** (-arange(0, dim, 2) / dim)`` in float32.""" + return 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + + @dataclass class EmbeddingGemma2TextConfig: """Text-tower config parsed from the HF ``text_config`` block. @@ -294,17 +299,11 @@ def __init__(self, config: EmbeddingGemma2TextConfig) -> None: full_hd = next(iter(full_dims), config.head_dim) self.register_buffer( "sliding_inv_freq", - self._inv_freq(sliding_hd, config.rope_theta_sliding), + rope_inv_freq(sliding_hd, config.rope_theta_sliding), persistent=False, ) self.register_buffer( - "full_inv_freq", self._inv_freq(full_hd, config.rope_theta_full), persistent=False - ) - - @staticmethod - def _inv_freq(head_dim: int, base: float, device: torch.device | None = None) -> torch.Tensor: - return 1.0 / ( - base ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim) + "full_inv_freq", rope_inv_freq(full_hd, config.rope_theta_full), persistent=False ) @staticmethod @@ -446,4 +445,5 @@ def load_embedding_gemma2_text_encoder( "extract_text_state_dict", "load_embedding_gemma2_text_encoder", "load_text_state_dict", + "rope_inv_freq", ] diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index 55d4704b..3466b55c 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -26,6 +26,7 @@ from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaDenseMLP from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaRMSNorm as RMSNorm +from coreai_models.models.macos.embedding_gemma2 import rope_inv_freq from coreai_models.primitives.ios.rope import apply_rope from coreai_models.primitives.macos.sdpa import SDPA @@ -96,14 +97,6 @@ def from_hf_config( ) -def axial_inv_freq(cfg: EmbeddingGemma2VisionConfig) -> torch.Tensor: - """Per-axis RoPE inverse frequencies (``head_dim // 2`` rotary dims per spatial axis).""" - spatial_dim = cfg.head_dim // 2 - return 1.0 / ( - cfg.rope_theta ** (torch.arange(0, spatial_dim, 2, dtype=torch.float32) / spatial_dim) - ) - - class VisionAttention(nn.Module): """Bidirectional MHA with q/k RMSNorm, scale-free v RMSNorm and axial RoPE; scale 1.0.""" @@ -235,7 +228,10 @@ def __init__(self, config: EmbeddingGemma2VisionConfig | None = None) -> None: self.patch_embedder = VisionPatchEmbedder(cfg) self.encoder = VisionEncoder(cfg) self.embed_vision = MultimodalEmbedder(cfg) - self.register_buffer("rope_inv_freq", axial_inv_freq(cfg), persistent=False) + # Axial RoPE: head_dim // 2 rotary dims per spatial axis. + self.register_buffer( + "rope_inv_freq", rope_inv_freq(cfg.head_dim // 2, cfg.rope_theta), persistent=False + ) self.register_buffer( "pool_slots", torch.arange(cfg.max_soft_tokens, dtype=torch.int64), persistent=False ) From 531775ff088b082a118f5d5b4a4cc7b92be80dcb Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 23/38] Read the EmbeddingGemma 2 checkpoint in one place and load strictly read_checkpoint_tensors serves both macOS loaders; the iOS loaders build on them, and load_conv_weights replaces the hand-written key checks. --- .../models/ios/embedding_gemma2.py | 42 +++++--------- .../models/ios/embedding_gemma2_vision.py | 44 +++++---------- .../models/macos/embedding_gemma2.py | 56 ++++++++----------- .../models/macos/embedding_gemma2_vision.py | 29 +++------- .../test_embedding_gemma2.py | 19 +------ 5 files changed, 64 insertions(+), 126 deletions(-) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index 45f6fb3f..4f2fa763 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -18,11 +18,12 @@ import torch.nn.functional as F from coreai_models.models.macos.embedding_gemma2 import ( - HF_TEXT_PREFIX, EmbeddingGemma2TextConfig, - _resolve_model_dir, rope_inv_freq, ) +from coreai_models.models.macos.embedding_gemma2 import ( + load_embedding_gemma2_text_encoder as load_macos_text_encoder, +) from coreai_models.primitives.ios.sdpa import SDPA # Smallest positive normal float16; floors the row max-abs and the squared norm so all-zero rows @@ -213,6 +214,14 @@ def forward( return ys * torch.rsqrt((ys * ys).sum(-1, keepdim=True).clamp(min=_FP16_TINY)) +def load_conv_weights(model: nn.Module, state_dict: dict[str, torch.Tensor]) -> None: + """Strict ``load_state_dict`` with each tensor reshaped to the module's shape (e.g. conv).""" + own = model.state_dict() + model.load_state_dict( + {k: v.reshape(own[k].shape) if k in own else v for k, v in state_dict.items()} + ) + + def text_encoder_from_state_dict( config: EmbeddingGemma2TextConfig, state_dict: dict[str, torch.Tensor], @@ -231,17 +240,7 @@ def text_encoder_from_state_dict( k2 = k2.replace("mlp.", "").replace("ple_block.", "") mapped[k2] = v model = EmbeddingGemma2TextEncoder(config) - own = model.state_dict() - for k, v in mapped.items(): - if k not in own: - raise RuntimeError(f"unexpected key {k}") - if own[k].dim() == 4 and v.dim() == 2: - v = v[:, :, None, None] - mapped[k] = v.reshape(own[k].shape) - missing = set(own) - set(mapped) - if missing: - raise RuntimeError(f"missing keys {sorted(missing)[:10]}") - model.load_state_dict(mapped, strict=True) + load_conv_weights(model, mapped) return model.to(dtype).eval(), table @@ -250,20 +249,8 @@ def load_embedding_gemma2_text_encoder( dtype: torch.dtype = torch.float16, ) -> tuple[EmbeddingGemma2TextEncoder, torch.Tensor]: """Load from a local snapshot dir or HF id; returns ``(encoder, embed_table_fp32)``.""" - from safetensors import safe_open - - model_dir = _resolve_model_dir(model_dir_or_id) - config = EmbeddingGemma2TextConfig.from_pretrained(model_dir) - sd: dict[str, torch.Tensor] = {} - files = sorted(model_dir.glob("*.safetensors")) - if not files: - raise FileNotFoundError(f"No .safetensors in {model_dir}") - for file in files: - with safe_open(str(file), framework="pt") as handle: - for key in handle.keys(): # noqa: SIM118 - if key.startswith(HF_TEXT_PREFIX): - sd[key[len(HF_TEXT_PREFIX) :]] = handle.get_tensor(key) - return text_encoder_from_state_dict(config, sd, dtype) + mac = load_macos_text_encoder(model_dir_or_id, torch.float32) + return text_encoder_from_state_dict(mac.config, mac.state_dict(), dtype) __all__ = [ @@ -271,6 +258,7 @@ def load_embedding_gemma2_text_encoder( "ScaledRMSNorm", "EmbeddingGemma2TextEncoder", "conv1x1", + "load_conv_weights", "load_embedding_gemma2_text_encoder", "text_encoder_from_state_dict", ] diff --git a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py index 44658df5..1d3d3023 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -11,16 +11,20 @@ from __future__ import annotations -import json -import os - import torch import torch.nn as nn import torch.nn.functional as F -from safetensors import safe_open -from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, ScaledRMSNorm, conv1x1 +from coreai_models.models.ios.embedding_gemma2 import ( + MASK_NEG, + ScaledRMSNorm, + conv1x1, + load_conv_weights, +) from coreai_models.models.macos.embedding_gemma2 import rope_inv_freq +from coreai_models.models.macos.embedding_gemma2_vision import ( + EmbeddingGemma2ImageEncoder as MacEmbeddingGemma2ImageEncoder, +) from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2VisionConfig @@ -185,43 +189,25 @@ def load_image_encoder( max_soft_tokens: int | None = None, ) -> EmbeddingGemma2ImageEncoder: """Load from a local EG2 snapshot directory (``config.json`` + ``*.safetensors``).""" - with open(os.path.join(model_dir, "config.json")) as f: - cfg = EmbeddingGemma2VisionConfig.from_hf_config(json.load(f), max_soft_tokens) - model = EmbeddingGemma2ImageEncoder(cfg) - sd: dict[str, torch.Tensor] = {} - files = sorted(x for x in os.listdir(model_dir) if x.endswith(".safetensors")) - if not files: - raise FileNotFoundError(f"No .safetensors in {model_dir}") - for fn in files: - with safe_open(os.path.join(model_dir, fn), framework="pt") as h: - for key in h.keys(): # noqa: SIM118 - if key.startswith(("vision_tower.", "embed_vision.")): - sd[key] = h.get_tensor(key).float() - load_hf_state_dict(model, sd) + mac = MacEmbeddingGemma2ImageEncoder.from_pretrained(model_dir, torch.float32, max_soft_tokens) + model = load_hf_state_dict(EmbeddingGemma2ImageEncoder(mac.config), mac.state_dict()) return model.to(dtype).eval() def load_hf_state_dict( model: EmbeddingGemma2ImageEncoder, sd: dict[str, torch.Tensor] ) -> EmbeddingGemma2ImageEncoder: - """Load checkpoint-named ``vision_tower.*`` / ``embed_vision.*`` tensors (fail closed).""" + """Load checkpoint-named (``vision_tower.*`` / ``embed_vision.*``) or macOS-encoder-named + tensors (fail closed).""" mapped: dict[str, torch.Tensor] = {} for key, v in sd.items(): k2 = key.replace(".linear.weight", ".weight") k2 = k2.removeprefix("vision_tower.").removeprefix("embed_vision.") k2 = k2.replace("patch_embedder.input_proj", "input_proj") k2 = k2.replace("encoder.layers.", "layers.").replace("mlp.", "") - mapped[k2] = v.float() + mapped[k2] = v pos_table = mapped.pop("patch_embedder.position_embedding_table") - own = model.state_dict() - for key, v in list(mapped.items()): - if key not in own: - raise RuntimeError(f"unexpected key {key}") - mapped[key] = v.reshape(own[key].shape) - missing = set(own) - set(mapped) - if missing: - raise RuntimeError(f"missing keys {sorted(missing)[:10]}") - model.load_state_dict(mapped, strict=True) + load_conv_weights(model, mapped) model.bake_tables(pos_table) return model diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index 93211b88..ccb89dea 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -382,23 +382,26 @@ def forward( return sentence_embedding -def extract_text_state_dict( - state_dict: dict[str, torch.Tensor], prefix: str = HF_TEXT_PREFIX +def read_checkpoint_tensors( + model_dir: str | Path, prefixes: tuple[str, ...], dtype: torch.dtype | None = None ) -> dict[str, torch.Tensor]: - """Keep only ``prefix*`` keys (vision/audio dropped) and strip the prefix.""" - return {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} - - -def load_text_state_dict( - model: EmbeddingGemma2TextEncoder, state_dict: dict[str, torch.Tensor] -) -> None: - """Load an (already prefix-stripped) text state dict; fail closed on any missing/extra key.""" - info = model.load_state_dict(state_dict, strict=False) - if info.missing_keys or info.unexpected_keys: - raise RuntimeError( - f"EG2 text load mismatch: missing={info.missing_keys[:10]} " - f"unexpected={info.unexpected_keys[:10]}" - ) + """Tensors from ``model_dir/*.safetensors`` whose names start with one of ``prefixes``. + + Names are kept as in the checkpoint; each tensor is cast to ``dtype`` when given. + """ + from safetensors import safe_open + + files = sorted(Path(model_dir).glob("*.safetensors")) + if not files: + raise FileNotFoundError(f"No .safetensors in {model_dir}") + tensors: dict[str, torch.Tensor] = {} + for file in files: + with safe_open(str(file), framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + if key.startswith(prefixes): + tensor = handle.get_tensor(key) + tensors[key] = tensor if dtype is None else tensor.to(dtype) + return tensors def _resolve_model_dir(model_dir_or_id: str | Path) -> Path: @@ -419,22 +422,10 @@ def load_embedding_gemma2_text_encoder( dtype: torch.dtype = torch.float32, ) -> EmbeddingGemma2TextEncoder: """Build the text encoder from a local snapshot dir or HF id, loading ``language_model.*``.""" - from safetensors import safe_open - model_dir = _resolve_model_dir(model_dir_or_id) - config = EmbeddingGemma2TextConfig.from_pretrained(model_dir) - model = EmbeddingGemma2TextEncoder(config) - - state_dict: dict[str, torch.Tensor] = {} - files = sorted(model_dir.glob("*.safetensors")) - if not files: - raise FileNotFoundError(f"No .safetensors in {model_dir}") - for file in files: - with safe_open(str(file), framework="pt") as handle: - for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping - if key.startswith(HF_TEXT_PREFIX): - state_dict[key[len(HF_TEXT_PREFIX) :]] = handle.get_tensor(key).to(dtype) - load_text_state_dict(model, state_dict) + model = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_pretrained(model_dir)) + tensors = read_checkpoint_tensors(model_dir, (HF_TEXT_PREFIX,), dtype) + model.load_state_dict({k.removeprefix(HF_TEXT_PREFIX): v for k, v in tensors.items()}) return model.to(dtype).eval() @@ -442,8 +433,7 @@ def load_embedding_gemma2_text_encoder( "EmbeddingGemma2TextConfig", "EmbeddingGemma2TextEncoder", "build_attention_masks", - "extract_text_state_dict", "load_embedding_gemma2_text_encoder", - "load_text_state_dict", + "read_checkpoint_tensors", "rope_inv_freq", ] diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index 3466b55c..9e383864 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -15,18 +15,17 @@ import json import math -import os from collections.abc import Callable, Sequence from dataclasses import dataclass +from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F -from safetensors import safe_open from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaDenseMLP from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaRMSNorm as RMSNorm -from coreai_models.models.macos.embedding_gemma2 import rope_inv_freq +from coreai_models.models.macos.embedding_gemma2 import read_checkpoint_tensors, rope_inv_freq from coreai_models.primitives.ios.rope import apply_rope from coreai_models.primitives.macos.sdpa import SDPA @@ -310,23 +309,13 @@ def from_pretrained( dtype: torch.dtype = torch.float32, max_soft_tokens: int | None = None, ) -> EmbeddingGemma2ImageEncoder: - """Load from a local EG2 snapshot directory (``config.json`` + ``model.safetensors``).""" - with open(os.path.join(model_dir, "config.json")) as f: - config = EmbeddingGemma2VisionConfig.from_hf_config(json.load(f), max_soft_tokens) - model = cls(config) - state: dict[str, torch.Tensor] = {} - files = sorted(fn for fn in os.listdir(model_dir) if fn.endswith(".safetensors")) - if not files: - raise FileNotFoundError(f"No .safetensors in {model_dir}") - for fn in files: - with safe_open(os.path.join(model_dir, fn), framework="pt", device="cpu") as f: - for key in f.keys(): # noqa: SIM118 - local = cls.map_hf_key(key) - if local is not None: - state[local] = f.get_tensor(key).to(dtype) - missing, unexpected = model.load_state_dict(state, strict=False) - if missing or unexpected: - raise RuntimeError(f"weight mismatch: missing={missing} unexpected={unexpected}") + """Load from a local EG2 snapshot directory (``config.json`` + ``*.safetensors``).""" + config = json.loads((Path(model_dir) / "config.json").read_text()) + model = cls(EmbeddingGemma2VisionConfig.from_hf_config(config, max_soft_tokens)) + tensors = read_checkpoint_tensors(model_dir, ("vision_tower.", "embed_vision."), dtype) + model.load_state_dict( + {local: v for k, v in tensors.items() if (local := cls.map_hf_key(k)) is not None} + ) return model.to(dtype).eval() diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py index ed7bcdda..f5e4fea0 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -33,9 +33,7 @@ EmbeddingGemma2TextConfig, EmbeddingGemma2TextEncoder, build_attention_masks, - extract_text_state_dict, load_embedding_gemma2_text_encoder, - load_text_state_dict, ) from tests.test_model_units.test_models._embedding_gemma2_common import ( TINY_TEXT_CONFIG, @@ -164,19 +162,6 @@ def test_inputs_embeds_path_matches_input_ids(): model(input_ids=ids, inputs_embeds=model.embed_tokens(ids)) -def test_load_text_state_dict_fails_closed(): - model = _tiny_model() - prefixed = {f"language_model.{k}": v for k, v in model.state_dict().items()} - prefixed["vision_tower.encoder.foo"] = torch.zeros(1) - stripped = extract_text_state_dict(prefixed) - assert "encoder.foo" not in stripped and len(stripped) == len(model.state_dict()) - fresh = EmbeddingGemma2TextEncoder(model.config) - load_text_state_dict(fresh, stripped) - stripped.pop("layers.0.layer_scalar") - with pytest.raises(RuntimeError, match="missing"): - load_text_state_dict(EmbeddingGemma2TextEncoder(model.config), stripped) - - def test_torch_export_static_bucket(): """torch.export at a static (B=2, L=128) bucket with int32 ids/mask; outputs match eager.""" model = _tiny_model() @@ -228,7 +213,7 @@ def test_tiny_parity_vs_hf(attn_impl): hf = _tiny_hf(attn_impl) ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG)).eval() - load_text_state_dict(ours, dict(hf.state_dict())) + ours.load_state_dict(hf.state_dict()) # Long enough (L=24 > 2*window+1) that the sliding window actually bites, with padding. ids, mask = _ragged_batch([24, 13, 5], seq_len=24, vocab=128, seed=7) @@ -248,7 +233,7 @@ def test_tiny_sliding_window_matters_vs_hf(): hf = _tiny_hf("eager") wide = dict(TINY_TEXT_CONFIG, sliding_window=1000) ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(wide)).eval() - load_text_state_dict(ours, dict(hf.state_dict())) + ours.load_state_dict(hf.state_dict()) ids, mask = _ragged_batch([24], seq_len=24, vocab=128, seed=7) with torch.no_grad(): hf_emb = _st_mean_pool_normalize( From a57aacdf2fe677553dfc1309ef695dae00124d2b Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 24/38] Call add_function directly in the embedding exporters Drops TraceSpec, add_traces, export_text_encoder_program and the loader dispatchers; the token-table trace is shared, and export_stateless uses new_converter. --- python/src/coreai_models/diffusion/gpu.py | 8 +- .../src/coreai_models/embedding/export_ios.py | 38 ++--- .../coreai_models/embedding/export_macos.py | 143 +++++------------- .../coreai_models/embedding/export_utils.py | 76 +++------- .../test_export/test_embedding_export.py | 2 +- 5 files changed, 82 insertions(+), 185 deletions(-) diff --git a/python/src/coreai_models/diffusion/gpu.py b/python/src/coreai_models/diffusion/gpu.py index 61db98d9..add41125 100644 --- a/python/src/coreai_models/diffusion/gpu.py +++ b/python/src/coreai_models/diffusion/gpu.py @@ -82,13 +82,7 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram: decomposed: torch.export.ExportedProgram = exported.run_decompositions(coreai_decomp_table) return decomposed - converter = coreai_torch.TorchConverter( - mode=( - coreai_torch.TorchConverter.Mode.DEBUG - if include_debug_info - else coreai_torch.TorchConverter.Mode.RELEASE - ) - ) + converter = new_converter(include_debug_info) converter.add_pytorch_module( wrapper, export_fn=export_fn, diff --git a/python/src/coreai_models/embedding/export_ios.py b/python/src/coreai_models/embedding/export_ios.py index 0a090587..62c4c9f3 100644 --- a/python/src/coreai_models/embedding/export_ios.py +++ b/python/src/coreai_models/embedding/export_ios.py @@ -10,6 +10,7 @@ from __future__ import annotations +from collections.abc import Sequence from pathlib import Path from typing import TYPE_CHECKING @@ -31,8 +32,6 @@ encoder_functions, ) from coreai_models.embedding.export_utils import ( - TraceSpec, - add_traces, export_token_table_program, image_reference_inputs, save_program, @@ -82,20 +81,27 @@ def ane_text_reference_inputs( def convert_ane_functions( - traces: list[tuple[nn.Module, TraceSpec]], + functions: Sequence[ + tuple[str, nn.Module, tuple[torch.Tensor, ...], tuple[str, ...], tuple[str, ...]] + ], *, include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, ) -> AIProgram: - """Convert every ``(module, trace)`` into one multi-function Neural Engine program. + """Convert each ``(name, module, args, input_names, output_names)`` into one Neural Engine + program. Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. """ - from coreai_models.diffusion.gpu import new_converter + from coreai_models.diffusion.gpu import add_function, new_converter from coreai_models.export.mlir_ops import register_custom_torch_lowering converter = new_converter(include_debug_info) register_custom_torch_lowering(converter) - add_traces(converter, traces, torch.export.default_decompositions()) + decompositions = torch.export.default_decompositions() + for name, module, args, input_names, output_names in functions: + add_function( + converter, name, module, args, input_names, output_names, decompositions=decompositions + ) return converter.to_coreai() @@ -109,13 +115,11 @@ def export_ane_encoder_program( hidden = int(encoder.config.hidden_size) traces = [ ( + spec.name, EncoderBandWrapper(encoder, spec.length, sliding_window), - TraceSpec( - spec.name, - ane_text_reference_inputs(spec.length, hidden), - ANE_ENCODER_INPUT_NAMES, - (OUTPUT_NAME,), - ), + ane_text_reference_inputs(spec.length, hidden), + ANE_ENCODER_INPUT_NAMES, + (OUTPUT_NAME,), ) for spec in functions ] @@ -128,13 +132,11 @@ def export_ane_image_encoder_program( include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, ) -> AIProgram: """Neural Engine ``image_encoder``: function ``main`` at the static budget.""" - trace = TraceSpec( - "main", - image_reference_inputs(image, torch.float16, torch.float16), - IMAGE_INPUT_NAMES, - (IMAGE_OUTPUT_NAME,), + args = image_reference_inputs(image, torch.float16, torch.float16) + return convert_ane_functions( + [("main", image_encoder, args, IMAGE_INPUT_NAMES, (IMAGE_OUTPUT_NAME,))], + include_debug_info=include_debug_info, ) - return convert_ane_functions([(image_encoder, trace)], include_debug_info=include_debug_info) def export_ane_programs( diff --git a/python/src/coreai_models/embedding/export_macos.py b/python/src/coreai_models/embedding/export_macos.py index ecb82a8f..73bfdc10 100644 --- a/python/src/coreai_models/embedding/export_macos.py +++ b/python/src/coreai_models/embedding/export_macos.py @@ -37,10 +37,9 @@ ) from coreai_models.embedding.export_utils import ( TokenTableExportWrapper, - TraceSpec, - add_traces, image_reference_inputs, save_program, + token_table_trace, ) if TYPE_CHECKING: @@ -116,24 +115,6 @@ def embeds_reference_inputs(spec: FunctionSpec, hidden: int) -> tuple[torch.Tens return embeds, torch.ones((spec.batch, spec.length), dtype=torch.int32) -def export_text_encoder_program( - wrapper: nn.Module, - functions: list[FunctionSpec], - vocab_size: int, - include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, -) -> AIProgram: - """GPU ``text_encoder``: one static (batch, length) function per spec, weights shared.""" - from coreai_models.diffusion.gpu import export_multifunction - - # Static (batch, length) traces; the host pads to the next listed batch. - return export_multifunction( - [(spec.name, wrapper, reference_inputs(spec, vocab_size)) for spec in functions], - TEXT_INPUT_NAMES, - (OUTPUT_NAME,), - include_debug_info, - ) - - def export_shared_text_program( encoder: nn.Module, text_fns: list[FunctionSpec], @@ -141,58 +122,37 @@ def export_shared_text_program( max_length: int, include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, ) -> AIProgram: - """GPU ``text_encoder`` that also holds the ``token_table`` and ``encoder`` functions. + """GPU ``text_encoder``; with ``encoder_fns`` it also holds the ``token_table`` and ``encoder``. - Text functions take ids; ``main`` maps ids to ``inputs_embeds`` (dynamic length up to - ``max_length``); ``embeds_*`` functions take ``inputs_embeds``. The functions share one copy - of the table and the transformer. + Text functions take ids at a static (batch, length); the host pads to the next listed batch. + ``main`` maps ids to ``inputs_embeds`` (dynamic length up to ``max_length``); ``embeds_*`` + functions take ``inputs_embeds``. The functions share one copy of the table and the transformer. """ - from coreai_models.diffusion.gpu import new_converter + from coreai_models.diffusion.gpu import add_function, new_converter - vocab_size = int(encoder.config.vocab_size) - hidden = int(encoder.config.hidden_size) - length = torch.export.Dim("length", min=2, max=max_length) + converter = new_converter(include_debug_info) text = TextEmbeddingExportWrapper(encoder) - traces = [ - ( - text, - TraceSpec( - spec.name, reference_inputs(spec, vocab_size), TEXT_INPUT_NAMES, (OUTPUT_NAME,) - ), - ) - for spec in text_fns - ] - traces.append( - ( + for spec in text_fns: + args = reference_inputs(spec, int(encoder.config.vocab_size)) + add_function(converter, spec.name, text, args, TEXT_INPUT_NAMES, (OUTPUT_NAME,)) + if encoder_fns: + ids, dynamic_shapes = token_table_trace(max_length) + add_function( + converter, + "main", TokenTableExportWrapper(encoder.embed_tokens), - TraceSpec( - "main", - (torch.zeros((1, 64), dtype=torch.int32),), - TOKEN_TABLE_INPUT_NAMES, - (TOKEN_TABLE_OUTPUT_NAME,), - ({1: length},), - ), + ids, + TOKEN_TABLE_INPUT_NAMES, + (TOKEN_TABLE_OUTPUT_NAME,), + dynamic_shapes=dynamic_shapes, ) - ) - converter = new_converter(include_debug_info) - add_traces(converter, traces) - with without_token_table(encoder) as inner: - embeds = InputsEmbedsEncoderWrapper(inner) - add_traces( - converter, - [ - ( - embeds, - TraceSpec( - spec.name, - embeds_reference_inputs(spec, hidden), - ENCODER_INPUT_NAMES, - (OUTPUT_NAME,), - ), + with without_token_table(encoder) as inner: + embeds = InputsEmbedsEncoderWrapper(inner) + for spec in encoder_fns: + args = embeds_reference_inputs(spec, int(encoder.config.hidden_size)) + add_function( + converter, spec.name, embeds, args, ENCODER_INPUT_NAMES, (OUTPUT_NAME,) ) - for spec in encoder_fns - ], - ) return converter.to_coreai() @@ -213,26 +173,6 @@ def export_image_encoder_program( ) -def load_text_encoder(family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype) -> nn.Module: - if family.key == "embeddinggemma-2": - from coreai_models.models.macos.embedding_gemma2 import load_embedding_gemma2_text_encoder - - return load_embedding_gemma2_text_encoder(snapshot, dtype=dtype) - raise NotImplementedError(f"No text-encoder loader for family {family.key!r}") - - -def load_image_encoder( - family: EmbeddingFamily, snapshot: Path, dtype: torch.dtype, max_soft_tokens: int -) -> nn.Module: - if family.key == "embeddinggemma-2": - from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder - - return EmbeddingGemma2ImageEncoder.from_pretrained( - str(snapshot), dtype=dtype, max_soft_tokens=max_soft_tokens - ) - raise NotImplementedError(f"No image-encoder loader for family {family.key!r}") - - def export_gpu_programs( family: EmbeddingFamily, snapshot: Path, @@ -242,29 +182,28 @@ def export_gpu_programs( compute_precision: str, include_debug_info: bool, ) -> None: + if family.key != "embeddinggemma-2": + raise NotImplementedError(f"No GPU export for family {family.key!r}") + from coreai_models.models.macos.embedding_gemma2 import load_embedding_gemma2_text_encoder + from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder + dtype = _DTYPES[compute_precision] - encoder = load_text_encoder(family, snapshot, dtype) - if with_image: - program = export_shared_text_program( - encoder, - text_functions(family, "macOS"), - encoder_functions(family, "macOS"), - family.model_max_tokens, - include_debug_info, - ) - else: - program = export_text_encoder_program( - TextEmbeddingExportWrapper(encoder), - text_functions(family, "macOS"), - int(encoder.config.vocab_size), - include_debug_info, - ) + encoder = load_embedding_gemma2_text_encoder(snapshot, dtype=dtype) + program = export_shared_text_program( + encoder, + text_functions(family, "macOS"), + encoder_functions(family, "macOS") if with_image else [], + family.model_max_tokens, + include_debug_info, + ) save_program(program, out, family, TEXT_COMPONENT) del program, encoder if not with_image: return assert family.image is not None - image_encoder = load_image_encoder(family, snapshot, dtype, family.image.budget) + image_encoder = EmbeddingGemma2ImageEncoder.from_pretrained( + str(snapshot), dtype=dtype, max_soft_tokens=family.image.budget + ) save_program( export_image_encoder_program(image_encoder, family.image, include_debug_info), out, diff --git a/python/src/coreai_models/embedding/export_utils.py b/python/src/coreai_models/embedding/export_utils.py index bc5d558a..2064f320 100644 --- a/python/src/coreai_models/embedding/export_utils.py +++ b/python/src/coreai_models/embedding/export_utils.py @@ -13,8 +13,6 @@ import json import logging import shutil -import time -from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -146,59 +144,24 @@ def image_reference_inputs( The GPU program takes float32 pixels and int32 positions; the Neural Engine takes float16 for both. """ + from coreai_models.models.macos.embedding_gemma2_vision import pad_to_budget, patch_position_ids + k = image.pooling_kernel_size - n = image.budget - # Square-ish grid of n soft tokens (blocks) -> (bw*k) x (bh*k) patches, row-major. - bw = max(1, int(n**0.5)) - bh = max(1, n // bw) - pw, ph = bw * k, bh * k - ys, xs = torch.meshgrid(torch.arange(ph), torch.arange(pw), indexing="ij") - real = torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1) - positions = torch.full((1, image.max_patches, 2), -1, dtype=torch.int32) - positions[0, : real.shape[0]] = real.to(torch.int32) + # Square-ish grid of budget soft tokens (blocks) -> (bh*k) x (bw*k) patches, row-major. + bw = max(1, int(image.budget**0.5)) + bh = max(1, image.budget // bw) generator = torch.Generator().manual_seed(image.budget) - pixels = torch.rand((1, image.max_patches, image.patch_dim), generator=generator) - pixels[0, real.shape[0] :] = 0.0 - return pixels.to(dtype), positions.to(position_dtype) - - -@dataclass -class TraceSpec: - """One function of a multi-function program: example args, IO names, dynamic dims.""" - - name: str - args: tuple[torch.Tensor, ...] - input_names: tuple[str, ...] - output_names: tuple[str, ...] - dynamic_shapes: tuple[dict[int, Any] | None, ...] | None = None - - -def add_traces( - converter: Any, - traces: list[tuple[nn.Module, TraceSpec]], - decompositions: dict[Any, Any] | None = None, -) -> None: - """Add each ``(module, trace)`` to ``converter`` as a function (Core AI's table by default).""" - from coreai_models.diffusion.gpu import add_function - - for module, trace in traces: - start = time.time() - add_function( - converter, - trace.name, - module, - trace.args, - list(trace.input_names), - list(trace.output_names), - dynamic_shapes=trace.dynamic_shapes, - decompositions=decompositions, - ) - logger.info( - "traced %s %s in %.1fs", - trace.name, - [tuple(a.shape) for a in trace.args], - time.time() - start, - ) + pixels = torch.rand((image.max_patches, image.patch_dim), generator=generator) + pixels, positions = pad_to_budget( + pixels[: bh * k * bw * k], patch_position_ids(bh * k, bw * k), image.max_patches + ) + return pixels[None].to(dtype), positions[None].to(position_dtype) + + +def token_table_trace(max_length: int) -> tuple[tuple[torch.Tensor], tuple[dict[int, Any]]]: + """Token-table trace args and dynamic shapes: batch-1 ids, length 2..``max_length``.""" + length = torch.export.Dim("length", min=2, max=max_length) + return (torch.zeros((1, 64), dtype=torch.int32),), ({1: length},) def export_token_table_program( @@ -211,14 +174,13 @@ def export_token_table_program( from coreai_models.diffusion.gpu import export_stateless # Placeholder ids look up like any id; the host overwrites those rows with soft tokens. - length = torch.export.Dim("length", min=2, max=max_length) - ids = torch.zeros((1, 64), dtype=torch.int32) + args, dynamic_shapes = token_table_trace(max_length) return export_stateless( TokenTableExportWrapper(embed_tokens, dtype), - (ids,), + args, TOKEN_TABLE_INPUT_NAMES, (TOKEN_TABLE_OUTPUT_NAME,), - dynamic_shapes=({1: length},), + dynamic_shapes=dynamic_shapes, include_debug_info=include_debug_info, ) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index f05dc6d6..ac1ff98d 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -516,7 +516,7 @@ def test_tiny_multifunction_conversion(): out = wrapper(ids, mask) assert out.dtype == torch.float32 and out.shape == (4, 24) torch.testing.assert_close(out.norm(dim=-1), torch.ones(4)) - assert export_macos.export_text_encoder_program(wrapper, plan, 64) is not None + assert export_macos.export_shared_text_program(wrapper.encoder, plan, [], 64) is not None def _tiny_text_encoder(): From 55ae6ad56f78e475d3278d4a033944629c5a5dda Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 25/38] Look up the RoPE head dims through the text config rope_head_dim serves the macOS and iOS text encoders; also drops an empty TYPE_CHECKING block. --- python/src/coreai_models/embedding/export.py | 5 +--- .../models/ios/embedding_gemma2.py | 5 +--- .../models/macos/embedding_gemma2.py | 27 ++++++++++--------- 3 files changed, 17 insertions(+), 20 deletions(-) diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index 894f6165..710e86cc 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -21,7 +21,7 @@ import shutil import sys from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO from coreai_models.embedding.contract import ( @@ -54,9 +54,6 @@ verify_pooling_config, ) -if TYPE_CHECKING: - pass - logger = logging.getLogger(__name__) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index 4f2fa763..6e457c34 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -167,10 +167,7 @@ def __init__(self, config: EmbeddingGemma2TextConfig) -> None: self.layers = nn.ModuleList([EmbeddingGemma2EncoderLayer(config, i) for i in range(L)]) self.norm = ScaledRMSNorm(dim, eps) self.embedding_projection = conv1x1(dim, config.embedding_dim) - self._rope_hd = { - "sliding": config.layer_head_dim(next(i for i in range(L) if config.is_sliding(i))), - "full": config.layer_head_dim(next(i for i in range(L) if not config.is_sliding(i))), - } + self._rope_hd = {"sliding": config.rope_head_dim(True), "full": config.rope_head_dim(False)} def rope( self, length: int, dtype: torch.dtype, device diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index ccb89dea..c351cb8d 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -90,6 +90,17 @@ def layer_num_key_value_heads(self, layer_idx: int) -> int: def is_sliding(self, layer_idx: int) -> bool: return self.layer_types[layer_idx] == "sliding_attention" + def rope_head_dim(self, sliding: bool) -> int: + """Head dim of the sliding (or full) layers; ``head_dim`` when there are none.""" + dims = { + self.layer_head_dim(i) + for i in range(self.num_hidden_layers) + if self.is_sliding(i) == sliding + } + if len(dims) > 1: + raise ValueError("RoPE tables assume one head_dim per layer type") + return next(iter(dims), self.head_dim) + @classmethod def from_dict(cls, text_config: dict[str, Any]) -> EmbeddingGemma2TextConfig: """Build from an HF ``text_config`` dict (``config.json["text_config"]``). @@ -287,23 +298,15 @@ def __init__(self, config: EmbeddingGemma2TextConfig) -> None: self.norm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.embedding_projection = nn.Linear(config.hidden_size, config.embedding_dim, bias=False) - sliding_dims = { - config.layer_head_dim(i) for i in range(len(self.layers)) if config.is_sliding(i) - } - full_dims = { - config.layer_head_dim(i) for i in range(len(self.layers)) if not config.is_sliding(i) - } - if len(sliding_dims) > 1 or len(full_dims) > 1: - raise ValueError("RoPE tables assume one head_dim per layer type") - sliding_hd = next(iter(sliding_dims), config.head_dim) - full_hd = next(iter(full_dims), config.head_dim) self.register_buffer( "sliding_inv_freq", - rope_inv_freq(sliding_hd, config.rope_theta_sliding), + rope_inv_freq(config.rope_head_dim(True), config.rope_theta_sliding), persistent=False, ) self.register_buffer( - "full_inv_freq", rope_inv_freq(full_hd, config.rope_theta_full), persistent=False + "full_inv_freq", + rope_inv_freq(config.rope_head_dim(False), config.rope_theta_full), + persistent=False, ) @staticmethod From b1f93de37da2b21ac8eb4f3f405287dc99adf588 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 26/38] Promote the Gemma 4 max-abs RMSNorm to primitives/ios MaxAbsRMSNorm and MaxAbsRMSNormNoScale take a reduction axis and replace ScaledRMSNorm in the EmbeddingGemma 2 Neural Engine encoders; the magnitude sweep moves with them. --- .../models/ios/embedding_gemma2.py | 59 ++++---------- .../models/ios/embedding_gemma2_vision.py | 28 +++---- .../coreai_models/models/ios/gemma4_vision.py | 78 +++---------------- .../coreai_models/primitives/ios/rms_norm.py | 66 ++++++++++++++++ .../test_ios_layers/test_embedding_gemma2.py | 14 ---- .../test_primitives/test_ios/test_rms_norm.py | 36 +++++++++ 6 files changed, 137 insertions(+), 144 deletions(-) create mode 100644 python/tests/test_model_units/test_primitives/test_ios/test_rms_norm.py diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index 6e457c34..f36e134b 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -5,8 +5,7 @@ """Neural Engine float16 EmbeddingGemma 2 text encoder, channels-first (BC1S) conv layout. -What: the macOS encoder math as 1x1 Conv2d projections with static shapes. -Why: ScaledRMSNorm divides by the row max-abs before squaring, keeping float16 in range. +The macOS encoder math as 1x1 Conv2d projections with static shapes and max-abs RMSNorms. """ from __future__ import annotations @@ -24,44 +23,15 @@ from coreai_models.models.macos.embedding_gemma2 import ( load_embedding_gemma2_text_encoder as load_macos_text_encoder, ) +from coreai_models.primitives.ios.rms_norm import MaxAbsRMSNorm, MaxAbsRMSNormNoScale from coreai_models.primitives.ios.sdpa import SDPA -# Smallest positive normal float16; floors the row max-abs and the squared norm so all-zero rows -# stay finite. +# Smallest positive normal float16; floors the output L2's row max-abs and squared norm so an +# all-zero row stays finite. _FP16_TINY = 6.103515625e-05 MASK_NEG = -30000.0 -class ScaledRMSNorm(nn.Module): - """RMSNorm over ``dim`` that normalizes by the row max-abs first, so squares stay at most 1.""" - - def __init__( - self, channels: int, eps: float = 1e-6, with_scale: bool = True, dim: int = 1 - ) -> None: - super().__init__() - self.eps = eps - self.channels = channels - self.dim = dim - self.with_scale = with_scale - if with_scale: - self.weight = nn.Parameter(torch.ones(channels)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - d = self.dim - m = x.abs().amax(d, keepdim=True).clamp(min=_FP16_TINY) - r = torch.reciprocal(m) - xs = x * r - # eps * r^2 as (sqrt(eps) * r)^2. r goes up to 1 / _FP16_TINY, and r * r overflows - # float16 once r > 256, which zeroes the row. - e = self.eps**0.5 * r - y = xs * torch.rsqrt((xs * xs).mean(d, keepdim=True) + e * e) - if self.with_scale: - shape = [1] * x.dim() - shape[d] = self.channels - y = y * self.weight.to(y.dtype).reshape(shape) - return y - - def rotate_half_channels(x: torch.Tensor, head_dim: int) -> torch.Tensor: """NeoX rotate-half over the per-head channel axis of ``(B, H*hd, 1, L)``.""" b, c, one, length = x.shape @@ -99,9 +69,9 @@ def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: self.k_proj = conv1x1(dim, self.n_kv * hd) self.v_proj = conv1x1(dim, self.n_kv * hd) self.o_proj = conv1x1(self.n_heads * hd, dim) - self.q_norm = ScaledRMSNorm(hd, eps, dim=2) - self.k_norm = ScaledRMSNorm(hd, eps, dim=2) - self.v_norm = ScaledRMSNorm(hd, eps, with_scale=False, dim=2) + self.q_norm = MaxAbsRMSNorm(hd, eps, axis=-2) + self.k_norm = MaxAbsRMSNorm(hd, eps, axis=-2) + self.v_norm = MaxAbsRMSNormNoScale(eps, axis=-2) self.sdpa = SDPA(head_dim=hd, scale=1.0) def _per_head(self, x: torch.Tensor, n: int, norm: nn.Module) -> torch.Tensor: @@ -131,11 +101,11 @@ def __init__(self, config: EmbeddingGemma2TextConfig, i: int) -> None: self.down_proj = conv1x1(config.intermediate_size, dim) self.per_layer_input_gate = conv1x1(dim, ple) self.per_layer_projection = conv1x1(ple, dim) - self.input_layernorm = ScaledRMSNorm(dim, eps) - self.post_attention_layernorm = ScaledRMSNorm(dim, eps) - self.pre_feedforward_layernorm = ScaledRMSNorm(dim, eps) - self.post_feedforward_layernorm = ScaledRMSNorm(dim, eps) - self.post_per_layer_input_norm = ScaledRMSNorm(dim, eps) + self.input_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) + self.post_attention_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) + self.pre_feedforward_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) + self.post_feedforward_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) + self.post_per_layer_input_norm = MaxAbsRMSNorm(dim, eps, axis=-3) self.register_buffer("layer_scalar", torch.ones(1)) def forward(self, x, ple_i, cos, sin, bias): @@ -163,9 +133,9 @@ def __init__(self, config: EmbeddingGemma2TextConfig) -> None: dim, eps, L = config.hidden_size, config.rms_norm_eps, config.num_hidden_layers self.ple_dim = config.hidden_size_per_layer_input self.per_layer_model_projection = conv1x1(dim, L * self.ple_dim) - self.per_layer_projection_norm = ScaledRMSNorm(self.ple_dim, eps, dim=2) + self.per_layer_projection_norm = MaxAbsRMSNorm(self.ple_dim, eps, axis=-2) self.layers = nn.ModuleList([EmbeddingGemma2EncoderLayer(config, i) for i in range(L)]) - self.norm = ScaledRMSNorm(dim, eps) + self.norm = MaxAbsRMSNorm(dim, eps, axis=-3) self.embedding_projection = conv1x1(dim, config.embedding_dim) self._rope_hd = {"sliding": config.rope_head_dim(True), "full": config.rope_head_dim(False)} @@ -252,7 +222,6 @@ def load_embedding_gemma2_text_encoder( __all__ = [ "MASK_NEG", - "ScaledRMSNorm", "EmbeddingGemma2TextEncoder", "conv1x1", "load_conv_weights", diff --git a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py index 1d3d3023..51cd2c93 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -5,8 +5,8 @@ """Neural Engine float16 EmbeddingGemma 2 image encoder, channels-first (BC1S) conv layout. -What: the vision tower as 1x1 Conv2d projections with a static patch budget and in-graph lookups. -Why: ScaledRMSNorm divides by the row max-abs before squaring, keeping float16 in range. +The vision tower as 1x1 Conv2d projections with max-abs RMSNorms, a static patch budget and +in-graph lookups. """ from __future__ import annotations @@ -15,17 +15,13 @@ import torch.nn as nn import torch.nn.functional as F -from coreai_models.models.ios.embedding_gemma2 import ( - MASK_NEG, - ScaledRMSNorm, - conv1x1, - load_conv_weights, -) +from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, conv1x1, load_conv_weights from coreai_models.models.macos.embedding_gemma2 import rope_inv_freq from coreai_models.models.macos.embedding_gemma2_vision import ( EmbeddingGemma2ImageEncoder as MacEmbeddingGemma2ImageEncoder, ) from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2VisionConfig +from coreai_models.primitives.ios.rms_norm import MaxAbsRMSNorm, MaxAbsRMSNormNoScale def one_hot_f16(coord: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: @@ -52,9 +48,9 @@ def __init__(self, cfg: EmbeddingGemma2VisionConfig, q_chunk: int = 630) -> None self.k_proj = conv1x1(dim, self.n_heads * hd) self.v_proj = conv1x1(dim, self.n_heads * hd) self.o_proj = conv1x1(self.n_heads * hd, dim) - self.q_norm = ScaledRMSNorm(hd, eps, dim=2) - self.k_norm = ScaledRMSNorm(hd, eps, dim=2) - self.v_norm = ScaledRMSNorm(hd, eps, with_scale=False, dim=2) + self.q_norm = MaxAbsRMSNorm(hd, eps, axis=-2) + self.k_norm = MaxAbsRMSNorm(hd, eps, axis=-2) + self.v_norm = MaxAbsRMSNormNoScale(eps, axis=-2) def _heads(self, x: torch.Tensor) -> torch.Tensor: b, _, _, p = x.shape @@ -88,10 +84,10 @@ def __init__(self, cfg: EmbeddingGemma2VisionConfig) -> None: self.gate_proj = conv1x1(dim, cfg.intermediate_size) self.up_proj = conv1x1(dim, cfg.intermediate_size) self.down_proj = conv1x1(cfg.intermediate_size, dim) - self.input_layernorm = ScaledRMSNorm(dim, eps) - self.post_attention_layernorm = ScaledRMSNorm(dim, eps) - self.pre_feedforward_layernorm = ScaledRMSNorm(dim, eps) - self.post_feedforward_layernorm = ScaledRMSNorm(dim, eps) + self.input_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) + self.post_attention_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) + self.pre_feedforward_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) + self.post_feedforward_layernorm = MaxAbsRMSNorm(dim, eps, axis=-3) def forward(self, x, cos, sin, bias): x = x + self.post_attention_layernorm( @@ -113,7 +109,7 @@ def __init__(self, config: EmbeddingGemma2VisionConfig | None = None): dim = cfg.hidden_size self.input_proj = conv1x1(cfg.patch_dim, dim) self.layers = nn.ModuleList([VisionEncoderLayer(cfg) for _ in range(cfg.num_hidden_layers)]) - self.embedding_pre_projection_norm = ScaledRMSNorm(dim, cfg.rms_norm_eps, with_scale=False) + self.embedding_pre_projection_norm = MaxAbsRMSNormNoScale(cfg.rms_norm_eps, axis=-3) self.embedding_projection = conv1x1(dim, cfg.text_hidden_size) # Per-axis lookup tables, filled by ``bake_tables`` (fp32 math, stored in model dtype): # [G, hidden] position embedding | [G, hd/4] cos | [G, hd/4] sin. diff --git a/python/src/coreai_models/models/ios/gemma4_vision.py b/python/src/coreai_models/models/ios/gemma4_vision.py index 04dc74e7..fd790286 100644 --- a/python/src/coreai_models/models/ios/gemma4_vision.py +++ b/python/src/coreai_models/models/ios/gemma4_vision.py @@ -49,70 +49,10 @@ import torch.nn as nn import torch.nn.functional as F +from coreai_models.primitives.ios.rms_norm import MaxAbsRMSNorm, MaxAbsRMSNormNoScale from coreai_models.primitives.ios.rope import apply_rope from coreai_models.primitives.ios.sdpa import SDPA -# fp16-overflow-safe RMSNorm. RMSNorm is scale-invariant, so we normalize each row -# by its own max magnitude *before* squaring: ``m = amax(|x|)`` (a data-dependent, -# per-row reduction), ``xn = x/m``. Since ``|xn| <= 1``, ``xn^2`` can never overflow -# fp16 (whereas raw ``x^2`` overflows once ``|x|>~256``: ``256^2~=65504``, the fp16 max -# -- the tower's residual stream peaks at ``|x|~=500`` on real images, so raw ``x^2`` -# overflows to ``inf`` -> ``rsqrt(inf)=0`` -> zeroed features -> garbage soft tokens). -# -# The result is the *exact* RMSNorm because ``eps`` is folded back at the right scale:: -# -# xn*rsqrt(mean(xn^2) + (sqrt(eps)/m)^2) * w -# = (x/m)*rsqrt(mean(x^2)/m^2 + eps/m^2) * w -# = (x/m)*m*rsqrt(mean(x^2) + eps) * w -# = x*rsqrt(mean(x^2) + eps) * w -# -# ``(sqrt(eps)/m)^2`` computes ``eps/m^2`` without ever forming ``m^2`` (which would -# itself overflow fp16 for ``m~=500``). Crucially this survives the ANE graph -# optimizer: a constant prescale (``x*0.125``) gets reassociated back into ``x^2*c^2`` -# (recomputing the overflowing ``x^2``) or stripped as scale-invariant, but -# ``x/amax(|x|)`` is data-dependent and cannot be folded away. All ops -# (``abs``/``amax``/``div``/``mean``/``rsqrt``) are fp16 and ANE-resident -- no fp32 -# upcast. - - -class RMSNormNoScale(nn.Module): - """RMSNorm without a learnable scale (v_norm and the projector pre-norm). - - fp16-overflow-safe via per-row max-normalization before squaring (see the comment - directly above, and :class:`RMSNorm`, which adds the scale). Exact and ANE-resident (no fp32 - upcast). - """ - - def __init__(self, eps: float = 1e-6) -> None: - super().__init__() - self._eps = eps - # Registered as a buffer (not a plain attribute) so it follows the module's - # device on `.to(device)` -- a plain `torch.tensor(...)` created in forward() - # defaults to CPU and raises on a GPU/MPS input (e.g. during parity tests). - self.register_buffer("_eps_sqrt", torch.tensor(eps**0.5), persistent=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # Per-row max magnitude; clamp floor avoids 0/0 on an all-zero row (real rows - # have m >> floor, so this is a no-op for them). - m = torch.clamp(x.abs().amax(-1, keepdim=True), min=1e-4) - xn = x / m - eps_sqrt = self._eps_sqrt.to(dtype=x.dtype) - eps_term = (eps_sqrt / m) * (eps_sqrt / m) # eps/m^2, no m^2 overflow - mean_square = (xn * xn).mean(-1, keepdim=True) - return xn * torch.rsqrt(mean_square + eps_term) - - -class RMSNorm(RMSNormNoScale): - """:class:`RMSNormNoScale` followed by a learnable per-channel scale (last-dim reduction).""" - - def __init__(self, dim: int, eps: float = 1e-6) -> None: - super().__init__(eps) - with torch.device("cpu"): - self.weight = nn.Parameter(torch.ones(dim)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return super().forward(x) * self.weight - class ClippedConvLinear(nn.Module): """Bias-free Linear (as a 1x1 Conv2d) with scalar input/output clamping. @@ -157,9 +97,9 @@ def __init__(self, config) -> None: self.v_proj = ClippedConvLinear(dim, self.n_heads * self.head_dim) self.o_proj = ClippedConvLinear(self.n_heads * self.head_dim, dim) - self.q_norm = RMSNorm(self.head_dim, eps=eps) - self.k_norm = RMSNorm(self.head_dim, eps=eps) - self.v_norm = RMSNormNoScale(eps=eps) + self.q_norm = MaxAbsRMSNorm(self.head_dim, eps=eps) + self.k_norm = MaxAbsRMSNorm(self.head_dim, eps=eps) + self.v_norm = MaxAbsRMSNormNoScale(eps=eps) # scale 1.0: Q/K are RMS-normed, matching HF Gemma4VisionAttention.scaling. self.sdpa = SDPA(head_dim=self.head_dim, scale=1.0) @@ -220,10 +160,10 @@ def __init__(self, config) -> None: dim = config.hidden_size self.self_attn = VisionAttention(config) self.mlp = VisionMLP(config) - self.input_layernorm = RMSNorm(dim, eps=eps) - self.post_attention_layernorm = RMSNorm(dim, eps=eps) - self.pre_feedforward_layernorm = RMSNorm(dim, eps=eps) - self.post_feedforward_layernorm = RMSNorm(dim, eps=eps) + self.input_layernorm = MaxAbsRMSNorm(dim, eps=eps) + self.post_attention_layernorm = MaxAbsRMSNorm(dim, eps=eps) + self.pre_feedforward_layernorm = MaxAbsRMSNorm(dim, eps=eps) + self.post_feedforward_layernorm = MaxAbsRMSNorm(dim, eps=eps) def forward( self, @@ -281,7 +221,7 @@ def __init__(self, config, text_hidden_size: int, image_size: int) -> None: ) # Projector (Gemma4MultimodalEmbedder): scale-less RMSNorm -> bias-free Linear. - self.embedding_pre_projection_norm = RMSNormNoScale(eps=config.rms_norm_eps) + self.embedding_pre_projection_norm = MaxAbsRMSNormNoScale(eps=config.rms_norm_eps) self.embedding_projection = nn.Conv2d( self.hidden_size, text_hidden_size, kernel_size=1, bias=False ) diff --git a/python/src/coreai_models/primitives/ios/rms_norm.py b/python/src/coreai_models/primitives/ios/rms_norm.py index 8eb90847..e649cce5 100644 --- a/python/src/coreai_models/primitives/ios/rms_norm.py +++ b/python/src/coreai_models/primitives/ios/rms_norm.py @@ -51,3 +51,69 @@ def forward( x_2normalized = x_2normalized.to(input_dtype) return x_2normalized * self.weight + + +# fp16-overflow-safe RMSNorm. RMSNorm is scale-invariant, so we normalize each row +# by its own max magnitude *before* squaring: ``m = amax(|x|)`` (a data-dependent, +# per-row reduction), ``xn = x/m``. Since ``|xn| <= 1``, ``xn^2`` can never overflow +# fp16 (whereas raw ``x^2`` overflows once ``|x|>~256``: ``256^2~=65504``, the fp16 max +# -- the tower's residual stream peaks at ``|x|~=500`` on real images, so raw ``x^2`` +# overflows to ``inf`` -> ``rsqrt(inf)=0`` -> zeroed features -> garbage soft tokens). +# +# The result is the *exact* RMSNorm because ``eps`` is folded back at the right scale:: +# +# xn*rsqrt(mean(xn^2) + (sqrt(eps)/m)^2) * w +# = (x/m)*rsqrt(mean(x^2)/m^2 + eps/m^2) * w +# = (x/m)*m*rsqrt(mean(x^2) + eps) * w +# = x*rsqrt(mean(x^2) + eps) * w +# +# ``(sqrt(eps)/m)^2`` computes ``eps/m^2`` without ever forming ``m^2`` (which would +# itself overflow fp16 for ``m~=500``). Crucially this survives the ANE graph +# optimizer: a constant prescale (``x*0.125``) gets reassociated back into ``x^2*c^2`` +# (recomputing the overflowing ``x^2``) or stripped as scale-invariant, but +# ``x/amax(|x|)`` is data-dependent and cannot be folded away. All ops +# (``abs``/``amax``/``div``/``mean``/``rsqrt``) are fp16 and ANE-resident -- no fp32 +# upcast. + + +class MaxAbsRMSNormNoScale(nn.Module): + """RMSNorm without a learnable scale, reducing over ``axis`` (counted from the end). + + fp16-overflow-safe via per-row max-normalization before squaring (see the comment + directly above, and :class:`MaxAbsRMSNorm`, which adds the scale). Exact and ANE-resident (no + fp32 upcast). + """ + + def __init__(self, eps: float = 1e-6, axis: int = -1) -> None: + super().__init__() + self._eps = eps + self.axis = axis + # Registered as a buffer (not a plain attribute) so it follows the module's + # device on `.to(device)` -- a plain `torch.tensor(...)` created in forward() + # defaults to CPU and raises on a GPU/MPS input (e.g. during parity tests). + self.register_buffer("_eps_sqrt", torch.tensor(eps**0.5), persistent=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Per-row max magnitude; clamp floor avoids 0/0 on an all-zero row (real rows + # have m >> floor, so this is a no-op for them). + m = torch.clamp(x.abs().amax(self.axis, keepdim=True), min=1e-4) + xn = x / m + eps_sqrt = self._eps_sqrt.to(dtype=x.dtype) + eps_term = (eps_sqrt / m) * (eps_sqrt / m) # eps/m^2, no m^2 overflow + mean_square = (xn * xn).mean(self.axis, keepdim=True) + return xn * torch.rsqrt(mean_square + eps_term) + + +class MaxAbsRMSNorm(MaxAbsRMSNormNoScale): + """:class:`MaxAbsRMSNormNoScale` followed by a learnable per-channel scale. + + The scale has shape ``(dim, 1, ..., 1)``, broadcasting over the dims after ``axis``. + """ + + def __init__(self, dim: int, eps: float = 1e-6, axis: int = -1) -> None: + super().__init__(eps, axis) + with torch.device("cpu"): + self.weight = nn.Parameter(torch.ones((dim,) + (1,) * (-axis - 1))) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return super().forward(x) * self.weight diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py index 7451957a..67c008ff 100644 --- a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -12,7 +12,6 @@ from coreai_models.models.ios.embedding_gemma2 import ( MASK_NEG, - ScaledRMSNorm, load_embedding_gemma2_text_encoder, text_encoder_from_state_dict, ) @@ -112,19 +111,6 @@ def test_float16_stays_finite_with_large_activations() -> None: assert float(_cos(got.float(), want)) >= 0.99 -@pytest.mark.parametrize("peak", [0.0, 1e-7, 1e-5, 1e-3, 3.9e-3, 4e-3, 1.0, 1e2, 1e4, 6e4]) -def test_float16_scaled_rms_norm_matches_float32_across_magnitudes(peak: float) -> None: - eps = 1e-6 - x = torch.randn(8, 256, generator=torch.Generator().manual_seed(0)) - x = x / x.abs().amax(1, keepdim=True) * peak - x16 = x.half() - want = x16.float() * torch.rsqrt((x16.float() ** 2).mean(1, keepdim=True) + eps) - with torch.no_grad(): - got = ScaledRMSNorm(256, eps, with_scale=False).half()(x16).float() - assert torch.isfinite(got).all() - torch.testing.assert_close(got, want, rtol=2e-3, atol=2e-3) - - @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) def test_all_zero_pool_weights_give_a_zero_vector(dtype: torch.dtype) -> None: cfg, _, state = _reference_and_state_dict() diff --git a/python/tests/test_model_units/test_primitives/test_ios/test_rms_norm.py b/python/tests/test_model_units/test_primitives/test_ios/test_rms_norm.py new file mode 100644 index 00000000..a6298fa8 --- /dev/null +++ b/python/tests/test_model_units/test_primitives/test_ios/test_rms_norm.py @@ -0,0 +1,36 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for the max-abs RMSNorm (``primitives/ios/rms_norm.py``).""" + +import pytest +import torch + +from coreai_models.primitives.ios.rms_norm import MaxAbsRMSNorm, MaxAbsRMSNormNoScale + + +@pytest.mark.parametrize("peak", [0.0, 1e-7, 1e-5, 1e-3, 3.9e-3, 4e-3, 1.0, 1e2, 1e4, 6e4]) +def test_float16_max_abs_rms_norm_matches_float32_across_magnitudes(peak: float) -> None: + eps = 1e-6 + x = torch.randn(8, 256, generator=torch.Generator().manual_seed(0)) + x = x / x.abs().amax(1, keepdim=True) * peak + x16 = x.half() + want = x16.float() * torch.rsqrt((x16.float() ** 2).mean(1, keepdim=True) + eps) + with torch.no_grad(): + got = MaxAbsRMSNormNoScale(eps).half()(x16).float() + assert torch.isfinite(got).all() + torch.testing.assert_close(got, want, rtol=2e-3, atol=2e-3) + + +def test_scale_over_a_leading_axis_matches_the_last_axis() -> None: + x = torch.randn(2, 6, 1, 5, generator=torch.Generator().manual_seed(1)) + channels_first, channels_last = MaxAbsRMSNorm(6, axis=-3), MaxAbsRMSNorm(6) + assert channels_first.weight.shape == (6, 1, 1) + with torch.no_grad(): + channels_first.weight.copy_(torch.arange(6.0).reshape(6, 1, 1)) + channels_last.weight.copy_(torch.arange(6.0)) + got = channels_first(x) + want = channels_last(x.transpose(1, 3)).transpose(1, 3) + torch.testing.assert_close(got, want) From ad81f5e1cb4c07fc69e2c56bba3f71e7cd891b02 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 27/38] Reshape each iOS text projection to per-head layout once Norms and RoPE run on (B, n, hd, L); drops rotate_half_channels, apply_rope_bc1s and _per_head. --- .../models/ios/embedding_gemma2.py | 39 +++++-------------- 1 file changed, 10 insertions(+), 29 deletions(-) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index f36e134b..fa5a6f39 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -32,25 +32,6 @@ MASK_NEG = -30000.0 -def rotate_half_channels(x: torch.Tensor, head_dim: int) -> torch.Tensor: - """NeoX rotate-half over the per-head channel axis of ``(B, H*hd, 1, L)``.""" - b, c, one, length = x.shape - x = x.reshape(b, c // head_dim, head_dim, length) - x1, x2 = x[:, :, : head_dim // 2], x[:, :, head_dim // 2 :] - return torch.cat((-x2, x1), dim=2).reshape(b, c, one, length) - - -def apply_rope_bc1s(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, head_dim: int): - """Apply rotary embedding; ``cos``/``sin``: ``(1, 1, hd, L)``; ``x``: ``(B, H*hd, 1, L)``.""" - b, c, one, length = x.shape - xr = x.reshape(b, c // head_dim, head_dim, length) - out = ( - xr * cos - + rotate_half_channels(x, head_dim).reshape(b, c // head_dim, head_dim, length) * sin - ) - return out.reshape(b, c, one, length) - - def conv1x1(cin: int, cout: int) -> nn.Conv2d: """Bias-free 1x1 ``Conv2d``, the BC1S form of a bias-free ``Linear``.""" return nn.Conv2d(cin, cout, kernel_size=1, bias=False) @@ -74,17 +55,17 @@ def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: self.v_norm = MaxAbsRMSNormNoScale(eps, axis=-2) self.sdpa = SDPA(head_dim=hd, scale=1.0) - def _per_head(self, x: torch.Tensor, n: int, norm: nn.Module) -> torch.Tensor: - b, c, one, length = x.shape - return norm(x.reshape(b, n, self.head_dim, length)).reshape(b, c, one, length) - def forward(self, x, cos, sin, bias): - hd = self.head_dim - q = self._per_head(self.q_proj(x), self.n_heads, self.q_norm) - k = self._per_head(self.k_proj(x), self.n_kv, self.k_norm) - v = self._per_head(self.v_proj(x), self.n_kv, self.v_norm) - q = apply_rope_bc1s(q, cos, sin, hd) - k = apply_rope_bc1s(k, cos, sin, hd) + """``x``: ``(B, C, 1, L)``; ``cos``/``sin``: ``(1, 1, hd, L)``; norms and RoPE run per head + on ``(B, n, hd, L)``.""" + b, _, _, length = x.shape + hd, half = self.head_dim, self.head_dim // 2 + q = self.q_norm(self.q_proj(x).reshape(b, self.n_heads, hd, length)) + k = self.k_norm(self.k_proj(x).reshape(b, self.n_kv, hd, length)) + v = self.v_norm(self.v_proj(x).reshape(b, self.n_kv, hd, length)) + q = q * cos + torch.cat((-q[:, :, half:], q[:, :, :half]), dim=2) * sin + k = k * cos + torch.cat((-k[:, :, half:], k[:, :, :half]), dim=2) * sin + q, k, v = (t.reshape(b, -1, 1, length) for t in (q, k, v)) return self.o_proj(self.sdpa(q, k, v, bias)) From 2f2e0869cf88409dee9666cd107c3865c0cdbc58 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 28/38] Share the macOS attention and sandwich layer between text and image VisionAttention overrides only the RoPE; state_dict names are unchanged. --- .../models/macos/embedding_gemma2.py | 90 +++++++++++++------ .../models/macos/embedding_gemma2_vision.py | 85 ++++-------------- 2 files changed, 83 insertions(+), 92 deletions(-) diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index c351cb8d..8ff24302 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -167,27 +167,38 @@ def build_attention_masks( class EmbeddingGemma2Attention(nn.Module): - """Bidirectional attention; q/k RMSNorm + RoPE, scale-free v_norm, scaling 1.0.""" + """Bidirectional attention; q/k RMSNorm + RoPE, scale-free v_norm, scaling 1.0. - def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: + ``_rope`` applies the rotary embedding to ``[B, L, heads, head_dim]``; subclasses override it. + """ + + def __init__( + self, + dim: int, + n_heads: int, + n_kv_heads: int, + head_dim: int, + eps: float, + bias: bool = False, + ) -> None: super().__init__() - self.n_heads = config.num_attention_heads - self.n_kv_heads = config.layer_num_key_value_heads(layer_idx) - self.head_dim = config.layer_head_dim(layer_idx) - dim, hd = config.hidden_size, self.head_dim - bias = config.attention_bias + self.n_heads, self.n_kv_heads, self.head_dim = n_heads, n_kv_heads, head_dim + hd = head_dim self.q_proj = nn.Linear(dim, self.n_heads * hd, bias=bias) self.k_proj = nn.Linear(dim, self.n_kv_heads * hd, bias=bias) self.v_proj = nn.Linear(dim, self.n_kv_heads * hd, bias=bias) self.o_proj = nn.Linear(self.n_heads * hd, dim, bias=bias) - self.q_norm = EmbeddingGemma2RMSNorm(hd, eps=config.rms_norm_eps) - self.k_norm = EmbeddingGemma2RMSNorm(hd, eps=config.rms_norm_eps) - self.v_norm = EmbeddingGemma2RMSNorm(hd, eps=config.rms_norm_eps, with_scale=False) + self.q_norm = EmbeddingGemma2RMSNorm(hd, eps=eps) + self.k_norm = EmbeddingGemma2RMSNorm(hd, eps=eps) + self.v_norm = EmbeddingGemma2RMSNorm(hd, eps=eps, with_scale=False) # The mask carries both the window and the key padding; the composite itself stays # windowless and non-causal. self.sdpa = SDPA(scale=1.0, window_size=0, is_causal=False) + def _rope(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + return apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim=2) + def forward( self, x: torch.Tensor, @@ -199,9 +210,9 @@ def forward( n_h, n_kv, hd = self.n_heads, self.n_kv_heads, self.head_dim query = self.q_norm(self.q_proj(x).reshape(batch, seq_len, n_h, hd)) - query = apply_rotary_pos_emb(query, cos, sin, unsqueeze_dim=2).transpose(1, 2) + query = self._rope(query, cos, sin).transpose(1, 2) key = self.k_norm(self.k_proj(x).reshape(batch, seq_len, n_kv, hd)) - key = apply_rotary_pos_emb(key, cos, sin, unsqueeze_dim=2).transpose(1, 2) + key = self._rope(key, cos, sin).transpose(1, 2) value = self.v_norm(self.v_proj(x).reshape(batch, seq_len, n_kv, hd)).transpose(1, 2) out = self.sdpa(query=query, key=key, value=value, attn_mask=attn_mask) @@ -248,24 +259,25 @@ def forward(self, x: torch.Tensor, per_layer_input: torch.Tensor) -> torch.Tenso return x + self.post_per_layer_input_norm(self.per_layer_projection(h)) -class EmbeddingGemma2EncoderLayer(nn.Module): - def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: +class EmbeddingGemma2SandwichLayer(nn.Module): + """Sandwich-norm block: ``x + post_norm(attn(pre_norm(x)))``, same for the MLP. + + ``config`` needs ``hidden_size``, ``intermediate_size`` and ``rms_norm_eps``. + """ + + def __init__(self, config: Any, self_attn: EmbeddingGemma2Attention) -> None: super().__init__() - eps = config.rms_norm_eps - self.is_sliding = config.is_sliding(layer_idx) - self.self_attn = EmbeddingGemma2Attention(config, layer_idx) + h, eps = config.hidden_size, config.rms_norm_eps + self.self_attn = self_attn self.mlp = EmbeddingGemma2MLP(config) - self.input_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) - self.post_attention_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) - self.pre_feedforward_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) - self.post_feedforward_layernorm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=eps) - self.ple_block = EmbeddingGemma2PLEBlock(config) - self.register_buffer("layer_scalar", torch.ones(1)) + self.input_layernorm = EmbeddingGemma2RMSNorm(h, eps=eps) + self.post_attention_layernorm = EmbeddingGemma2RMSNorm(h, eps=eps) + self.pre_feedforward_layernorm = EmbeddingGemma2RMSNorm(h, eps=eps) + self.post_feedforward_layernorm = EmbeddingGemma2RMSNorm(h, eps=eps) def forward( self, x: torch.Tensor, - per_layer_input: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attn_mask: torch.Tensor, @@ -273,7 +285,35 @@ def forward( h = self.self_attn(self.input_layernorm(x), cos, sin, attn_mask) x = x + self.post_attention_layernorm(h) h = self.mlp(self.pre_feedforward_layernorm(x)) - x = x + self.post_feedforward_layernorm(h) + return x + self.post_feedforward_layernorm(h) + + +class EmbeddingGemma2EncoderLayer(EmbeddingGemma2SandwichLayer): + """Sandwich-norm block followed by the per-layer-input block and ``layer_scalar``.""" + + def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: + attn = EmbeddingGemma2Attention( + config.hidden_size, + config.num_attention_heads, + config.layer_num_key_value_heads(layer_idx), + config.layer_head_dim(layer_idx), + config.rms_norm_eps, + config.attention_bias, + ) + super().__init__(config, attn) + self.is_sliding = config.is_sliding(layer_idx) + self.ple_block = EmbeddingGemma2PLEBlock(config) + self.register_buffer("layer_scalar", torch.ones(1)) + + def forward( # type: ignore[override] + self, + x: torch.Tensor, + per_layer_input: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + attn_mask: torch.Tensor, + ) -> torch.Tensor: + x = super().forward(x, cos, sin, attn_mask) x = self.ple_block(x, per_layer_input) return x * self.layer_scalar diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index 9e383864..000667b0 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -23,11 +23,14 @@ import torch.nn as nn import torch.nn.functional as F -from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaDenseMLP from coreai_models.models.macos.diffusion_gemma import DiffusionGemmaRMSNorm as RMSNorm -from coreai_models.models.macos.embedding_gemma2 import read_checkpoint_tensors, rope_inv_freq +from coreai_models.models.macos.embedding_gemma2 import ( + EmbeddingGemma2Attention, + EmbeddingGemma2SandwichLayer, + read_checkpoint_tensors, + rope_inv_freq, +) from coreai_models.primitives.ios.rope import apply_rope -from coreai_models.primitives.macos.sdpa import SDPA # Special token ids from the EG2 tokenizer / config.json. EG2_BOS_TOKEN_ID = 2 @@ -96,72 +99,17 @@ def from_hf_config( ) -class VisionAttention(nn.Module): - """Bidirectional MHA with q/k RMSNorm, scale-free v RMSNorm and axial RoPE; scale 1.0.""" +class VisionAttention(EmbeddingGemma2Attention): + """Text attention with axial 2-D RoPE: half of head_dim rotates by x, half by y (cos/sin from + ``axial_rope_cos_sin``, laid out ``[fx, fx, fy, fy]``).""" def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: - super().__init__() - self.num_heads = config.num_attention_heads - self.num_kv_heads = config.num_key_value_heads - self.head_dim = config.head_dim - h = config.hidden_size - self.q_proj = nn.Linear(h, self.num_heads * self.head_dim, bias=False) - self.k_proj = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False) - self.v_proj = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False) - self.o_proj = nn.Linear(self.num_heads * self.head_dim, h, bias=False) - self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) - self.v_norm = RMSNorm(self.head_dim, config.rms_norm_eps, with_scale=False) - self.sdpa = SDPA(scale=1.0, is_causal=False) + c = config + heads, kv_heads = c.num_attention_heads, c.num_key_value_heads + super().__init__(c.hidden_size, heads, kv_heads, c.head_dim, c.rms_norm_eps) - def forward( - self, - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - attn_mask: torch.Tensor, - ) -> torch.Tensor: - b, n, _ = x.shape - q = self.q_norm(self.q_proj(x).view(b, n, self.num_heads, self.head_dim)) - k = self.k_norm(self.k_proj(x).view(b, n, self.num_kv_heads, self.head_dim)) - v = self.v_norm(self.v_proj(x).view(b, n, self.num_kv_heads, self.head_dim)) - # Axial 2-D RoPE: half of head_dim rotates by x, half by y (cos/sin from - # axial_rope_cos_sin, laid out [fx, fx, fy, fy]). - q = apply_rope(q, cos, sin, head_axis=2, ndim=2).transpose(1, 2) - k = apply_rope(k, cos, sin, head_axis=2, ndim=2).transpose(1, 2) - v = v.transpose(1, 2) - if self.num_kv_heads != self.num_heads: - rep = self.num_heads // self.num_kv_heads - k = k.repeat_interleave(rep, dim=1) - v = v.repeat_interleave(rep, dim=1) - out = self.sdpa(q, k, v, attn_mask=attn_mask) - return self.o_proj(out.transpose(1, 2).reshape(b, n, self.num_heads * self.head_dim)) - - -class VisionEncoderLayer(nn.Module): - """Sandwich-norm block: ``x + post_norm(attn(pre_norm(x)))``, same for the MLP.""" - - def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: - super().__init__() - h, eps = config.hidden_size, config.rms_norm_eps - self.self_attn = VisionAttention(config) - self.mlp = DiffusionGemmaDenseMLP(config) - self.input_layernorm = RMSNorm(h, eps) - self.post_attention_layernorm = RMSNorm(h, eps) - self.pre_feedforward_layernorm = RMSNorm(h, eps) - self.post_feedforward_layernorm = RMSNorm(h, eps) - - def forward( - self, - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - attn_mask: torch.Tensor, - ) -> torch.Tensor: - h = self.self_attn(self.input_layernorm(x), cos, sin, attn_mask) - x = x + self.post_attention_layernorm(h) - h = self.mlp(self.pre_feedforward_layernorm(x)) - return x + self.post_feedforward_layernorm(h) + def _rope(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + return apply_rope(x, cos, sin, head_axis=2, ndim=2) class VisionPatchEmbedder(nn.Module): @@ -193,7 +141,10 @@ class VisionEncoder(nn.Module): def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: super().__init__() self.layers = nn.ModuleList( - [VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)] + [ + EmbeddingGemma2SandwichLayer(config, VisionAttention(config)) + for _ in range(config.num_hidden_layers) + ] ) From 377f5035a6b9ad16c223e309cabd5281f45efb05 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 29/38] Drop a rationale from a section comment in the macOS image encoder --- .../src/coreai_models/models/macos/embedding_gemma2_vision.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index 000667b0..493b86cb 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -186,7 +186,7 @@ def __init__(self, config: EmbeddingGemma2VisionConfig | None = None) -> None: "pool_slots", torch.arange(cfg.max_soft_tokens, dtype=torch.int64), persistent=False ) - # -- pieces (public so tests / activation probes can hook them) --------------------- + # -- pieces ------------------------------------------------------------------------- def axial_rope_cos_sin( self, position_ids: torch.Tensor, dtype: torch.dtype From d7a56f24a91237a5c2d5f2dcfbe16d36f4558195 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 30/38] Fold overlapping embedding schema tests and reuse the shared tiny configs The macOS text, envelope and asset-sharing asserts move into the matrix and macOS image tests. --- .../test_export/test_embedding_export.py | 244 ++++++------------ 1 file changed, 76 insertions(+), 168 deletions(-) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index ac1ff98d..c3d16c83 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -6,7 +6,7 @@ """Tests for the ``kind: "embedding"`` bundle contract and exporter (``coreai_models.embedding``). Metadata, precision and CLI tests need no hardware or weights. The conversion tests use tiny random -EmbeddingGemma 2 configs and skip when ``coreai_torch`` is unavailable. +EmbeddingGemma 2 configs. """ from __future__ import annotations @@ -36,6 +36,10 @@ family_for_model, text_functions, ) +from tests.test_model_units.test_models._embedding_gemma2_common import ( + TINY_TEXT_CONFIG, + TINY_VISION_CONFIG, +) EG2 = family_for_model("embeddinggemma-2") IMAGE = EG2.image @@ -67,35 +71,6 @@ "FactChecking": PROMPTS["fact_checking"], **dict.fromkeys(["CodeRetrieval", "InstructionRetrieval"], PROMPTS["code_retrieval"]), } -REMOVED_KEYS = { - "pooling", - "attention", - "granularity", - "budgets", - "max_sequence_length", - "prompt_applies_to", - "side", - "strategy", - "wrap", - "soft_tokens_scaled", - "floor", - "norm", - "dynamic_batch", - "dynamic_length", - "hidden_size", - "note", - "max_patches", - "soft_token_dim", - "convert_rgb", - "exif_transpose", - "resize", - "resize_policy", - "patch_pixel_order", - "position_order", - "padding_position", - "prompt", - "truncation_side", -} def _block(platform="macOS", with_image=False, **overrides): @@ -104,16 +79,6 @@ def _block(platform="macOS", with_image=False, **overrides): return embedding_block(EG2, platform=platform, with_image=with_image, **kwargs) -def _keys(node): - if isinstance(node, dict): - for key, value in node.items(): - yield key - yield from _keys(value) - elif isinstance(node, list): - for value in node: - yield from _keys(value) - - # --- Precision guard --- @@ -201,49 +166,6 @@ def test_function_plans_per_platform(): # --- embedding_block schema --- -def test_embedding_block_macos_text(): - block = _block() - assert list(block) == [ - "space", - "dimension", - "matryoshka_dimensions", - "renormalize_after_truncation", - "normalized", - "metric", - "precision", - "components", - "towers", - "text", - ] - assert block["space"] == {"id": "google/embeddinggemma-2", "revision": REVISION} - assert block["matryoshka_dimensions"] == [768, 512, 256, 128] - assert block["precision"] == { - "compute": "float32", - "forbidden": ["float16"], - "compute_units": ["gpu"], - } - assert set(block["components"]) == {"text_encoder"} - component = block["components"]["text_encoder"] - assert component["output"] == "embedding" and component["compute_units"] == ["gpu"] - assert {"function": "len2048_b8", "batch": 8, "length": 2048} in component["functions"] - (tower,) = block["towers"] - assert tower["pipeline"] == ["text_encoder"] - assert tower["batch_sizes"] == [1, 8, 32] - assert tower["length_buckets"] == [128, 512, 2048] - assert tower["max_batch_tokens"] == 16384 - text = block["text"] - assert text["truncation"] == {"keep_eos": True, "max_tokens": 2048} - assert text["special_tokens"] == { - "bos": True, - "eos": True, - "bos_token_id": 2, - "eos_token_id": 1, - "pad_token_id": 0, - } - assert text["prompts"] == PROMPTS and text["default_prompt_name"] is None - json.dumps(block) - - @pytest.mark.parametrize( ("platform", "with_image", "components", "pipelines", "units"), [ @@ -287,10 +209,16 @@ def test_embedding_block_platform_matrix(platform, with_image, components, pipel assert {k: v["compute_units"] for k, v in block["components"].items()} == units assert ("merge" in block) == with_image and ("preprocessing" in block) == with_image assert set(bundle_assets(platform, with_image)) == set(components) - assert not REMOVED_KEYS & set(_keys(block)), REMOVED_KEYS & set(_keys(block)) - assert all("batch" not in tower and "output" in tower for tower in block["towers"]) - assert "output" not in block["precision"] - json.dumps(block) + meta = build_bundle_metadata( + EG2, + block, + name=bundle_name(EG2, platform, with_image), + assets=bundle_assets(platform, with_image), + hf_revision="914f7f89142e33e77833254d9c9b90c3cef7303b", + ) + assert meta["metadata_version"] == "0.2" and meta["kind"] == "embedding" + assert {c for t in meta["embedding"]["towers"] for c in t["pipeline"]} == set(meta["assets"]) + assert meta["source"]["hf_model_id"] == "google/embeddinggemma-2" def test_embedding_block_ios_details(): @@ -315,6 +243,49 @@ def test_embedding_block_ios_details(): def test_embedding_block_macos_image_details(): block = _block("macOS", True) + assert list(block) == [ + "space", + "dimension", + "matryoshka_dimensions", + "renormalize_after_truncation", + "normalized", + "metric", + "precision", + "components", + "towers", + "text", + "merge", + "preprocessing", + ] + assert block["space"] == {"id": "google/embeddinggemma-2", "revision": REVISION} + assert block["matryoshka_dimensions"] == [768, 512, 256, 128] + assert block["precision"] == { + "compute": "float32", + "forbidden": ["float16"], + "compute_units": ["gpu"], + } + text_encoder = block["components"]["text_encoder"] + assert text_encoder["output"] == "embedding" + assert {"function": "len2048_b8", "batch": 8, "length": 2048} in text_encoder["functions"] + text_tower = block["towers"][0] + assert text_tower["batch_sizes"] == [1, 8, 32] + assert text_tower["length_buckets"] == [128, 512, 2048] + assert text_tower["max_batch_tokens"] == 16384 + text = block["text"] + assert text["truncation"] == {"keep_eos": True, "max_tokens": 2048} + assert text["special_tokens"] == { + "bos": True, + "eos": True, + "bos_token_id": 2, + "eos_token_id": 1, + "pad_token_id": 0, + } + assert text["prompts"] == PROMPTS and text["default_prompt_name"] is None + assets = bundle_assets("macOS", True) + assert assets["token_table"] == assets["encoder"] == assets["text_encoder"] + assert assets["image_encoder"] == "image_encoder.aimodel" + names = [f.name for f in text_functions(EG2, "macOS") + encoder_functions(EG2, "macOS")] + assert len(names) == len(set(names)) and "main" not in names enc = block["components"]["encoder"] assert enc["inputs"] == ["inputs_embeds", "attention_mask"] assert enc["mask_encoding"] == "attention_mask" @@ -352,8 +323,9 @@ def test_embedding_block_records_the_compute_precision(): _block(compute_precision="float16") -def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch): - seen = {} +@pytest.fixture +def stub_snapshot(tmp_path, monkeypatch): + """Snapshot helpers stubbed so ``export_embedding_bundle`` runs without weights.""" monkeypatch.setattr( embedding_export, "resolve_snapshot", lambda *a: (tmp_path, "0123456789abcdef0123") ) @@ -361,6 +333,10 @@ def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch monkeypatch.setattr(embedding_export, "read_prompts", lambda *a: (PROMPTS, None)) monkeypatch.setattr(embedding_export, "consolidate_prompts", lambda family, prompts: prompts) monkeypatch.setattr(embedding_export, "copy_tokenizer", lambda *a: ["tokenizer.json"]) + + +def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch, stub_snapshot): + seen = {} monkeypatch.setattr(embedding_export, "export_gpu_programs", lambda *a, **kw: seen.update(kw)) meta = embedding_export.export_embedding_bundle( EG2, tmp_path / "out", with_image=False, compute_precision="bfloat16" @@ -371,17 +347,10 @@ def test_export_writes_the_requested_precision_to_metadata(tmp_path, monkeypatch assert written["embedding"]["precision"]["compute"] == "bfloat16" -def test_failed_export_removes_the_bundle_dir(tmp_path, monkeypatch): +def test_failed_export_removes_the_bundle_dir(tmp_path, monkeypatch, stub_snapshot): def fail(*args, **kwargs): raise RuntimeError("conversion failed") - monkeypatch.setattr( - embedding_export, "resolve_snapshot", lambda *a: (tmp_path, "0123456789abcdef0123") - ) - monkeypatch.setattr(embedding_export, "verify_pooling_config", lambda *a: None) - monkeypatch.setattr(embedding_export, "read_prompts", lambda *a: (PROMPTS, None)) - monkeypatch.setattr(embedding_export, "consolidate_prompts", lambda family, prompts: prompts) - monkeypatch.setattr(embedding_export, "copy_tokenizer", lambda *a: ["tokenizer.json"]) monkeypatch.setattr(embedding_export, "export_gpu_programs", fail) with pytest.raises(RuntimeError, match="conversion failed"): embedding_export.export_embedding_bundle(EG2, tmp_path / "out", with_image=False) @@ -396,21 +365,6 @@ def test_forbidden_precision_message_lists_the_allowed_ones(): check_compute_precision(no_floor, "float16") -def test_bundle_metadata_envelope(): - meta = build_bundle_metadata( - EG2, - _block("iOS", True), - name=bundle_name(EG2, "iOS", True), - assets=bundle_assets("iOS", True), - hf_revision="914f7f89142e33e77833254d9c9b90c3cef7303b", - ) - assert meta["metadata_version"] == "0.2" and meta["kind"] == "embedding" - assert meta["name"] == "embeddinggemma-2-ios" - pipelines = {c for t in meta["embedding"]["towers"] for c in t["pipeline"]} - assert pipelines == set(meta["assets"]) - assert meta["source"]["hf_model_id"] == "google/embeddinggemma-2" - - def test_family_registry(): assert "embeddinggemma-2" in EMBEDDING_FAMILIES assert family_for_model("Google/EmbeddingGemma-2") is EG2 @@ -487,34 +441,14 @@ def test_verify_image_config(tmp_path): # --- Conversion (tiny random config) --- -TINY_CONFIG = { - "vocab_size": 64, - "hidden_size": 32, - "intermediate_size": 48, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "head_dim": 16, - "sliding_window": 3, - "layer_types": ["sliding_attention", "full_attention"], - "per_layer_config": {"01": {"head_dim": 32, "num_key_value_heads": 1}}, - "rope_parameters": { - "full_attention": {"rope_theta": 1_000_000.0, "rope_type": "default"}, - "sliding_attention": {"rope_theta": 10_000.0, "rope_type": "default"}, - }, - "hidden_size_per_layer_input": 8, - "embedding_dim": 24, -} - def test_tiny_multifunction_conversion(): - pytest.importorskip("coreai_torch") wrapper = export_macos.TextEmbeddingExportWrapper(_tiny_text_encoder()) plan = bucket_plan([16], [1, 4], 1024) - ids, mask = export_macos.reference_inputs(plan[1], 64) + ids, mask = export_macos.reference_inputs(plan[1], 128) assert ids.dtype == torch.int32 and mask.dtype == torch.int32 and ids.shape == (4, 16) out = wrapper(ids, mask) - assert out.dtype == torch.float32 and out.shape == (4, 24) + assert out.dtype == torch.float32 and out.shape == (4, 48) torch.testing.assert_close(out.norm(dim=-1), torch.ones(4)) assert export_macos.export_shared_text_program(wrapper.encoder, plan, [], 64) is not None @@ -525,32 +459,16 @@ def _tiny_text_encoder(): EmbeddingGemma2TextEncoder, ) - config = EmbeddingGemma2TextConfig.from_dict(TINY_CONFIG) + config = EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG) torch.manual_seed(0) return EmbeddingGemma2TextEncoder(config).eval() def _tiny_image_encoder(): - from coreai_models.models.macos.embedding_gemma2_vision import ( - EmbeddingGemma2ImageEncoder, - EmbeddingGemma2VisionConfig, - ) + from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder - config = EmbeddingGemma2VisionConfig( - hidden_size=64, - intermediate_size=128, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=4, - head_dim=16, - patch_size=4, - pooling_kernel_size=3, - position_embedding_size=64, - text_hidden_size=32, - max_soft_tokens=12, - ) torch.manual_seed(1) - return EmbeddingGemma2ImageEncoder(config).eval() + return EmbeddingGemma2ImageEncoder(TINY_VISION_CONFIG).eval() TINY_IMAGE = dataclasses.replace(IMAGE, patch_size=4, budget=12, soft_token_dim=32) @@ -559,7 +477,7 @@ def _tiny_image_encoder(): def test_inputs_embeds_path_equals_ids_path(): """token_table(ids) -> encoder(inputs_embeds) == text_encoder(ids) (the merge is lossless).""" text = _tiny_text_encoder() - ids = torch.randint(3, 64, (1, 16), dtype=torch.int32) + ids = torch.randint(3, 128, (1, 16), dtype=torch.int32) mask = torch.ones(1, 16, dtype=torch.int32) mask[0, 11:] = 0 with torch.no_grad(): @@ -567,7 +485,7 @@ def test_inputs_embeds_path_equals_ids_path(): embeds = export_utils.TokenTableExportWrapper(text.embed_tokens)(ids) with export_macos.without_token_table(text) as encoder: merged = export_macos.InputsEmbedsEncoderWrapper(encoder)(embeds, mask) - assert embeds.dtype == torch.float32 and embeds.shape == (1, 16, 32) + assert embeds.dtype == torch.float32 and embeds.shape == (1, 16, 64) torch.testing.assert_close(merged, direct, atol=1e-6, rtol=0) assert isinstance(text.embed_tokens, torch.nn.Embedding) # restored @@ -575,7 +493,7 @@ def test_inputs_embeds_path_equals_ids_path(): def test_float16_scaled_token_table_equals_the_float32_table_cast(): """The float16 token table gives the encoder the same inputs as the float32 one.""" text = _tiny_text_encoder() - ids = torch.randint(3, 64, (1, 16), dtype=torch.int32) + ids = torch.randint(3, 128, (1, 16), dtype=torch.int32) table = export_utils.scaled_token_table( text.embed_tokens.weight, text.config.hidden_size**0.5, torch.float16 ) @@ -587,7 +505,6 @@ def test_float16_scaled_token_table_equals_the_float32_table_cast(): def test_tiny_float16_token_table_conversion(): - pytest.importorskip("coreai_torch") text = _tiny_text_encoder() table = export_utils.scaled_token_table( text.embed_tokens.weight, text.config.hidden_size**0.5, torch.float16 @@ -610,7 +527,6 @@ def test_image_reference_inputs_are_a_valid_padded_grid(): def test_tiny_shared_text_program_conversion(): - pytest.importorskip("coreai_torch") program = export_macos.export_shared_text_program( _tiny_text_encoder(), [FunctionSpec("len16_b1", 1, 16), FunctionSpec("len32_b1", 1, 32)], @@ -622,14 +538,6 @@ def test_tiny_shared_text_program_conversion(): assert image is not None -def test_macos_image_bundle_shares_the_text_encoder_asset(): - assets = bundle_assets("macOS", True) - assert assets["token_table"] == assets["encoder"] == assets["text_encoder"] - assert assets["image_encoder"] == "image_encoder.aimodel" - names = [f.name for f in text_functions(EG2, "macOS") + encoder_functions(EG2, "macOS")] - assert len(names) == len(set(names)) and "main" not in names - - def test_ane_reference_inputs_shapes(): embeds, bias, pool = export_ios.ane_text_reference_inputs(128, 512) assert embeds.shape == (1, 128, 512) and bias.shape == (1, 128, 1, 128) @@ -646,7 +554,7 @@ def test_ane_band_wrapper_matches_explicit_sliding_bias(): from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, EmbeddingGemma2TextEncoder from coreai_models.models.macos.embedding_gemma2 import EmbeddingGemma2TextConfig - config = EmbeddingGemma2TextConfig.from_dict(TINY_CONFIG) + config = EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG) torch.manual_seed(0) model = EmbeddingGemma2TextEncoder(config).eval() length, valid = 16, 12 @@ -656,7 +564,7 @@ def test_ane_band_wrapper_matches_explicit_sliding_bias(): band = (pos[:, None] - pos[None, :]).abs() <= 3 sliding = torch.where(key_ok & band, 0.0, MASK_NEG) bias, sliding = (t.reshape(1, length, 1, length) for t in (bias, sliding)) - embeds = torch.randn(1, length, 32) + embeds = torch.randn(1, length, 64) pool = (pos < valid).float().reshape(1, length) / valid wrapper = export_ios.EncoderBandWrapper(model, length, 3) with torch.no_grad(): From 3559d6279f865ffe98c29821bee58e2c94fcf3ae Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:20:48 -0700 Subject: [PATCH 31/38] Trim the EmbeddingGemma 2 model tests to one path per check 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. --- .../test_export/test_embedding_export.py | 25 --- .../test_ios_layers/test_embedding_gemma2.py | 61 +++--- .../test_embedding_gemma2.py | 177 +++--------------- .../test_embedding_gemma2_vision.py | 65 ++----- 4 files changed, 71 insertions(+), 257 deletions(-) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index c3d16c83..e5364fdd 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -549,31 +549,6 @@ def test_ane_reference_inputs_shapes(): assert int(real.sum()) % 9 == 0 and torch.all(positions[0, ~real] == -1) -def test_ane_band_wrapper_matches_explicit_sliding_bias(): - """The in-graph band equals an explicit sliding-window bias built on the host.""" - from coreai_models.models.ios.embedding_gemma2 import MASK_NEG, EmbeddingGemma2TextEncoder - from coreai_models.models.macos.embedding_gemma2 import EmbeddingGemma2TextConfig - - config = EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG) - torch.manual_seed(0) - model = EmbeddingGemma2TextEncoder(config).eval() - length, valid = 16, 12 - pos = torch.arange(length) - key_ok = (pos < valid)[:, None] # (key, query) - bias = torch.where(key_ok, 0.0, MASK_NEG).expand(length, length) - band = (pos[:, None] - pos[None, :]).abs() <= 3 - sliding = torch.where(key_ok & band, 0.0, MASK_NEG) - bias, sliding = (t.reshape(1, length, 1, length) for t in (bias, sliding)) - embeds = torch.randn(1, length, 64) - pool = (pos < valid).float().reshape(1, length) / valid - wrapper = export_ios.EncoderBandWrapper(model, length, 3) - with torch.no_grad(): - want = model(embeds, bias, pool, sliding) - got = wrapper(embeds, bias, pool) - torch.testing.assert_close(got, want) - assert not export_ios.EncoderBandWrapper(model, 4, 3).banded - - # --- CLI --- diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py index 67c008ff..184618df 100644 --- a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -10,6 +10,8 @@ import pytest import torch +from coreai_models.embedding.export_ios import EncoderBandWrapper +from coreai_models.embedding.export_utils import TokenTableExportWrapper, scaled_token_table from coreai_models.models.ios.embedding_gemma2 import ( MASK_NEG, load_embedding_gemma2_text_encoder, @@ -33,28 +35,19 @@ def _host_inputs( ids: torch.Tensor, attention_mask: torch.Tensor, embed_table: torch.Tensor, - sliding_window: int, dtype: torch.dtype = torch.float16, -) -> dict[str, torch.Tensor]: - """Graph inputs for one sequence (``ids [1, L]``): scaled embeddings, bias, pooling weights.""" - b, length = ids.shape - emb = embed_table[ids.long()].float() * math.sqrt(embed_table.shape[1]) - valid = attention_mask.bool() - key_ok = valid[0].reshape(length, 1, 1).expand(length, 1, length) # (Lk, 1, Lq) +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Graph inputs for one sequence (``ids [1, L]``): token-table embeddings, bias, pooling.""" + length = ids.shape[1] + table = scaled_token_table(embed_table, math.sqrt(embed_table.shape[1]), dtype) + key_ok = attention_mask[0].bool().reshape(length, 1, 1).expand(length, 1, length) bias = torch.where(key_ok, 0.0, MASK_NEG).reshape(1, length, 1, length) - out = { - "inputs_embeds": emb.to(dtype), - "bias": bias.to(dtype), - "pool_weights": (attention_mask.float() / attention_mask.float().sum(-1, keepdim=True)).to( - dtype - ), - } - if length > sliding_window + 1: - pos = torch.arange(length) - band = (pos[:, None] - pos[None, :]).abs() <= sliding_window # (Lk, Lq) symmetric - sb = torch.where(key_ok & band.reshape(length, 1, length), 0.0, MASK_NEG) - out["sliding_bias"] = sb.reshape(1, length, 1, length).to(dtype) - return out + pool = attention_mask.float() / attention_mask.float().sum(-1, keepdim=True) + return TokenTableExportWrapper(table, dtype)(ids), bias.to(dtype), pool.to(dtype) + + +def _banded(ios: torch.nn.Module, length: int) -> EncoderBandWrapper: + return EncoderBandWrapper(ios, length, ios.config.sliding_window) def _perturb(model: torch.nn.Module, seed: int) -> None: @@ -84,15 +77,14 @@ def test_float32_matches_macos_reference_per_row() -> None: ids, mask = _ragged_batch([12, 7, 4], seq_len=12, vocab=cfg.vocab_size) with torch.no_grad(): want = ref(input_ids=ids, attention_mask=mask) - for row in range(ids.shape[0]): - n = int(mask[row].sum()) - inputs = _host_inputs( - ids[row : row + 1], mask[row : row + 1], table, cfg.sliding_window, torch.float32 - ) - got = ios(**inputs) + # Length 12 takes the banded path (> sliding_window + 1); length 4 the unbanded one. + for row, length in ((0, 12), (1, 12), (2, 12), (2, 4)): + r = slice(row, row + 1) + inputs = _host_inputs(ids[r, :length], mask[r, :length], table, torch.float32) + got = _banded(ios, length)(*inputs) assert got.shape == (1, cfg.embedding_dim) - cos = _cos(got, want[row : row + 1]) - assert float(cos) >= 0.9999, (row, n, float(cos)) + cos = _cos(got, want[r]) + assert float(cos) >= 0.9999, (row, length, float(cos)) def test_float16_stays_finite_with_large_activations() -> None: @@ -105,8 +97,8 @@ def test_float16_stays_finite_with_large_activations() -> None: ios32, _ = text_encoder_from_state_dict(cfg, state, dtype=torch.float32) ios16, _ = text_encoder_from_state_dict(cfg, state, dtype=torch.float16) with torch.no_grad(): - want = ios32(**_host_inputs(ids, mask, table, cfg.sliding_window, torch.float32)) - got = ios16(**_host_inputs(ids, mask, table, cfg.sliding_window, torch.float16)) + want = _banded(ios32, 12)(*_host_inputs(ids, mask, table, torch.float32)) + got = _banded(ios16, 12)(*_host_inputs(ids, mask, table, torch.float16)) assert torch.isfinite(got.float()).all() assert float(_cos(got.float(), want)) >= 0.99 @@ -116,14 +108,13 @@ def test_all_zero_pool_weights_give_a_zero_vector(dtype: torch.dtype) -> None: cfg, _, state = _reference_and_state_dict() ios, table = text_encoder_from_state_dict(cfg, state, dtype=dtype) ids, mask = _ragged_batch([6], seq_len=6, vocab=cfg.vocab_size) - inputs = _host_inputs(ids, mask, table, cfg.sliding_window, dtype) - inputs["pool_weights"] = torch.zeros_like(inputs["pool_weights"]) + embeds, bias, pool = _host_inputs(ids, mask, table, dtype) with torch.no_grad(): - got = ios(**inputs) + got = _banded(ios, 6)(embeds, bias, torch.zeros_like(pool)) assert torch.equal(got, torch.zeros_like(got)) -@pytest.mark.parametrize(("length", "valid"), [(8, 8), (64, 40), (128, 128)]) +@pytest.mark.parametrize(("length", "valid"), [(64, 40), (1024, 700)]) def test_real_weight_float16_matches_float32(embeddinggemma2_snapshot, length, valid) -> None: reference = load_macos_text_encoder(embeddinggemma2_snapshot, dtype=torch.float32).eval() ios16, table = load_embedding_gemma2_text_encoder(embeddinggemma2_snapshot, torch.float16) @@ -133,6 +124,6 @@ def test_real_weight_float16_matches_float32(embeddinggemma2_snapshot, length, v mask = (torch.arange(length) < valid).to(torch.int32)[None] with torch.no_grad(): want = reference(input_ids=ids, attention_mask=mask) - got = ios16(**_host_inputs(ids, mask, table, ios16.config.sliding_window, torch.float16)) + got = _banded(ios16, length)(*_host_inputs(ids, mask, table, torch.float16)) assert torch.isfinite(got).all() assert float(_cos(got.float(), want)) >= 0.9999 diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py index f5e4fea0..e1b17b2e 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -7,24 +7,26 @@ Tiers ----- -1. Self-contained (always run): config parsing, mask semantics, padding invariance, pooling - mask negative test, inputs_embeds path, torch.export at a static bucket. +1. Self-contained (always run): config parsing, padding invariance, pooling mask negative test, + inputs_embeds path. 2. HF reference (tiny random config): needs ``transformers>=5.19`` (``models/embedding_gemma2``). The repo pins an older transformers, so run in an overlay:: - uv run --no-sync --with "transformers>=5.19" --with "sentence-transformers>=6.1" \\ + uv run --no-sync --with "transformers>=5.19" \\ python -m pytest -q -p no:rerunfailures -p no:xdist -o addopts="" 3. Real weights (CPU fp32): additionally needs the ``google/embeddinggemma-2`` snapshot in the - local HF cache (the ``embeddinggemma2_snapshot`` fixture skips otherwise). Compared against - sentence-transformers>=6.1 when present; otherwise against HF ``EmbeddingGemma2Model`` + the - mean-pool + L2-normalize re-implementation in ``_st_mean_pool_normalize``. + local HF cache (the ``embeddinggemma2_snapshot`` fixture skips otherwise). Compared against HF + ``EmbeddingGemma2Model`` + the sentence-transformers mean-pool + L2-normalize + re-implementation in ``_st_mean_pool_normalize``. Tiers 2/3 skip cleanly when their dependencies are unavailable. """ from __future__ import annotations +import json + import pytest import torch import torch.nn.functional as F @@ -32,7 +34,6 @@ from coreai_models.models.macos.embedding_gemma2 import ( EmbeddingGemma2TextConfig, EmbeddingGemma2TextEncoder, - build_attention_masks, load_embedding_gemma2_text_encoder, ) from tests.test_model_units.test_models._embedding_gemma2_common import ( @@ -87,28 +88,6 @@ def test_config_requires_both_rope_thetas(): EmbeddingGemma2TextConfig.from_dict({**TINY_TEXT_CONFIG, "rope_parameters": rope}) -def test_config_from_snapshot_matches_published_layout(embeddinggemma2_snapshot): - cfg = EmbeddingGemma2TextConfig.from_pretrained(embeddinggemma2_snapshot) - full = [i for i in range(cfg.num_hidden_layers) if not cfg.is_sliding(i)] - assert full == [5, 11, 17, 23] - assert all(cfg.layer_head_dim(i) == 512 and cfg.layer_num_key_value_heads(i) == 1 for i in full) - assert cfg.layer_head_dim(0) == 256 and cfg.layer_num_key_value_heads(0) == 2 - assert (cfg.hidden_size, cfg.embedding_dim, cfg.sliding_window) == (512, 768, 512) - - -def test_masks_bidirectional_window_and_key_padding(): - attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0, 0]], dtype=torch.int32) - full, sliding = build_attention_masks(attention_mask, sliding_window=2) - assert full.shape == sliding.shape == (1, 1, 7, 7) - i, j = torch.meshgrid(torch.arange(7), torch.arange(7), indexing="ij") - key_valid = (j < 5).unsqueeze(0).unsqueeze(0) - assert torch.equal(full, key_valid.expand_as(full)) - assert torch.equal(sliding, ((i - j).abs() <= 2).unsqueeze(0).unsqueeze(0) & key_valid) - # Symmetric (bidirectional) on valid positions. - valid = sliding[0, 0, :5, :5] - assert torch.equal(valid, valid.T) - - def test_forward_shapes_and_unit_norm(): model = _tiny_model() ids, mask = _ragged_batch([9, 5], seq_len=12, vocab=128) @@ -162,27 +141,6 @@ def test_inputs_embeds_path_matches_input_ids(): model(input_ids=ids, inputs_embeds=model.embed_tokens(ids)) -def test_torch_export_static_bucket(): - """torch.export at a static (B=2, L=128) bucket with int32 ids/mask; outputs match eager.""" - model = _tiny_model() - ids, mask = _ragged_batch([128, 37], seq_len=128, vocab=128) - with torch.no_grad(): - exported = torch.export.export( - model, args=(), kwargs={"input_ids": ids, "attention_mask": mask} - ) - eager = model(input_ids=ids, attention_mask=mask) - traced = exported.module()(input_ids=ids, attention_mask=mask) - assert traced.shape == (2, 48) - torch.testing.assert_close(traced, eager, atol=1e-6, rtol=0) - placeholders = [ - n - for n in exported.graph.nodes - if n.op == "placeholder" and n.name in ("input_ids", "attention_mask") - ] - assert len(placeholders) == 2 - assert all(n.meta["val"].dtype == torch.int32 for n in placeholders) - - # --- Tier 2: tiny random config vs HF EmbeddingGemma2TextModel --- @@ -193,8 +151,8 @@ def _hf_modeling(): ) -def _tiny_hf(attn_impl: str): - """Tiny random HF ``EmbeddingGemma2TextModel`` with perturbed weights.""" +def _tiny_hf(): + """Tiny random HF ``EmbeddingGemma2TextModel`` (eager attention) with perturbed weights.""" modeling = _hf_modeling() from transformers.models.embedding_gemma2.configuration_embedding_gemma2 import ( EmbeddingGemma2Config, @@ -202,15 +160,14 @@ def _tiny_hf(attn_impl: str): torch.manual_seed(0) hf_cfg = EmbeddingGemma2Config(text_config=dict(TINY_TEXT_CONFIG)).text_config - hf_cfg._attn_implementation = attn_impl + hf_cfg._attn_implementation = "eager" hf = modeling.EmbeddingGemma2TextModel(hf_cfg).eval() _randomize_(hf, seed=3) return hf -@pytest.mark.parametrize("attn_impl", ["eager", "sdpa"]) -def test_tiny_parity_vs_hf(attn_impl): - hf = _tiny_hf(attn_impl) +def test_tiny_parity_vs_hf(): + hf = _tiny_hf() ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG)).eval() ours.load_state_dict(hf.state_dict()) @@ -228,21 +185,6 @@ def test_tiny_parity_vs_hf(attn_impl): torch.testing.assert_close(emb, hf_emb, atol=1e-5, rtol=0) -def test_tiny_sliding_window_matters_vs_hf(): - """Guard against a silently-ignored window: widening ours must break parity with HF.""" - hf = _tiny_hf("eager") - wide = dict(TINY_TEXT_CONFIG, sliding_window=1000) - ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(wide)).eval() - ours.load_state_dict(hf.state_dict()) - ids, mask = _ragged_batch([24], seq_len=24, vocab=128, seed=7) - with torch.no_grad(): - hf_emb = _st_mean_pool_normalize( - hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state, mask - ) - emb = ours(input_ids=ids, attention_mask=mask) - assert (emb - hf_emb).abs().max().item() > 1e-3 - - # --- Tier 3: real weights (CPU fp32) --- REAL_TEXTS = [ @@ -264,22 +206,7 @@ def real_setup(embeddinggemma2_snapshot): from transformers import AutoConfig, AutoTokenizer from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import EmbeddingGemma2Model - prompts = None - st_model = None - try: - from sentence_transformers import SentenceTransformer - - st_model = SentenceTransformer( - str(snapshot), device="cpu", model_kwargs={"dtype": torch.float32} - ) - prompts = st_model.prompts - except Exception: # noqa: BLE001 - optional reference - import json - - prompts = json.loads((snapshot / "config_sentence_transformers.json").read_text())[ - "prompts" - ] - + prompts = json.loads((snapshot / "config_sentence_transformers.json").read_text())["prompts"] cfg = AutoConfig.from_pretrained(str(snapshot)) cfg.vision_config = None cfg.audio_config = None @@ -287,41 +214,28 @@ def real_setup(embeddinggemma2_snapshot): tokenizer = AutoTokenizer.from_pretrained(str(snapshot)) ours = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float32) texts = [prompts[p] + t for p, t in REAL_TEXTS] - if st_model is not None: - tokenize = lambda batch: st_model.tokenize(batch) # noqa: E731 - else: - tokenize = lambda batch: tokenizer(batch, padding=True, return_tensors="pt") # noqa: E731 - return {"st": st_model, "hf": hf, "ours": ours, "texts": texts, "tokenize": tokenize} - - -def test_real_weight_key_coverage(real_setup, embeddinggemma2_snapshot): - from safetensors import safe_open - - with safe_open(str(embeddinggemma2_snapshot / "model.safetensors"), framework="pt") as handle: - all_keys = list(handle.keys()) # noqa: SIM118 - safe_open is not a Mapping - hf_text_keys = { - k[len("language_model.") :] for k in all_keys if k.startswith("language_model.") - } - ours = set(real_setup["ours"].state_dict().keys()) - assert ours == hf_text_keys, (sorted(ours - hf_text_keys)[:5], sorted(hf_text_keys - ours)[:5]) + tokenize = lambda batch: tokenizer(batch, padding=True, return_tensors="pt") # noqa: E731 + return {"hf": hf, "ours": ours, "texts": texts, "tokenize": tokenize} def test_real_weight_parity_and_padding_invariance(real_setup): - ours, hf, texts = real_setup["ours"], real_setup["hf"], real_setup["texts"] + """Includes a row over 2 * 512 + 1 tokens, so the sliding layers' window bites for real.""" + ours, hf = real_setup["ours"], real_setup["hf"] + paragraph = ( + "On-device semantic search embeds documents once and compares vectors at query time. " + ) + texts = real_setup["texts"] + [real_setup["texts"][1] + paragraph * 90] features = real_setup["tokenize"](texts) ids = features["input_ids"].to(torch.int32) mask = features["attention_mask"].to(torch.int32) assert int(mask.sum(1).min()) < int(mask.sum(1).max()), "batch must be mixed-length" + assert ids.shape[1] > 1100 with torch.inference_mode(): hf.language_model.config._attn_implementation = "sdpa" hf_tokens = hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state ref = _st_mean_pool_normalize(hf_tokens, mask) - if real_setup["st"] is not None: - st_emb = real_setup["st"].encode(texts, batch_size=len(texts), convert_to_tensor=True) - assert _cos(st_emb, ref).min().item() > 0.99999 # our pooling re-impl == ST - ref = st_emb - emb, tokens = ours(input_ids=ids, attention_mask=mask, return_token_states=True) + emb = ours(input_ids=ids, attention_mask=mask) alone = torch.cat( [ ours( @@ -332,17 +246,8 @@ def test_real_weight_parity_and_padding_invariance(real_setup): ] ) - cos = _cos(emb, ref) - emb_max_abs = (emb - ref).abs().max().item() - pad_cos = _cos(alone, emb) - print( - f"\nEG2 real fp32 CPU: n={len(texts)} lens={mask.sum(1).tolist()} " - f"cos min={cos.min().item():.9f} median={cos.median().item():.9f} " - f"emb max_abs={emb_max_abs:.3e} " - f"padding-invariance cos min={pad_cos.min().item():.9f}" - ) - assert cos.min().item() >= 0.99999 - assert pad_cos.min().item() >= 0.99999 + assert _cos(emb, ref).min().item() >= 0.99999 + assert _cos(alone, emb).min().item() >= 0.99999 def test_real_weight_excluding_prompt_changes_embedding(real_setup): @@ -360,32 +265,4 @@ def test_real_weight_excluding_prompt_changes_embedding(real_setup): with torch.inference_mode(): a = ours(input_ids=ids, attention_mask=mask) b = ours(input_ids=ids, attention_mask=mask, pooling_mask=no_prompt) - cos = _cos(a, b).item() - print(f"\nEG2 include_prompt vs exclude_prompt cos={cos:.6f} (prompt_len={prompt_len})") - assert cos < 0.999 - - -def test_real_weight_long_sequence_window_bites(real_setup): - """> 2 * 512 + 1 tokens so the sliding layers' |i - j| <= 512 window is exercised for real.""" - ours, hf = real_setup["ours"], real_setup["hf"] - paragraph = ( - "On-device semantic search embeds documents once and compares vectors at query time. " - ) - features = real_setup["tokenize"]( - [real_setup["texts"][1] + paragraph * 90, real_setup["texts"][0]] - ) - ids = features["input_ids"].to(torch.int32) - mask = features["attention_mask"].to(torch.int32) - assert ids.shape[1] > 1100 - with torch.inference_mode(): - hf.language_model.config._attn_implementation = "sdpa" - ref = _st_mean_pool_normalize( - hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state, mask - ) - emb = ours(input_ids=ids, attention_mask=mask) - cos = _cos(emb, ref) - print( - f"\nEG2 long L={ids.shape[1]} lens={mask.sum(1).tolist()} cos={cos.tolist()} " - f"max_abs={(emb - ref).abs().max().item():.3e}" - ) - assert cos.min().item() >= 0.99999 + assert _cos(a, b).item() < 0.999 diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py index 076b0112..2d9e31ea 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py @@ -5,7 +5,7 @@ """Parity tests for the EmbeddingGemma 2 image path (``embedding_gemma2_vision.py``). -* Tiny-config parity against HF ``Gemma4VisionModel`` / ``Gemma4MultimodalEmbedder`` +* Tiny-config parity against HF ``Gemma4VisionModel`` + ``Gemma4MultimodalEmbedder`` (present in the repo's transformers; no checkpoint needed). * Host preprocessing helpers against HF ``Gemma4ImageProcessor``. * Scatter / sequence assembly against HF ``masked_scatter`` semantics. @@ -19,7 +19,6 @@ """ import math -import os import pytest import torch @@ -89,21 +88,15 @@ def _hf_tiny_vision(cfg: EmbeddingGemma2VisionConfig): return hf_cfg, model -def _load_from_hf_vision(ours: EmbeddingGemma2ImageEncoder, hf_vision, hf_embedder=None): +def _load_from_hf_vision(ours: EmbeddingGemma2ImageEncoder, hf_vision, hf_embedder): state = {} for k, v in hf_vision.state_dict().items(): local = EmbeddingGemma2ImageEncoder.map_hf_key("vision_tower." + k) if local is not None and local in dict(ours.named_parameters()): state[local] = v - if hf_embedder is not None: - for k, v in hf_embedder.state_dict().items(): - state["embed_vision." + k] = v - missing, unexpected = ours.load_state_dict(state, strict=False) - expected_missing = ( - [] if hf_embedder is not None else ["embed_vision.embedding_projection.weight"] - ) - assert sorted(missing) == expected_missing - assert not unexpected + for k, v in hf_embedder.state_dict().items(): + state["embed_vision." + k] = v + ours.load_state_dict(state) def _random_inputs(cfg: EmbeddingGemma2VisionConfig, grids: list[tuple[int, int]], seed: int = 0): @@ -126,44 +119,28 @@ def _random_inputs(cfg: EmbeddingGemma2VisionConfig, grids: list[tuple[int, int] # --- Tiny-config parity vs HF --- -@pytest.mark.parametrize("use_hf_sdpa", ["true", "false"]) -@pytest.mark.parametrize("grids", [[(3, 4)], [(2, 3)], [(1, 5), (4, 3)]]) -def test_tiny_vision_tower_matches_hf(grids, use_hf_sdpa, monkeypatch) -> None: - # The directory conftest forces USE_HF_IMPL=true; also cover the Core AI lowering path - # of the SDPA primitive (bool-mask -> -1e4 bias), which is what gets exported. - monkeypatch.setenv("USE_HF_IMPL", use_hf_sdpa) - _, hf = _hf_tiny_vision(TINY) - ours = EmbeddingGemma2ImageEncoder(TINY).eval() - _load_from_hf_vision(ours, hf) - pixel_values, position_ids, counts = _random_inputs(TINY, grids) - with torch.no_grad(): - ref = hf(pixel_values=pixel_values, pixel_position_ids=position_ids).last_hidden_state - got = ours.encode_pooled(pixel_values, position_ids) - # HF strips padding slots and concatenates images; ours keeps the static budget. - got_valid = torch.cat([got[i, :n] for i, n in enumerate(counts)]) - assert got_valid.shape == ref.shape - torch.testing.assert_close(got_valid, ref.float(), rtol=1e-5, atol=1e-5) - # Padding slots are exactly zero. - for i, n in enumerate(counts): - assert torch.count_nonzero(got[i, n:]) == 0 - - -def test_tiny_soft_tokens_match_hf_embedder() -> None: +def test_tiny_soft_tokens_match_hf(monkeypatch) -> None: + # The directory conftest forces USE_HF_IMPL=true; run the Core AI lowering path of the SDPA + # primitive (bool-mask -> -1e4 bias), which is what gets exported. + monkeypatch.setenv("USE_HF_IMPL", "false") hf_cfg, hf = _hf_tiny_vision(TINY) text_cfg = gemma4_config.Gemma4TextConfig(hidden_size=TINY.text_hidden_size) torch.manual_seed(1) embedder = gemma4_modeling.Gemma4MultimodalEmbedder(hf_cfg, text_cfg).eval() ours = EmbeddingGemma2ImageEncoder(TINY).eval() _load_from_hf_vision(ours, hf, embedder) - pixel_values, position_ids, counts = _random_inputs(TINY, [(4, 3)], seed=3) + pixel_values, position_ids, counts = _random_inputs(TINY, [(1, 5), (4, 3)]) with torch.no_grad(): ref = embedder( hf(pixel_values=pixel_values, pixel_position_ids=position_ids).last_hidden_state ) got = ours(pixel_values, position_ids) - assert got.shape == (1, TINY.max_soft_tokens, TINY.text_hidden_size) - torch.testing.assert_close(got[0, : counts[0]], ref, rtol=1e-5, atol=1e-5) - assert torch.count_nonzero(got[0, counts[0] :]) == 0 + assert got.shape == (2, TINY.max_soft_tokens, TINY.text_hidden_size) + # HF strips padding slots and concatenates images; ours keeps the static budget. + got_valid = torch.cat([got[i, :n] for i, n in enumerate(counts)]) + torch.testing.assert_close(got_valid, ref, rtol=1e-5, atol=1e-5) + for i, n in enumerate(counts): + assert torch.count_nonzero(got[i, n:]) == 0 def test_rejects_wrong_patch_count() -> None: @@ -336,15 +313,9 @@ def eg2_hf(embeddinggemma2_snapshot): snapshot, dtype=torch.float32, audio_config=None ).eval() # Build ours on the SDPA lowering path that gets exported (not the HF-eager shim). - saved = os.environ.get("USE_HF_IMPL") - os.environ["USE_HF_IMPL"] = "false" - try: + with pytest.MonkeyPatch.context() as patch: + patch.setenv("USE_HF_IMPL", "false") ours = EmbeddingGemma2ImageEncoder.from_pretrained(snapshot, dtype=torch.float32) - finally: - if saved is None: - os.environ.pop("USE_HF_IMPL", None) - else: - os.environ["USE_HF_IMPL"] = saved return processor, model, ours From 5b3db09223f5f94f7940a92847d40dfa2c3a5b90 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 17:44:02 -0700 Subject: [PATCH 32/38] Clarify embedding helper names and derive bundle assets from the pipelines 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. --- python/src/coreai_models/diffusion/gpu.py | 8 +-- .../src/coreai_models/embedding/contract.py | 57 ++++++++++--------- python/src/coreai_models/embedding/export.py | 6 +- .../src/coreai_models/embedding/export_ios.py | 4 +- .../coreai_models/embedding/export_macos.py | 4 +- .../models/ios/embedding_gemma2.py | 30 ++++------ .../test_export/test_embedding_export.py | 18 +++--- 7 files changed, 62 insertions(+), 65 deletions(-) diff --git a/python/src/coreai_models/diffusion/gpu.py b/python/src/coreai_models/diffusion/gpu.py index add41125..6cfdea64 100644 --- a/python/src/coreai_models/diffusion/gpu.py +++ b/python/src/coreai_models/diffusion/gpu.py @@ -82,7 +82,7 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram: decomposed: torch.export.ExportedProgram = exported.run_decompositions(coreai_decomp_table) return decomposed - converter = new_converter(include_debug_info) + converter = make_converter(include_debug_info) converter.add_pytorch_module( wrapper, export_fn=export_fn, @@ -93,10 +93,10 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram: return program -def new_converter( +def make_converter( include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, ) -> coreai_torch.TorchConverter: - """A converter in ``DEBUG`` mode (full debug info) or ``RELEASE`` mode (smaller asset).""" + """Return a ``TorchConverter``; ``include_debug_info`` keeps debug info in the asset.""" mode = coreai_torch.TorchConverter.Mode return coreai_torch.TorchConverter(mode=mode.DEBUG if include_debug_info else mode.RELEASE) @@ -152,7 +152,7 @@ def export_multifunction( Returns: An optimized AIProgram with multiple named functions. """ - converter = new_converter(include_debug_info) + converter = make_converter(include_debug_info) for name, wrapper, dummy_inputs in functions: logger.info(f"Tracing function '{name}' (seq dims: {[t.shape for t in dummy_inputs[:2]]})") add_function(converter, name, wrapper, dummy_inputs, input_names, output_names) diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py index 37706b26..5b98cfaf 100644 --- a/python/src/coreai_models/embedding/contract.py +++ b/python/src/coreai_models/embedding/contract.py @@ -41,10 +41,10 @@ class EmbeddingPrecisionError(ValueError): @dataclass(frozen=True) class ImageTowerFacts: - """One family's image-tower facts for the host (merge ids, preprocessing, soft-token budget). + """Image-tower facts the host needs: merge ids, preprocessing and the soft-token budget. The exporter checks them against the snapshot ``config.json`` / ``processor_config.json`` - and fails closed on a mismatch. + and raises on a mismatch. """ placeholder_id: int @@ -168,7 +168,7 @@ def default_precision(family: EmbeddingFamily, platform: str) -> str: return "float16" if platform == "iOS" else family.default_compute_precision -def check_compute_precision( +def validate_compute_precision( family: EmbeddingFamily, precision: str, *, platform: str = "macOS" ) -> str: """Validate ``precision`` for ``family`` on ``platform`` and return it. @@ -282,25 +282,24 @@ def bundle_name(family: EmbeddingFamily, platform: str, with_image: bool) -> str return family.key + ("-ios" if platform == "iOS" else "") + ("" if with_image else "-text") +IMAGE_PIPELINE = (IMAGE_COMPONENT, TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT) + + +def text_pipeline(platform: str) -> tuple[str, ...]: + """Components the text tower runs, in order.""" + return (TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT) if platform == "iOS" else (TEXT_COMPONENT,) + + def bundle_assets(platform: str, with_image: bool) -> dict[str, str]: - """Component name -> asset file for the bundle. + """Component name -> asset file for every component of the bundle's pipelines. - In the macOS image bundle, ``token_table`` and ``encoder`` are functions of - ``text_encoder.aimodel``, so the table and the transformer are stored once. + On macOS, ``token_table`` and ``encoder`` are functions of ``text_encoder.aimodel``. """ - if platform == "iOS": - components = [TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT] - if with_image: - components.append(IMAGE_COMPONENT) - return {name: f"{name}.aimodel" for name in components} - text = f"{TEXT_COMPONENT}.aimodel" - if not with_image: - return {TEXT_COMPONENT: text} + components = text_pipeline(platform) + (IMAGE_PIPELINE if with_image else ()) + shared = () if platform == "iOS" else (TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT) return { - TEXT_COMPONENT: text, - IMAGE_COMPONENT: f"{IMAGE_COMPONENT}.aimodel", - TOKEN_TABLE_COMPONENT: text, - ENCODER_COMPONENT: text, + name: f"{TEXT_COMPONENT if name in shared else name}.aimodel" + for name in dict.fromkeys(components) } @@ -316,17 +315,21 @@ def embedding_block( ) -> dict[str, Any]: """Build the ``embedding`` block of ``metadata.json``. - macOS: a ``text_encoder`` graph (ids in, vector out), plus the image tower - ``image_encoder -> token_table -> encoder`` when ``with_image``; ``token_table`` and - ``encoder`` are functions of ``text_encoder.aimodel``. iOS: the text tower runs - ``token_table -> encoder`` on the Neural Engine. The host merges soft tokens into the table - output at the placeholder rows before the encoder. ``compute_precision`` defaults to the - platform's precision. + Text runs ``text_encoder`` on macOS and ``token_table -> encoder`` on iOS. With + ``with_image``, images run ``image_encoder -> token_table -> encoder``, and the host writes + the soft tokens into the ``token_table`` output at the placeholder rows. ``token_table`` runs + on the GPU; on iOS the encoders run on the Neural Engine. ``compute_precision`` defaults to + the platform's precision. + + Raises: + ValueError: on an unknown platform, a missing revision or prompts, an unknown + ``default_prompt_name``, or family facts that do not fit the function plan. + EmbeddingPrecisionError: if ``compute_precision`` is not allowed on ``platform``. """ ios = platform == "iOS" if platform not in PLATFORMS: raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") - precision = check_compute_precision( + precision = validate_compute_precision( family, compute_precision or default_precision(family, platform), platform=platform ) if not revision: @@ -383,7 +386,7 @@ def embedding_block( text_tower: dict[str, Any] = { "modality": "text", - "pipeline": [TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT] if ios else [TEXT_COMPONENT], + "pipeline": list(text_pipeline(platform)), "inputs": {"input_ids": "int32[B,L]", "attention_mask": "int32[B,L]"}, "output": OUTPUT_NAME, "batch_sizes": sorted({f.batch for f in text_fns}), @@ -402,7 +405,7 @@ def embedding_block( towers.append( { "modality": "image", - "pipeline": [IMAGE_COMPONENT, TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT], + "pipeline": list(IMAGE_PIPELINE), "output": OUTPUT_NAME, "batch_sizes": [1], "length_buckets": image_lengths, diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index 710e86cc..72a0d654 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -32,13 +32,13 @@ build_bundle_metadata, bundle_assets, bundle_name, - check_compute_precision, consolidate_prompts, default_precision, embedding_block, encoder_functions, family_for_model, text_functions, + validate_compute_precision, ) from coreai_models.embedding.export_ios import ( export_ane_programs, @@ -90,7 +90,7 @@ def export_embedding_bundle( """ if platform not in PLATFORMS: raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") - precision = check_compute_precision( + precision = validate_compute_precision( family, compute_precision or default_precision(family, platform), platform=platform ) # before touching disk if with_image and family.image is None: @@ -231,7 +231,7 @@ def main(argv: list[str] | None = None) -> int: platform = args.platform with_image = not args.skip_vision try: - precision = check_compute_precision( + precision = validate_compute_precision( family, args.compute_precision or default_precision(family, platform), platform=platform, diff --git a/python/src/coreai_models/embedding/export_ios.py b/python/src/coreai_models/embedding/export_ios.py index 62c4c9f3..3575771e 100644 --- a/python/src/coreai_models/embedding/export_ios.py +++ b/python/src/coreai_models/embedding/export_ios.py @@ -92,10 +92,10 @@ def convert_ane_functions( Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. """ - from coreai_models.diffusion.gpu import add_function, new_converter + from coreai_models.diffusion.gpu import add_function, make_converter from coreai_models.export.mlir_ops import register_custom_torch_lowering - converter = new_converter(include_debug_info) + converter = make_converter(include_debug_info) register_custom_torch_lowering(converter) decompositions = torch.export.default_decompositions() for name, module, args, input_names, output_names in functions: diff --git a/python/src/coreai_models/embedding/export_macos.py b/python/src/coreai_models/embedding/export_macos.py index 73bfdc10..e09eab84 100644 --- a/python/src/coreai_models/embedding/export_macos.py +++ b/python/src/coreai_models/embedding/export_macos.py @@ -128,9 +128,9 @@ def export_shared_text_program( ``main`` maps ids to ``inputs_embeds`` (dynamic length up to ``max_length``); ``embeds_*`` functions take ``inputs_embeds``. The functions share one copy of the table and the transformer. """ - from coreai_models.diffusion.gpu import add_function, new_converter + from coreai_models.diffusion.gpu import add_function, make_converter - converter = new_converter(include_debug_info) + converter = make_converter(include_debug_info) text = TextEmbeddingExportWrapper(encoder) for spec in text_fns: args = reference_inputs(spec, int(encoder.config.vocab_size)) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index fa5a6f39..3cdc5fc2 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -118,22 +118,16 @@ def __init__(self, config: EmbeddingGemma2TextConfig) -> None: self.layers = nn.ModuleList([EmbeddingGemma2EncoderLayer(config, i) for i in range(L)]) self.norm = MaxAbsRMSNorm(dim, eps, axis=-3) self.embedding_projection = conv1x1(dim, config.embedding_dim) - self._rope_hd = {"sliding": config.rope_head_dim(True), "full": config.rope_head_dim(False)} - - def rope( - self, length: int, dtype: torch.dtype, device - ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: - out = {} - for kind, theta in ( - ("sliding", self.config.rope_theta_sliding), - ("full", self.config.rope_theta_full), - ): - hd = self._rope_hd[kind] - inv = rope_inv_freq(hd, theta, device) - fr = inv[:, None] * torch.arange(length, dtype=torch.float32, device=device)[None, :] - e = torch.cat([fr, fr], 0).reshape(1, 1, hd, length) - out[kind] = (e.cos().to(dtype), e.sin().to(dtype)) - return out + + def _cos_sin(self, sliding: bool, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """cos/sin ``[1, 1, head_dim, L]`` in ``x``'s dtype for the sliding or full layers.""" + head_dim = self.config.rope_head_dim(sliding) + theta = self.config.rope_theta_sliding if sliding else self.config.rope_theta_full + length = x.shape[-1] + inv_freq = rope_inv_freq(head_dim, theta, x.device)[:, None] + freqs = inv_freq * torch.arange(length, dtype=torch.float32, device=x.device)[None, :] + emb = torch.cat([freqs, freqs], 0).reshape(1, 1, head_dim, length) + return emb.cos().to(x.dtype), emb.sin().to(x.dtype) def forward( self, @@ -148,10 +142,10 @@ def forward( ple = self.per_layer_model_projection(x) * (dim**-0.5) ple = self.per_layer_projection_norm(ple.reshape(b, nl, self.ple_dim, length)) ple = ple.reshape(b, nl * self.ple_dim, 1, length) - ropes = self.rope(length, x.dtype, x.device) + ropes = {sliding: self._cos_sin(sliding, x) for sliding in (True, False)} sb = bias if sliding_bias is None else sliding_bias for i, layer in enumerate(self.layers): - cos, sin = ropes["sliding" if layer.is_sliding else "full"] + cos, sin = ropes[layer.is_sliding] ple_i = ple[:, i * self.ple_dim : (i + 1) * self.ple_dim] x = layer(x, ple_i, cos, sin, sb if layer.is_sliding else bias) x = self.norm(x) # (B, H, 1, L) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index e5364fdd..d001249c 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -29,12 +29,12 @@ build_bundle_metadata, bundle_assets, bundle_name, - check_compute_precision, consolidate_prompts, embedding_block, encoder_functions, family_for_model, text_functions, + validate_compute_precision, ) from tests.test_model_units.test_models._embedding_gemma2_common import ( TINY_TEXT_CONFIG, @@ -83,22 +83,22 @@ def _block(platform="macOS", with_image=False, **overrides): def test_ios_requires_float16(): - assert check_compute_precision(EG2, "float16", platform="iOS") == "float16" + assert validate_compute_precision(EG2, "float16", platform="iOS") == "float16" for precision in ("float32", "bfloat16"): with pytest.raises(EmbeddingPrecisionError, match="float16"): - check_compute_precision(EG2, precision, platform="iOS") + validate_compute_precision(EG2, precision, platform="iOS") def test_macos_refuses_float16_and_allows_bfloat16(): with pytest.raises(EmbeddingPrecisionError, match="forbids float16"): - check_compute_precision(EG2, "float16") - assert check_compute_precision(EG2, "bfloat16") == "bfloat16" - assert check_compute_precision(EG2, "float32") == "float32" + validate_compute_precision(EG2, "float16") + assert validate_compute_precision(EG2, "bfloat16") == "bfloat16" + assert validate_compute_precision(EG2, "float32") == "float32" def test_unknown_precision_is_refused(): with pytest.raises(EmbeddingPrecisionError, match="Unknown compute precision"): - check_compute_precision(EG2, "int8") + validate_compute_precision(EG2, "int8") def test_export_bundle_refuses_float16_before_touching_out(tmp_path): @@ -359,10 +359,10 @@ def fail(*args, **kwargs): def test_forbidden_precision_message_lists_the_allowed_ones(): with pytest.raises(EmbeddingPrecisionError, match=r"Use float32 \(default\) or bfloat16\."): - check_compute_precision(EG2, "float16") + validate_compute_precision(EG2, "float16") no_floor = dataclasses.replace(EG2, precision_floor=None) with pytest.raises(EmbeddingPrecisionError, match=r"Use float32 \(default\)\.$"): - check_compute_precision(no_floor, "float16") + validate_compute_precision(no_floor, "float16") def test_family_registry(): From ac7fe2a5a076abef33ed7d46f4a5f80c28b87954 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 22:56:38 -0700 Subject: [PATCH 33/38] Make the base.py safetensors helpers public Rename _resolve_safetensors_files and _load_tensors_for_keys to resolve_safetensors_files and load_tensors_for_keys and switch every caller. --- python/src/coreai_models/models/base.py | 10 +++++----- .../models/macos/diffusion_gemma.py | 10 +++++----- python/src/coreai_models/models/macos/gemma3n.py | 8 ++++---- .../coreai_models/models/macos/muse_glimmer.py | 16 ++++++++-------- .../models/macos/muse_glimmer_drafter_ring.py | 12 ++++++------ .../models/macos/muse_glimmer_vision.py | 10 +++++----- .../test_macos_layers/test_diffusion_gemma.py | 4 ++-- 7 files changed, 35 insertions(+), 35 deletions(-) diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index f583b208..cbc12a1e 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -231,7 +231,7 @@ def _save_and_mmap_safetensors( module.load_state_dict(new_sd, assign=True, strict=False) -def _resolve_safetensors_files(model_dir: str) -> list[str]: +def resolve_safetensors_files(model_dir: str) -> list[str]: """Resolve the safetensors file paths in a local HuggingFace model directory.""" single_path = os.path.join(model_dir, "model.safetensors") index_path = os.path.join(model_dir, "model.safetensors.index.json") @@ -306,7 +306,7 @@ def _build_safetensors_key_index( return per_layer, shared -def _load_tensors_for_keys( +def load_tensors_for_keys( key_to_file: dict[str, str], target_dtype: torch.dtype, ) -> dict[str, torch.Tensor]: @@ -876,7 +876,7 @@ def from_hf_memory_efficient( model = cls(**init_kwargs) model.to(dtype=target_dtype) - safetensors_files = _resolve_safetensors_files(model_dir) + safetensors_files = resolve_safetensors_files(model_dir) per_layer_index, shared_index = _build_safetensors_key_index( safetensors_files, num_layers=num_layers, @@ -885,7 +885,7 @@ def from_hf_memory_efficient( ) # Shared params first (embeddings, norm, lm_head, ...). - shared_dict = _load_tensors_for_keys(shared_index, target_dtype) + shared_dict = load_tensors_for_keys(shared_index, target_dtype) shared_dict = {k.removeprefix(hf_state_dict_prefix): v for k, v in shared_dict.items()} del shared_index @@ -903,7 +903,7 @@ def from_hf_memory_efficient( exclude_buffers = {KVCache.HF_K_BUFFER_NAME, KVCache.HF_V_BUFFER_NAME} for layer_idx in sorted(per_layer_index.keys()): layer_key_to_file = per_layer_index.pop(layer_idx) - layer_sd = _load_tensors_for_keys(layer_key_to_file, target_dtype) + layer_sd = load_tensors_for_keys(layer_key_to_file, target_dtype) layer_sd = {k.removeprefix(hf_state_dict_prefix): v for k, v in layer_sd.items()} del layer_key_to_file diff --git a/python/src/coreai_models/models/macos/diffusion_gemma.py b/python/src/coreai_models/models/macos/diffusion_gemma.py index b9f52905..82e60da0 100644 --- a/python/src/coreai_models/models/macos/diffusion_gemma.py +++ b/python/src/coreai_models/models/macos/diffusion_gemma.py @@ -823,20 +823,20 @@ def _load_state_dict_from_hub( from coreai_models.models.base import ( _build_safetensors_key_index, - _load_tensors_for_keys, - _resolve_safetensors_files, + load_tensors_for_keys, + resolve_safetensors_files, ) model_dir = snapshot_download( hf_model_id, allow_patterns=["*.safetensors", "*.safetensors.index.json"] ) - sf_files = _resolve_safetensors_files(model_dir) + sf_files = resolve_safetensors_files(model_dir) per_layer, shared = _build_safetensors_key_index( sf_files, num_layers=num_layers, hf_state_dict_prefix="" ) - sd: dict[str, torch.Tensor] = _load_tensors_for_keys(shared, target_dtype) + sd: dict[str, torch.Tensor] = load_tensors_for_keys(shared, target_dtype) for layer_idx in sorted(per_layer.keys()): - sd.update(_load_tensors_for_keys(per_layer[layer_idx], target_dtype)) + sd.update(load_tensors_for_keys(per_layer[layer_idx], target_dtype)) return sd diff --git a/python/src/coreai_models/models/macos/gemma3n.py b/python/src/coreai_models/models/macos/gemma3n.py index 0dbd6153..c2a0d5df 100644 --- a/python/src/coreai_models/models/macos/gemma3n.py +++ b/python/src/coreai_models/models/macos/gemma3n.py @@ -495,7 +495,7 @@ def from_hf_memory_efficient( from huggingface_hub import snapshot_download from safetensors import safe_open - from coreai_models.models.base import _load_tensors_for_keys, _resolve_safetensors_files + from coreai_models.models.base import load_tensors_for_keys, resolve_safetensors_files model_dir = snapshot_download( huggingface_model_id, @@ -511,7 +511,7 @@ def from_hf_memory_efficient( model = cls(config=config, model_device="meta") model.to(dtype=target_dtype) - safetensors_files = _resolve_safetensors_files(model_dir) + safetensors_files = resolve_safetensors_files(model_dir) prefix = hf_state_dict_prefix layer_pattern = re.compile(re.escape(prefix) + r"layers\.(\d+)\.") @@ -532,7 +532,7 @@ def from_hf_memory_efficient( else: shared[key] = path - shared_dict = _load_tensors_for_keys(shared, target_dtype) + shared_dict = load_tensors_for_keys(shared, target_dtype) normalized: dict[str, torch.Tensor] = {} for k, v in shared_dict.items(): if k.startswith(prefix): @@ -547,7 +547,7 @@ def from_hf_memory_efficient( for layer_idx in sorted(per_layer.keys()): layer_key_to_file = per_layer.pop(layer_idx) - layer_sd = _load_tensors_for_keys(layer_key_to_file, target_dtype) + layer_sd = load_tensors_for_keys(layer_key_to_file, target_dtype) del layer_key_to_file remapped: dict[str, torch.Tensor] = {} for k, v in layer_sd.items(): diff --git a/python/src/coreai_models/models/macos/muse_glimmer.py b/python/src/coreai_models/models/macos/muse_glimmer.py index f2221fde..0e42c739 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer.py +++ b/python/src/coreai_models/models/macos/muse_glimmer.py @@ -46,8 +46,8 @@ from coreai_models.models.base import ( BaseForCausalLM, TraceSpec, - _load_tensors_for_keys, - _resolve_safetensors_files, + load_tensors_for_keys, + resolve_safetensors_files, ) from coreai_models.primitives.macos.cache import ( KVCache, @@ -381,7 +381,7 @@ def from_hf_memory_efficient( model = cls(config=config, model_device="meta") model.to(dtype=target_dtype) - safetensors_files = _resolve_safetensors_files(model_dir) + safetensors_files = resolve_safetensors_files(model_dir) # Build key index with Muse Glimmer's actual key layout: # model.language_model.layers.N.* → per-layer @@ -408,7 +408,7 @@ def from_hf_memory_efficient( shared[key] = path # Load shared params (embed_tokens, norm, lm_head) - shared_dict = _load_tensors_for_keys(shared, target_dtype) + shared_dict = load_tensors_for_keys(shared, target_dtype) # Normalize keys: strip "model.language_model." prefix where present normalized: dict[str, torch.Tensor] = {} prefix = "model.language_model." @@ -425,7 +425,7 @@ def from_hf_memory_efficient( # Load one layer at a time for layer_idx in sorted(per_layer.keys()): layer_key_to_file = per_layer.pop(layer_idx) - layer_sd = _load_tensors_for_keys(layer_key_to_file, target_dtype) + layer_sd = load_tensors_for_keys(layer_key_to_file, target_dtype) del layer_key_to_file # Strip prefix → "layers.N.*", then add "model." → "model.layers.N.*" remapped: dict[str, torch.Tensor] = {} @@ -855,7 +855,7 @@ def from_hf_memory_efficient( model = cls(config=config, model_device="meta") model.to(dtype=target_dtype) - safetensors_files = _resolve_safetensors_files(model_dir) + safetensors_files = resolve_safetensors_files(model_dir) layer_pattern = re.compile(r"model\.language_model\.layers\.(\d+)\.") from safetensors import safe_open @@ -877,7 +877,7 @@ def from_hf_memory_efficient( shared[key] = path # Load shared params (norm, lm_head — embed_tokens dropped below) - shared_dict = _load_tensors_for_keys(shared, target_dtype) + shared_dict = load_tensors_for_keys(shared, target_dtype) normalized: dict[str, torch.Tensor] = {} prefix = "model.language_model." for k, v in shared_dict.items(): @@ -902,7 +902,7 @@ def from_hf_memory_efficient( # Load one layer at a time for layer_idx in sorted(per_layer.keys()): layer_key_to_file = per_layer.pop(layer_idx) - layer_sd = _load_tensors_for_keys(layer_key_to_file, target_dtype) + layer_sd = load_tensors_for_keys(layer_key_to_file, target_dtype) del layer_key_to_file remapped: dict[str, torch.Tensor] = {} for k, v in layer_sd.items(): diff --git a/python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py b/python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py index 8b460d76..9d0b93cf 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py +++ b/python/src/coreai_models/models/macos/muse_glimmer_drafter_ring.py @@ -38,8 +38,8 @@ from coreai_models.models.base import ( BaseForCausalLM, TraceSpec, - _load_tensors_for_keys, - _resolve_safetensors_files, + load_tensors_for_keys, + resolve_safetensors_files, ) from coreai_models.primitives.macos.cache import RingKVCache, concat_window_causal_mask from coreai_models.primitives.macos.mlp import MLP @@ -296,7 +296,7 @@ def from_hf( model.to(dtype=target_dtype) # ---- 4. Load drafter weights (skip encoder.*, add model. prefix) ----- - drafter_files = _resolve_safetensors_files(drafter_dir) + drafter_files = resolve_safetensors_files(drafter_dir) drafter_keys: dict[str, str] = {} for path in drafter_files: with safe_open(path, framework="pt", device="cpu") as f: @@ -309,7 +309,7 @@ def from_hf( continue drafter_keys[key] = path - drafter_sd = _load_tensors_for_keys(drafter_keys, target_dtype) + drafter_sd = load_tensors_for_keys(drafter_keys, target_dtype) # "layers.0.*" → "model.layers.0.*", "norm.weight" → "model.norm.weight" remapped: dict[str, torch.Tensor] = {} for k, v in drafter_sd.items(): @@ -320,7 +320,7 @@ def from_hf( gc.collect() # ---- 5. Load embed_tokens and lm_head from target (2 tensors) ------- - target_files = _resolve_safetensors_files(target_dir) + target_files = resolve_safetensors_files(target_dir) embed_hf_key = "model.language_model.embed_tokens.weight" lm_head_hf_key = "lm_head.weight" target_keys: dict[str, str] = {} @@ -337,7 +337,7 @@ def from_hf( f"found: {found}" ) - target_sd = _load_tensors_for_keys(target_keys, target_dtype) + target_sd = load_tensors_for_keys(target_keys, target_dtype) shared: dict[str, torch.Tensor] = { "model.embed_tokens.weight": target_sd[embed_hf_key], "lm_head.weight": target_sd[lm_head_hf_key], diff --git a/python/src/coreai_models/models/macos/muse_glimmer_vision.py b/python/src/coreai_models/models/macos/muse_glimmer_vision.py index 679dc9f5..6b15aaca 100644 --- a/python/src/coreai_models/models/macos/muse_glimmer_vision.py +++ b/python/src/coreai_models/models/macos/muse_glimmer_vision.py @@ -55,8 +55,8 @@ from safetensors import safe_open from coreai_models.models.base import ( - _load_tensors_for_keys, - _resolve_safetensors_files, + load_tensors_for_keys, + resolve_safetensors_files, ) # --------------------------------------------------------------------------- @@ -721,7 +721,7 @@ def from_pretrained( model = model.to(dtype=dtype) # --- 4. Load vision weights from safetensors --- - safetensors_files = _resolve_safetensors_files(model_dir) + safetensors_files = resolve_safetensors_files(model_dir) # Build key->file index for vision-only keys vision_key_to_file: dict[str, str] = {} @@ -736,7 +736,7 @@ def from_pretrained( vision_key_to_file[key] = path # Load all vision tensors (~ 1.8B params, fits in RAM) - raw_sd = _load_tensors_for_keys(vision_key_to_file, dtype) + raw_sd = load_tensors_for_keys(vision_key_to_file, dtype) # --- 5. Remap keys to our module structure --- mapped_sd: dict[str, torch.Tensor] = {} @@ -759,7 +759,7 @@ def from_pretrained( # _rope_inv_freq is non-persistent, not in the checkpoint model.load_state_dict(mapped_sd, strict=False, assign=True) - # _load_tensors_for_keys skips casting for "embedding_table" keys, + # load_tensors_for_keys skips casting for "embedding_table" keys, # so force everything to the target dtype now. model = model.to(dtype=dtype) diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_diffusion_gemma.py b/python/tests/test_model_units/test_models/test_macos_layers/test_diffusion_gemma.py index b85a9c73..44e2e39a 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_diffusion_gemma.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_diffusion_gemma.py @@ -373,9 +373,9 @@ def test_load_state_dict_from_hub_mocked() -> None: with ( mock.patch("huggingface_hub.snapshot_download", return_value="/tmp/x"), - mock.patch.object(base, "_resolve_safetensors_files", return_value=[]), + mock.patch.object(base, "resolve_safetensors_files", return_value=[]), mock.patch.object(base, "_build_safetensors_key_index", return_value=({}, {})), - mock.patch.object(base, "_load_tensors_for_keys", return_value={}), + mock.patch.object(base, "load_tensors_for_keys", return_value={}), ): sd = dgm._load_state_dict_from_hub("some/model", torch.float32, num_layers=2) assert sd == {} From d979627130b570a678dda4f8be36147cd0f94d29 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 22:56:38 -0700 Subject: [PATCH 34/38] Read EmbeddingGemma 2 checkpoints through the base.py safetensors helpers read_checkpoint_tensors resolves model.safetensors or the index.json shards with resolve_safetensors_files and loads with load_tensors_for_keys. --- .../models/ios/embedding_gemma2_vision.py | 2 +- .../models/macos/embedding_gemma2.py | 26 ++++++++++--------- .../models/macos/embedding_gemma2_vision.py | 2 +- .../test_embedding_gemma2.py | 24 +++++++++++++++++ 4 files changed, 40 insertions(+), 14 deletions(-) diff --git a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py index 51cd2c93..df08108d 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -184,7 +184,7 @@ def load_image_encoder( dtype: torch.dtype = torch.float16, max_soft_tokens: int | None = None, ) -> EmbeddingGemma2ImageEncoder: - """Load from a local EG2 snapshot directory (``config.json`` + ``*.safetensors``).""" + """Load from a local EG2 snapshot directory (``config.json`` + ``model.safetensors``).""" mac = MacEmbeddingGemma2ImageEncoder.from_pretrained(model_dir, torch.float32, max_soft_tokens) model = load_hf_state_dict(EmbeddingGemma2ImageEncoder(mac.config), mac.state_dict()) return model.to(dtype).eval() diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index 8ff24302..b09b7628 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -22,6 +22,7 @@ import torch.nn as nn import torch.nn.functional as F +from coreai_models.models.base import load_tensors_for_keys, resolve_safetensors_files from coreai_models.models.macos.diffusion_gemma import ( DiffusionGemmaDenseMLP as EmbeddingGemma2MLP, ) @@ -426,25 +427,26 @@ def forward( def read_checkpoint_tensors( - model_dir: str | Path, prefixes: tuple[str, ...], dtype: torch.dtype | None = None + model_dir: str | Path, prefixes: tuple[str, ...], dtype: torch.dtype ) -> dict[str, torch.Tensor]: - """Tensors from ``model_dir/*.safetensors`` whose names start with one of ``prefixes``. + """Tensors of the ``model_dir`` checkpoint whose names start with one of ``prefixes``. - Names are kept as in the checkpoint; each tensor is cast to ``dtype`` when given. + Reads ``model.safetensors`` or the shards listed in ``model.safetensors.index.json``. Names + are kept as in the checkpoint; tensors are cast to ``dtype`` except embedding tables and + quantization zero-points, as in ``load_tensors_for_keys``. + + Raises: + FileNotFoundError: if ``model_dir`` has no safetensors weights. """ from safetensors import safe_open - files = sorted(Path(model_dir).glob("*.safetensors")) - if not files: - raise FileNotFoundError(f"No .safetensors in {model_dir}") - tensors: dict[str, torch.Tensor] = {} - for file in files: - with safe_open(str(file), framework="pt") as handle: + key_to_file: dict[str, str] = {} + for file in resolve_safetensors_files(str(model_dir)): + with safe_open(file, framework="pt") as handle: for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping if key.startswith(prefixes): - tensor = handle.get_tensor(key) - tensors[key] = tensor if dtype is None else tensor.to(dtype) - return tensors + key_to_file[key] = file + return load_tensors_for_keys(key_to_file, dtype) def _resolve_model_dir(model_dir_or_id: str | Path) -> Path: diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index 493b86cb..f80f00c1 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -260,7 +260,7 @@ def from_pretrained( dtype: torch.dtype = torch.float32, max_soft_tokens: int | None = None, ) -> EmbeddingGemma2ImageEncoder: - """Load from a local EG2 snapshot directory (``config.json`` + ``*.safetensors``).""" + """Load from a local EG2 snapshot directory (``config.json`` + ``model.safetensors``).""" config = json.loads((Path(model_dir) / "config.json").read_text()) model = cls(EmbeddingGemma2VisionConfig.from_hf_config(config, max_soft_tokens)) tensors = read_checkpoint_tensors(model_dir, ("vision_tower.", "embed_vision."), dtype) diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py index e1b17b2e..fcc8a02b 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -35,6 +35,7 @@ EmbeddingGemma2TextConfig, EmbeddingGemma2TextEncoder, load_embedding_gemma2_text_encoder, + read_checkpoint_tensors, ) from tests.test_model_units.test_models._embedding_gemma2_common import ( TINY_TEXT_CONFIG, @@ -141,6 +142,29 @@ def test_inputs_embeds_path_matches_input_ids(): model(input_ids=ids, inputs_embeds=model.embed_tokens(ids)) +def test_read_checkpoint_tensors_reads_index_shards(tmp_path): + from safetensors.torch import save_file + + with pytest.raises(FileNotFoundError): + read_checkpoint_tensors(tmp_path, ("language_model.",), torch.float32) + shards = { + "model-00001-of-00002.safetensors": { + "language_model.a": torch.ones(2, dtype=torch.bfloat16) + }, + "model-00002-of-00002.safetensors": { + "language_model.b": torch.zeros(3), + "vision_tower.c": torch.ones(1), + }, + } + for name, tensors in shards.items(): + save_file(tensors, str(tmp_path / name)) + weight_map = {key: name for name, tensors in shards.items() for key in tensors} + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + tensors = read_checkpoint_tensors(tmp_path, ("language_model.",), torch.float32) + assert sorted(tensors) == ["language_model.a", "language_model.b"] + assert {t.dtype for t in tensors.values()} == {torch.float32} + + # --- Tier 2: tiny random config vs HF EmbeddingGemma2TextModel --- From 6f29e9347d0cea44f742f838e4541c8b14152cee Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 22:56:38 -0700 Subject: [PATCH 35/38] Restore the EmbeddingGemma 2 test guards dropped by the trims 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. --- .../test_export/test_embedding_export.py | 42 +++++++++++ .../test_embedding_gemma2.py | 73 +++++++++++++++++-- .../test_embedding_gemma2_vision.py | 9 ++- 3 files changed, 112 insertions(+), 12 deletions(-) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index d001249c..4706d206 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -72,6 +72,37 @@ **dict.fromkeys(["CodeRetrieval", "InstructionRetrieval"], PROMPTS["code_retrieval"]), } +# Keys of earlier contract drafts that must not come back. +REMOVED_KEYS = { + "pooling", + "attention", + "granularity", + "budgets", + "max_sequence_length", + "prompt_applies_to", + "side", + "strategy", + "wrap", + "soft_tokens_scaled", + "floor", + "norm", + "dynamic_batch", + "dynamic_length", + "hidden_size", + "note", + "max_patches", + "soft_token_dim", + "convert_rgb", + "exif_transpose", + "resize", + "resize_policy", + "patch_pixel_order", + "position_order", + "padding_position", + "prompt", + "truncation_side", +} + def _block(platform="macOS", with_image=False, **overrides): kwargs = dict(revision=REVISION, prompts=PROMPTS) @@ -79,6 +110,16 @@ def _block(platform="macOS", with_image=False, **overrides): return embedding_block(EG2, platform=platform, with_image=with_image, **kwargs) +def _keys(node): + if isinstance(node, dict): + for key, value in node.items(): + yield key + yield from _keys(value) + elif isinstance(node, list): + for value in node: + yield from _keys(value) + + # --- Precision guard --- @@ -209,6 +250,7 @@ def test_embedding_block_platform_matrix(platform, with_image, components, pipel assert {k: v["compute_units"] for k, v in block["components"].items()} == units assert ("merge" in block) == with_image and ("preprocessing" in block) == with_image assert set(bundle_assets(platform, with_image)) == set(components) + assert not REMOVED_KEYS & set(_keys(block)), REMOVED_KEYS & set(_keys(block)) meta = build_bundle_metadata( EG2, block, diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py index fcc8a02b..1eb12c66 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -7,8 +7,8 @@ Tiers ----- -1. Self-contained (always run): config parsing, padding invariance, pooling mask negative test, - inputs_embeds path. +1. Self-contained (always run): config parsing, mask semantics, padding invariance, pooling + mask negative test, inputs_embeds path. 2. HF reference (tiny random config): needs ``transformers>=5.19`` (``models/embedding_gemma2``). The repo pins an older transformers, so run in an overlay:: @@ -18,7 +18,8 @@ 3. Real weights (CPU fp32): additionally needs the ``google/embeddinggemma-2`` snapshot in the local HF cache (the ``embeddinggemma2_snapshot`` fixture skips otherwise). Compared against HF ``EmbeddingGemma2Model`` + the sentence-transformers mean-pool + L2-normalize - re-implementation in ``_st_mean_pool_normalize``. + re-implementation in ``_st_mean_pool_normalize``, and against sentence-transformers>=6.1 + when it is installed (``--with "sentence-transformers>=6.1"``). Tiers 2/3 skip cleanly when their dependencies are unavailable. """ @@ -34,6 +35,7 @@ from coreai_models.models.macos.embedding_gemma2 import ( EmbeddingGemma2TextConfig, EmbeddingGemma2TextEncoder, + build_attention_masks, load_embedding_gemma2_text_encoder, read_checkpoint_tensors, ) @@ -89,6 +91,28 @@ def test_config_requires_both_rope_thetas(): EmbeddingGemma2TextConfig.from_dict({**TINY_TEXT_CONFIG, "rope_parameters": rope}) +def test_config_from_snapshot_matches_published_layout(embeddinggemma2_snapshot): + cfg = EmbeddingGemma2TextConfig.from_pretrained(embeddinggemma2_snapshot) + full = [i for i in range(cfg.num_hidden_layers) if not cfg.is_sliding(i)] + assert full == [5, 11, 17, 23] + assert all(cfg.layer_head_dim(i) == 512 and cfg.layer_num_key_value_heads(i) == 1 for i in full) + assert cfg.layer_head_dim(0) == 256 and cfg.layer_num_key_value_heads(0) == 2 + assert (cfg.hidden_size, cfg.embedding_dim, cfg.sliding_window) == (512, 768, 512) + + +def test_masks_bidirectional_window_and_key_padding(): + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0, 0]], dtype=torch.int32) + full, sliding = build_attention_masks(attention_mask, sliding_window=2) + assert full.shape == sliding.shape == (1, 1, 7, 7) + i, j = torch.meshgrid(torch.arange(7), torch.arange(7), indexing="ij") + key_valid = (j < 5).unsqueeze(0).unsqueeze(0) + assert torch.equal(full, key_valid.expand_as(full)) + assert torch.equal(sliding, ((i - j).abs() <= 2).unsqueeze(0).unsqueeze(0) & key_valid) + # Symmetric (bidirectional) on valid positions. + valid = sliding[0, 0, :5, :5] + assert torch.equal(valid, valid.T) + + def test_forward_shapes_and_unit_norm(): model = _tiny_model() ids, mask = _ragged_batch([9, 5], seq_len=12, vocab=128) @@ -175,8 +199,8 @@ def _hf_modeling(): ) -def _tiny_hf(): - """Tiny random HF ``EmbeddingGemma2TextModel`` (eager attention) with perturbed weights.""" +def _tiny_hf(attn_impl: str = "eager"): + """Tiny random HF ``EmbeddingGemma2TextModel`` with perturbed weights.""" modeling = _hf_modeling() from transformers.models.embedding_gemma2.configuration_embedding_gemma2 import ( EmbeddingGemma2Config, @@ -184,14 +208,15 @@ def _tiny_hf(): torch.manual_seed(0) hf_cfg = EmbeddingGemma2Config(text_config=dict(TINY_TEXT_CONFIG)).text_config - hf_cfg._attn_implementation = "eager" + hf_cfg._attn_implementation = attn_impl hf = modeling.EmbeddingGemma2TextModel(hf_cfg).eval() _randomize_(hf, seed=3) return hf -def test_tiny_parity_vs_hf(): - hf = _tiny_hf() +@pytest.mark.parametrize("attn_impl", ["eager", "sdpa"]) +def test_tiny_parity_vs_hf(attn_impl): + hf = _tiny_hf(attn_impl) ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(TINY_TEXT_CONFIG)).eval() ours.load_state_dict(hf.state_dict()) @@ -209,6 +234,21 @@ def test_tiny_parity_vs_hf(): torch.testing.assert_close(emb, hf_emb, atol=1e-5, rtol=0) +def test_tiny_sliding_window_matters_vs_hf(): + """With the window widened from 3 to 1000, ours no longer matches HF.""" + hf = _tiny_hf() + wide = dict(TINY_TEXT_CONFIG, sliding_window=1000) + ours = EmbeddingGemma2TextEncoder(EmbeddingGemma2TextConfig.from_dict(wide)).eval() + ours.load_state_dict(hf.state_dict()) + ids, mask = _ragged_batch([24], seq_len=24, vocab=128, seed=7) + with torch.no_grad(): + hf_emb = _st_mean_pool_normalize( + hf(input_ids=ids.long(), attention_mask=mask.long()).last_hidden_state, mask + ) + emb = ours(input_ids=ids, attention_mask=mask) + assert (emb - hf_emb).abs().max().item() > 1e-3 + + # --- Tier 3: real weights (CPU fp32) --- REAL_TEXTS = [ @@ -274,6 +314,23 @@ def test_real_weight_parity_and_padding_invariance(real_setup): assert _cos(alone, emb).min().item() >= 0.99999 +def test_real_weight_matches_sentence_transformers(real_setup, embeddinggemma2_snapshot): + """The pooling re-implementation and the prompt handling agree with sentence-transformers.""" + st = pytest.importorskip("sentence_transformers", minversion="6.1") + model = st.SentenceTransformer( + str(embeddinggemma2_snapshot), device="cpu", model_kwargs={"dtype": torch.float32} + ) + texts = real_setup["texts"] + features = real_setup["tokenize"](texts) + with torch.inference_mode(): + st_emb = model.encode(texts, batch_size=len(texts), convert_to_tensor=True) + emb = real_setup["ours"]( + input_ids=features["input_ids"].to(torch.int32), + attention_mask=features["attention_mask"].to(torch.int32), + ) + assert _cos(emb, st_emb).min().item() >= 0.99999 + + def test_real_weight_excluding_prompt_changes_embedding(real_setup): ours, texts = real_setup["ours"], real_setup["texts"] features = real_setup["tokenize"](texts[:1]) diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py index 2d9e31ea..93c7e984 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py @@ -119,10 +119,11 @@ def _random_inputs(cfg: EmbeddingGemma2VisionConfig, grids: list[tuple[int, int] # --- Tiny-config parity vs HF --- -def test_tiny_soft_tokens_match_hf(monkeypatch) -> None: - # The directory conftest forces USE_HF_IMPL=true; run the Core AI lowering path of the SDPA - # primitive (bool-mask -> -1e4 bias), which is what gets exported. - monkeypatch.setenv("USE_HF_IMPL", "false") +@pytest.mark.parametrize("use_hf_impl", ["true", "false"]) +def test_tiny_soft_tokens_match_hf(use_hf_impl, monkeypatch) -> None: + # "false" runs the Core AI lowering path of the SDPA primitive (bool-mask -> -1e4 bias), + # which is what gets exported; "true" is the directory conftest's default. + monkeypatch.setenv("USE_HF_IMPL", use_hf_impl) hf_cfg, hf = _hf_tiny_vision(TINY) text_cfg = gemma4_config.Gemma4TextConfig(hidden_size=TINY.text_hidden_size) torch.manual_seed(1) From 49f73366ccc3dcf581c417aa32bb5df04bc7a83b Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 22:56:38 -0700 Subject: [PATCH 36/38] Test the Neural Engine embedding conversion on tiny configs 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. --- .../test_export/test_embedding_export.py | 135 +++++++++++++++++- 1 file changed, 134 insertions(+), 1 deletion(-) diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 4706d206..9e0cec3d 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -6,13 +6,16 @@ """Tests for the ``kind: "embedding"`` bundle contract and exporter (``coreai_models.embedding``). Metadata, precision and CLI tests need no hardware or weights. The conversion tests use tiny random -EmbeddingGemma 2 configs. +EmbeddingGemma 2 configs; the Neural Engine programs also run on the Core AI runtime (CPU and GPU) +when it is available. """ from __future__ import annotations +import asyncio import dataclasses import json +import tempfile from pathlib import Path import pytest @@ -591,6 +594,136 @@ def test_ane_reference_inputs_shapes(): assert int(real.sum()) % 9 == 0 and torch.all(positions[0, ~real] == -1) +# --- Neural Engine conversion (tiny random config) --- + +# Encoder lengths 4 (unbanded) and 16 (banded: > sliding_window + 1); 16 also fits the image +# sequence (budget 12 + 4). +TINY_ANE_FAMILY = dataclasses.replace( + EG2, image=TINY_IMAGE, ane_length_buckets=(4, 16), model_max_tokens=64 +) + + +@pytest.fixture(scope="module") +def tiny_ane_programs(): + """Component name -> (PyTorch module, program converted through ``export_ios``).""" + from coreai_models.models.ios import embedding_gemma2_vision as ios_vision + from coreai_models.models.ios.embedding_gemma2 import text_encoder_from_state_dict + + text = _tiny_text_encoder() + encoder, table = text_encoder_from_state_dict( + text.config, dict(text.state_dict()), torch.float16 + ) + encoder.eval() + token_table = export_utils.scaled_token_table( + table, text.config.hidden_size**0.5, torch.float16 + ) + torch.manual_seed(1) + image_encoder = ios_vision.EmbeddingGemma2ImageEncoder(TINY_VISION_CONFIG).half().eval() + functions = encoder_functions(TINY_ANE_FAMILY, "iOS") + window = int(text.config.sliding_window) + return { + "encoder": ( + encoder, + export_ios.export_ane_encoder_program(encoder, functions, window), + ), + "token_table": ( + token_table, + export_utils.export_token_table_program(token_table, 64, dtype=torch.float16), + ), + "image_encoder": ( + image_encoder, + export_ios.export_ane_image_encoder_program(image_encoder, TINY_IMAGE), + ), + } + + +def _function_names(program) -> list[str]: + with tempfile.TemporaryDirectory() as tmp: + asset = program.save_asset(Path(tmp) / "program.aimodel") + return sorted(asset.summary(include_statistics=False).function_names) + + +async def _run_function(program, name: str, inputs: dict, options) -> dict: + """Outputs of function ``name`` of ``program`` on ``inputs``, run by the Core AI runtime.""" + import coreai.runtime as rt + + with tempfile.TemporaryDirectory() as tmp: + asset = program.save_asset(Path(tmp) / "program.aimodel") + async with asset.executable(options) as model: + outputs = await model.load_function(name)( + {key: rt.NDArray(data=value.contiguous()) for key, value in inputs.items()} + ) + return {key: torch.from_numpy(value.numpy().copy()) for key, value in outputs.items()} + + +def test_tiny_ane_programs_have_the_contract_functions(tiny_ane_programs): + block = embedding_block( + TINY_ANE_FAMILY, platform="iOS", with_image=True, revision=REVISION, prompts=PROMPTS + ) + encoder = block["components"]["encoder"] + assert [f["function"] for f in encoder["functions"]] == ["embeds_len4_b1", "embeds_len16_b1"] + for name, spec in block["components"].items(): + if "functions" in spec: + functions = [f["function"] for f in spec["functions"]] + else: + functions = [spec["function"]] + assert _function_names(tiny_ane_programs[name][1]) == sorted(functions), name + + +@pytest.mark.parametrize("compute_unit", ["cpu", "gpu"]) +def test_tiny_ane_programs_match_pytorch(tiny_ane_programs, compute_unit): + """Each converted program, run by the Core AI runtime, matches its float16 module. + + The tolerance is the float16 one of ``test_macos_prefill_graph.py`` (atol = rtol = 1e-2). + """ + import coreai.runtime as rt + + from coreai_models.models.ios.embedding_gemma2 import MASK_NEG + from coreai_models.models.macos.embedding_gemma2_vision import pad_to_budget, patch_position_ids + + if not rt.SpecializationOptions.is_supported(): + pytest.skip("the Core AI runtime cannot run programs here") + options = ( + rt.SpecializationOptions.cpu_only() + if compute_unit == "cpu" + else rt.SpecializationOptions.from_preferred_compute_unit_kind(rt.ComputeUnitKind.gpu()) + ) + + encoder, program = tiny_ane_programs["encoder"] + for spec in encoder_functions(TINY_ANE_FAMILY, "iOS"): + n = spec.length + valid = torch.arange(n) < n - 1 + embeds, _, _ = export_ios.ane_text_reference_inputs(n, 64) + bias = torch.where(valid.reshape(1, n, 1, 1), 0.0, MASK_NEG).expand(1, n, 1, n) + pool = (valid.float() / valid.sum()).reshape(1, n) + inputs = {"inputs_embeds": embeds, "bias": bias.half(), "pool_weights": pool.half()} + with torch.no_grad(): + wrapper = export_ios.EncoderBandWrapper(encoder, n, encoder.config.sliding_window) + want = wrapper(*inputs.values()) + got = asyncio.run(_run_function(program, spec.name, inputs, options))["embedding"] + assert got.dtype == torch.float16 and got.shape == want.shape == (1, 48) + torch.testing.assert_close(got, want, atol=1e-2, rtol=1e-2) + + table, program = tiny_ane_programs["token_table"] + ids = torch.randint(3, 128, (1, 9), dtype=torch.int32) + got = asyncio.run(_run_function(program, "main", {"input_ids": ids}, options)) + assert got["inputs_embeds"].dtype == torch.float16 + assert torch.equal( + got["inputs_embeds"], export_utils.TokenTableExportWrapper(table, torch.float16)(ids) + ) + + image_encoder, program = tiny_ane_programs["image_encoder"] + pixels = torch.rand((54, TINY_IMAGE.patch_dim), generator=torch.Generator().manual_seed(2)) + pixels, positions = pad_to_budget(pixels, patch_position_ids(6, 9), TINY_IMAGE.max_patches) + inputs = {"pixel_values": pixels[None].half(), "position_ids": positions[None].half()} + with torch.no_grad(): + want = image_encoder(*inputs.values()) + got = asyncio.run(_run_function(program, "main", inputs, options))["soft_tokens"] + assert got.dtype == torch.float16 and got.shape == want.shape == (1, 12, 32) + torch.testing.assert_close(got, want, atol=1e-2, rtol=1e-2) + assert torch.all(got[0, 6:] == 0) + + # --- CLI --- From 191533451472e9d867aef3a58b67af6eefd1cebb Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 29 Oct 2026 22:56:38 -0700 Subject: [PATCH 37/38] Record each embedding asset once and the component output dtypes 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. --- models/embedding_gemma2/README.md | 13 ++++ .../src/coreai_models/embedding/contract.py | 51 +++++++++++---- python/src/coreai_models/embedding/export.py | 9 +-- .../src/coreai_models/embedding/export_ios.py | 24 ++++--- .../coreai_models/embedding/export_macos.py | 26 +++++--- .../coreai_models/embedding/export_utils.py | 3 + .../test_export/test_embedding_export.py | 64 +++++++++++++------ 7 files changed, 136 insertions(+), 54 deletions(-) diff --git a/models/embedding_gemma2/README.md b/models/embedding_gemma2/README.md index 3a4b2e70..7e57f138 100644 --- a/models/embedding_gemma2/README.md +++ b/models/embedding_gemma2/README.md @@ -39,4 +39,17 @@ uv run coreai.embedding.export --help | ---------------- | ----------------------- | ---------------------------- | | embeddinggemma-2 | google/embeddinggemma-2 | 439M (text 271M, image 168M) | +## Bundle + +`metadata.json` lists each `.aimodel` file once under `assets`. Each entry of `embedding.components` names the asset it lives in (`asset`), its functions and its `output_dtype`: + +| Component | macOS asset (GPU) | iOS asset | Functions | Output dtype | +| --------------- | ---------------------------- | ------------------------------- | ------------------ | ------------------------------ | +| `text_encoder` | `text_encoder` | — | `len{L}_b{B}` | `float32` | +| `token_table` | `text_encoder` (images only) | `token_table` (GPU) | `main` | `float32` macOS, `float16` iOS | +| `encoder` | `text_encoder` (images only) | `encoder` (Neural Engine) | `embeds_len{L}_b1` | `float32` macOS, `float16` iOS | +| `image_encoder` | `image_encoder` | `image_encoder` (Neural Engine) | `main` | `float32` macOS, `float16` iOS | + +The encoder functions are named `embeds_len{L}_b1` on both platforms; read the names from `components.encoder.functions`. + [^1]: [Documentation](https://ai.google.dev/gemma/docs/embeddinggemma) · [HuggingFace](https://huggingface.co/google/embeddinggemma-2) diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py index 5b98cfaf..c2854b97 100644 --- a/python/src/coreai_models/embedding/contract.py +++ b/python/src/coreai_models/embedding/contract.py @@ -268,7 +268,11 @@ def text_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec] def encoder_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec]: - """Batch-1 ``inputs_embeds`` functions of the ``encoder`` component, named ``embeds_len*``.""" + """Batch-1 ``inputs_embeds`` functions of the ``encoder`` component. + + Named ``embeds_len{L}_b1`` on both platforms; the host reads the names from + ``components.encoder.functions``. + """ if platform == "iOS": lengths = sorted(set(family.ane_length_buckets)) else: @@ -290,17 +294,29 @@ def text_pipeline(platform: str) -> tuple[str, ...]: return (TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT) if platform == "iOS" else (TEXT_COMPONENT,) -def bundle_assets(platform: str, with_image: bool) -> dict[str, str]: - """Component name -> asset file for every component of the bundle's pipelines. +def bundle_components(platform: str, with_image: bool) -> tuple[str, ...]: + """Components of the bundle's pipelines, text first.""" + return tuple(dict.fromkeys(text_pipeline(platform) + (IMAGE_PIPELINE if with_image else ()))) + - On macOS, ``token_table`` and ``encoder`` are functions of ``text_encoder.aimodel``. +def component_asset(platform: str, component: str) -> str: + """Name of the asset holding ``component``. + + On macOS, ``token_table`` and ``encoder`` are functions of the ``text_encoder`` asset. """ - components = text_pipeline(platform) + (IMAGE_PIPELINE if with_image else ()) - shared = () if platform == "iOS" else (TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT) - return { - name: f"{TEXT_COMPONENT if name in shared else name}.aimodel" - for name in dict.fromkeys(components) - } + shared = platform != "iOS" and component in (TOKEN_TABLE_COMPONENT, ENCODER_COMPONENT) + return TEXT_COMPONENT if shared else component + + +def bundle_assets(platform: str, with_image: bool) -> dict[str, str]: + """Asset name -> ``.aimodel`` file, one entry per file; components name theirs in ``asset``.""" + names = (component_asset(platform, c) for c in bundle_components(platform, with_image)) + return {name: f"{name}.aimodel" for name in dict.fromkeys(names)} + + +def output_dtype(platform: str) -> str: + """Dtype of every component output: float16 on iOS, float32 on macOS.""" + return "float16" if platform == "iOS" else "float32" def embedding_block( @@ -318,8 +334,9 @@ def embedding_block( Text runs ``text_encoder`` on macOS and ``token_table -> encoder`` on iOS. With ``with_image``, images run ``image_encoder -> token_table -> encoder``, and the host writes the soft tokens into the ``token_table`` output at the placeholder rows. ``token_table`` runs - on the GPU; on iOS the encoders run on the Neural Engine. ``compute_precision`` defaults to - the platform's precision. + on the GPU; on iOS the encoders run on the Neural Engine. Each component names its + ``bundle_assets`` entry (``asset``), its functions and its ``output_dtype``. + ``compute_precision`` defaults to the platform's precision. Raises: ValueError: on an unknown platform, a missing revision or prompts, an unknown @@ -354,35 +371,43 @@ def embedding_block( if with_image or ios: enc_fns = encoder_functions(family, platform) functions_json = [f.as_dict() for f in enc_fns] + out_dtype = output_dtype(platform) specs: dict[str, Any] = { TEXT_COMPONENT: { "functions": [f.as_dict() for f in text_fns], "inputs": list(TEXT_INPUT_NAMES), "output": OUTPUT_NAME, + "output_dtype": out_dtype, "compute_units": ["gpu"], }, IMAGE_COMPONENT: { "function": "main", "inputs": list(IMAGE_INPUT_NAMES), "output": IMAGE_OUTPUT_NAME, + "output_dtype": out_dtype, "compute_units": ["neural_engine"] if ios else ["gpu"], }, TOKEN_TABLE_COMPONENT: { "function": "main", "inputs": list(TOKEN_TABLE_INPUT_NAMES), "output": TOKEN_TABLE_OUTPUT_NAME, + "output_dtype": out_dtype, "compute_units": ["gpu"], }, ENCODER_COMPONENT: { "functions": functions_json, "inputs": list(ANE_ENCODER_INPUT_NAMES if ios else ENCODER_INPUT_NAMES), "output": OUTPUT_NAME, + "output_dtype": out_dtype, "mask_encoding": "additive_bias" if ios else "attention_mask", "compute_units": ["neural_engine"] if ios else ["gpu"], }, } - components = {name: specs[name] for name in bundle_assets(platform, with_image)} + components = { + name: {"asset": component_asset(platform, name), **specs[name]} + for name in bundle_components(platform, with_image) + } text_tower: dict[str, Any] = { "modality": "text", diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index 72a0d654..34399dc9 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -84,9 +84,10 @@ def export_embedding_bundle( ) -> dict[str, Any]: """Export one bundle to ``output_dir`` and return its ``metadata.json`` content. - macOS writes a GPU bundle (``text_encoder``, plus ``image_encoder`` / ``token_table`` / - ``encoder`` with images); iOS writes the Neural Engine bundle (``token_table``, ``encoder``, - plus ``image_encoder``). Both carry the tokenizer and a ``metadata.json``. + macOS writes a GPU bundle: ``text_encoder.aimodel``; with images also + ``image_encoder.aimodel``, and ``text_encoder`` then holds the ``token_table`` and ``encoder`` + functions. iOS writes the Neural Engine bundle: ``token_table.aimodel``, ``encoder.aimodel``, + plus ``image_encoder.aimodel``. Both carry the tokenizer and a ``metadata.json``. """ if platform not in PLATFORMS: raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") @@ -245,7 +246,7 @@ def main(argv: list[str] | None = None) -> int: print(f"platform: {platform}") print(f"precision: {precision}") print(f"towers: {towers}") - print(f"assets: {', '.join(dict.fromkeys(assets.values()))}") + print(f"assets: {', '.join(assets.values())}") if platform == "macOS": print(f"text functions: {[f.name for f in text_functions(family, platform)]}") if with_image or platform == "iOS": diff --git a/python/src/coreai_models/embedding/export_ios.py b/python/src/coreai_models/embedding/export_ios.py index 3575771e..ee2bed3b 100644 --- a/python/src/coreai_models/embedding/export_ios.py +++ b/python/src/coreai_models/embedding/export_ios.py @@ -30,8 +30,10 @@ FunctionSpec, ImageTowerFacts, encoder_functions, + output_dtype, ) from coreai_models.embedding.export_utils import ( + TORCH_DTYPES, export_token_table_program, image_reference_inputs, save_program, @@ -42,6 +44,10 @@ from coreai.authoring import AIProgram +# The Neural Engine programs compute in this dtype and output it. +DTYPE = TORCH_DTYPES[output_dtype("iOS")] + + class EncoderBandWrapper(nn.Module): """``(inputs_embeds, bias, pool_weights) -> embedding`` at one static length (Neural Engine). @@ -74,9 +80,9 @@ def ane_text_reference_inputs( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Neural Engine encoder trace inputs (all valid tokens): embeds, zero bias, uniform pooling.""" generator = torch.Generator().manual_seed(length) - embeds = torch.randn((1, length, hidden), generator=generator).to(torch.float16) - bias = torch.zeros((1, length, 1, length), dtype=torch.float16) - pool = torch.full((1, length), 1.0 / length, dtype=torch.float16) + embeds = torch.randn((1, length, hidden), generator=generator).to(DTYPE) + bias = torch.zeros((1, length, 1, length), dtype=DTYPE) + pool = torch.full((1, length), 1.0 / length, dtype=DTYPE) return embeds, bias, pool @@ -132,7 +138,7 @@ def export_ane_image_encoder_program( include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, ) -> AIProgram: """Neural Engine ``image_encoder``: function ``main`` at the static budget.""" - args = image_reference_inputs(image, torch.float16, torch.float16) + args = image_reference_inputs(image, DTYPE, DTYPE) return convert_ane_functions( [("main", image_encoder, args, IMAGE_INPUT_NAMES, (IMAGE_OUTPUT_NAME,))], include_debug_info=include_debug_info, @@ -156,13 +162,11 @@ def export_ane_programs( # TODO: write the table as a safetensors sidecar that the host gathers from, as the Gemma 4 # per-layer embeddings do (SafeTensorsReader), instead of a GPU program; lands with the Swift # embedding runtime. - encoder, table = load_embedding_gemma2_text_encoder(snapshot, dtype=torch.float16) - token_table = scaled_token_table(table, encoder.config.hidden_size**0.5, torch.float16) + encoder, table = load_embedding_gemma2_text_encoder(snapshot, dtype=DTYPE) + token_table = scaled_token_table(table, encoder.config.hidden_size**0.5, DTYPE) del table save_program( - export_token_table_program( - token_table, family.model_max_tokens, include_debug_info, dtype=torch.float16 - ), + export_token_table_program(token_table, family.model_max_tokens, include_debug_info, DTYPE), out, family, TOKEN_TABLE_COMPONENT, @@ -185,7 +189,7 @@ def export_ane_programs( from coreai_models.models.ios.embedding_gemma2_vision import load_image_encoder image_encoder = load_image_encoder( - str(snapshot), torch.float16, max_soft_tokens=family.image.budget + str(snapshot), DTYPE, max_soft_tokens=family.image.budget ) save_program( export_ane_image_encoder_program(image_encoder, family.image, include_debug_info), diff --git a/python/src/coreai_models/embedding/export_macos.py b/python/src/coreai_models/embedding/export_macos.py index e09eab84..cd146317 100644 --- a/python/src/coreai_models/embedding/export_macos.py +++ b/python/src/coreai_models/embedding/export_macos.py @@ -33,9 +33,11 @@ FunctionSpec, ImageTowerFacts, encoder_functions, + output_dtype, text_functions, ) from coreai_models.embedding.export_utils import ( + TORCH_DTYPES, TokenTableExportWrapper, image_reference_inputs, save_program, @@ -46,11 +48,11 @@ from coreai.authoring import AIProgram -_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16} +OUTPUT_DTYPE = TORCH_DTYPES[output_dtype("macOS")] class TextEmbeddingExportWrapper(nn.Module): - """``(input_ids, attention_mask) -> embedding`` with float32 output. + """``(input_ids, attention_mask) -> embedding`` as ``OUTPUT_DTYPE``. The wrapped encoder owns pooling (masked mean incl. prompt), projection and L2 normalization. """ @@ -60,11 +62,14 @@ def __init__(self, encoder: nn.Module) -> None: self.encoder = encoder def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.encoder(input_ids=input_ids, attention_mask=attention_mask).to(torch.float32) + return self.encoder(input_ids=input_ids, attention_mask=attention_mask).to(OUTPUT_DTYPE) class InputsEmbedsEncoderWrapper(nn.Module): - """``(inputs_embeds, attention_mask) -> embedding`` for merged (e.g. image) inputs.""" + """``(inputs_embeds, attention_mask) -> embedding`` for merged (e.g. image) inputs. + + Outputs ``OUTPUT_DTYPE``. + """ def __init__(self, encoder: nn.Module) -> None: super().__init__() @@ -74,11 +79,14 @@ def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> dtype = next(self.encoder.parameters()).dtype return self.encoder( inputs_embeds=inputs_embeds.to(dtype), attention_mask=attention_mask - ).to(torch.float32) + ).to(OUTPUT_DTYPE) class ImageEncoderExportWrapper(nn.Module): - """``(pixel_values f32, position_ids int32) -> soft_tokens f32`` at the module's dtype.""" + """``(pixel_values f32, position_ids int32) -> soft_tokens`` as ``OUTPUT_DTYPE``. + + The encoder computes at the module's dtype. + """ def __init__(self, image_encoder: nn.Module) -> None: super().__init__() @@ -86,7 +94,7 @@ def __init__(self, image_encoder: nn.Module) -> None: def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: dtype = next(self.image_encoder.parameters()).dtype - return self.image_encoder(pixel_values.to(dtype), position_ids).to(torch.float32) + return self.image_encoder(pixel_values.to(dtype), position_ids).to(OUTPUT_DTYPE) @contextmanager @@ -140,7 +148,7 @@ def export_shared_text_program( add_function( converter, "main", - TokenTableExportWrapper(encoder.embed_tokens), + TokenTableExportWrapper(encoder.embed_tokens, OUTPUT_DTYPE), ids, TOKEN_TABLE_INPUT_NAMES, (TOKEN_TABLE_OUTPUT_NAME,), @@ -187,7 +195,7 @@ def export_gpu_programs( from coreai_models.models.macos.embedding_gemma2 import load_embedding_gemma2_text_encoder from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder - dtype = _DTYPES[compute_precision] + dtype = TORCH_DTYPES[compute_precision] encoder = load_embedding_gemma2_text_encoder(snapshot, dtype=dtype) program = export_shared_text_program( encoder, diff --git a/python/src/coreai_models/embedding/export_utils.py b/python/src/coreai_models/embedding/export_utils.py index 2064f320..7f6bc415 100644 --- a/python/src/coreai_models/embedding/export_utils.py +++ b/python/src/coreai_models/embedding/export_utils.py @@ -33,6 +33,9 @@ logger = logging.getLogger(__name__) +TORCH_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16} + + TOKENIZER_FILES = ( "tokenizer.json", "tokenizer_config.json", diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 9e0cec3d..6ded27a9 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -252,17 +252,21 @@ def test_embedding_block_platform_matrix(platform, with_image, components, pipel assert [t["pipeline"] for t in block["towers"]] == pipelines assert {k: v["compute_units"] for k, v in block["components"].items()} == units assert ("merge" in block) == with_image and ("preprocessing" in block) == with_image - assert set(bundle_assets(platform, with_image)) == set(components) + assets = bundle_assets(platform, with_image) + assert {spec["asset"] for spec in block["components"].values()} == set(assets) + assert len(set(assets.values())) == len(assets) + want_dtype = "float16" if platform == "iOS" else "float32" + assert {spec["output_dtype"] for spec in block["components"].values()} == {want_dtype} assert not REMOVED_KEYS & set(_keys(block)), REMOVED_KEYS & set(_keys(block)) meta = build_bundle_metadata( EG2, block, name=bundle_name(EG2, platform, with_image), - assets=bundle_assets(platform, with_image), + assets=assets, hf_revision="914f7f89142e33e77833254d9c9b90c3cef7303b", ) assert meta["metadata_version"] == "0.2" and meta["kind"] == "embedding" - assert {c for t in meta["embedding"]["towers"] for c in t["pipeline"]} == set(meta["assets"]) + assert {c for t in meta["embedding"]["towers"] for c in t["pipeline"]} == set(components) assert meta["source"]["hf_model_id"] == "google/embeddinggemma-2" @@ -284,6 +288,16 @@ def test_embedding_block_ios_details(): assert image_tower["batch_sizes"] == [1] assert block["text"]["truncation"]["max_tokens"] == 1024 assert block["space"] == _block("macOS")["space"] + assert bundle_assets("iOS", True) == { + "token_table": "token_table.aimodel", + "encoder": "encoder.aimodel", + "image_encoder": "image_encoder.aimodel", + } + assert {k: v["asset"] for k, v in block["components"].items()} == { + "token_table": "token_table", + "encoder": "encoder", + "image_encoder": "image_encoder", + } def test_embedding_block_macos_image_details(): @@ -326,9 +340,16 @@ def test_embedding_block_macos_image_details(): "pad_token_id": 0, } assert text["prompts"] == PROMPTS and text["default_prompt_name"] is None - assets = bundle_assets("macOS", True) - assert assets["token_table"] == assets["encoder"] == assets["text_encoder"] - assert assets["image_encoder"] == "image_encoder.aimodel" + assert bundle_assets("macOS", True) == { + "text_encoder": "text_encoder.aimodel", + "image_encoder": "image_encoder.aimodel", + } + assert {k: v["asset"] for k, v in block["components"].items()} == { + "text_encoder": "text_encoder", + "image_encoder": "image_encoder", + "token_table": "text_encoder", + "encoder": "text_encoder", + } names = [f.name for f in text_functions(EG2, "macOS") + encoder_functions(EG2, "macOS")] assert len(names) == len(set(names)) and "main" not in names enc = block["components"]["encoder"] @@ -362,7 +383,9 @@ def test_embedding_block_validation(): def test_embedding_block_records_the_compute_precision(): assert _block()["precision"]["compute"] == "float32" - assert _block(compute_precision="bfloat16")["precision"]["compute"] == "bfloat16" + bf16 = _block(with_image=True, compute_precision="bfloat16") + assert bf16["precision"]["compute"] == "bfloat16" + assert {spec["output_dtype"] for spec in bf16["components"].values()} == {"float32"} assert _block("iOS")["precision"]["compute"] == "float16" with pytest.raises(EmbeddingPrecisionError, match="forbids float16"): _block(compute_precision="float16") @@ -571,6 +594,12 @@ def test_image_reference_inputs_are_a_valid_padded_grid(): assert torch.all(soft[0, n:] == 0) and torch.all(soft[0, :n].abs().sum(-1) > 0) +def _function_names(program) -> list[str]: + with tempfile.TemporaryDirectory() as tmp: + asset = program.save_asset(Path(tmp) / "program.aimodel") + return sorted(asset.summary(include_statistics=False).function_names) + + def test_tiny_shared_text_program_conversion(): program = export_macos.export_shared_text_program( _tiny_text_encoder(), @@ -578,7 +607,7 @@ def test_tiny_shared_text_program_conversion(): [FunctionSpec("embeds_len32_b1", 1, 32)], max_length=64, ) - assert program is not None + assert _function_names(program) == ["embeds_len32_b1", "len16_b1", "len32_b1", "main"] image = export_macos.export_image_encoder_program(_tiny_image_encoder(), TINY_IMAGE) assert image is not None @@ -637,12 +666,6 @@ def tiny_ane_programs(): } -def _function_names(program) -> list[str]: - with tempfile.TemporaryDirectory() as tmp: - asset = program.save_asset(Path(tmp) / "program.aimodel") - return sorted(asset.summary(include_statistics=False).function_names) - - async def _run_function(program, name: str, inputs: dict, options) -> dict: """Outputs of function ``name`` of ``program`` on ``inputs``, run by the Core AI runtime.""" import coreai.runtime as rt @@ -672,7 +695,8 @@ def test_tiny_ane_programs_have_the_contract_functions(tiny_ane_programs): @pytest.mark.parametrize("compute_unit", ["cpu", "gpu"]) def test_tiny_ane_programs_match_pytorch(tiny_ane_programs, compute_unit): - """Each converted program, run by the Core AI runtime, matches its float16 module. + """Each converted program, run by the Core AI runtime, matches its float16 module and + outputs the metadata's ``output_dtype``. The tolerance is the float16 one of ``test_macos_prefill_graph.py`` (atol = rtol = 1e-2). """ @@ -689,6 +713,10 @@ def test_tiny_ane_programs_match_pytorch(tiny_ane_programs, compute_unit): else rt.SpecializationOptions.from_preferred_compute_unit_kind(rt.ComputeUnitKind.gpu()) ) + dtypes = { + name: getattr(torch, spec["output_dtype"]) + for name, spec in _block("iOS", True)["components"].items() + } encoder, program = tiny_ane_programs["encoder"] for spec in encoder_functions(TINY_ANE_FAMILY, "iOS"): n = spec.length @@ -701,13 +729,13 @@ def test_tiny_ane_programs_match_pytorch(tiny_ane_programs, compute_unit): wrapper = export_ios.EncoderBandWrapper(encoder, n, encoder.config.sliding_window) want = wrapper(*inputs.values()) got = asyncio.run(_run_function(program, spec.name, inputs, options))["embedding"] - assert got.dtype == torch.float16 and got.shape == want.shape == (1, 48) + assert got.dtype == dtypes["encoder"] and got.shape == want.shape == (1, 48) torch.testing.assert_close(got, want, atol=1e-2, rtol=1e-2) table, program = tiny_ane_programs["token_table"] ids = torch.randint(3, 128, (1, 9), dtype=torch.int32) got = asyncio.run(_run_function(program, "main", {"input_ids": ids}, options)) - assert got["inputs_embeds"].dtype == torch.float16 + assert got["inputs_embeds"].dtype == dtypes["token_table"] assert torch.equal( got["inputs_embeds"], export_utils.TokenTableExportWrapper(table, torch.float16)(ids) ) @@ -719,7 +747,7 @@ def test_tiny_ane_programs_match_pytorch(tiny_ane_programs, compute_unit): with torch.no_grad(): want = image_encoder(*inputs.values()) got = asyncio.run(_run_function(program, "main", inputs, options))["soft_tokens"] - assert got.dtype == torch.float16 and got.shape == want.shape == (1, 12, 32) + assert got.dtype == dtypes["image_encoder"] and got.shape == want.shape == (1, 12, 32) torch.testing.assert_close(got, want, atol=1e-2, rtol=1e-2) assert torch.all(got[0, 6:] == 0) From 41b294b8d211f0ab8984ce5178cd6f0008ce95a7 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Fri, 30 Oct 2026 10:31:43 -0700 Subject: [PATCH 38/38] Tighten the EmbeddingGemma 2 comments and docstrings State what each helper does, returns and raises; drop rationale and stale references. Rename the two 'fails closed' tests after what they assert. --- models/README.md | 2 +- models/embedding_gemma2/README.md | 24 +++--- .../src/coreai_models/embedding/contract.py | 39 ++++----- python/src/coreai_models/embedding/export.py | 19 ++++- .../src/coreai_models/embedding/export_ios.py | 19 +++-- .../coreai_models/embedding/export_macos.py | 9 +- .../coreai_models/embedding/export_utils.py | 26 ++++-- python/src/coreai_models/models/base.py | 16 +++- .../models/ios/embedding_gemma2.py | 19 +++-- .../models/ios/embedding_gemma2_vision.py | 30 +++++-- .../models/macos/embedding_gemma2.py | 62 +++++++++----- .../models/macos/embedding_gemma2_vision.py | 83 +++++++++++-------- .../coreai_models/primitives/ios/rms_norm.py | 40 +++------ .../test_export/test_embedding_export.py | 20 +++-- .../test_model_units/test_model_registry.py | 2 +- .../test_models/_embedding_gemma2_common.py | 2 +- .../test_ios_layers/test_embedding_gemma2.py | 6 +- .../test_embedding_gemma2.py | 8 +- .../test_embedding_gemma2_vision.py | 15 ++-- 19 files changed, 267 insertions(+), 174 deletions(-) diff --git a/models/README.md b/models/README.md index 27d913df..85173a2b 100644 --- a/models/README.md +++ b/models/README.md @@ -261,7 +261,7 @@ uv run models//export.py --include-debug-info # embed debug information ## Adding a Model -To make a new model exportable via short-name, add a `ModelPreset(...)` entry to `LLM_PRESETS` or `DIFFUSION_PRESETS` in [`python/src/coreai_models/model_registry.py`](../python/src/coreai_models/model_registry.py). Set the short name, HuggingFace ID, family, variant, and the export defaults (compression, compute precision, max context length). Embedding models are registered as `embedding` utility models in the registry and as a family in `coreai_models/embedding/contract.py`. +To make a new model exportable via short-name, add a `ModelPreset(...)` entry to `LLM_PRESETS` or `DIFFUSION_PRESETS` in [`python/src/coreai_models/model_registry.py`](../python/src/coreai_models/model_registry.py). Set the short name, HuggingFace ID, family, variant, and the export defaults (compression, compute precision, max context length). Models exported by `coreai.embedding.export` need an `EmbeddingFamily` in `EMBEDDING_FAMILIES` ([`python/src/coreai_models/embedding/contract.py`](../python/src/coreai_models/embedding/contract.py)) and a `UtilityModel` entry with task `embedding` in `UTILITY_PRESETS`. For models with bespoke export logic that doesn't fit the standard `coreai.llm.export` / `coreai.diffusion.export` flow, write a standalone recipe under `models//export.py` — see existing recipes for the [PEP 723](https://peps.python.org/pep-0723/) pattern and `models/README.md` for the contribution checklist. diff --git a/models/embedding_gemma2/README.md b/models/embedding_gemma2/README.md index 7e57f138..e316939b 100644 --- a/models/embedding_gemma2/README.md +++ b/models/embedding_gemma2/README.md @@ -1,6 +1,6 @@ # EmbeddingGemma 2 -Multimodal retrieval embedding model that maps text and images into one shared space: each input becomes a single 768-dimensional, L2-normalized vector, and its leading 512, 256 or 128 dimensions are valid smaller embeddings.[^1] +Multimodal retrieval embedding model that maps text and images into one shared space: each input becomes a single 768-dimensional, L2-normalized vector, and its leading 512, 256 or 128 dimensions, re-normalized, are valid smaller embeddings.[^1] ## Setup @@ -16,7 +16,7 @@ brew install uv uv run coreai.embedding.export embeddinggemma-2 ``` -Saves to `/exports/embeddinggemma-2/` (text and image towers, GPU). Pass `--output-dir ` to override the destination. +Saves to `/exports/embeddinggemma-2/` (text and image towers, GPU). `--platform iOS` writes `embeddinggemma-2-ios/` and `--skip-vision` appends `-text`. Pass `--output-dir ` to change the parent directory. ```sh uv run coreai.embedding.export --help @@ -24,14 +24,16 @@ uv run coreai.embedding.export --help **Options:** -| Flag | Description | Default | -| --------------------- | ---------------------------------------------------- | ------------------------------------ | -| `--platform` | `macOS` (GPU bundle) or `iOS` (Neural Engine bundle) | `macOS` | -| `--skip-vision` | Export the text tower only | — | -| `--compute-precision` | `float32`, `bfloat16` on macOS; `float16` on iOS | `float32` on macOS, `float16` on iOS | -| `--output-dir` | Parent directory for the bundle | `/exports/` | -| `--overwrite` | Overwrite an existing bundle | — | -| `--dry-run` | Print the export plan without writing | — | +| Flag | Description | Default | +| ---------------------- | ---------------------------------------------------- | ------------------------------------ | +| `--platform` | `macOS` (GPU bundle) or `iOS` (Neural Engine bundle) | `macOS` | +| `--skip-vision` | Export the text tower only | — | +| `--compute-precision` | `float32`, `bfloat16` on macOS; `float16` on iOS | `float32` on macOS, `float16` on iOS | +| `--output-dir` | Parent directory for the bundle | `/exports/` | +| `--overwrite` | Overwrite an existing bundle | — | +| `--include-debug-info` | Embed conversion debug info in the `.aimodel` files | — | +| `--dry-run` | Print the export plan without writing | — | +| `--list-models` | List the supported models and exit | — | **Supported models:** @@ -50,6 +52,6 @@ uv run coreai.embedding.export --help | `encoder` | `text_encoder` (images only) | `encoder` (Neural Engine) | `embeds_len{L}_b1` | `float32` macOS, `float16` iOS | | `image_encoder` | `image_encoder` | `image_encoder` (Neural Engine) | `main` | `float32` macOS, `float16` iOS | -The encoder functions are named `embeds_len{L}_b1` on both platforms; read the names from `components.encoder.functions`. +The lengths `{L}` differ by platform; read the function names from `embedding.components..functions`. [^1]: [Documentation](https://ai.google.dev/gemma/docs/embeddinggemma) · [HuggingFace](https://huggingface.co/google/embeddinggemma-2) diff --git a/python/src/coreai_models/embedding/contract.py b/python/src/coreai_models/embedding/contract.py index c2854b97..15e9b997 100644 --- a/python/src/coreai_models/embedding/contract.py +++ b/python/src/coreai_models/embedding/contract.py @@ -5,9 +5,9 @@ """Builds the ``embedding`` block and ``metadata.json`` for ``kind: "embedding"`` bundles. -Also holds the per-family export facts the exporter and the Swift runtime both read. -EmbeddingGemma 2 overflows in float16 on the GPU, so macOS bundles compute in float32 or -bfloat16, while Neural Engine (iOS) bundles use float16 with scale-invariant norms. +Also defines the per-family export facts (``EmbeddingFamily``). iOS (Neural Engine) bundles +compute in float16 with scale-invariant norms. EmbeddingGemma 2 macOS bundles compute in float32 +or bfloat16; float16 overflows on the GPU. """ from __future__ import annotations @@ -74,11 +74,7 @@ def sequence_length(self) -> int: @dataclass(frozen=True) class EmbeddingFamily: - """Facts about one retrieval-embedding family that the bundle must carry. - - Values come from the HF snapshot files and model card; the space *revision* is the snapshot - SHA of the weights actually exported. - """ + """Facts about one retrieval-embedding family that the bundle must carry.""" key: str hf_model_id: str @@ -91,8 +87,7 @@ class EmbeddingFamily: renormalize_after_truncation: bool = True pooling_mode: str = "mean" include_prompt: bool = True - # Precision of the GPU encoder, which runs the reference math. The Neural Engine encoder - # rescales its norms to compute in float16. + # Default macOS (GPU) compute precision; iOS bundles always compute in float16. default_compute_precision: str = "float32" precision_floor: str | None = None forbidden_precisions: tuple[str, ...] = () @@ -205,8 +200,9 @@ def consolidate_prompts( ) -> dict[str, str]: """Collapse the snapshot's prompt table to canonical name -> prefix, in family order. - Fails closed: a mapped snapshot key that is missing, or a snapshot prefix that no canonical - prompt carries (a new upstream task), raises instead of being dropped. + Raises: + ValueError: if a snapshot key mapped in ``family.prompt_names`` is missing, or a + snapshot prompt uses a prefix that no canonical prompt carries. """ missing = [key for key in family.prompt_names.values() if key not in snapshot_prompts] if missing: @@ -243,8 +239,11 @@ def bucket_plan( ) -> list[FunctionSpec]: """Cross length buckets with batch sizes, dropping pairs above ``max_batch_tokens``. - Batch 1 is always kept for every length (the unbatched path). Attention cost is O(B * L^2), so - the token cap is what keeps the large buckets small. + Batch 1 is always kept for every length. Functions are named ``len{L}_b{B}``. + + Raises: + ValueError: if ``length_buckets`` is empty or holds a length below 2, or a batch size is + below 1. """ lengths = sorted(set(int(x) for x in length_buckets)) batches = sorted(set(int(x) for x in batch_sizes) | {1}) @@ -268,10 +267,11 @@ def text_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec] def encoder_functions(family: EmbeddingFamily, platform: str) -> list[FunctionSpec]: - """Batch-1 ``inputs_embeds`` functions of the ``encoder`` component. + """Batch-1 ``inputs_embeds`` functions ``embeds_len{L}_b1`` of the ``encoder`` component. - Named ``embeds_len{L}_b1`` on both platforms; the host reads the names from - ``components.encoder.functions``. + Lengths are ``ane_length_buckets`` on iOS. On macOS (image tower required) they are the + longest image sequence rounded up to a multiple of 32, plus every longer ``length_buckets`` + entry. The host reads the names from ``components.encoder.functions``. """ if platform == "iOS": lengths = sorted(set(family.ane_length_buckets)) @@ -340,7 +340,8 @@ def embedding_block( Raises: ValueError: on an unknown platform, a missing revision or prompts, an unknown - ``default_prompt_name``, or family facts that do not fit the function plan. + ``default_prompt_name``, ``with_image`` on a family without an image tower, or family + facts that do not fit the function plan. EmbeddingPrecisionError: if ``compute_precision`` is not allowed on ``platform``. """ ios = platform == "iOS" @@ -500,7 +501,7 @@ def build_bundle_metadata( assets: dict[str, str], hf_revision: str | None, ) -> dict[str, Any]: - """Wrap an :func:`embedding_block` in the 0.2 ``metadata.json`` envelope.""" + """Wrap an :func:`embedding_block` in the ``metadata.json`` envelope.""" return { "metadata_version": METADATA_VERSION, "kind": BUNDLE_KIND, diff --git a/python/src/coreai_models/embedding/export.py b/python/src/coreai_models/embedding/export.py index 34399dc9..661abe84 100644 --- a/python/src/coreai_models/embedding/export.py +++ b/python/src/coreai_models/embedding/export.py @@ -84,10 +84,23 @@ def export_embedding_bundle( ) -> dict[str, Any]: """Export one bundle to ``output_dir`` and return its ``metadata.json`` content. - macOS writes a GPU bundle: ``text_encoder.aimodel``; with images also + macOS writes a GPU bundle: ``text_encoder.aimodel``; if ``with_image``, also ``image_encoder.aimodel``, and ``text_encoder`` then holds the ``token_table`` and ``encoder`` - functions. iOS writes the Neural Engine bundle: ``token_table.aimodel``, ``encoder.aimodel``, - plus ``image_encoder.aimodel``. Both carry the tokenizer and a ``metadata.json``. + functions. iOS writes a Neural Engine bundle: ``token_table.aimodel`` and ``encoder.aimodel``; + if ``with_image``, also ``image_encoder.aimodel``. Both carry the tokenizer and a + ``metadata.json``. ``compute_precision`` defaults to the family's default on macOS and + float16 on iOS. ``model_dir`` is a local HF cache snapshot (``.../snapshots/``); without + it the snapshot is downloaded. A failed export removes ``output_dir``. + + Raises: + ValueError: on an unknown platform, ``with_image`` for a family without an image tower, + a ``model_dir`` outside the HF cache, or snapshot configs or prompts that disagree + with the family. + EmbeddingPrecisionError: if ``compute_precision`` is unknown or not allowed for the + family on ``platform``. + FileExistsError: if ``output_dir`` exists and ``overwrite`` is False. + FileNotFoundError: if ``model_dir`` is not a directory or the snapshot has no + ``tokenizer.json``. """ if platform not in PLATFORMS: raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") diff --git a/python/src/coreai_models/embedding/export_ios.py b/python/src/coreai_models/embedding/export_ios.py index ee2bed3b..907910f9 100644 --- a/python/src/coreai_models/embedding/export_ios.py +++ b/python/src/coreai_models/embedding/export_ios.py @@ -3,9 +3,9 @@ # Use of this source code is governed by a BSD-3-clause license that can # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause -"""Neural Engine (iOS) programs of an embedding bundle. +"""iOS programs of an embedding bundle. -The static-length encoder, the float16 token table, and the image encoder. +The float16 token table (GPU), the static-length encoder and the image encoder (Neural Engine). """ from __future__ import annotations @@ -44,15 +44,15 @@ from coreai.authoring import AIProgram -# The Neural Engine programs compute in this dtype and output it. +# The iOS programs compute in this dtype and output it. DTYPE = TORCH_DTYPES[output_dtype("iOS")] class EncoderBandWrapper(nn.Module): """``(inputs_embeds, bias, pool_weights) -> embedding`` at one static length (Neural Engine). - When ``length > sliding_window + 1`` the sliding layers also need the ``|i - j| <= window`` - band; it is a per-length constant folded in here (``min(bias, band)``), so every length keeps + When ``length > sliding_window + 1``, the sliding-window layers use ``min(bias, band)``, with + ``band`` a constant ``|i - j| <= sliding_window`` mask for that length. Every length takes the same three inputs. """ @@ -93,8 +93,8 @@ def convert_ane_functions( *, include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, ) -> AIProgram: - """Convert each ``(name, module, args, input_names, output_names)`` into one Neural Engine - program. + """Convert ``(name, module, args, input_names, output_names)`` entries into one Neural Engine + program with one function per entry. Uses the default decompositions plus the repo's custom lowerings, as ``export/ios.py`` does. """ @@ -153,6 +153,11 @@ def export_ane_programs( with_image: bool, include_debug_info: bool, ) -> None: + """Write the ``token_table``, ``encoder`` and, if ``with_image``, ``image_encoder`` assets to + ``out``. + + Raises ``NotImplementedError`` for a family other than ``embeddinggemma-2``. + """ if family.key != "embeddinggemma-2": raise NotImplementedError(f"No Neural Engine export for family {family.key!r}") from coreai_models.models.ios.embedding_gemma2 import load_embedding_gemma2_text_encoder diff --git a/python/src/coreai_models/embedding/export_macos.py b/python/src/coreai_models/embedding/export_macos.py index cd146317..412eb9c6 100644 --- a/python/src/coreai_models/embedding/export_macos.py +++ b/python/src/coreai_models/embedding/export_macos.py @@ -5,7 +5,8 @@ """GPU (macOS) programs of an embedding bundle. -The text encoder, the shared text / token-table / encoder program, and the image encoder. +The ``text_encoder`` program (with images it also holds the ``token_table`` and ``encoder`` +functions) and the ``image_encoder`` program. """ from __future__ import annotations @@ -99,7 +100,7 @@ def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> tor @contextmanager def without_token_table(encoder: nn.Module) -> Iterator[nn.Module]: - """Temporarily drop ``encoder.embed_tokens`` so an ``inputs_embeds`` trace carries no table.""" + """Swap ``encoder.embed_tokens`` for ``nn.Identity`` inside the block; restore it on exit.""" table = encoder.embed_tokens encoder.embed_tokens = nn.Identity() try: @@ -190,6 +191,10 @@ def export_gpu_programs( compute_precision: str, include_debug_info: bool, ) -> None: + """Write the ``text_encoder`` and, if ``with_image``, ``image_encoder`` assets to ``out``. + + Raises ``NotImplementedError`` for a family other than ``embeddinggemma-2``. + """ if family.key != "embeddinggemma-2": raise NotImplementedError(f"No GPU export for family {family.key!r}") from coreai_models.models.macos.embedding_gemma2 import load_embedding_gemma2_text_encoder diff --git a/python/src/coreai_models/embedding/export_utils.py b/python/src/coreai_models/embedding/export_utils.py index 7f6bc415..ce4d0431 100644 --- a/python/src/coreai_models/embedding/export_utils.py +++ b/python/src/coreai_models/embedding/export_utils.py @@ -58,7 +58,8 @@ def resolve_snapshot(hf_model_id: str, model_dir: str | None = None) -> tuple[Pa """Return ``(snapshot_dir, revision_sha)`` for a local dir or a hub download. The SHA is the HF cache snapshot directory name (``.../snapshots/``); a plain directory - outside the cache yields ``""``. + outside the cache yields ``""``. Raises ``FileNotFoundError`` if ``model_dir`` is not a + directory. """ if model_dir is not None: path = Path(model_dir).expanduser().resolve() @@ -82,7 +83,12 @@ def read_prompts(snapshot: Path) -> tuple[dict[str, str], str | None]: def verify_pooling_config(snapshot: Path, family: EmbeddingFamily) -> None: - """Fail closed if the snapshot's sentence-transformers pooling disagrees with the family.""" + """Raise ``ValueError`` if the snapshot's pooling config disagrees with the family. + + Checks the pooling mode, ``include_prompt``, the embedding dimension and, for a normalized + family, the ``2_Normalize`` module; a snapshot without ``1_Pooling/config.json`` only logs a + warning. + """ pooling = snapshot / "1_Pooling" / "config.json" if not pooling.exists(): logger.warning("No 1_Pooling/config.json in %s; trusting the family defaults", snapshot) @@ -107,7 +113,10 @@ def verify_pooling_config(snapshot: Path, family: EmbeddingFamily) -> None: def copy_tokenizer(snapshot: Path, dest: Path) -> list[str]: - """Copy the tokenizer files verbatim (no re-serialization through a tokenizer class).""" + """Copy the snapshot's tokenizer files verbatim into ``dest`` and return the copied names. + + Raises ``FileNotFoundError`` if the snapshot has no ``tokenizer.json``. + """ dest.mkdir(parents=True, exist_ok=True) copied = [] for name in TOKENIZER_FILES: @@ -133,7 +142,7 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor: def scaled_token_table(table: torch.Tensor, scale: float, dtype: torch.dtype) -> nn.Embedding: - """``table * scale`` stored as ``dtype``. Scaling in float32 first keeps the cast exact.""" + """Return an ``nn.Embedding`` of ``table * scale`` (computed in float32) as ``dtype``.""" return nn.Embedding.from_pretrained((table.float() * scale).to(dtype)) @@ -150,7 +159,7 @@ def image_reference_inputs( from coreai_models.models.macos.embedding_gemma2_vision import pad_to_budget, patch_position_ids k = image.pooling_kernel_size - # Square-ish grid of budget soft tokens (blocks) -> (bh*k) x (bw*k) patches, row-major. + # Square-ish grid of at most budget blocks (soft tokens) -> (bh*k) x (bw*k) patches, row-major. bw = max(1, int(image.budget**0.5)) bh = max(1, image.budget // bw) generator = torch.Generator().manual_seed(image.budget) @@ -189,7 +198,12 @@ def export_token_table_program( def verify_image_config(snapshot: Path, family: EmbeddingFamily) -> None: - """Fail closed if the snapshot's merge ids / image processor disagree with the family facts.""" + """Raise ``ValueError`` if the snapshot's image config disagrees with ``family.image``. + + Checks the placeholder, begin-image and end-image token ids and the text hidden size in + ``config.json``, and the image processor settings in ``processor_config.json``. Also raises + if the family has no image tower. + """ image = family.image if image is None: raise ValueError(f"{family.hf_model_id} has no image tower") diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index cbc12a1e..982d8524 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -232,7 +232,15 @@ def _save_and_mmap_safetensors( def resolve_safetensors_files(model_dir: str) -> list[str]: - """Resolve the safetensors file paths in a local HuggingFace model directory.""" + """Return the safetensors paths of a local HuggingFace model directory. + + Returns the shards listed in ``model.safetensors.index.json`` if present, else + ``model.safetensors``. + + Raises: + FileNotFoundError: if neither file exists or an indexed shard is missing. + RuntimeError: if the index has no ``weight_map``. + """ single_path = os.path.join(model_dir, "model.safetensors") index_path = os.path.join(model_dir, "model.safetensors.index.json") @@ -310,10 +318,10 @@ def load_tensors_for_keys( key_to_file: dict[str, str], target_dtype: torch.dtype, ) -> dict[str, torch.Tensor]: - """Load specific tensors from safetensors files by key. + """Load the tensors in ``key_to_file`` (tensor name -> safetensors path), keyed by name. - Each safetensors file is opened at most once. Tensors are cast to - ``target_dtype`` except embedding tables and quantization zero-points. + Tensors are cast to ``target_dtype`` except those whose names contain ``embedding_table`` or + ``zero_point``. """ file_to_keys: dict[str, list[str]] = {} for key, path in key_to_file.items(): diff --git a/python/src/coreai_models/models/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py index 3cdc5fc2..15315a39 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -26,8 +26,8 @@ from coreai_models.primitives.ios.rms_norm import MaxAbsRMSNorm, MaxAbsRMSNormNoScale from coreai_models.primitives.ios.sdpa import SDPA -# Smallest positive normal float16; floors the output L2's row max-abs and squared norm so an -# all-zero row stays finite. +# Smallest positive normal float16; floor for the row max-abs and squared norm in the output L2 +# normalization. _FP16_TINY = 6.103515625e-05 MASK_NEG = -30000.0 @@ -56,8 +56,8 @@ def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int) -> None: self.sdpa = SDPA(head_dim=hd, scale=1.0) def forward(self, x, cos, sin, bias): - """``x``: ``(B, C, 1, L)``; ``cos``/``sin``: ``(1, 1, hd, L)``; norms and RoPE run per head - on ``(B, n, hd, L)``.""" + """``x``: ``(B, C, 1, L)``; ``cos``/``sin``: ``(1, 1, hd, L)``; ``bias``: additive + ``(1, L, 1, L)``. Returns ``(B, C, 1, L)``.""" b, _, _, length = x.shape hd, half = self.head_dim, self.head_dim // 2 q = self.q_norm(self.q_proj(x).reshape(b, self.n_heads, hd, length)) @@ -102,10 +102,12 @@ def forward(self, x, ple_i, cos, sin, bias): class EmbeddingGemma2TextEncoder(nn.Module): - """``(inputs_embeds [B,L,H], bias [1,L,1,L], pool_weights [B,L]) -> embedding [B, E]``. + """``(inputs_embeds [B, L, H], bias [1, L, 1, L], pool_weights [B, L]) -> embedding [B, E]``. - ``sliding_bias`` optionally replaces ``bias`` on sliding layers. ``pool_weights`` is the - attention mask divided by its row sum, so pooling is a masked mean. + ``H`` is ``hidden_size`` and ``E`` is ``embedding_dim``. ``inputs_embeds`` are token + embeddings scaled by ``sqrt(hidden_size)``. ``bias`` is additive (0 or ``MASK_NEG``); + ``sliding_bias`` optionally replaces it on sliding layers. ``pool_weights`` is the attention + mask divided by its row sum, so pooling is a masked mean. The output has unit L2 norm. """ def __init__(self, config: EmbeddingGemma2TextConfig) -> None: @@ -171,7 +173,8 @@ def text_encoder_from_state_dict( ) -> tuple[EmbeddingGemma2TextEncoder, torch.Tensor]: """Build the encoder from HF text weights (``language_model.`` prefix stripped). - Returns ``(encoder, embed_table_fp32)``; the embedding table stays on the host. + Returns ``(encoder, embed_table_fp32)``. The table is the unscaled ``[vocab, hidden_size]`` + token embedding, not part of the encoder. """ sd = {k: v.float() for k, v in state_dict.items()} table = sd.pop("embed_tokens.weight") diff --git a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py index df08108d..09bb1c86 100644 --- a/python/src/coreai_models/models/ios/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -99,7 +99,11 @@ def forward(self, x, cos, sin, bias): class EmbeddingGemma2ImageEncoder(nn.Module): - """``pixel_values [1, P, 768]`` + ``position_ids [1, P, 2]`` (float16) -> ``[1, L, 512]``.""" + """``(pixel_values [B, P, patch_dim], position_ids [B, P, 2]) -> [B, L, text_hidden_size]``. + + Padding patches have position ``-1``. ``L`` is ``max_soft_tokens``; the image's soft tokens + fill the first slots and the rest are zero. Inputs are cast to the model dtype. + """ def __init__(self, config: EmbeddingGemma2VisionConfig | None = None): super().__init__() @@ -128,6 +132,12 @@ def __init__(self, config: EmbeddingGemma2VisionConfig | None = None): @torch.no_grad() def bake_tables(self, position_embedding_table: torch.Tensor) -> None: + """Fill the position tables from ``position_embedding_table [2, N, hidden_size]`` and the + RoPE tables from the config. + + Raises: + ValueError: if ``N`` is smaller than ``self.grid``. + """ cfg = self.config g = self.grid table = position_embedding_table.float() @@ -165,8 +175,8 @@ def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> tor for layer in self.layers: x = layer(x, cos, sin, bias) - # 3x3 pooling as one matmul with exact one-hot slot weights (sum; the scale cancels in - # the scale-free norm below, like the dropped x sqrt(hidden)). + # ``k``-by-``k`` block pooling as one matmul with exact one-hot slot weights. It sums where + # HF takes the mean times ``sqrt(hidden_size)``; the scale-free norm below cancels this. bx = ox @ self.block_tab.to(dt) # [B, P, 1] by = oy @ self.block_tab.to(dt) max_x = pos[..., 0:1].amax(1, keepdim=True) # [B, 1, 1] @@ -175,8 +185,8 @@ def forward(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) -> tor w = one_hot_f16(slot, self.slot_ar.to(dt)) * valid # [B, P, L] pooled = torch.matmul(x.squeeze(2), w) # (B, C, L) y = self.embedding_pre_projection_norm(pooled.unsqueeze(2)) # (B, C, 1, L) - y = self.embedding_projection(y) # (B, 512, 1, L) - return y.squeeze(2).transpose(1, 2).contiguous() # [B, L, 512] + y = self.embedding_projection(y) # (B, text_hidden_size, 1, L) + return y.squeeze(2).transpose(1, 2).contiguous() # [B, L, text_hidden_size] def load_image_encoder( @@ -184,7 +194,7 @@ def load_image_encoder( dtype: torch.dtype = torch.float16, max_soft_tokens: int | None = None, ) -> EmbeddingGemma2ImageEncoder: - """Load from a local EG2 snapshot directory (``config.json`` + ``model.safetensors``).""" + """Load from a local EG2 snapshot directory (``config.json`` and its safetensors weights).""" mac = MacEmbeddingGemma2ImageEncoder.from_pretrained(model_dir, torch.float32, max_soft_tokens) model = load_hf_state_dict(EmbeddingGemma2ImageEncoder(mac.config), mac.state_dict()) return model.to(dtype).eval() @@ -194,7 +204,13 @@ def load_hf_state_dict( model: EmbeddingGemma2ImageEncoder, sd: dict[str, torch.Tensor] ) -> EmbeddingGemma2ImageEncoder: """Load checkpoint-named (``vision_tower.*`` / ``embed_vision.*``) or macOS-encoder-named - tensors (fail closed).""" + tensors into ``model``, bake its lookup tables, and return it. + + Raises: + RuntimeError: on missing, unexpected, or mis-shaped tensors. + KeyError: if ``patch_embedder.position_embedding_table`` is absent. + ValueError: if the position table is shorter than ``model.grid``. + """ mapped: dict[str, torch.Tensor] = {} for key, v in sd.items(): k2 = key.replace(".linear.weight", ".weight") diff --git a/python/src/coreai_models/models/macos/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py index b09b7628..2e5631be 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -6,9 +6,9 @@ """EmbeddingGemma 2 (EG2) text encoder for Core AI export. Bidirectional encoder with per-layer inputs and a sentence-embedding head: masked mean over the -valid tokens (prompt included), bias-free projection to 768, then L2 normalization. Masks derive -from the static sequence length and ``attention_mask``, so the graph has no data-dependent control -flow. The GPU path runs in float32 or bfloat16; the float16 variant lives in the iOS modules. +valid tokens (prompt included), bias-free projection to ``embedding_dim``, then L2 normalization. +Masks are built from the static sequence length and ``attention_mask``. macOS bundles compute in +float32 or bfloat16; the float16 Neural Engine variant is ``models.ios.embedding_gemma2``. """ from __future__ import annotations @@ -39,16 +39,15 @@ def rope_inv_freq(dim: int, base: float, device: torch.device | None = None) -> torch.Tensor: - """RoPE inverse frequencies ``base ** (-arange(0, dim, 2) / dim)`` in float32.""" + """Return RoPE inverse frequencies ``base ** (-arange(0, dim, 2) / dim)`` in float32.""" return 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) @dataclass class EmbeddingGemma2TextConfig: - """Text-tower config parsed from the HF ``text_config`` block. + """Text-tower config parsed from the ``text_config`` block of an HF ``config.json``. - Reads ``config.json`` directly, so export runs on the repo's pinned transformers; the HF - EmbeddingGemma 2 classes need transformers>=5.19. + Parsed without transformers; the HF EmbeddingGemma 2 classes need transformers>=5.19. """ vocab_size: int = 262_144 @@ -92,7 +91,11 @@ def is_sliding(self, layer_idx: int) -> bool: return self.layer_types[layer_idx] == "sliding_attention" def rope_head_dim(self, sliding: bool) -> int: - """Head dim of the sliding (or full) layers; ``head_dim`` when there are none.""" + """Return the head dim of the sliding (or full) layers; ``head_dim`` when there are none. + + Raises: + ValueError: if those layers have more than one head dim. + """ dims = { self.layer_head_dim(i) for i in range(self.num_hidden_layers) @@ -106,8 +109,14 @@ def rope_head_dim(self, sliding: bool) -> int: def from_dict(cls, text_config: dict[str, Any]) -> EmbeddingGemma2TextConfig: """Build from an HF ``text_config`` dict (``config.json["text_config"]``). - ``layer_types``, ``per_layer_config`` and both ``rope_parameters`` entries define the layer - layout, so they are required; the other fields default to EmbeddingGemma 2's values. + ``layer_types``, ``per_layer_config`` and ``rope_parameters`` (a ``rope_theta`` for both + ``sliding_attention`` and ``full_attention``) are required; other missing or null fields + take EmbeddingGemma 2's values. + + Raises: + ValueError: if a required field is missing, a ``rope_type`` is not ``"default"``, + ``layer_types`` does not have ``num_hidden_layers`` entries, or + ``hidden_activation`` is not ``"gelu_pytorch_tanh"``. """ missing = [ key @@ -153,9 +162,8 @@ def build_attention_masks( Matches ``create_bidirectional_mask`` / ``create_bidirectional_sliding_window_mask`` in transformers (``masking_utils.sliding_window_bidirectional_overlay`` + - ``padding_mask_function``). Rows of padded queries may be fully masked; the SDPA composite - converts the boolean mask into a finite additive bias, so those rows stay finite and they are - excluded from pooling anyway. + ``padding_mask_function``). Rows of padded queries may be fully masked; their outputs stay + finite and are excluded from pooling. """ batch, seq_len = attention_mask.shape key_valid = (attention_mask != 0).reshape(batch, 1, 1, seq_len) @@ -366,7 +374,15 @@ def encode_tokens( attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor: - """Run the tower; returns final-normed hidden states ``[B, L, hidden_size]``.""" + """Run the tower; returns final-normed hidden states ``[B, L, hidden_size]``. + + Pass exactly one of ``input_ids`` ``[B, L]`` or ``inputs_embeds`` ``[B, L, hidden_size]`` + (already scaled, as ``embed_tokens`` returns). ``attention_mask`` ``[B, L]`` defaults to + all ones. + + Raises: + ValueError: if both or neither of ``input_ids`` and ``inputs_embeds`` are given. + """ if (input_ids is None) == (inputs_embeds is None): raise ValueError("Pass exactly one of input_ids or inputs_embeds") if inputs_embeds is None: @@ -411,9 +427,9 @@ def forward( """Return ``sentence_embedding`` ``[B, embedding_dim]`` (unit L2 norm, fp32). ``pooling_mask`` defaults to ``attention_mask`` (include_prompt=True); pass a narrower - mask to exclude e.g. prompt tokens from the mean. With ``return_token_states`` the - per-token projected states ``[B, L, embedding_dim]`` (HF ``last_hidden_state`` - equivalent) are returned as a second output for parity checks. + mask to exclude e.g. prompt tokens from the mean. With ``return_token_states``, also + returns the per-token projected states ``[B, L, embedding_dim]`` (HF ``last_hidden_state``) + as a second output, unnormalized and in the model dtype. """ hidden = self.encode_tokens(input_ids, attention_mask, inputs_embeds) if attention_mask is None: @@ -429,14 +445,16 @@ def forward( def read_checkpoint_tensors( model_dir: str | Path, prefixes: tuple[str, ...], dtype: torch.dtype ) -> dict[str, torch.Tensor]: - """Tensors of the ``model_dir`` checkpoint whose names start with one of ``prefixes``. + """Load the tensors of the ``model_dir`` checkpoint whose names start with one of ``prefixes``. - Reads ``model.safetensors`` or the shards listed in ``model.safetensors.index.json``. Names - are kept as in the checkpoint; tensors are cast to ``dtype`` except embedding tables and - quantization zero-points, as in ``load_tensors_for_keys``. + Reads the shards listed in ``model.safetensors.index.json``, else ``model.safetensors``. Names + are kept as in the checkpoint; tensors are cast to ``dtype`` except those whose names contain + ``embedding_table`` or ``zero_point``, as in ``load_tensors_for_keys``. Raises: - FileNotFoundError: if ``model_dir`` has no safetensors weights. + FileNotFoundError: if ``model_dir`` has no safetensors weights or an indexed shard is + missing. + RuntimeError: if ``model.safetensors.index.json`` has no ``weight_map``. """ from safetensors import safe_open diff --git a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py index f80f00c1..f2a44fa4 100644 --- a/python/src/coreai_models/models/macos/embedding_gemma2_vision.py +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -5,10 +5,10 @@ """EmbeddingGemma 2 image encoder for Core AI export (vision tower + ``embed_vision``). -Host-prepared patches ``[B, P, 768]`` and ``position_ids [B, P, 2]`` (-1 for padding) become soft -tokens ``[B, L, 512]`` with P and L static. Padding and the 3x3 pooling slots are derived in-graph -from ``position_ids``, with no data-dependent control flow. The GPU path runs in float32 or -bfloat16; the float16 variant lives in the iOS modules. +Host-prepared patches ``[B, P, patch_dim]`` and ``position_ids [B, P, 2]`` (-1 for padding) become +soft tokens ``[B, L, text_hidden_size]``, with P and L static. Padding and the ``k x k`` pooling +slots come from ``position_ids`` inside the graph. Runs in float32 or bfloat16; the float16 +variant is in ``models/ios/embedding_gemma2_vision.py``. """ from __future__ import annotations @@ -43,7 +43,7 @@ @dataclass class EmbeddingGemma2VisionConfig: - """Vision tower + ``embed_vision`` hyper-parameters (defaults = google/embeddinggemma-2).""" + """Vision tower and ``embed_vision`` hyper-parameters (defaults: google/embeddinggemma-2).""" hidden_size: int = 768 intermediate_size: int = 3072 @@ -71,7 +71,11 @@ def patch_dim(self) -> int: def from_hf_config( cls, config: dict, max_soft_tokens: int | None = None ) -> EmbeddingGemma2VisionConfig: - """Build from an EG2 ``config.json`` dict (top level, with ``vision_config``).""" + """Build from an EG2 ``config.json`` dict (top level, with ``vision_config``). + + ``max_soft_tokens`` overrides the snapshot's ``vision_soft_tokens_per_image``. Raises + ``NotImplementedError`` for clipped linears, standardization or a non-axial RoPE. + """ vc = config["vision_config"] if vc.get("use_clipped_linears", False): raise NotImplementedError("clipped vision linears are not supported") @@ -100,8 +104,11 @@ def from_hf_config( class VisionAttention(EmbeddingGemma2Attention): - """Text attention with axial 2-D RoPE: half of head_dim rotates by x, half by y (cos/sin from - ``axial_rope_cos_sin``, laid out ``[fx, fx, fy, fy]``).""" + """``EmbeddingGemma2Attention`` with axial 2-D RoPE. + + Half of ``head_dim`` rotates by x and half by y, with cos/sin from ``axial_rope_cos_sin`` + laid out ``[fx, fx, fy, fy]``. + """ def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: c = config @@ -113,7 +120,10 @@ def _rope(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch. class VisionPatchEmbedder(nn.Module): - """``Linear(2*(p-0.5))`` plus separable learned x/y position embeddings.""" + """Projection of ``2 * (pixels - 0.5)`` plus learned x and y position embeddings. + + Padding rows get no position embedding. + """ def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: super().__init__() @@ -136,7 +146,7 @@ def forward( class VisionEncoder(nn.Module): - """Holds the layer stack (module path ``encoder.layers.N`` mirrors the checkpoint).""" + """Layer stack at module path ``encoder.layers.N``, matching the checkpoint keys.""" def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: super().__init__() @@ -167,7 +177,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class EmbeddingGemma2ImageEncoder(nn.Module): - """gemma4_vision tower + ``embed_vision``: patches -> soft tokens ``[B, L, text_hidden]``.""" + """Gemma 4 vision tower and ``embed_vision``: patches -> soft tokens ``[B, L, text_hidden]``.""" def __init__(self, config: EmbeddingGemma2VisionConfig | None = None) -> None: super().__init__() @@ -193,20 +203,23 @@ def axial_rope_cos_sin( ) -> tuple[torch.Tensor, torch.Tensor]: """``[B, P, 2]`` positions -> cos/sin ``[B, P, head_dim]`` laid out ``[fx, fx, fy, fy]``. - Padding rows keep their raw ``-1`` coordinates, as in HF; they are masked as keys - and dropped by the pooler, so their values never reach a real soft token. + Padding rows keep their raw ``-1`` coordinates, as in HF; attention masks them and the + pooler drops them. """ freqs = position_ids.unsqueeze(-1).float() * self.rope_inv_freq # [B, P, 2, D/4] fx, fy = freqs[..., 0, :], freqs[..., 1, :] angles = torch.cat([fx, fx, fy, fy], dim=-1) return angles.cos().to(dtype), angles.sin().to(dtype) - def pool(self, hidden: torch.Tensor, position_ids: torch.Tensor, padding: torch.Tensor): - """Mean over each ``k x k`` patch block, then ``x sqrt(hidden)``; fp32 output. + def pool( + self, hidden: torch.Tensor, position_ids: torch.Tensor, padding: torch.Tensor + ) -> torch.Tensor: + """Mean over each ``k x k`` patch block, scaled by ``sqrt(hidden_size)``. - Block index = ``x//k + (max_x+1)//k * (y//k)``. With the host contract (a full - ``(W/k) x (H/k)`` block grid, row-major, padded at the end) the real soft tokens are - exactly the first ``num_soft_tokens`` slots. + Returns ``[B, L, hidden]`` in fp32. Block ``(x, y)`` lands in slot + ``x//k + (W//k) * (y//k)``, where ``W`` is the patch-grid width. For a full + ``(W/k) x (H/k)`` block grid, row-major and padded at the end, the real soft tokens are + the first ``num_soft_tokens`` slots. """ k = self.config.pooling_kernel_size h = hidden.float().masked_fill(padding.unsqueeze(-1), 0.0) @@ -229,9 +242,8 @@ def encode_pooled(self, pixel_values: torch.Tensor, position_ids: torch.Tensor) padding = position_ids[..., 0] < 0 x = self.patch_embedder(pixel_values, position_ids, padding) cos, sin = self.axial_rope_cos_sin(position_ids, x.dtype) - # Boolean key mask (True = attend), the convention of the repo SDPA primitive: its - # Core AI lowering turns ``~mask`` into a -1e4 additive bias. Do not pass an additive - # float mask here; that path would read it as a bool and invert it. + # Boolean key mask, True = attend. The SDPA primitive reads any mask as boolean; do not + # pass an additive float mask. attn_mask = torch.logical_not(padding)[:, None, None, :] for layer in self.encoder.layers: x = layer(x, cos, sin, attn_mask) @@ -260,7 +272,7 @@ def from_pretrained( dtype: torch.dtype = torch.float32, max_soft_tokens: int | None = None, ) -> EmbeddingGemma2ImageEncoder: - """Load from a local EG2 snapshot directory (``config.json`` + ``model.safetensors``).""" + """Load from a local EG2 snapshot directory: ``config.json`` and its safetensors weights.""" config = json.loads((Path(model_dir) / "config.json").read_text()) model = cls(EmbeddingGemma2VisionConfig.from_hf_config(config, max_soft_tokens)) tensors = read_checkpoint_tensors(model_dir, ("vision_tower.", "embed_vision."), dtype) @@ -277,8 +289,11 @@ def aspect_ratio_preserving_size( max_soft_tokens: int = 280, pooling_kernel_size: int = 3, ) -> tuple[int, int]: - """Target ``(height, width)``: the largest size with sides divisible by ``k*patch`` and at - most ``max_soft_tokens * k^2`` patches, preserving aspect ratio (Gemma4ImageProcessor).""" + """Aspect-preserving resize target ``(height, width)``, as Gemma4ImageProcessor computes it. + + It is the largest size with sides divisible by ``pooling_kernel_size * patch_size`` and at + most ``max_soft_tokens * pooling_kernel_size**2`` patches. + """ max_patches = max_soft_tokens * pooling_kernel_size**2 target_px = max_patches * patch_size**2 factor = math.sqrt(target_px / (height * width)) @@ -336,8 +351,7 @@ def preprocess_resized_image( """Already-resized ``[3, H, W]`` RGB (uint8 or 0..255 float) -> model inputs. Returns ``(pixel_values [1, P, patch_dim] fp32, position_ids [1, P, 2] int64, - num_soft_tokens)``. Rescale only (no mean/std normalization: the tower does - ``2*(x-0.5)`` itself). + num_soft_tokens)``. Applies ``rescale_factor`` only; the tower applies ``2 * (x - 0.5)``. """ _, h, w = image.shape side = patch_size * pooling_kernel_size @@ -360,8 +374,10 @@ def build_image_input_ids( eoi_token_id: int = EG2_EOI_TOKEN_ID, image_token_id: int = EG2_IMAGE_TOKEN_ID, ) -> list[int]: - """Ids for an image-only input, as the HF processor emits them: - `` * num_soft_tokens ``.""" + """Ids for an image-only input, as the HF processor emits them. + + `` * num_soft_tokens ``; ``None`` drops ```` or ````. + """ ids = [boi_token_id] + [image_token_id] * num_soft_tokens + [eoi_token_id] if bos_token_id is not None: ids = [bos_token_id] + ids @@ -375,9 +391,8 @@ def scatter_soft_tokens( ) -> torch.Tensor: """Replace placeholder rows with soft tokens, in order (``masked_scatter`` semantics). - Gather-based so it exports at a static length: the ``j``-th placeholder in row ``b`` - takes ``soft_tokens[b, j]``. ``soft_tokens`` may carry extra trailing rows (e.g. the - zero padding slots of a fixed budget); they are never read. + Exports at a static length. The ``j``-th placeholder in row ``b`` takes ``soft_tokens[b, j]``; + extra trailing rows of ``soft_tokens``, such as budget padding, are ignored. Args: inputs_embeds: ``[B, S, D]`` text embeddings. @@ -403,9 +418,9 @@ def build_image_inputs_embeds( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Full text-encoder input for one image per row. - ``embed_tokens`` maps ids -> *scaled* text embeddings (EG2 multiplies the table by - ``sqrt(hidden)``; HF ``language_model.embed_tokens`` already does). Placeholders are - looked up as ```` (like HF) and then overwritten by the soft tokens. + ``embed_tokens`` maps ids to text embeddings already scaled by ``sqrt(hidden_size)``, as HF + ``language_model.embed_tokens`` does. Placeholders are looked up as ````, as in HF, then + replaced by the soft tokens. Returns ``(inputs_embeds [B, S, D], attention_mask [B, S] int64, input_ids [B, S])``, right-padded to the longest row. diff --git a/python/src/coreai_models/primitives/ios/rms_norm.py b/python/src/coreai_models/primitives/ios/rms_norm.py index e649cce5..ed34d32f 100644 --- a/python/src/coreai_models/primitives/ios/rms_norm.py +++ b/python/src/coreai_models/primitives/ios/rms_norm.py @@ -53,49 +53,32 @@ def forward( return x_2normalized * self.weight -# fp16-overflow-safe RMSNorm. RMSNorm is scale-invariant, so we normalize each row -# by its own max magnitude *before* squaring: ``m = amax(|x|)`` (a data-dependent, -# per-row reduction), ``xn = x/m``. Since ``|xn| <= 1``, ``xn^2`` can never overflow -# fp16 (whereas raw ``x^2`` overflows once ``|x|>~256``: ``256^2~=65504``, the fp16 max -# -- the tower's residual stream peaks at ``|x|~=500`` on real images, so raw ``x^2`` -# overflows to ``inf`` -> ``rsqrt(inf)=0`` -> zeroed features -> garbage soft tokens). +# Max-abs RMSNorm: a float16-overflow-safe RMSNorm. The input is divided by its max magnitude +# along ``axis``, ``m = amax(|x|)``, before squaring, and ``eps`` is rescaled to match; the +# result equals plain RMSNorm:: # -# The result is the *exact* RMSNorm because ``eps`` is folded back at the right scale:: +# xn = x / m +# xn * rsqrt(mean(xn^2) + (sqrt(eps) / m)^2) * w == x * rsqrt(mean(x^2) + eps) * w # -# xn*rsqrt(mean(xn^2) + (sqrt(eps)/m)^2) * w -# = (x/m)*rsqrt(mean(x^2)/m^2 + eps/m^2) * w -# = (x/m)*m*rsqrt(mean(x^2) + eps) * w -# = x*rsqrt(mean(x^2) + eps) * w -# -# ``(sqrt(eps)/m)^2`` computes ``eps/m^2`` without ever forming ``m^2`` (which would -# itself overflow fp16 for ``m~=500``). Crucially this survives the ANE graph -# optimizer: a constant prescale (``x*0.125``) gets reassociated back into ``x^2*c^2`` -# (recomputing the overflowing ``x^2``) or stripped as scale-invariant, but -# ``x/amax(|x|)`` is data-dependent and cannot be folded away. All ops -# (``abs``/``amax``/``div``/``mean``/``rsqrt``) are fp16 and ANE-resident -- no fp32 -# upcast. +# Plain RMSNorm overflows float16 once ``|x| > ~256`` (``x^2 > 65504``). ``eps / m^2`` is computed +# as ``(sqrt(eps) / m)^2`` to avoid forming ``m^2``. A constant prescale does not work instead: +# the Neural Engine compiler folds it back into ``x^2``. All ops run in the input dtype. class MaxAbsRMSNormNoScale(nn.Module): """RMSNorm without a learnable scale, reducing over ``axis`` (counted from the end). - fp16-overflow-safe via per-row max-normalization before squaring (see the comment - directly above, and :class:`MaxAbsRMSNorm`, which adds the scale). Exact and ANE-resident (no - fp32 upcast). + Float16-overflow-safe and exact; see the comment above. :class:`MaxAbsRMSNorm` adds the scale. """ def __init__(self, eps: float = 1e-6, axis: int = -1) -> None: super().__init__() self._eps = eps self.axis = axis - # Registered as a buffer (not a plain attribute) so it follows the module's - # device on `.to(device)` -- a plain `torch.tensor(...)` created in forward() - # defaults to CPU and raises on a GPU/MPS input (e.g. during parity tests). self.register_buffer("_eps_sqrt", torch.tensor(eps**0.5), persistent=False) def forward(self, x: torch.Tensor) -> torch.Tensor: - # Per-row max magnitude; clamp floor avoids 0/0 on an all-zero row (real rows - # have m >> floor, so this is a no-op for them). + # Max magnitude along ``axis``; the floor keeps an all-zero row at zero instead of NaN. m = torch.clamp(x.abs().amax(self.axis, keepdim=True), min=1e-4) xn = x / m eps_sqrt = self._eps_sqrt.to(dtype=x.dtype) @@ -107,7 +90,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class MaxAbsRMSNorm(MaxAbsRMSNormNoScale): """:class:`MaxAbsRMSNormNoScale` followed by a learnable per-channel scale. - The scale has shape ``(dim, 1, ..., 1)``, broadcasting over the dims after ``axis``. + ``axis`` must be negative. The scale has shape ``(dim, 1, ..., 1)``, broadcasting over the dims + after ``axis``. """ def __init__(self, dim: int, eps: float = 1e-6, axis: int = -1) -> None: diff --git a/python/tests/test_model_units/test_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py index 6ded27a9..a86ce087 100644 --- a/python/tests/test_model_units/test_export/test_embedding_export.py +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -162,7 +162,7 @@ def test_consolidate_prompts_maps_the_snapshot_table_to_eight_names(): assert list(prompts) == list(EG2.prompt_names) -def test_consolidate_prompts_fails_closed(): +def test_consolidate_prompts_rejects_unknown_and_missing_keys(): with pytest.raises(ValueError, match="canonical"): consolidate_prompts(EG2, {**SNAPSHOT_PROMPTS, "Translation": "task: translate | query: "}) missing = {k: v for k, v in SNAPSHOT_PROMPTS.items() if k != "FactChecking"} @@ -468,7 +468,7 @@ def test_snapshot_helpers(tmp_path): assert copied == ["tokenizer.json", "tokenizer_config.json"] -def test_verify_pooling_config_fails_closed(tmp_path): +def test_verify_pooling_config_rejects_include_prompt_false(tmp_path): snap = _fake_snapshot( tmp_path, {"embedding_dimension": 768, "pooling_mode": "mean", "include_prompt": False} ) @@ -543,7 +543,9 @@ def _tiny_image_encoder(): def test_inputs_embeds_path_equals_ids_path(): - """token_table(ids) -> encoder(inputs_embeds) == text_encoder(ids) (the merge is lossless).""" + """``encoder(token_table(ids))`` matches ``text_encoder(ids)`` within atol 1e-6, and + ``without_token_table`` restores ``embed_tokens`` afterwards. + """ text = _tiny_text_encoder() ids = torch.randint(3, 128, (1, 16), dtype=torch.int32) mask = torch.ones(1, 16, dtype=torch.int32) @@ -559,7 +561,9 @@ def test_inputs_embeds_path_equals_ids_path(): def test_float16_scaled_token_table_equals_the_float32_table_cast(): - """The float16 token table gives the encoder the same inputs as the float32 one.""" + """Lookups in the float16 ``scaled_token_table`` equal the float32 ``embed_tokens`` lookups + cast to float16. + """ text = _tiny_text_encoder() ids = torch.randint(3, 128, (1, 16), dtype=torch.int32) table = export_utils.scaled_token_table( @@ -634,7 +638,7 @@ def test_ane_reference_inputs_shapes(): @pytest.fixture(scope="module") def tiny_ane_programs(): - """Component name -> (PyTorch module, program converted through ``export_ios``).""" + """Component name -> (float16 PyTorch module, its converted program) for the tiny iOS bundle.""" from coreai_models.models.ios import embedding_gemma2_vision as ios_vision from coreai_models.models.ios.embedding_gemma2 import text_encoder_from_state_dict @@ -695,10 +699,8 @@ def test_tiny_ane_programs_have_the_contract_functions(tiny_ane_programs): @pytest.mark.parametrize("compute_unit", ["cpu", "gpu"]) def test_tiny_ane_programs_match_pytorch(tiny_ane_programs, compute_unit): - """Each converted program, run by the Core AI runtime, matches its float16 module and - outputs the metadata's ``output_dtype``. - - The tolerance is the float16 one of ``test_macos_prefill_graph.py`` (atol = rtol = 1e-2). + """Each converted program, run by the Core AI runtime, outputs the metadata's ``output_dtype`` + and matches its float16 module: the token table exactly, the encoders within atol = rtol = 1e-2. """ import coreai.runtime as rt diff --git a/python/tests/test_model_units/test_model_registry.py b/python/tests/test_model_units/test_model_registry.py index 6c55cd3b..673648e8 100644 --- a/python/tests/test_model_units/test_model_registry.py +++ b/python/tests/test_model_units/test_model_registry.py @@ -239,7 +239,7 @@ def test_embedding_export_args_use_the_console_script() -> None: assert _utility_to_export_args(model, "macOS")[-1] == "embeddinggemma-2" -def test_script_utility_export_args_are_unchanged() -> None: +def test_script_utility_export_args_run_the_script() -> None: model = lookup_utility_model("clip-vit-b32") assert model is not None assert _utility_to_export_args(model) == [ diff --git a/python/tests/test_model_units/test_models/_embedding_gemma2_common.py b/python/tests/test_model_units/test_models/_embedding_gemma2_common.py index 38630a95..bf540925 100644 --- a/python/tests/test_model_units/test_models/_embedding_gemma2_common.py +++ b/python/tests/test_model_units/test_models/_embedding_gemma2_common.py @@ -3,7 +3,7 @@ # Use of this source code is governed by a BSD-3-clause license that can # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause -"""Tiny EmbeddingGemma 2 configs and helpers shared by the macOS and iOS test modules.""" +"""Tiny EmbeddingGemma 2 configs and helpers shared by the EmbeddingGemma 2 test modules.""" import torch import torch.nn.functional as F diff --git a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py index 184618df..4faeb4d0 100644 --- a/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -37,7 +37,7 @@ def _host_inputs( embed_table: torch.Tensor, dtype: torch.dtype = torch.float16, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Graph inputs for one sequence (``ids [1, L]``): token-table embeddings, bias, pooling.""" + """Encoder inputs for ``ids [1, L]``: scaled embeddings, bias ``[1, L, 1, L]``, pool weights.""" length = ids.shape[1] table = scaled_token_table(embed_table, math.sqrt(embed_table.shape[1]), dtype) key_ok = attention_mask[0].bool().reshape(length, 1, 1).expand(length, 1, length) @@ -51,7 +51,7 @@ def _banded(ios: torch.nn.Module, length: int) -> EncoderBandWrapper: def _perturb(model: torch.nn.Module, seed: int) -> None: - """Randomize norm weights and layer scalars so the comparison exercises them.""" + """Randomize the norm weights and layer scalars of ``model``.""" gen = torch.Generator().manual_seed(seed) with torch.no_grad(): for name, tensor in sorted(model.state_dict().items()): @@ -91,7 +91,7 @@ def test_float16_stays_finite_with_large_activations() -> None: cfg, _, state = _reference_and_state_dict() _, table = text_encoder_from_state_dict(cfg, state, dtype=torch.float32) ids, mask = _ragged_batch([12], seq_len=12, vocab=cfg.vocab_size) - # Scale the token embeddings so the hidden stream reaches about 3e3, past sqrt(fp16 max). + # Scale the token table to a peak scaled embedding of 3000, past sqrt(float16 max). 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) diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py index 1eb12c66..340a17e8 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -48,7 +48,7 @@ def _randomize_(module: torch.nn.Module, seed: int = 0) -> None: - """Perturb every parameter/buffer (norm weights, layer_scalar included) so bugs are visible.""" + """Overwrite every floating-point parameter and buffer with seeded random values.""" gen = torch.Generator().manual_seed(seed) with torch.no_grad(): for name, tensor in sorted(module.state_dict().items()): @@ -123,7 +123,7 @@ def test_forward_shapes_and_unit_norm(): def test_padding_invariance_tiny(): - """A sequence embedded alone (no padding) == the same sequence inside a padded batch.""" + """A sequence embedded alone (no padding) matches its row of a padded batch (atol 2e-5).""" model = _tiny_model() ids, mask = _ragged_batch([11, 4, 7], seq_len=16, vocab=128) with torch.no_grad(): @@ -283,7 +283,9 @@ def real_setup(embeddinggemma2_snapshot): def test_real_weight_parity_and_padding_invariance(real_setup): - """Includes a row over 2 * 512 + 1 tokens, so the sliding layers' window bites for real.""" + """Ours matches HF hidden states mean-pooled and normalized, and each row embedded alone + matches its padded-batch row (cosine >= 0.99999). One row exceeds 2 * sliding_window + 1 tokens. + """ ours, hf = real_setup["ours"], real_setup["hf"] paragraph = ( "On-device semantic search embeds documents once and compares vectors at query time. " diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py index 93c7e984..0c6d8c35 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py @@ -3,7 +3,7 @@ # Use of this source code is governed by a BSD-3-clause license that can # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause -"""Parity tests for the EmbeddingGemma 2 image path (``embedding_gemma2_vision.py``). +"""Parity tests for the EmbeddingGemma 2 image path (``models/macos/embedding_gemma2_vision.py``). * Tiny-config parity against HF ``Gemma4VisionModel`` + ``Gemma4MultimodalEmbedder`` (present in the repo's transformers; no checkpoint needed). @@ -78,7 +78,7 @@ def _hf_tiny_vision(cfg: EmbeddingGemma2VisionConfig): hf_cfg._attn_implementation = "eager" torch.manual_seed(0) model = gemma4_modeling.Gemma4VisionModel(hf_cfg).eval() - # HF init leaves norms at 1 / zeros; randomize every float tensor so the test bites. + # HF init leaves norms at 1 and biases at 0; overwrite every parameter with random values. with torch.no_grad(): for name, p in model.named_parameters(): if "norm" in name: @@ -100,7 +100,10 @@ def _load_from_hf_vision(ours: EmbeddingGemma2ImageEncoder, hf_vision, hf_embedd def _random_inputs(cfg: EmbeddingGemma2VisionConfig, grids: list[tuple[int, int]], seed: int = 0): - """Random patches for images whose block grid is ``(bh, bw)`` blocks, padded to budget.""" + """Random patches, positions and soft-token counts for images of ``(bh, bw)`` pooling blocks. + + Patches and positions are padded to ``cfg.max_patches``. + """ g = torch.Generator().manual_seed(seed) k = cfg.pooling_kernel_size pv, pos, counts = [], [], [] @@ -179,7 +182,9 @@ def test_aspect_ratio_size_matches_hf(budget: int) -> None: def test_preprocess_matches_hf_processor_without_resize() -> None: - """At the target size HF skips the resize, so patching/positions/padding must be exact.""" + """At sizes the resize policy keeps, ``preprocess_resized_image`` matches the HF processor: + pixels within atol 1e-7, positions and soft-token count exactly. + """ np = pytest.importorskip("numpy") try: from transformers.models.gemma4.image_processing_pil_gemma4 import ( @@ -265,7 +270,7 @@ def test_torch_export_default_budget_static_shape(monkeypatch) -> None: exported = ep.module()(pixel_values, position_ids) assert exported.shape == (1, 280, 512) torch.testing.assert_close(exported, eager, rtol=1e-5, atol=1e-5) - # No dynamic dims and no data-dependent ops (nonzero / masked_select / item) in the graph. + # No data-dependent ops (nonzero / masked_select / item) in the graph. targets = {str(node.target) for node in ep.graph.nodes if node.op == "call_function"} data_dependent = ("aten.nonzero", "aten.masked_select", "aten.item", "aten._local_scalar_dense") assert not any(op in t for t in targets for op in data_dependent), sorted(targets)