diff --git a/models/README.md b/models/README.md index 51e975eb..85173a2b 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). 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 new file mode 100644 index 00000000..e316939b --- /dev/null +++ b/models/embedding_gemma2/README.md @@ -0,0 +1,57 @@ +# 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, re-normalized, 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). `--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 +``` + +**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 | — | +| `--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:** + +| Model | HuggingFace ID | Parameters | +| ---------------- | ----------------------- | ---------------------------- | +| 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 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/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/diffusion/gpu.py b/python/src/coreai_models/diffusion/gpu.py index b61751af..6cfdea64 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 = make_converter(include_debug_info) converter.add_pytorch_module( wrapper, export_fn=export_fn, @@ -99,6 +93,42 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram: return program +def make_converter( + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> coreai_torch.TorchConverter: + """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) + + +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 +152,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 = make_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/__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..15e9b997 --- /dev/null +++ b/python/src/coreai_models/embedding/contract.py @@ -0,0 +1,518 @@ +# 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 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 + +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: + """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 raises on a mismatch. + """ + + placeholder_id: int + boi_token_id: int + eoi_token_id: int + 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: + 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.""" + + 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 macOS (GPU) compute precision; iOS bundles always compute in float16. + 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, + 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} + + +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 validate_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: + 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 {alternatives}." + ) + 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. + + 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: + 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, prefix: str = "") -> str: + return f"{prefix}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. 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}) + 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 ``embeds_len{L}_b1`` of the ``encoder`` component. + + 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)) + 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, "embeds_"), 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") + + +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_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 ()))) + + +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. + """ + 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( + family: EmbeddingFamily, + *, + platform: str, + with_image: bool, + 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``. + + 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. 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 + ``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" + if platform not in PLATFORMS: + raise ValueError(f"platform must be one of {list(PLATFORMS)}, got {platform!r}") + precision = validate_compute_precision( + 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") + 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] + 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: {"asset": component_asset(platform, name), **specs[name]} + for name in bundle_components(platform, with_image) + } + + text_tower: dict[str, Any] = { + "modality": "text", + "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}), + "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": list(IMAGE_PIPELINE), + "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, + # The Neural Engine graphs compute in float16. + "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 ``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..661abe84 --- /dev/null +++ b/python/src/coreai_models/embedding/export.py @@ -0,0 +1,286 @@ +# 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 +from pathlib import Path +from typing import Any + +from coreai_models._constants import DEFAULT_INCLUDE_DEBUG_INFO +from coreai_models.embedding.contract import ( + EMBEDDING_FAMILIES, + PLATFORMS, + EmbeddingFamily, + EmbeddingPrecisionError, + build_bundle_metadata, + bundle_assets, + bundle_name, + 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, +) +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, +) + +logger = logging.getLogger(__name__) + + +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 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.aimodel``; if ``with_image``, also + ``image_encoder.aimodel``, and ``text_encoder`` then holds the ``token_table`` and ``encoder`` + 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}") + 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: + 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) + 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, + ) + + 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") + except BaseException: + shutil.rmtree(out, ignore_errors=True) # a failed export leaves no partial bundle + raise + 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="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", + 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 + with_image = not args.skip_vision + try: + precision = validate_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/embedding/export_ios.py b/python/src/coreai_models/embedding/export_ios.py new file mode 100644 index 00000000..907910f9 --- /dev/null +++ b/python/src/coreai_models/embedding/export_ios.py @@ -0,0 +1,204 @@ +# 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 + +"""iOS programs of an embedding bundle. + +The float16 token table (GPU), the static-length encoder and the image encoder (Neural Engine). +""" + +from __future__ import annotations + +from collections.abc import Sequence +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, + output_dtype, +) +from coreai_models.embedding.export_utils import ( + TORCH_DTYPES, + export_token_table_program, + image_reference_inputs, + save_program, + scaled_token_table, +) + +if TYPE_CHECKING: + from coreai.authoring import AIProgram + + +# 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-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. + """ + + 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(DTYPE) + bias = torch.zeros((1, length, 1, length), dtype=DTYPE) + pool = torch.full((1, length), 1.0 / length, dtype=DTYPE) + return embeds, bias, pool + + +def convert_ane_functions( + functions: Sequence[ + tuple[str, nn.Module, tuple[torch.Tensor, ...], tuple[str, ...], tuple[str, ...]] + ], + *, + include_debug_info: bool = DEFAULT_INCLUDE_DEBUG_INFO, +) -> AIProgram: + """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. + """ + from coreai_models.diffusion.gpu import add_function, make_converter + from coreai_models.export.mlir_ops import register_custom_torch_lowering + + 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: + add_function( + converter, name, module, args, input_names, output_names, decompositions=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 = [ + ( + spec.name, + EncoderBandWrapper(encoder, spec.length, sliding_window), + 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.""" + 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, + ) + + +def export_ane_programs( + family: EmbeddingFamily, + snapshot: Path, + out: Path, + *, + 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 + + # 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=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), + 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), DTYPE, 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..412eb9c6 --- /dev/null +++ b/python/src/coreai_models/embedding/export_macos.py @@ -0,0 +1,225 @@ +# 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`` program (with images it also holds the ``token_table`` and ``encoder`` +functions) and the ``image_encoder`` program. +""" + +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, + output_dtype, + text_functions, +) +from coreai_models.embedding.export_utils import ( + TORCH_DTYPES, + TokenTableExportWrapper, + image_reference_inputs, + save_program, + token_table_trace, +) + +if TYPE_CHECKING: + from coreai.authoring import AIProgram + + +OUTPUT_DTYPE = TORCH_DTYPES[output_dtype("macOS")] + + +class TextEmbeddingExportWrapper(nn.Module): + """``(input_ids, attention_mask) -> embedding`` as ``OUTPUT_DTYPE``. + + 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(OUTPUT_DTYPE) + + +class InputsEmbedsEncoderWrapper(nn.Module): + """``(inputs_embeds, attention_mask) -> embedding`` for merged (e.g. image) inputs. + + Outputs ``OUTPUT_DTYPE``. + """ + + 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(OUTPUT_DTYPE) + + +class ImageEncoderExportWrapper(nn.Module): + """``(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__() + 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(OUTPUT_DTYPE) + + +@contextmanager +def without_token_table(encoder: nn.Module) -> Iterator[nn.Module]: + """Swap ``encoder.embed_tokens`` for ``nn.Identity`` inside the block; restore it on exit.""" + 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_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``; with ``encoder_fns`` it also holds the ``token_table`` and ``encoder``. + + 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 add_function, make_converter + + converter = make_converter(include_debug_info) + text = TextEmbeddingExportWrapper(encoder) + 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, OUTPUT_DTYPE), + ids, + TOKEN_TABLE_INPUT_NAMES, + (TOKEN_TABLE_OUTPUT_NAME,), + dynamic_shapes=dynamic_shapes, + ) + 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,) + ) + 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_gpu_programs( + family: EmbeddingFamily, + snapshot: Path, + out: Path, + *, + with_image: bool, + 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 + from coreai_models.models.macos.embedding_gemma2_vision import EmbeddingGemma2ImageEncoder + + dtype = TORCH_DTYPES[compute_precision] + 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 = 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, + 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..ce4d0431 --- /dev/null +++ b/python/src/coreai_models/embedding/export_utils.py @@ -0,0 +1,249 @@ +# 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 +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__) + + +TORCH_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16} + + +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 ``""``. Raises ``FileNotFoundError`` if ``model_dir`` is not a + directory. + """ + 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: + """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) + 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 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: + 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: + """Return an ``nn.Embedding`` of ``table * scale`` (computed in float32) as ``dtype``.""" + 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. + """ + from coreai_models.models.macos.embedding_gemma2_vision import pad_to_budget, patch_position_ids + + k = image.pooling_kernel_size + # 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) + 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( + 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. + args, dynamic_shapes = token_table_trace(max_length) + return export_stateless( + TokenTableExportWrapper(embed_tokens, dtype), + args, + TOKEN_TABLE_INPUT_NAMES, + (TOKEN_TABLE_OUTPUT_NAME,), + dynamic_shapes=dynamic_shapes, + include_debug_info=include_debug_info, + ) + + +def verify_image_config(snapshot: Path, family: EmbeddingFamily) -> None: + """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") + 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/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/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index f583b208..982d8524 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -231,8 +231,16 @@ def _save_and_mmap_safetensors( module.load_state_dict(new_sd, assign=True, strict=False) -def _resolve_safetensors_files(model_dir: str) -> list[str]: - """Resolve the safetensors file paths in a local HuggingFace model directory.""" +def resolve_safetensors_files(model_dir: str) -> list[str]: + """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") @@ -306,14 +314,14 @@ 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]: - """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(): @@ -876,7 +884,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 +893,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 +911,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/ios/embedding_gemma2.py b/python/src/coreai_models/models/ios/embedding_gemma2.py new file mode 100644 index 00000000..15315a39 --- /dev/null +++ b/python/src/coreai_models/models/ios/embedding_gemma2.py @@ -0,0 +1,208 @@ +# 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. + +The macOS encoder math as 1x1 Conv2d projections with static shapes and max-abs RMSNorms. +""" + +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 ( + EmbeddingGemma2TextConfig, + 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.rms_norm import MaxAbsRMSNorm, MaxAbsRMSNormNoScale +from coreai_models.primitives.ios.sdpa import SDPA + +# 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 + + +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) + + +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 = 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 = 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 forward(self, x, cos, sin, bias): + """``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)) + 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)) + + +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 = 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 = 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): + 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]``. + + ``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: + 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 = conv1x1(dim, L * self.ple_dim) + 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 = MaxAbsRMSNorm(dim, eps, axis=-3) + self.embedding_projection = conv1x1(dim, config.embedding_dim) + + 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, + 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 = {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[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) + 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).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], + 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 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") + 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) + load_conv_weights(model, mapped) + 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)``.""" + mac = load_macos_text_encoder(model_dir_or_id, torch.float32) + return text_encoder_from_state_dict(mac.config, mac.state_dict(), dtype) + + +__all__ = [ + "MASK_NEG", + "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 new file mode 100644 index 00000000..09bb1c86 --- /dev/null +++ b/python/src/coreai_models/models/ios/embedding_gemma2_vision.py @@ -0,0 +1,232 @@ +# 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. + +The vision tower as 1x1 Conv2d projections with max-abs RMSNorms, a static patch budget and +in-graph lookups. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +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: + """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 = 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 = 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 + 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 = conv1x1(dim, cfg.intermediate_size) + self.up_proj = conv1x1(dim, cfg.intermediate_size) + self.down_proj = conv1x1(cfg.intermediate_size, dim) + 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( + 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 [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__() + 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 = conv1x1(cfg.patch_dim, dim) + self.layers = nn.ModuleList([VisionEncoderLayer(cfg) for _ in range(cfg.num_hidden_layers)]) + 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. + 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: + """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() + 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 = 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()) + + 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 (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) + + # ``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] + 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, text_hidden_size, 1, L) + return y.squeeze(2).transpose(1, 2).contiguous() # [B, L, text_hidden_size] + + +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`` 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() + + +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 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") + 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 + pos_table = mapped.pop("patch_embedder.position_embedding_table") + load_conv_weights(model, mapped) + 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/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/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/embedding_gemma2.py b/python/src/coreai_models/models/macos/embedding_gemma2.py new file mode 100644 index 00000000..2e5631be --- /dev/null +++ b/python/src/coreai_models/models/macos/embedding_gemma2.py @@ -0,0 +1,502 @@ +# 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 ``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 + +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.base import load_tensors_for_keys, resolve_safetensors_files +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." + + +def rope_inv_freq(dim: int, base: float, device: torch.device | None = None) -> torch.Tensor: + """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 ``text_config`` block of an HF ``config.json``. + + Parsed without transformers; the HF EmbeddingGemma 2 classes need 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 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" + + def rope_head_dim(self, sliding: bool) -> int: + """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) + 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"]``). + + ``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 + 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"): + 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"] = 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() + } + 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; 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) + 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. + + ``_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, 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=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, + 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 = self._rope(query, cos, sin).transpose(1, 2) + key = self.k_norm(self.k_proj(x).reshape(batch, seq_len, n_kv, hd)) + 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) + 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 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__() + h, eps = config.hidden_size, config.rms_norm_eps + self.self_attn = self_attn + self.mlp = EmbeddingGemma2MLP(config) + 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, + 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 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 + + +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) + + self.register_buffer( + "sliding_inv_freq", + rope_inv_freq(config.rope_head_dim(True), config.rope_theta_sliding), + persistent=False, + ) + self.register_buffer( + "full_inv_freq", + rope_inv_freq(config.rope_head_dim(False), config.rope_theta_full), + persistent=False, + ) + + @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]``. + + 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: + 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``, 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: + 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 read_checkpoint_tensors( + model_dir: str | Path, prefixes: tuple[str, ...], dtype: torch.dtype +) -> dict[str, torch.Tensor]: + """Load the tensors of the ``model_dir`` checkpoint whose names start with one of ``prefixes``. + + 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 or an indexed shard is + missing. + RuntimeError: if ``model.safetensors.index.json`` has no ``weight_map``. + """ + from safetensors import safe_open + + 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): + 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: + 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.*``.""" + model_dir = _resolve_model_dir(model_dir_or_id) + 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() + + +__all__ = [ + "EmbeddingGemma2TextConfig", + "EmbeddingGemma2TextEncoder", + "build_attention_masks", + "load_embedding_gemma2_text_encoder", + "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 new file mode 100644 index 00000000..f2a44fa4 --- /dev/null +++ b/python/src/coreai_models/models/macos/embedding_gemma2_vision.py @@ -0,0 +1,442 @@ +# 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, 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 + +import json +import math +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 coreai_models.models.macos.diffusion_gemma import DiffusionGemmaRMSNorm as RMSNorm +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 + +# 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 and ``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``). + + ``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") + 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), + ) + + +class VisionAttention(EmbeddingGemma2Attention): + """``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 + 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 _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): + """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__() + 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): + """Layer stack at module path ``encoder.layers.N``, matching the checkpoint keys.""" + + def __init__(self, config: EmbeddingGemma2VisionConfig) -> None: + super().__init__() + self.layers = nn.ModuleList( + [ + EmbeddingGemma2SandwichLayer(config, VisionAttention(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): + """Gemma 4 vision tower and ``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) + # 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 + ) + + # -- pieces ------------------------------------------------------------------------- + + 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; 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 + ) -> torch.Tensor: + """Mean over each ``k x k`` patch block, scaled by ``sqrt(hidden_size)``. + + 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) + 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 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) + 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`` 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) + 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() + + +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]: + """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)) + 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)``. Applies ``rescale_factor`` only; the tower applies ``2 * (x - 0.5)``. + """ + _, 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 ``; ``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 + 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). + + 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. + 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 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. + """ + 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/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/src/coreai_models/primitives/ios/rms_norm.py b/python/src/coreai_models/primitives/ios/rms_norm.py index 8eb90847..ed34d32f 100644 --- a/python/src/coreai_models/primitives/ios/rms_norm.py +++ b/python/src/coreai_models/primitives/ios/rms_norm.py @@ -51,3 +51,53 @@ def forward( x_2normalized = x_2normalized.to(input_dtype) return x_2normalized * self.weight + + +# 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:: +# +# xn = x / m +# xn * rsqrt(mean(xn^2) + (sqrt(eps) / m)^2) * w == x * rsqrt(mean(x^2) + eps) * w +# +# 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). + + 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 + self.register_buffer("_eps_sqrt", torch.tensor(eps**0.5), persistent=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # 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) + 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. + + ``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: + 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_export/test_embedding_export.py b/python/tests/test_model_units/test_export/test_embedding_export.py new file mode 100644 index 00000000..a86ce087 --- /dev/null +++ b/python/tests/test_model_units/test_export/test_embedding_export.py @@ -0,0 +1,845 @@ +# 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; 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 +import tomllib +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, + FunctionSpec, + bucket_plan, + build_bundle_metadata, + bundle_assets, + bundle_name, + 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, + TINY_VISION_CONFIG, +) + +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"]), +} + +# 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) + 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 validate_compute_precision(EG2, "float16", platform="iOS") == "float16" + for precision in ("float32", "bfloat16"): + with pytest.raises(EmbeddingPrecisionError, match="float16"): + validate_compute_precision(EG2, precision, platform="iOS") + + +def test_macos_refuses_float16_and_allows_bfloat16(): + with pytest.raises(EmbeddingPrecisionError, match="forbids float16"): + 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"): + validate_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_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"} + 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 --- + + +@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 + 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=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(components) + assert meta["source"]["hf_model_id"] == "google/embeddinggemma-2" + + +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"] + 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(): + 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 + 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"] + 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_embedding_block_records_the_compute_precision(): + assert _block()["precision"]["compute"] == "float32" + 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") + + +@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") + ) + 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"]) + + +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" + ) + 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_failed_export_removes_the_bundle_dir(tmp_path, monkeypatch, stub_snapshot): + def fail(*args, **kwargs): + raise RuntimeError("conversion failed") + + 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\."): + validate_compute_precision(EG2, "float16") + no_floor = dataclasses.replace(EG2, precision_floor=None) + with pytest.raises(EmbeddingPrecisionError, match=r"Use float32 \(default\)\.$"): + validate_compute_precision(no_floor, "float16") + + +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 = export_utils.resolve_snapshot("google/embeddinggemma-2", str(snap)) + assert sha == "a" * 40 and path == snap.resolve() + 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"] + + +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} + ) + with pytest.raises(ValueError, match="include_prompt"): + export_utils.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): + export_utils.verify_image_config(_image_snapshot(tmp_path / "ok"), EG2) + with pytest.raises(ValueError, match="boi_token_id"): + export_utils.verify_image_config(_image_snapshot(tmp_path / "bad", boi_token_id=1), EG2) + + +# --- Conversion (tiny random config) --- + + +def test_tiny_multifunction_conversion(): + wrapper = export_macos.TextEmbeddingExportWrapper(_tiny_text_encoder()) + plan = bucket_plan([16], [1, 4], 1024) + 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, 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 + + +def _tiny_text_encoder(): + from coreai_models.models.macos.embedding_gemma2 import ( + EmbeddingGemma2TextConfig, + EmbeddingGemma2TextEncoder, + ) + + 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 + + torch.manual_seed(1) + return EmbeddingGemma2ImageEncoder(TINY_VISION_CONFIG).eval() + + +TINY_IMAGE = dataclasses.replace(IMAGE, patch_size=4, budget=12, soft_token_dim=32) + + +def test_inputs_embeds_path_equals_ids_path(): + """``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) + mask[0, 11:] = 0 + with torch.no_grad(): + 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, 64) + torch.testing.assert_close(merged, direct, atol=1e-6, rtol=0) + assert isinstance(text.embed_tokens, torch.nn.Embedding) # restored + + +def test_float16_scaled_token_table_equals_the_float32_table_cast(): + """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( + text.embed_tokens.weight, text.config.hidden_size**0.5, torch.float16 + ) + with torch.no_grad(): + 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) + + +def test_tiny_float16_token_table_conversion(): + text = _tiny_text_encoder() + table = export_utils.scaled_token_table( + text.embed_tokens.weight, text.config.hidden_size**0.5, 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 = 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 = 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) + + +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(), + [FunctionSpec("len16_b1", 1, 16), FunctionSpec("len32_b1", 1, 32)], + [FunctionSpec("embeds_len32_b1", 1, 32)], + max_length=64, + ) + 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 + + +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) + assert pool.shape == (1, 128) and {t.dtype for t in (embeds, bias, pool)} == {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 + 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 -> (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 + + 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), + ), + } + + +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, 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 + + 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()) + ) + + 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 + 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 == 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 == dtypes["token_table"] + 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 == 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) + + +# --- 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 "embeds_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..673648e8 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_run_the_script() -> 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", + ] 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..bf540925 --- /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 EmbeddingGemma 2 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..4faeb4d0 --- /dev/null +++ b/python/tests/test_model_units/test_models/test_ios_layers/test_embedding_gemma2.py @@ -0,0 +1,129 @@ +# 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 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, + 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, + _ragged_batch, +) + + +def _host_inputs( + ids: torch.Tensor, + attention_mask: torch.Tensor, + embed_table: torch.Tensor, + dtype: torch.dtype = torch.float16, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """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) + bias = torch.where(key_ok, 0.0, MASK_NEG).reshape(1, length, 1, length) + 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: + """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()): + 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) + # 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[r]) + assert float(cos) >= 0.9999, (row, length, 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 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) + ios16, _ = text_encoder_from_state_dict(cfg, state, dtype=torch.float16) + with torch.no_grad(): + 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 + + +@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) + embeds, bias, pool = _host_inputs(ids, mask, table, dtype) + with torch.no_grad(): + got = _banded(ios, 6)(embeds, bias, torch.zeros_like(pool)) + assert torch.equal(got, torch.zeros_like(got)) + + +@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) + 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 = _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_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_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 == {} 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..340a17e8 --- /dev/null +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2.py @@ -0,0 +1,351 @@ +# 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. +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" \\ + 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 HF + ``EmbeddingGemma2Model`` + the sentence-transformers mean-pool + L2-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. +""" + +from __future__ import annotations + +import json + +import pytest +import torch +import torch.nn.functional as F + +from coreai_models.models.macos.embedding_gemma2 import ( + EmbeddingGemma2TextConfig, + EmbeddingGemma2TextEncoder, + build_attention_masks, + load_embedding_gemma2_text_encoder, + read_checkpoint_tensors, +) +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: + """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()): + 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) + + +@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)] + 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) 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(): + 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_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 --- + + +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 = "eager"): + """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() + 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) + 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(): + """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 = [ + ("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 = 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] + 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 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. " + ) + 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) + emb = ours(input_ids=ids, attention_mask=mask) + 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)) + ] + ) + + assert _cos(emb, ref).min().item() >= 0.99999 + 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]) + 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) + 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 new file mode 100644 index 00000000..0c6d8c35 --- /dev/null +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_embedding_gemma2_vision.py @@ -0,0 +1,377 @@ +# 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 (``models/macos/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 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 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: + 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): + 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 + 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): + """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 = [], [], [] + 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_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) + 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, [(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 == (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: + 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 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 ( + 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 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) + + +# --- 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). + with pytest.MonkeyPatch.context() as patch: + patch.setenv("USE_HF_IMPL", "false") + ours = EmbeddingGemma2ImageEncoder.from_pretrained(snapshot, dtype=torch.float32) + 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) 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)