diff --git a/python/src/coreai_models/_constants.py b/python/src/coreai_models/_constants.py index 4f3d4642..43240036 100644 --- a/python/src/coreai_models/_constants.py +++ b/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 diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 1b19272a..bdcc8a96 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/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": diff --git a/python/src/coreai_models/models/macos/gemma3n.py b/python/src/coreai_models/models/macos/gemma3n.py index e380c21a..0dbd6153 100644 --- a/python/src/coreai_models/models/macos/gemma3n.py +++ b/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, ) diff --git a/python/tests/test_model_units/test_export/test_gather_table_guard.py b/python/tests/test_model_units/test_export/test_gather_table_guard.py new file mode 100644 index 00000000..91ff0a15 --- /dev/null +++ b/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 == "" diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_gemma3n.py b/python/tests/test_model_units/test_models/test_macos_layers/test_gemma3n.py index b4290610..a5cfd8d6 100644 --- a/python/tests/test_model_units/test_models/test_macos_layers/test_gemma3n.py +++ b/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)