Issue 313 Fix: ChunkedEmbedding for Gemma-3n-E4B

python/src/coreai_models/_constants.py

@@ -84,3 +84,7 @@
 METADATA_VERSION = "0.2"
 # Batched LLM bundles use 0.3 (adds the max_batch_size field).
 METADATA_VERSION_BATCHED = "0.3"
+
+# Core AI GPU gathers from a single constant return wrong/zero rows for rows
+# starting at or past this element (apple/coreai-models#313). Tables this large are split.
+MAX_GATHER_CONSTANT_ELEMENTS = 2**31

python/src/coreai_models/export/pipeline.py

@@ -28,6 +28,7 @@
     DEFAULT_DYNAMIC_MAX_BATCH_SIZE,
     DEFAULT_INCLUDE_DEBUG_INFO,
     IOS_DEFAULT_MAX_CONTEXT_LENGTH,
+    MAX_GATHER_CONSTANT_ELEMENTS,
     TRACE_KV_CACHE_SEQ_LEN,
 )
 from coreai_models.export.bundle import bundle_llm_asset
@@ -319,6 +320,19 @@ def resolve_model(config: ExportConfig) -> ResolvedModel:
     )
 
 
+def warn_on_oversized_gather_tables(model: torch.nn.Module) -> None:
+    """Warn about embedding tables a GPU gather would read incorrectly (#313)."""
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Embedding):
+            numel = module.weight.numel()
+            if numel >= MAX_GATHER_CONSTANT_ELEMENTS:
+                logger.warning(
+                    f"Embedding '{name}' has {numel} elements (>= 2**31); GPU gathers "
+                    "from it return wrong rows for rows starting past element 2**31. "
+                    "Split it along the feature axis into smaller constants."
+                )
+
+
 def build_model(
     resolved: ResolvedModel, temp_dir: str | None = None, export_backend: object | None = None
 ) -> BuiltModel:
@@ -516,6 +530,7 @@ async def _async_export_model(config: ExportConfig) -> str:
 
     with temp_dir_ctx as temp_dir:
         model, externalized_model = build_model(resolved, temp_dir)
+        warn_on_oversized_gather_tables(model)
 
         # ---- 4. Variant-specific export ----
         if config.variant == "macOS":

python/src/coreai_models/models/macos/gemma3n.py

@@ -25,6 +25,7 @@
 )
 from typing_extensions import Self, override
 
+from coreai_models._constants import MAX_GATHER_CONSTANT_ELEMENTS
 from coreai_models.models.base import BaseForCausalLM
 from coreai_models.primitives.macos.cache import KVCache
 from coreai_models.primitives.macos.rms_norm import RMSNorm
@@ -297,6 +298,55 @@ def forward(
         return corrected
 
 
+class ChunkedEmbedding(nn.Module):
+    """Embedding whose table is split along the feature axis into ``nn.Embedding``
+    chunks below ``max_elements``, since GPU gathers from larger constants return
+    wrong rows (apple/coreai-models#313). Chunks stay ``nn.Embedding`` so compression
+    configs match them like the unsplit table.
+
+    Checkpoints still store one ``weight``; it is split into ``chunks.{i}.weight`` on load.
+    """
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        max_elements: int = MAX_GATHER_CONSTANT_ELEMENTS,
+        align: int = 32,
+    ) -> None:
+        super().__init__()
+        # Widths are multiples of ``align`` because the int4 presets quantize in 32-column
+        # blocks and silently skip tables whose width is not a multiple of the block.
+        max_width = (max_elements - 1) // num_embeddings // align * align
+        if max_width < 1:
+            raise ValueError(
+                f"{align} columns of {num_embeddings} rows already reach {max_elements} "
+                "elements; the table cannot be split along the feature axis."
+            )
+        width = math.ceil(embedding_dim / math.ceil(embedding_dim / max_width))
+        width = math.ceil(width / align) * align
+        self.widths = [min(width, embedding_dim - i) for i in range(0, embedding_dim, width)]
+        self.chunks = nn.ModuleList([nn.Embedding(num_embeddings, w) for w in self.widths])
+        self._register_load_state_dict_pre_hook(self._split_weight)
+
+    def _split_weight(self, state_dict, prefix, *args) -> None:
+        full = state_dict.pop(prefix + "weight", None)
+        if full is None:
+            return
+        # Column slices are strided, so contiguous() already copies each one.
+        for i, chunk in enumerate(full.split(self.widths, dim=-1)):
+            state_dict[f"{prefix}chunks.{i}.weight"] = chunk.contiguous()
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return torch.cat([chunk(input_ids) for chunk in self.chunks], dim=-1)
+
+
+def _per_layer_embedding(num_embeddings: int, embedding_dim: int) -> nn.Module:
+    if num_embeddings * embedding_dim >= MAX_GATHER_CONSTANT_ELEMENTS:
+        return ChunkedEmbedding(num_embeddings, embedding_dim, MAX_GATHER_CONSTANT_ELEMENTS)
+    return nn.Embedding(num_embeddings, embedding_dim)
+
+
 class Gemma3nModel(nn.Module):
     def __init__(self, config: Gemma3nTextConfig) -> None:
         super().__init__()
