From 664526a3b01036b2f593794a4f8965a26a8034de Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Tue, 25 Aug 2026 18:37:10 -0700 Subject: [PATCH] Add Gemma 3n support (E2B, macOS) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Google Gemma 3n E2B — a novel on-device architecture with 7 features not found in other supported models: - AltUp: 4-copy hidden state with learned predict/correct routing - LAUREL: low-rank residual branch (rank 64) per layer - Per-layer input embeddings with gated injection - KV cache sharing: last 10/30 layers reuse K/V (33% cache savings) - Gaussian TopK: 95% activation sparsity in first 10 MLP layers - Dual RoPE: local (theta=10K, window 512) + global (theta=1M) - QKV norms: RMSNorm on Q/K (with scale), V (without scale) Evaluation: WikiText PPL 61.53 (float16), tinyMMLU 58.75%. High wikitext PPL is expected for instruction-tuned models; tinyMMLU accuracy is the appropriate quality metric. --- models/gemma3n/README.md | 38 ++ python/src/coreai_models/export/metadata.py | 9 + python/src/coreai_models/model_registry.py | 10 + .../src/coreai_models/models/macos/gemma3n.py | 473 ++++++++++++++++++ python/src/coreai_models/models/registry.py | 7 + .../test_macos_layers/test_gemma3n.py | 235 +++++++++ 6 files changed, 772 insertions(+) create mode 100644 models/gemma3n/README.md create mode 100644 python/src/coreai_models/models/macos/gemma3n.py create mode 100644 python/tests/test_model_units/test_models/test_macos_layers/test_gemma3n.py diff --git a/models/gemma3n/README.md b/models/gemma3n/README.md new file mode 100644 index 00000000..2d30d57d --- /dev/null +++ b/models/gemma3n/README.md @@ -0,0 +1,38 @@ +# Gemma 3n + +Google's Gemma 3n for on-device inference via Core AI. + +## Supported Models + +| Model | Parameters | Context | macOS | iOS | +| --------------- | ------------------- | ------- | ----- | --- | +| gemma-3n-E2B-it | ~5B (E2B effective) | 32768 | Yes | No | + +## Export + +```bash +uv run coreai.llm.export google/gemma-3n-E2B-it +``` + +## Run + +```bash +swift run -c release llm-runner --model path/to/exported_model --prompt "Hello" +``` + +## Benchmark + +```bash +swift run -c release llm-benchmark --model path/to/exported_model -p 512 -g 1024 -n 5 +``` + +## Evaluation + +Perplexity on WikiText-2 and accuracy on tinyMMLU, computed with the Core AI PyTorch models (4K context). + +| Compression | Wikitext PPL | tinyMMLU acc | +| ---------------- | ------------ | ------------ | +| none (float16) | 61.53 | 58.75% | +| 4-bit quantized | 72.13 | 59.27% | + +High wikitext PPL is expected for instruction-tuned models — the probability distribution is optimized for dialogue, not raw text. tinyMMLU accuracy is the appropriate quality metric. diff --git a/python/src/coreai_models/export/metadata.py b/python/src/coreai_models/export/metadata.py index ed3f9867..7552ff14 100644 --- a/python/src/coreai_models/export/metadata.py +++ b/python/src/coreai_models/export/metadata.py @@ -95,6 +95,15 @@ class AIModelMetadataFields: "Source: https://huggingface.co/google/gemma-3-12b-it" ), ), + "google/gemma-3n-E2B-it": AIModelMetadataFields( + author="Gemma Team", + license="Gemma Terms of Use", + model_description=( + "Gemma 3n E2B is a ~5B-parameter on-device language model from Google's " + "Gemma 3n family using AltUp for efficient inference. " + "Source: https://huggingface.co/google/gemma-3n-E2B-it" + ), + ), "mistralai/Mistral-7B-Instruct-v0.3": AIModelMetadataFields( author="Mistral AI", license="Apache-2.0", diff --git a/python/src/coreai_models/model_registry.py b/python/src/coreai_models/model_registry.py index e843499f..0596f113 100644 --- a/python/src/coreai_models/model_registry.py +++ b/python/src/coreai_models/model_registry.py @@ -174,6 +174,16 @@ class UtilityModel: "float16", 131072, ), + ModelPreset( + "gemma-3n-e2b-it", + "google/gemma-3n-E2B-it", + "gemma3n", + "llm", + "macOS", + "4bit", + "float16", + 32768, + ), # --- iOS (compression = palettized) --- ModelPreset( "qwen3-0.6b", diff --git a/python/src/coreai_models/models/macos/gemma3n.py b/python/src/coreai_models/models/macos/gemma3n.py new file mode 100644 index 00000000..017fbce5 --- /dev/null +++ b/python/src/coreai_models/models/macos/gemma3n.py @@ -0,0 +1,473 @@ +# 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 + +"""Gemma 3n text decoder for CoreAI model export. + +Novel features vs standard LLaMA/Mistral: +- AltUp: 4-copy hidden state with predict/correct routing per layer +- Per-layer input embeddings with gated injection +- LAUREL: low-rank residual branch +- KV sharing: last N layers reuse KV from earlier layers +- Gaussian TopK activation sparsity in early MLP layers +- Dual RoPE: local (θ=10K) + global (θ=1M) +- QKV norms with scaling=1.0 +""" + +import math + +import torch +import torch.nn as nn +from transformers.models.gemma3n.configuration_gemma3n import Gemma3nTextConfig +from transformers.models.gemma3n.modeling_gemma3n import ( + Gemma3nForCausalLM as HFGemma3nForCausalLM, +) +from typing_extensions import Self, override + +from coreai_models.models.base import BaseForCausalLM +from coreai_models.primitives.macos.cache import KVCache +from coreai_models.primitives.macos.rms_norm import RMSNorm +from coreai_models.primitives.macos.rope import initialize_rope +from coreai_models.primitives.macos.sdpa import SDPA + + +class AltUp(nn.Module): + def __init__(self, config: Gemma3nTextConfig) -> None: + super().__init__() + self.num_inputs = config.altup_num_inputs + self.active_idx = config.altup_active_idx + hidden = config.hidden_size + + self.correct_output_scale = nn.Parameter(torch.zeros(hidden)) + self.correction_coefs = nn.Linear(self.num_inputs, self.num_inputs, bias=False) + self.prediction_coefs = nn.Linear(self.num_inputs, self.num_inputs**2, bias=False) + self.modality_router = nn.Linear(hidden, self.num_inputs, bias=False) + self.router_norm = RMSNorm(hidden, eps=config.rms_norm_eps) + self.router_input_scale = hidden**-1.0 + + def _compute_modalities(self, x: torch.Tensor) -> torch.Tensor: + router_inputs = self.router_norm(x) * self.router_input_scale + return torch.tanh(self.modality_router(router_inputs).float()).to(x.dtype) + + def predict(self, hidden_states: torch.Tensor) -> torch.Tensor: + modalities = self._compute_modalities(hidden_states[self.active_idx]) + all_coefs = ( + self.prediction_coefs(modalities) + .reshape(*modalities.shape[:-1], self.num_inputs, self.num_inputs) + .permute(0, 1, 3, 2) + ) + predictions = torch.matmul(hidden_states.permute(1, 2, 3, 0), all_coefs) + predictions = predictions.permute(3, 0, 1, 2) + predictions = predictions + hidden_states + return predictions.contiguous() + + def correct(self, predictions: torch.Tensor, activated: torch.Tensor) -> torch.Tensor: + modalities = self._compute_modalities(activated) + innovation = activated - predictions[self.active_idx] + innovation = innovation.repeat(self.num_inputs, 1, 1, 1) + all_coefs = (self.correction_coefs(modalities) + 1.0).permute(2, 0, 1).unsqueeze(-1) + corrected = torch.mul(innovation, all_coefs) + predictions + return corrected.contiguous() + + def scale_output(self, x: torch.Tensor) -> torch.Tensor: + return x * self.correct_output_scale + + +class Laurel(nn.Module): + def __init__(self, config: Gemma3nTextConfig) -> None: + super().__init__() + self.linear_left = nn.Linear(config.hidden_size, config.laurel_rank, bias=False) + self.linear_right = nn.Linear(config.laurel_rank, config.hidden_size, bias=False) + self.post_laurel_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.post_laurel_norm(self.linear_right(self.linear_left(x))) + + +class Attention(nn.Module): + def __init__(self, config: Gemma3nTextConfig, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + + dim = config.hidden_size + self.n_heads = n_heads = config.num_attention_heads + self.n_kv_heads = n_kv_heads = config.num_key_value_heads + self.head_dim = head_dim = config.head_dim + + self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False) + self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False) + self.q_norm = RMSNorm(head_dim, eps=config.rms_norm_eps) + + layer_types = config.layer_types or [ + "full_attention" if (i + 1) % 5 == 0 else "sliding_attention" + for i in range(config.num_hidden_layers) + ] + self.is_sliding = layer_types[layer_idx] == "sliding_attention" + + first_kv_shared = config.num_hidden_layers - config.num_kv_shared_layers + self.is_kv_shared = layer_idx >= first_kv_shared > 0 + + if not self.is_kv_shared: + self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.k_norm = RMSNorm(head_dim, eps=config.rms_norm_eps) + + rope_params = getattr(config, "rope_parameters", {}) + if self.is_sliding: + self.sdpa = SDPA(is_causal=True, scale=1.0, window_size=config.sliding_window) + local_theta = rope_params.get("sliding_attention", {}).get( + "rope_theta", getattr(config, "rope_local_base_freq", 10000.0) + ) + self.rope = initialize_rope(base=local_theta) + else: + self.sdpa = SDPA(is_causal=True, scale=1.0) + global_theta = rope_params.get("full_attention", {}).get( + "rope_theta", getattr(config, "rope_theta", 1000000.0) + ) + self.rope = initialize_rope(base=global_theta) + + self._v_norm_eps = config.rms_norm_eps + + # Cache slot index: shared layers map to their source layer's slot. + # Built by Gemma3nModel and assigned after construction. + self.cache_slot: int = layer_idx + + def forward( + self, + x: torch.Tensor, + position_ids: torch.IntTensor, + cache: KVCache | None = None, + ) -> torch.Tensor: + batch_size, query_len, _ = x.shape + n_heads, n_kv_heads = self.n_heads, self.n_kv_heads + + query = ( + self.q_proj(x) + .reshape(batch_size, query_len, n_heads, self.head_dim) + .permute(0, 2, 1, 3) + ) + query = self.q_norm(query) + + seq_len = position_ids.shape[-1] + torch._check_is_size(query_len) + torch._check_is_size(seq_len) + offset = seq_len - query_len + torch._check_is_size(offset) + rope_positions = position_ids.narrow(-1, offset, query_len) + + query = self.rope(query, position_ids=rope_positions) + + if self.is_kv_shared and cache is not None: + # Read K/V from source layer's cache slot without writing + key = cache._k_cache.narrow(0, self.cache_slot, 1).narrow(-2, 0, seq_len).squeeze(0) + value = cache._v_cache.narrow(0, self.cache_slot, 1).narrow(-2, 0, seq_len).squeeze(0) + else: + key = ( + self.k_proj(x) + .reshape(batch_size, query_len, n_kv_heads, self.head_dim) + .permute(0, 2, 1, 3) + ) + key = self.k_norm(key) + key = self.rope(key, position_ids=rope_positions) + + value = ( + self.v_proj(x) + .reshape(batch_size, query_len, n_kv_heads, self.head_dim) + .permute(0, 2, 1, 3) + ) + value = value / torch.sqrt(value.pow(2).mean(-1, keepdim=True) + self._v_norm_eps) + + if cache is not None: + key, value = cache.update_and_fetch( + self.cache_slot, offset, key, value, seq_len=seq_len, query_len=query_len + ) + + output = ( + self.sdpa(query, key, value) + .permute(0, 2, 1, 3) + .reshape(batch_size, query_len, self.n_heads * self.head_dim) + ) + return self.o_proj(output) + + +class MLP(nn.Module): + def __init__(self, config: Gemma3nTextConfig, layer_idx: int) -> None: + super().__init__() + hidden = config.hidden_size + intermediate = config.intermediate_size[layer_idx] + self.gate_proj = nn.Linear(hidden, intermediate, bias=False) + self.up_proj = nn.Linear(hidden, intermediate, bias=False) + self.down_proj = nn.Linear(intermediate, hidden, bias=False) + self.act_fn = nn.GELU(approximate="tanh") + self.activation_sparsity = config.activation_sparsity_pattern[layer_idx] + + def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor: + target = torch.tensor(self.activation_sparsity, dtype=torch.float32, device=x.device) + std_mult = (torch.erfinv(2.0 * target - 1.0) * math.sqrt(2.0)).to(x.dtype) + mu = x.mean(dim=-1, keepdim=True) + std = x.std(dim=-1, keepdim=True, unbiased=False) + cutoff = mu + std * std_mult + return nn.functional.relu(x - cutoff) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = self.gate_proj(x) + if self.activation_sparsity > 0.0: + gate = self._gaussian_topk(gate) + return self.down_proj(self.act_fn(gate) * self.up_proj(x)) + + +class TransformerBlock(nn.Module): + def __init__(self, config: Gemma3nTextConfig, layer_idx: int) -> None: + super().__init__() + hidden = config.hidden_size + self.layer_idx = layer_idx + self.active_idx = config.altup_active_idx + self.hidden_size_per_layer = config.hidden_size_per_layer_input + + self.self_attn = Attention(config, layer_idx) + self.mlp = MLP(config, layer_idx) + self.altup = AltUp(config) + self.laurel = Laurel(config) + + self.input_layernorm = RMSNorm(hidden, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(hidden, eps=config.rms_norm_eps) + self.pre_feedforward_layernorm = RMSNorm(hidden, eps=config.rms_norm_eps) + self.post_feedforward_layernorm = RMSNorm(hidden, eps=config.rms_norm_eps) + + pli_dim = config.hidden_size_per_layer_input + self.per_layer_input_gate = nn.Linear(hidden, pli_dim, bias=False) + self.per_layer_projection = nn.Linear(pli_dim, hidden, bias=False) + self.post_per_layer_input_norm = RMSNorm(hidden, eps=config.rms_norm_eps) + + self.altup_correct_scale = getattr(config, "altup_correct_scale", True) + self.act_fn = nn.GELU(approximate="tanh") + + def forward( + self, + hidden_states: torch.Tensor, + per_layer_input: torch.Tensor, + position_ids: torch.IntTensor, + cache: KVCache | None = None, + ) -> torch.Tensor: + predictions = self.altup.predict(hidden_states) + active = predictions[self.active_idx] + + active_normed = self.input_layernorm(active) + laurel_output = self.laurel(active_normed) + + attn = self.self_attn(active_normed, position_ids, cache) + attn = self.post_attention_layernorm(attn) + + attn_gated = active + attn + attn_laurel = (attn_gated + laurel_output) / math.sqrt(2) + + ffw = self.mlp(self.pre_feedforward_layernorm(attn_laurel)) + ffw = self.post_feedforward_layernorm(ffw) + ffw_out = attn_laurel + ffw + + corrected = self.altup.correct(predictions, ffw_out) + + first = corrected[self.active_idx].clone() + if self.altup_correct_scale: + first = self.altup.scale_output(first) + + first = self.act_fn(self.per_layer_input_gate(first)) + first = torch.multiply(first, per_layer_input) + first = self.post_per_layer_input_norm(self.per_layer_projection(first)) + corrected[1:] += first + + return corrected + + +class Gemma3nModel(nn.Module): + def __init__(self, config: Gemma3nTextConfig) -> None: + super().__init__() + self.config = config + hidden = config.hidden_size + self.per_layer_projection_scale = hidden**-0.5 + 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( + config.vocab_size_per_layer_input, + config.num_hidden_layers * config.hidden_size_per_layer_input, + ) + self.embed_scale = hidden**0.5 + self.per_layer_embed_scale = config.hidden_size_per_layer_input**0.5 + self.per_layer_model_projection = nn.Linear( + hidden, config.num_hidden_layers * config.hidden_size_per_layer_input, bias=False + ) + self.per_layer_projection_norm = RMSNorm( + config.hidden_size_per_layer_input, eps=config.rms_norm_eps + ) + + self.altup_projections = nn.ModuleList( + [nn.Linear(hidden, hidden, bias=False) for _ in range(config.altup_num_inputs - 1)] + ) + self.altup_unembed_projections = nn.ModuleList( + [nn.Linear(hidden, hidden, bias=False) for _ in range(config.altup_num_inputs - 1)] + ) + + self.layers = nn.ModuleList( + [TransformerBlock(config, i) for i in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(hidden, eps=config.rms_norm_eps) + + layer_types = config.layer_types or [ + "full_attention" if (i + 1) % 5 == 0 else "sliding_attention" + for i in range(config.num_hidden_layers) + ] + first_kv_shared = max(0, config.num_hidden_layers - config.num_kv_shared_layers) + self._kv_share_map: dict[int, int] = {} + if first_kv_shared > 0 and config.num_kv_shared_layers > 0: + non_shared_types = layer_types[:first_kv_shared] + for i in range(first_kv_shared, config.num_hidden_layers): + lt = layer_types[i] + if lt in non_shared_types: + src = len(non_shared_types) - 1 - non_shared_types[::-1].index(lt) + self._kv_share_map[i] = src + + # Build compact cache slot mapping: non-shared layers get slots 0..N-1, + # shared layers point to their source layer's slot. + slot_for_layer: dict[int, int] = {} + next_slot = 0 + for i in range(config.num_hidden_layers): + if i not in self._kv_share_map: + slot_for_layer[i] = next_slot + next_slot += 1 + for i, src in self._kv_share_map.items(): + slot_for_layer[i] = slot_for_layer[src] + self.num_cache_slots = next_slot + + for layer in self.layers: + layer.self_attn.cache_slot = slot_for_layer[layer.layer_idx] + + def _get_per_layer_inputs( + self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor + ) -> torch.Tensor: + per_layer = self.embed_tokens_per_layer(input_ids) * self.per_layer_embed_scale + pli_dim = self.config.hidden_size_per_layer_input + per_layer = per_layer.reshape(*input_ids.shape, self.config.num_hidden_layers, pli_dim) + proj = self.per_layer_model_projection(inputs_embeds) * self.per_layer_projection_scale + proj = proj.reshape(*inputs_embeds.shape[:-1], self.config.num_hidden_layers, pli_dim) + proj = self.per_layer_projection_norm(proj) + return (proj + per_layer) * self.per_layer_input_scale + + def _altup_expand(self, h: torch.Tensor) -> torch.Tensor: + target_mag = torch.mean(h**2, dim=-1, keepdim=True) ** 0.5 + eps = torch.tensor(1e-5, device=h.device, dtype=h.dtype) + copies = [h] + for proj in self.altup_projections: + p = proj(h).to(h.dtype) + mag = torch.sqrt(torch.maximum(torch.mean(p**2, dim=-1, keepdim=True), eps)) + copies.append(p * target_mag / mag) + return torch.stack(copies, dim=0) + + def _altup_collapse(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_mag = torch.mean(hidden_states[0] ** 2, dim=-1, keepdim=True) ** 0.5 + eps = torch.tensor(1e-5, device=hidden_states.device, dtype=hidden_states.dtype) + copies = [hidden_states[0]] + for i, proj in enumerate(self.altup_unembed_projections): + p = proj(hidden_states[i + 1]).to(hidden_states.dtype) + mag = torch.sqrt(torch.maximum(torch.mean(p**2, dim=-1, keepdim=True), eps)) + copies.append(p * target_mag / mag) + return torch.mean(torch.stack(copies), dim=0) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.IntTensor, + cache: KVCache | None = None, + ) -> torch.Tensor: + inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale + per_layer_inputs = self._get_per_layer_inputs(input_ids, inputs_embeds) + + hidden_states = self._altup_expand(inputs_embeds) + + for layer in self.layers: + idx = layer.layer_idx + per_layer_input = per_layer_inputs[:, :, idx, :] + hidden_states = layer(hidden_states, per_layer_input, position_ids, cache) + + hidden_states = self._altup_collapse(hidden_states) + return self.norm(hidden_states) + + +class Gemma3nForCausalLM(BaseForCausalLM): + _HF_MODEL_CLASS = HFGemma3nForCausalLM + + @classmethod + @override + def _get_reauthored_config(cls, hf_config, max_context_length=None, num_layers=None): + text_config = hf_config.text_config if hasattr(hf_config, "text_config") else hf_config + if max_context_length is not None: + text_config.max_position_embeddings = max_context_length + if num_layers is not None: + text_config.num_hidden_layers = num_layers + if text_config.rope_scaling is not None: + text_config.rope_scaling = None + return text_config + + @override + def _init_model(self, config: Gemma3nTextConfig) -> None: + self.model = Gemma3nModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + if config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + + @classmethod + def create_cache_tensors( + cls, config, dtype: torch.dtype = torch.float32 + ) -> tuple[torch.Tensor, torch.Tensor]: + """Create KV cache with compact slots (shared layers reuse source slots).""" + n_kv_heads = config.num_key_value_heads + max_seq_len = config.max_position_embeddings + head_dim = config.head_dim + + first_kv_shared = config.num_hidden_layers - config.num_kv_shared_layers + n_slots = first_kv_shared if config.num_kv_shared_layers > 0 else config.num_hidden_layers + + k_cache = torch.zeros(n_slots, 1, n_kv_heads, max_seq_len, head_dim, dtype=dtype) + v_cache = torch.zeros(n_slots, 1, n_kv_heads, max_seq_len, head_dim, dtype=dtype) + return k_cache, v_cache + + @BaseForCausalLM.cast_logits_bfloat16_to_float16 + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.IntTensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + ) -> torch.Tensor: + cache = KVCache(k_cache, v_cache) + out = self.model(input_ids, position_ids, cache) + logits = self.lm_head(out) + if hasattr(self.config, "final_logit_softcapping") and self.config.final_logit_softcapping: + cap = self.config.final_logit_softcapping + logits = logits / cap + logits = torch.tanh(logits) + logits = logits * cap + return logits + + @override + def _mutate_state_dict(self: Self, state_dict: dict[str, torch.Tensor]) -> None: + prefix = "model.language_model." + keys = list(state_dict.keys()) + for key in keys: + if key.startswith("model.visual."): + del state_dict[key] + elif key.startswith(prefix): + state_dict["model." + key[len(prefix) :]] = state_dict.pop(key) + elif key.startswith("language_model."): + state_dict["model." + key[len("language_model.") :]] = state_dict.pop(key) + + # Strip v_norm (scale-free norm, no learned weights in our impl) + for key in list(state_dict.keys()): + if "v_norm" in key: + del state_dict[key] + + def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False): + super().load_state_dict(state_dict, strict=strict, assign=assign) + if getattr(self.config, "tie_word_embeddings", False): + self.lm_head.weight = self.model.embed_tokens.weight diff --git a/python/src/coreai_models/models/registry.py b/python/src/coreai_models/models/registry.py index 85f0f06c..059b880d 100644 --- a/python/src/coreai_models/models/registry.py +++ b/python/src/coreai_models/models/registry.py @@ -90,6 +90,7 @@ def _get_registry() -> dict[str, ModelEntry]: from coreai_models.models.ios.qwen2 import Qwen2ForCausalLMForiOS from coreai_models.models.ios.qwen3 import Qwen3ForCausalLMForiOS from coreai_models.models.macos.gemma3_text import Gemma3ForCausalLM + from coreai_models.models.macos.gemma3n import Gemma3nForCausalLM from coreai_models.models.macos.gpt_oss import GptOssForCausalLM from coreai_models.models.macos.mistral import MistralForCausalLM from coreai_models.models.macos.mixtral import MixtralForCausalLM @@ -108,6 +109,11 @@ def _get_registry() -> dict[str, ModelEntry]: hf_config_attr="text_config", hf_state_dict_prefix="language_model.", ), + "gemma3n_text": ModelEntry( + macos_class=Gemma3nForCausalLM, + hf_config_attr="text_config", + hf_state_dict_prefix="language_model.", + ), "gpt_oss": ModelEntry( macos_class=GptOssForCausalLM, ), @@ -151,6 +157,7 @@ def _get_registry() -> dict[str, ModelEntry]: # Type alias for the remapping dict MODEL_TYPE_REMAPPING: dict[str, str] = { "gemma3": "gemma3_text", + "gemma3n": "gemma3n_text", "muse_glimmer": "muse_glimmer_text", "qwen2_5": "qwen2", } 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 new file mode 100644 index 00000000..c355356f --- /dev/null +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_gemma3n.py @@ -0,0 +1,235 @@ +# 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 macOS Gemma 3n model parity with HuggingFace.""" + +import torch +from transformers.models.gemma3n.configuration_gemma3n import Gemma3nTextConfig +from transformers.models.gemma3n.modeling_gemma3n import ( + Gemma3nForCausalLM as HFGemma3nForCausalLM, +) + +from coreai_models.models.macos.gemma3n import Gemma3nForCausalLM +from coreai_models.primitives.macos.cache import KVCache + + +def _make_gemma3n_config(**overrides) -> Gemma3nTextConfig: + n_layers = overrides.pop("num_hidden_layers", 10) + defaults = dict( + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=n_layers, + intermediate_size=[128] * n_layers, + vocab_size=200, + vocab_size_per_layer_input=200, + max_position_embeddings=32, + head_dim=16, + rms_norm_eps=1e-6, + rope_theta=1000000.0, + rope_local_base_freq=10000.0, + hidden_size_per_layer_input=16, + altup_num_inputs=4, + altup_active_idx=0, + altup_correct_scale=True, + laurel_rank=8, + num_kv_shared_layers=0, + activation_sparsity_pattern=[0.0] * n_layers, + sliding_window=8, + hidden_activation="gelu_pytorch_tanh", + tie_word_embeddings=True, + ) + defaults.update(overrides) + return Gemma3nTextConfig(**defaults) + + +def _build_models(config): + torch.manual_seed(42) + hf_model = HFGemma3nForCausalLM(config).to(torch.float32).eval() + our_model = Gemma3nForCausalLM(config, model_device="cpu") + our_model.to(torch.float32).eval() + sd = dict(hf_model.state_dict()) + our_model._mutate_state_dict(sd) + our_model.load_state_dict(sd, assign=True, strict=True) + return hf_model, our_model + + +class TestGemma3nForCausalLM: + """Test macOS Gemma3nForCausalLM against HuggingFace reference.""" + + def test_forward_parity_single_token(self): + config = _make_gemma3n_config(num_hidden_layers=5) + hf_model, our_model = _build_models(config) + + input_ids = torch.randint(0, 200, (1, 1)) + position_ids = torch.tensor([[0]], dtype=torch.int32) + 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_float16(self): + config = _make_gemma3n_config(num_hidden_layers=5) + torch.manual_seed(42) + hf_model = HFGemma3nForCausalLM(config).to(torch.float16).eval() + our_model = Gemma3nForCausalLM(config, model_device="cpu") + our_model.to(torch.float16).eval() + sd = dict(hf_model.state_dict()) + our_model._mutate_state_dict(sd) + our_model.load_state_dict(sd, assign=True, strict=True) + + input_ids = torch.randint(0, 200, (1, 4)) + position_ids = torch.arange(4, dtype=torch.int32).unsqueeze(0) + k_cache, v_cache = KVCache.create_cache_tensors(config, dtype=torch.float16) + + 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=5e-2, rtol=5e-2) + + def test_forward_parity_multi_token(self): + config = _make_gemma3n_config(num_hidden_layers=5) + hf_model, our_model = _build_models(config) + + 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) + hf_model, our_model = _build_models(config) + + input_ids = torch.randint(0, 200, (1, 4)) + position_ids = torch.arange(4, dtype=torch.int32).unsqueeze(0) + k_cache, v_cache = Gemma3nForCausalLM.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=True) + + torch.testing.assert_close(our_out, hf_out.logits, atol=1e-4, rtol=1e-4) + + def test_forward_parity_with_sparsity(self): + """Gaussian TopK activation sparsity in early layers.""" + sparsity = [0.95, 0.95, 0.0, 0.0, 0.0] + config = _make_gemma3n_config(num_hidden_layers=5, activation_sparsity_pattern=sparsity) + hf_model, our_model = _build_models(config) + + input_ids = torch.randint(0, 200, (1, 4)) + position_ids = torch.arange(4, 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_full_config(self): + """Full config: KV sharing + sparsity + 10 layers.""" + sparsity = [0.95, 0.95, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + config = _make_gemma3n_config( + num_hidden_layers=10, + num_kv_shared_layers=2, + activation_sparsity_pattern=sparsity, + ) + hf_model, our_model = _build_models(config) + + input_ids = torch.randint(0, 200, (1, 4)) + position_ids = torch.arange(4, dtype=torch.int32).unsqueeze(0) + k_cache, v_cache = Gemma3nForCausalLM.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=True) + + torch.testing.assert_close(our_out, hf_out.logits, atol=1e-4, rtol=1e-4) + + def test_output_shape(self): + config = _make_gemma3n_config(num_hidden_layers=5) + our_model = Gemma3nForCausalLM(config, model_device="cpu").to(torch.float32).eval() + + batch, seq_len = 1, 6 + input_ids = torch.randint(0, 200, (batch, seq_len)) + position_ids = torch.arange(seq_len, dtype=torch.int32).unsqueeze(0) + k_cache, v_cache = KVCache.create_cache_tensors(config, dtype=torch.float32) + + with torch.no_grad(): + out = our_model(input_ids, position_ids, k_cache, v_cache) + + assert out.shape == (batch, seq_len, config.vocab_size) + + def test_kv_sharing_cache_compaction(self): + """Cache should have fewer slots than layers when KV sharing is enabled.""" + config = _make_gemma3n_config(num_hidden_layers=10, num_kv_shared_layers=2) + model = Gemma3nForCausalLM(config, model_device="cpu") + + assert model.model.num_cache_slots == 8 + k_cache, v_cache = Gemma3nForCausalLM.create_cache_tensors(config, dtype=torch.float32) + assert k_cache.shape[0] == 8 + + def test_kv_sharing_slot_mapping(self): + """Shared layers should map to source layer's cache slot.""" + config = _make_gemma3n_config(num_hidden_layers=10, num_kv_shared_layers=2) + model = Gemma3nForCausalLM(config, model_device="cpu") + + layers = model.model.layers + # Layer 8 (sliding) should share with last non-shared sliding layer + assert layers[8].self_attn.is_kv_shared + assert layers[8].self_attn.cache_slot == layers[7].self_attn.cache_slot + # Layer 9 (full) should share with last non-shared full layer (layer 4) + assert layers[9].self_attn.is_kv_shared + assert layers[9].self_attn.cache_slot == layers[4].self_attn.cache_slot + + def test_altup_predict_correct_shape(self): + """AltUp should maintain [altup_num_inputs, B, S, D] shape through predict/correct.""" + config = _make_gemma3n_config(num_hidden_layers=5) + model = Gemma3nForCausalLM(config, model_device="cpu").to(torch.float32).eval() + altup = model.model.layers[0].altup + + hidden = torch.randn(4, 1, 3, 64) + predictions = altup.predict(hidden) + assert predictions.shape == (4, 1, 3, 64) + + activated = torch.randn(1, 3, 64) + corrected = altup.correct(predictions, activated) + assert corrected.shape == (4, 1, 3, 64) + + def test_sliding_window_pattern(self): + """Layers should alternate: 4 sliding + 1 full.""" + config = _make_gemma3n_config(num_hidden_layers=10) + model = Gemma3nForCausalLM(config, model_device="cpu") + pattern = [layer.self_attn.is_sliding for layer in model.model.layers] + expected = [True, True, True, True, False, True, True, True, True, False] + assert pattern == expected + + def test_tie_word_embeddings(self): + config = _make_gemma3n_config(num_hidden_layers=5, tie_word_embeddings=True) + hf_model = HFGemma3nForCausalLM(config).eval() + our_model = Gemma3nForCausalLM(config, model_device="cpu").eval() + sd = dict(hf_model.state_dict()) + our_model._mutate_state_dict(sd) + our_model.load_state_dict(sd, assign=True, strict=True) + + assert our_model.lm_head.weight is our_model.model.embed_tokens.weight