@@ -306,7 +356,7 @@ def __init__(self, config: Gemma3nTextConfig) -> None:
         self.per_layer_input_scale = 2.0**-0.5
 
         self.embed_tokens = nn.Embedding(config.vocab_size, hidden)
-        self.embed_tokens_per_layer = nn.Embedding(
+        self.embed_tokens_per_layer = _per_layer_embedding(
             config.vocab_size_per_layer_input,
             config.num_hidden_layers * config.hidden_size_per_layer_input,
         )

python/tests/test_model_units/test_export/test_gather_table_guard.py

@@ -0,0 +1,48 @@
+# 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
+
+"""Embedding tables a GPU gather reads incorrectly are flagged at export (#313)."""
+
+import logging
+
+import pytest
+import torch
+
+from coreai_models.export.pipeline import warn_on_oversized_gather_tables
+from coreai_models.models.macos.gemma3n import ChunkedEmbedding
+
+# Gemma 3n per-layer table shapes; meta tensors keep these free to build.
+E4B_SHAPE = (262144, 35 * 256)  # 2.35e9 elements, over 2**31
+E2B_SHAPE = (262144, 30 * 256)  # 2.01e9 elements, under 2**31
+
+
+def _model(**modules: torch.nn.Module) -> torch.nn.Module:
+    return torch.nn.ModuleDict(modules)
+
+
+def test_warns_on_oversized_embedding(caplog: pytest.LogCaptureFixture) -> None:
+    with torch.device("meta"):
+        model = _model(per_layer=torch.nn.Embedding(*E4B_SHAPE))
+    with caplog.at_level(logging.WARNING):
+        warn_on_oversized_gather_tables(model)
+    assert "'per_layer'" in caplog.text
+    assert str(E4B_SHAPE[0] * E4B_SHAPE[1]) in caplog.text
+
+
+def test_silent_below_limit(caplog: pytest.LogCaptureFixture) -> None:
+    with torch.device("meta"):
+        model = _model(per_layer=torch.nn.Embedding(*E2B_SHAPE))
+    with caplog.at_level(logging.WARNING):
+        warn_on_oversized_gather_tables(model)
+    assert caplog.text == ""
+
+
+def test_silent_for_chunked_embedding(caplog: pytest.LogCaptureFixture) -> None:
+    """The fix itself must not trip the guard."""
+    with torch.device("meta"):
+        model = _model(per_layer=ChunkedEmbedding(*E4B_SHAPE))
+    with caplog.at_level(logging.WARNING):
+        warn_on_oversized_gather_tables(model)
+    assert caplog.text == ""

python/tests/test_model_units/test_models/test_macos_layers/test_gemma3n.py

@@ -5,6 +5,9 @@
 
 """Tests for macOS Gemma 3n model parity with HuggingFace."""
 
+import copy
+
+import pytest
 import torch
 from transformers.models.gemma3n.configuration_gemma3n import Gemma3nTextConfig
 from transformers.models.gemma3n.modeling_gemma3n import (
@@ -113,6 +116,27 @@ def test_forward_parity_multi_token(self):
 
         torch.testing.assert_close(our_out, hf_out.logits, atol=1e-4, rtol=1e-4)
 
+    def test_forward_parity_with_chunked_per_layer_table(self, monkeypatch):
+        """A lowered limit splits the [200, 80] per-layer table; the HF checkpoint still loads."""
+        from coreai_models.models.macos import gemma3n
+
+        monkeypatch.setattr(gemma3n, "MAX_GATHER_CONSTANT_ELEMENTS", 200 * 32 + 1)
+        config = _make_gemma3n_config(num_hidden_layers=5)
+        hf_model, our_model = _build_models(config)
+        assert our_model.model.embed_tokens_per_layer.widths == [32, 32, 16]
+
+        input_ids = torch.randint(0, 200, (1, 6))
+        position_ids = torch.arange(6, dtype=torch.int32).unsqueeze(0)
+        k_cache, v_cache = KVCache.create_cache_tensors(config, dtype=torch.float32)
+
+        with torch.no_grad():
+            our_out = our_model(input_ids, position_ids, k_cache, v_cache)
+            hf_out = hf_model(
+                input_ids=input_ids, position_ids=position_ids.long(), use_cache=False
+            )
+
+        torch.testing.assert_close(our_out, hf_out.logits, atol=1e-4, rtol=1e-4)
+
     def test_forward_parity_with_kv_sharing(self):
         """KV sharing: last 2 layers reuse K/V from source layers."""
         config = _make_gemma3n_config(num_hidden_layers=10, num_kv_shared_layers=2)
@@ -239,3 +263,58 @@ def test_tie_word_embeddings(self):
         our_model.load_state_dict(sd, assign=True, strict=True)
 
         assert our_model.lm_head.weight is our_model.model.embed_tokens.weight
+
+
+class TestChunkedEmbedding:
+    """Per-layer table split to stay under the GPU gather element limit (#313)."""
+
+    def test_matches_unsplit_embedding(self):
+        from coreai_models.models.macos.gemma3n import ChunkedEmbedding
+
+        ref = torch.nn.Embedding(50, 21)
+        chunked = ChunkedEmbedding(50, 21, max_elements=400, align=1)
+        assert len(chunked.chunks) > 1
+        assert all(c.weight.numel() < 400 for c in chunked.chunks)
+        chunked.load_state_dict({"weight": ref.weight.detach()})
+        ids = torch.randint(0, 50, (2, 7))
+        assert torch.equal(chunked(ids), ref(ids))
+
+    def test_chunks_stay_under_limit_when_width_rounds_up(self):
+        """Splitting by element count alone gives 4-column chunks of 400 >= 351 here."""
+        from coreai_models.models.macos.gemma3n import ChunkedEmbedding
+
+        chunked = ChunkedEmbedding(100, 10, max_elements=351, align=1)
+        assert sum(chunked.widths) == 10
+        assert all(c.weight.numel() < 351 for c in chunked.chunks)
+
+    def test_single_chunk_below_limit(self):
+        from coreai_models.models.macos.gemma3n import ChunkedEmbedding
+
+        assert len(ChunkedEmbedding(50, 21, max_elements=2**31).chunks) == 1
+
+    def test_e4b_per_layer_table_is_split(self):
+        from coreai_models.models.macos.gemma3n import ChunkedEmbedding, _per_layer_embedding
+
+        with torch.device("meta"):
+            e4b = _per_layer_embedding(262144, 35 * 256)
+            e2b = _per_layer_embedding(262144, 30 * 256)
+        assert isinstance(e4b, ChunkedEmbedding)
+        assert e4b.widths == [4480, 4480]
+        assert all(c.weight.numel() < 2**31 for c in e4b.chunks)
+        assert isinstance(e2b, torch.nn.Embedding)
+
+    def test_4bit_preset_quantizes_every_chunk(self):
+        """Unaligned, this splits into two 48-wide chunks the 32-block int4 preset skips."""
+        quantization = pytest.importorskip("coreai_opt.quantization")
+        from coreai_models.export.presets import MACOS_PRESETS
+        from coreai_models.models.macos.gemma3n import ChunkedEmbedding
+
+        model = ChunkedEmbedding(64, 96, max_elements=64 * 48 + 1)
+        assert model.widths == [32, 32, 32]
+        qcfg = copy.deepcopy(MACOS_PRESETS["4bit"]["torch_quantization_config"])
+        quantizer = quantization.Quantizer(
+            model, quantization.QuantizerConfig.from_dict({"quantization_config": qcfg})
+        )
+        ids = torch.randint(0, 64, (1, 5), dtype=torch.int32)
+        prepared = quantizer.prepare(example_inputs=(ids,), dynamic_shapes=None)
+        assert all(hasattr(c, "parametrizations") for c in prepared.chunks)