From 6697b9d07d48f7a1ebc7a59251ffabb25e4735da Mon Sep 17 00:00:00 2001 From: none Date: Thu, 5 Nov 2026 16:27:45 -0500 Subject: [PATCH] feat(gemma4): serve the E2B and E4B releases E2B and E4B add three things to the gemma4 decoder: a per-layer input embedding table, KV-shared layers, and (E2B) a double-wide MLP on those layers. Both checkpoint formats are served; the GGUF path no longer refuses these files. - Per-layer inputs: embed_tokens_per_layer plus a per-layer projection and norm, one more residual step before layer_scalar. The table defaults to the GPU. --embed-weights host now moves it to pinned host RAM, which is how these models fit an 8 GB GPU, and CUDA graphs stay on. A table tied to the LM head stays on the GPU. - KV sharing: the last num_kv_shared_layers layers project q only and read k, v from the last earlier layer of their attention type, handed over in a dict that lives for one forward. Each shared layer still stores its own KV copy; reading the source's slab can be a follow-up. - MLP width: ModelConfig.layer_intermediate_sizes is a per-layer override, read from HF's use_double_wide_mlp or GGUF's feed_forward_length list. intermediate_size stays the base width. The E-series vision tower applies clipped linears, which the tower does not implement, so these checkpoints serve text only with a warning instead of refusing to start. An image request gets a 400. Co-Authored-By: Claude Opus 5.5 --- docs/cli.md | 2 +- docs/models.md | 4 +- python/freetoken/engine/config.py | 17 +- python/freetoken/engine/engine.py | 10 +- python/freetoken/models/blocks.py | 14 +- python/freetoken/models/config.py | 8 + python/freetoken/models/gemma4/attention.py | 81 +++-- python/freetoken/models/gemma4/config.py | 46 ++- python/freetoken/models/gemma4/gguf.py | 74 ++++- python/freetoken/models/gemma4/model.py | 71 ++++- python/freetoken/models/gemma4/moe.py | 18 +- python/freetoken/models/gemma4/weight.py | 23 +- python/freetoken/models/gguf/tensors.py | 12 +- python/freetoken/models/register.py | 12 +- python/freetoken/server/args.py | 6 +- tests/models/test_gemma4_e_series.py | 319 ++++++++++++++++++++ 16 files changed, 637 insertions(+), 80 deletions(-) create mode 100644 tests/models/test_gemma4_e_series.py diff --git a/docs/cli.md b/docs/cli.md index 8e7a2e0850..1449e7a919 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -71,7 +71,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b # the same card by UUID (a unique prefi |---|---|---| | `--memory-ratio` | 0.9 | Fraction of free VRAM the engine may use (weights + MoE cache + KV) | | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) | -| `--embed-weights` | gpu | Where the input token embedding table lives. `host` keeps it in pinned host RAM and each lookup reads its rows over PCIe, so its VRAM goes to the expert cache and KV. Refused when the table is tied to the LM head | +| `--embed-weights` | gpu | Where the input token embedding tables live. `host` keeps them in pinned host RAM and each lookup reads its rows over PCIe, so their VRAM goes to the expert cache and KV. A table tied to the LM head stays on the GPU, and `host` is refused when there is no other table; Gemma-4 E2B / E4B's per-layer table is the one that moves, the way those models fit an 8 GB GPU | | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 | | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` | | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU | diff --git a/docs/models.md b/docs/models.md index dbf321e23a..ed12ef66b0 100644 --- a/docs/models.md +++ b/docs/models.md @@ -16,7 +16,7 @@ for them; other checkpoints of the same architectures work too. | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | | Qwen3-VL | [Qwen/Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct), [Qwen/Qwen3-VL-30B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct) | | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | -| Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. | +| Gemma-4 | [google/gemma-4-E2B-it](https://huggingface.co/google/gemma-4-E2B-it), [google/gemma-4-E4B-it](https://huggingface.co/google/gemma-4-E4B-it), [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. | | MiniMax-M2.5 | [nvidia/MiniMax-M2.5-NVFP4](https://huggingface.co/nvidia/MiniMax-M2.5-NVFP4) | | MiniMax-M3 | [nvidia/MiniMax-M3-NVFP4](https://huggingface.co/nvidia/MiniMax-M3-NVFP4) | | Muse-Glimmer | [meta-models/Muse-Glimmer-30B](https://huggingface.co/meta-models/Muse-Glimmer-30B), [RedHatAI/Muse-Glimmer-30B-NVFP4](https://huggingface.co/RedHatAI/Muse-Glimmer-30B-NVFP4) | @@ -35,7 +35,7 @@ These llama.cpp architectures (`general.architecture`) are served: | `qwen3vlmoe`, `qwen3vl` | Qwen3-VL | | `qwen2` | Qwen2, Qwen2.5 | | `llama` | Llama, Mistral | -| `gemma4` | Gemma-4, except E2B / E4B | +| `gemma4` | Gemma-4 | | `glm4moe` | GLM-4.5, GLM-4.6, GLM-4.7 | | `glm5-next` | GLM-5.3-Flash | | `deepseek4` | DeepSeek-V4 Flash and Pro | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 531a4310c2..6796bbdcc3 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -144,12 +144,17 @@ def model_spec(self) -> ModelSpec: @cached_property def active_encoders(self) -> tuple[EncoderSpec, ...]: """The encoder towers this process builds: the family registers them, the checkpoint config carries their section, --mm-disable did not name them.""" - return tuple( - e - for e in self.model_spec.encoders - if getattr(self.hf_config, e.config_key, None) is not None - and e.kind not in self.mm.disabled_encoders - ) + active = [] + for e in self.model_spec.encoders: + section = getattr(self.hf_config, e.config_key, None) + if section is None or e.kind in self.mm.disabled_encoders: + continue + reason = e.unserved(section) if e.unserved is not None else None + if reason is not None: + logger.warning(f"{e.kind} input is not served on this checkpoint: {reason}") + continue + active.append(e) + return tuple(active) @cached_property def served_modalities(self) -> frozenset[str]: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 5365cd659f..ad8e83cc27 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -16,7 +16,7 @@ from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info from freetoken.gpu_select import gpu_identity from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint -from freetoken.layers import OffloadMoELayer, VocabParallelEmbedding, iter_moe_layers, iter_ops, set_rope_device +from freetoken.layers import OffloadMoELayer, ParallelLMHead, VocabParallelEmbedding, iter_moe_layers, iter_ops, set_rope_device from freetoken.layers.quantization import LayerKind, QuantBackend, QuantKind, finalize_quant, set_quant_backend from freetoken.moe.offload_cache import iter_offload_moe_layers from freetoken.mm.config import ENCODER_SECTIONS @@ -340,13 +340,15 @@ def _host_embedding_keys(model, config: EngineConfig) -> frozenset[str]: """State-dict keys of the tables --embed-weights host keeps in pinned host RAM.""" if config.embed_weights != "host": return frozenset() - if config.model_config.tie_word_embeddings: + # a tied LM head reads its whole table every step, so that table stays on the GPU + tied = {id(op.tied_embedding) for op in iter_ops(model) if isinstance(op, ParallelLMHead) and op.tied_embedding is not None} + tables = {id(op.weight) for op in iter_ops(model) if type(op) is VocabParallelEmbedding and id(op) not in tied} + keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables) + if not keys and tied: raise ValueError( "--embed-weights host: this model's embedding table is tied to its LM head, which reads the whole table " "every step; serve it with --embed-weights gpu" ) - tables = {id(op.weight) for op in iter_ops(model) if type(op) is VocabParallelEmbedding} - keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables) if not keys: raise ValueError( f"--embed-weights host: {type(model).__name__} has no plain embedding table to move (a DF11 table stays " diff --git a/python/freetoken/models/blocks.py b/python/freetoken/models/blocks.py index e51f3637ec..18ef79c6e9 100644 --- a/python/freetoken/models/blocks.py +++ b/python/freetoken/models/blocks.py @@ -76,10 +76,18 @@ def embed_input_ids(embed_tokens, input_ids: torch.Tensor, batch: Batch) -> torc class GatedMLP(BaseOP): - def __init__(self, config: ModelConfig, *, quant_config: QuantConfig | None = None, prefix: str = ""): + def __init__( + self, + config: ModelConfig, + *, + intermediate_size: int | None = None, + quant_config: QuantConfig | None = None, + prefix: str = "", + ): + intermediate_size = intermediate_size or config.intermediate_size self.gate_up_proj = LinearColParallelMerged( config.hidden_size, - [config.intermediate_size, config.intermediate_size], + [intermediate_size, intermediate_size], has_bias=False, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", @@ -91,7 +99,7 @@ def __init__(self, config: ModelConfig, *, quant_config: QuantConfig | None = No raise ValueError(f"Unsupported activation function: {config.hidden_act}") self.act_fn = act_fn self.down_proj = LinearRowParallel( - config.intermediate_size, + intermediate_size, config.hidden_size, has_bias=False, quant_config=quant_config, diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 7e954c06df..c2dbab5876 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -323,6 +323,8 @@ class ModelConfig: # the checkpoint's quant kind for the lm_head, detected the same way (only some NVFP4 exports quantize it) lm_head_quant: str = "none" shared_expert_intermediate_size: int = 0 + # The dense MLP's width per layer, for a family whose layers differ (Gemma-4 E2B's double-wide KV-shared tail); () -> intermediate_size on every layer. + layer_intermediate_sizes: Tuple[int, ...] = () use_qk_norm: bool = False # ----- DeepSeek/GLM-style MoE extensions (default keeps other models intact) ----- # The first ``first_k_dense_replace`` decoder layers use a dense MLP instead of the @@ -378,6 +380,9 @@ class ModelConfig: # n-gram embedding geometry and the QSA indexer scoring geometry the model module # needs. Opaque to model-agnostic engine code; None for every other model. qwen4_args: Any | None = None + # Gemma-4 E2B / E4B payload (Gemma4EArgs): per-layer input embeddings and KV-shared + # layers. Opaque to model-agnostic engine code; None for every other model. + gemma4_e_args: Any | None = None # Generic execution-path capability flags (set by a model's parse_config) so the engine and # factories stay model-agnostic instead of branching on dsv4_args: single_stream_only: bool = False # model runs one sequence at a time -> force bs=1 @@ -385,6 +390,9 @@ class ModelConfig: # () for models without any. Requires a linear-attention group to ride on. slot_states: Tuple[SlotStateSpec, ...] = () + def intermediate_size_for_layer(self, layer_id: int) -> int: + return self.layer_intermediate_sizes[layer_id] if self.layer_intermediate_sizes else self.intermediate_size + @property def is_moe(self) -> bool: return "moe" in self.model_type or self.moe_enabled diff --git a/python/freetoken/models/gemma4/attention.py b/python/freetoken/models/gemma4/attention.py index 3f34e1cd11..005dee369b 100644 --- a/python/freetoken/models/gemma4/attention.py +++ b/python/freetoken/models/gemma4/attention.py @@ -15,7 +15,13 @@ class Gemma4Attention(BaseOP): - """Gemma 4 attention for one full-context or SWA layer.""" + """Gemma 4 attention for one full-context or SWA layer. + + On E2B / E4B the KV-shared tail projects q only and attends over the k, v of the last + earlier layer of its attention type. That layer hands its rows over in ``shared_kv``, a + dict living for one forward, so a captured graph records the same tensors. Each shared + layer still stores them in its own KV slab; reading the source's slab would save that memory. + """ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): self.layer_id = layer_id @@ -28,25 +34,35 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): self.num_kv_heads = group.num_kv_heads self.num_qo_heads = config.num_qo_heads self.k_eq_v = isinstance(group, FullAttentionGroupConfig) and group.k_eq_v + e_args = config.gemma4_e_args + self.kv_shared = e_args is not None and e_args.is_kv_shared(layer_id) + self.kv_source = e_args is not None and e_args.is_kv_source(layer_id, group.layer_ids) + self._group_name = group.name self.q_dim = self.num_qo_heads * self.head_dim self.kv_dim = self.num_kv_heads * self.head_dim - self.qkv_proj = LinearQKVMerged( - config.hidden_size, - self.head_dim, - self.num_qo_heads, - self.num_kv_heads, - has_bias=False, - quant_config=config.quant, - prefix=f"{prefix}.qkv_proj", - ) + if self.kv_shared: + self.q_proj = LinearReplicated( + config.hidden_size, self.q_dim, has_bias=False, + quant_config=config.quant, prefix=f"{prefix}.q_proj", + ) + else: + self.qkv_proj = LinearQKVMerged( + config.hidden_size, + self.head_dim, + self.num_qo_heads, + self.num_kv_heads, + has_bias=False, + quant_config=config.quant, + prefix=f"{prefix}.qkv_proj", + ) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.v_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) self.o_proj = LinearReplicated( self.q_dim, config.hidden_size, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.o_proj", ) self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.v_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) self.attn_spec = AttentionSpec( sliding_window=group.sliding_window if self.is_swa else None, sm_scale=config.attn_sm_scale, @@ -79,29 +95,38 @@ def _apply_rope( return q_view.view_as(q), k_view.view_as(k) @nvtx_annotate("MHA") - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, shared_kv: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> torch.Tensor: ctx = get_global_ctx() positions = ctx.batch.positions T = x.shape[0] - qkv = self.qkv_proj.forward(x) - q_lin, k_lin, v_lin = qkv.split( - (self.q_dim, self.kv_dim, self.kv_dim), - dim=-1, - ) - del qkv - q = q_lin.view(T, self.num_qo_heads, self.head_dim) - k = k_lin.view(T, self.num_kv_heads, self.head_dim) - v = v_lin.view(T, self.num_kv_heads, self.head_dim) + if self.kv_shared: + q = self.q_proj.forward(x).view(T, self.num_qo_heads, self.head_dim) + q = self.q_norm.forward(q) + # the rope kernels rotate a key along with the query: give them one head of zeros + q, _ = self._apply_rope(positions, q, q.new_zeros(T, 1, self.head_dim)) + k, v = shared_kv[self._group_name] + else: + qkv = self.qkv_proj.forward(x) + q_lin, k_lin, v_lin = qkv.split( + (self.q_dim, self.kv_dim, self.kv_dim), + dim=-1, + ) + del qkv + q = q_lin.view(T, self.num_qo_heads, self.head_dim) + k = k_lin.view(T, self.num_kv_heads, self.head_dim) + v = v_lin.view(T, self.num_kv_heads, self.head_dim) - q = self.q_norm.forward(q) - k = self.k_norm.forward(k) - v = self.v_norm.forward(v) + q = self.q_norm.forward(q) + k = self.k_norm.forward(k) + v = self.v_norm.forward(v) - q, k = self._apply_rope(positions, q, k) + q, k = self._apply_rope(positions, q, k) - k = k.reshape(T, self.num_kv_heads * self.head_dim) - v = v.reshape(T, self.num_kv_heads * self.head_dim) + k = k.reshape(T, self.num_kv_heads * self.head_dim) + v = v.reshape(T, self.num_kv_heads * self.head_dim) + if self.kv_source: + shared_kv[self._group_name] = (k, v) o = ctx.attn_backend.forward( q.contiguous(), k.contiguous(), diff --git a/python/freetoken/models/gemma4/config.py b/python/freetoken/models/gemma4/config.py index 0f99fd77b2..a613f3aafc 100644 --- a/python/freetoken/models/gemma4/config.py +++ b/python/freetoken/models/gemma4/config.py @@ -43,6 +43,34 @@ class UnifiedVisionConfig: text_hidden_size: int +@dataclass(frozen=True) +class Gemma4EArgs: + """What the E2B / E4B releases add to the decoder.""" + + # width of one layer's per-layer input; the table holds num_layers of them per token + per_layer_hidden_size: int + per_layer_vocab_size: int + # layers from here on project q only and reuse k, v of the last earlier layer of their attention type + first_kv_shared_layer: int + + def is_kv_shared(self, layer_id: int) -> bool: + return layer_id >= self.first_kv_shared_layer + + def is_kv_source(self, layer_id: int, group_layer_ids: tuple[int, ...]) -> bool: + return layer_id == max((i for i in group_layer_ids if i < self.first_kv_shared_layer), default=None) + + +def gemma4_e_args( + num_layers: int, per_layer_hidden_size: int, per_layer_vocab_size: int, num_kv_shared_layers: int +) -> Gemma4EArgs | None: + """None for the releases without per-layer inputs: their configs set the width to 0.""" + if not per_layer_hidden_size: + if num_kv_shared_layers: + raise NotImplementedError("gemma4 with KV-shared layers but no per-layer inputs is not supported") + return None + return Gemma4EArgs(per_layer_hidden_size, per_layer_vocab_size, num_layers - num_kv_shared_layers) + + def _text_config(hf_config: Any) -> tuple[Any, list[str] | None, Any]: top_architectures = getattr(hf_config, "architectures", None) if hasattr(hf_config, "text_config") and hf_config.text_config is not None: @@ -162,6 +190,15 @@ def parse_config(hf_config: Any) -> ModelConfig: expert_quant = detect_expert_quant(top_cfg) dense_quant = "nvfp4" if (expert_quant == "nvfp4" and not num_experts > 0) else "none" + # E2B doubles the MLP of its KV-shared tail. Gemma4UnifiedTextConfig omits these fields. + num_kv_shared_layers = getattr(cfg, "num_kv_shared_layers", None) or 0 + first_kv_shared_layer = cfg.num_hidden_layers - num_kv_shared_layers + layer_intermediate_sizes = ( + tuple(cfg.intermediate_size * (2 if i >= first_kv_shared_layer else 1) for i in range(cfg.num_hidden_layers)) + if getattr(cfg, "use_double_wide_mlp", False) and num_kv_shared_layers + else () + ) + return ModelConfig( num_layers=cfg.num_hidden_layers, num_qo_heads=cfg.num_attention_heads, @@ -170,6 +207,7 @@ def parse_config(hf_config: Any) -> ModelConfig: hidden_size=cfg.hidden_size, vocab_size=cfg.vocab_size, intermediate_size=cfg.intermediate_size, + layer_intermediate_sizes=layer_intermediate_sizes, hidden_act="gelu_tanh", rms_norm_eps=cfg.rms_norm_eps, tie_word_embeddings=bool(getattr(cfg, "tie_word_embeddings", True)), @@ -189,6 +227,12 @@ def parse_config(hf_config: Any) -> ModelConfig: embedding_scale=float(cfg.hidden_size) ** 0.5, vision_config=_parse_vision_config(top_cfg, cfg.hidden_size), image_token_id=getattr(top_cfg, "image_token_id", None), + gemma4_e_args=gemma4_e_args( + cfg.num_hidden_layers, + getattr(cfg, "hidden_size_per_layer_input", None) or 0, + getattr(cfg, "vocab_size_per_layer_input", None) or 0, + num_kv_shared_layers, + ), attention_groups=( FullAttentionGroupConfig( name="full", @@ -211,4 +255,4 @@ def parse_config(hf_config: Any) -> ModelConfig: ) -__all__ = ["UnifiedVisionConfig", "VisionConfig", "parse_config"] +__all__ = ["Gemma4EArgs", "UnifiedVisionConfig", "VisionConfig", "gemma4_e_args", "parse_config"] diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py index ce806ba8f3..2438818747 100644 --- a/python/freetoken/models/gemma4/gguf.py +++ b/python/freetoken/models/gemma4/gguf.py @@ -26,9 +26,9 @@ ) from freetoken.models.gguf.mmproj import Mmproj, find_mmproj, log_mmproj, mmproj_tensor from freetoken.models.gguf.reader import gguf_tensor_index, load_gguf_metadata -from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix, weight_scale +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix, table_matrix, weight_scale -from .config import _parse_vision_config +from .config import _parse_vision_config, gemma4_e_args if TYPE_CHECKING: from freetoken.models.gguf.config import GgufConfigShim @@ -140,6 +140,7 @@ def kv(key: str): "rope_parameters": {"rope_theta": _ROPE_THETA}, "hidden_activation": "gelu_pytorch_tanh", "standardize": "v.std_bias" in index, + "use_clipped_linears": "v.blk.0.attn_q.input_max" in index, } processor_type = "Gemma4ImageProcessor" else: @@ -217,13 +218,17 @@ def g(key: str): raise KeyError(f"missing GGUF metadata key gemma4.{key}") return val - if int(m.get("gemma4.embedding_length_per_layer_input", 0)) or int(m.get("gemma4.attention.shared_kv_layers", 0)): - raise NotImplementedError("gemma4 E2B / E4B GGUF (per-layer input embeddings, shared KV layers) is not supported") num_layers = int(g("block_count")) hidden = int(g("embedding_length")) num_experts = int(m.get("gemma4.expert_count", 0)) num_qo_heads = int(g("attention.head_count")) - kv_per_layer = g("attention.head_count_kv") # per-layer list + # per-layer lists, or one value for every layer + kv_per_layer = g("attention.head_count_kv") + ffn_per_layer = g("feed_forward_length") + if not isinstance(kv_per_layer, (list, tuple)): + kv_per_layer = [kv_per_layer] * num_layers + layer_intermediate_sizes = tuple(int(x) for x in ffn_per_layer) if isinstance(ffn_per_layer, (list, tuple)) else () + assert len(layer_intermediate_sizes) in (0, num_layers), "feed_forward_length length != block_count" # True -> sliding-window (SWA) layer, False -> full attention. swa_pattern = [bool(x) for x in g("attention.sliding_window_pattern")] assert len(swa_pattern) == num_layers, "sliding_window_pattern length != block_count" @@ -264,7 +269,8 @@ def g(key: str): head_dim=full_head_dim, hidden_size=hidden, vocab_size=int(shim.vocab_size), - intermediate_size=int(g("feed_forward_length")), + intermediate_size=layer_intermediate_sizes[0] if layer_intermediate_sizes else int(ffn_per_layer), + layer_intermediate_sizes=layer_intermediate_sizes, hidden_act="gelu_tanh", rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), tie_word_embeddings=bool(shim.tie_word_embeddings), @@ -283,6 +289,12 @@ def g(key: str): embedding_scale=float(hidden) ** 0.5, vision_config=vision_config, image_token_id=getattr(shim, "image_token_id", None), + gemma4_e_args=gemma4_e_args( + num_layers, + int(m.get("gemma4.embedding_length_per_layer_input", 0)), + int(shim.vocab_size), + int(m.get("gemma4.attention.shared_kv_layers", 0)), + ), attention_groups=( FullAttentionGroupConfig( name="full", @@ -290,8 +302,8 @@ def g(key: str): num_kv_heads=full_kv, head_dim=full_head_dim, rotary_config=full_rotary, - # Full layers ship no v_proj in the GGUF (k reused as v). - k_eq_v=True, + # Full layers ship no v_proj in the GGUF (k reused as v), except on E2B / E4B. + k_eq_v="gemma4.attention.value_length" not in m, ), SWAAttentionGroupConfig( name="swa", @@ -331,8 +343,15 @@ def g(key: str): "ffn_down_exps.scale": "feed_forward.router.per_expert_scale", } +# E2B / E4B's per-layer input +_PER_LAYER_SCALAR_MAP = { + "post_norm.weight": "post_per_layer_input_norm.weight", +} + # checkpoint module leaf (gemma4 HF naming) -> GGUF tensor _LAYER_LEAVES = { + "per_layer_input_gate": "inp_gate", + "per_layer_projection": "proj", "self_attn.q_proj": "attn_q", "self_attn.k_proj": "attn_k", "self_attn.o_proj": "attn_output", @@ -350,6 +369,10 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" if name == "model.embed_tokens": return ("token_embd",) + if name == "model.embed_tokens_per_layer": + return ("per_layer_token_embd",) + if name == "model.per_layer_model_projection": + return ("per_layer_model_proj",) if name == "lm_head": return ("output",) if "output.weight" in tensors else () m = _LAYER_RE.match(name) @@ -410,27 +433,48 @@ def t(name: str): yield from fused("lm_head", [out], [matrix(out)], hidden) yield "model.norm.weight", float_tensor(t("output_norm")) - scalar_map = {**_LAYER_SCALAR_MAP, **_MOE_LAYER_SCALAR_MAP} if config.moe_enabled else _LAYER_SCALAR_MAP + e_args = config.gemma4_e_args + if e_args is not None: + table, projection = t("per_layer_token_embd"), t("per_layer_model_proj") + yield from fused("model.embed_tokens_per_layer", [table], [table_matrix(table)], int(table.ne[0])) + yield from fused("model.per_layer_model_projection", [projection], [matrix(projection)], hidden) + yield "model.per_layer_projection_norm.weight", float_tensor(t("per_layer_proj_norm")) + + scalar_map = {**_LAYER_SCALAR_MAP, **_MOE_LAYER_SCALAR_MAP} if config.moe_enabled else dict(_LAYER_SCALAR_MAP) + if e_args is not None: + scalar_map.update(_PER_LAYER_SCALAR_MAP) for layer in range(config.num_layers): blk, base = f"blk.{layer}", f"model.layers.{layer}" + kv_shared = e_args is not None and e_args.is_kv_shared(layer) for suffix, rel in scalar_map.items(): + # a KV-shared layer has no k_norm, whether or not the file still ships one + if kv_shared and suffix == "attn_k_norm.weight": + continue tensor = float_tensor(index[f"{blk}.{suffix}"]) if suffix == "ffn_down_exps.scale" and weight_scale(index[f"{blk}.ffn_down_exps.weight"]) is not None: # the converter folds per_expert_scale into the NVFP4 down weight_scale_2, which the expert kernel applies tensor = torch.ones_like(tensor) yield f"{base}.{rel}", tensor.reshape(1) if rel.endswith("layer_scalar") else tensor - if f"{blk}.attn_qkv.weight" in index: - # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix - qkv = [t(f"{blk}.attn_qkv")] + if kv_shared: + q = t(f"{blk}.attn_q") + yield from fused(f"{base}.self_attn.q_proj", [q], [matrix(q)], hidden) else: - q, k = t(f"{blk}.attn_q"), t(f"{blk}.attn_k") - qkv = [q, k, index.get(f"{blk}.attn_v.weight", k)] - yield from fused(f"{base}.self_attn.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + if f"{blk}.attn_qkv.weight" in index: + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] + else: + q, k = t(f"{blk}.attn_q"), t(f"{blk}.attn_k") + qkv = [q, k, index.get(f"{blk}.attn_v.weight", k)] + yield from fused(f"{base}.self_attn.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) o = t(f"{blk}.attn_output") yield from fused(f"{base}.self_attn.o_proj", [o], [matrix(o)], int(o.ne[0])) gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) yield from fused(f"{base}.feed_forward.shared_mlp.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) yield from fused(f"{base}.feed_forward.shared_mlp.down_proj", [down], [matrix(down)], int(down.ne[0])) + if e_args is not None: + gate, proj = t(f"{blk}.inp_gate"), t(f"{blk}.proj") + yield from fused(f"{base}.per_layer_input_gate", [gate], [matrix(gate)], hidden) + yield from fused(f"{base}.per_layer_projection", [proj], [matrix(proj)], int(proj.ne[0])) if vision: yield from iter_mmproj_weights(find_mmproj(model_path)) diff --git a/python/freetoken/models/gemma4/model.py b/python/freetoken/models/gemma4/model.py index 1b9d0da7d9..4b5c685c22 100644 --- a/python/freetoken/models/gemma4/model.py +++ b/python/freetoken/models/gemma4/model.py @@ -3,10 +3,12 @@ from typing import TYPE_CHECKING import torch +import torch.nn.functional as F from freetoken.core import get_global_ctx from freetoken.layers import ( BaseOP, GemmaRMSNorm, + LinearReplicated, OPList, ParallelLMHead, VocabParallelEmbedding, @@ -23,11 +25,15 @@ from freetoken.message import MMItem from freetoken.models.config import ModelConfig +# every release's pad token; the per-layer table row a multimodal row reads +_PAD_TOKEN_ID = 0 + class Gemma4DecoderLayer(BaseOP): """Gemma 4 decoder block: attention sandwich + feed-forward sandwich, scaled by a per-layer ``layer_scalar``. The feed-forward is the dual (shared MLP || routed MoE) - branch for MoE checkpoints, or a single dense MLP branch for dense checkpoints.""" + branch for MoE checkpoints, or a single dense MLP branch for dense checkpoints. + E2B / E4B add their per-layer input as one more residual step before the scalar.""" def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): self._layer_id = layer_id @@ -35,7 +41,7 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): self.feed_forward = ( Gemma4MLP(config, layer_id, prefix=f"{prefix}.feed_forward") if config.is_moe - else Gemma4DenseMLP(config, prefix=f"{prefix}.feed_forward") + else Gemma4DenseMLP(config, layer_id, prefix=f"{prefix}.feed_forward") ) eps = config.rms_norm_eps @@ -43,16 +49,35 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): self.input_layernorm = GemmaRMSNorm(H, eps=eps) self.post_attention_layernorm = GemmaRMSNorm(H, eps=eps) self.pre_feedforward_layernorm = GemmaRMSNorm(H, eps=eps) + if config.gemma4_e_args is not None: + P = config.gemma4_e_args.per_layer_hidden_size + self.per_layer_input_gate = LinearReplicated( + H, P, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.per_layer_input_gate" + ) + self.per_layer_projection = LinearReplicated( + P, H, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.per_layer_projection" + ) + self.post_per_layer_input_norm = GemmaRMSNorm(H, eps=eps) @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + per_layer_input: torch.Tensor | None, + shared_kv: dict[str, tuple[torch.Tensor, torch.Tensor]], + ) -> torch.Tensor: # --- attention sandwich --- residual = x h = self.input_layernorm.forward(x) - h = self.self_attn.forward(h) + h = self.self_attn.forward(h, shared_kv) h = self.post_attention_layernorm.forward(h) pre_ff, x = self.pre_feedforward_layernorm.forward_add_residual(h, residual) - return self.feed_forward.forward(pre_ff, x) + if per_layer_input is None: + return self.feed_forward.forward(pre_ff, x) + h = self.feed_forward.forward_unscaled(pre_ff, x) + g = F.gelu(self.per_layer_input_gate.forward(h), approximate="tanh") * per_layer_input + g = self.post_per_layer_input_norm.forward(self.per_layer_projection.forward(g)) + return (h + g) * self.feed_forward.layer_scalar class Gemma4Model(BaseOP): @@ -71,11 +96,43 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): ] ) self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self._has_per_layer_inputs = config.gemma4_e_args is not None + if self._has_per_layer_inputs: + P = config.gemma4_e_args.per_layer_hidden_size + self._per_layer_shape = (config.num_layers, P) + self.embed_tokens_per_layer = VocabParallelEmbedding( + num_embeddings=config.gemma4_e_args.per_layer_vocab_size, + embedding_dim=config.num_layers * P, + embed_scale=P**0.5, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens_per_layer", + ) + self.per_layer_model_projection = LinearReplicated( + config.hidden_size, + config.num_layers * P, + has_bias=False, + quant_config=config.quant, + prefix=f"{prefix}.per_layer_model_projection", + ) + self.per_layer_projection_norm = GemmaRMSNorm(P, eps=config.rms_norm_eps) + self._per_layer_projection_scale = config.hidden_size**-0.5 + + def _per_layer_inputs(self, input_ids: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """``[T, num_layers, P]``: each row's token embedding plus a projection of its input embedding.""" + table = self.embed_tokens_per_layer + # multimodal rows carry content pad ids above the vocab + ids = input_ids if get_global_ctx().batch.mm_embeds is None else input_ids.masked_fill(input_ids >= table.num_embeddings, _PAD_TOKEN_ID) + token = table.forward(ids).view(-1, *self._per_layer_shape) + context = self.per_layer_model_projection.forward(x) * self._per_layer_projection_scale + context = self.per_layer_projection_norm.forward(context.view(-1, *self._per_layer_shape)) + return (context + token) * 2.0**-0.5 def forward(self, input_ids: torch.Tensor) -> torch.Tensor: x = embed_input_ids(self.embed_tokens, input_ids, get_global_ctx().batch) - for layer in self.layers.op_list: - x = layer.forward(x) + per_layer = self._per_layer_inputs(input_ids, x) if self._has_per_layer_inputs else None + shared_kv: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for i, layer in enumerate(self.layers.op_list): + x = layer.forward(x, None if per_layer is None else per_layer[:, i], shared_kv) return self.norm.forward(x) diff --git a/python/freetoken/models/gemma4/moe.py b/python/freetoken/models/gemma4/moe.py index f5c8ca22c5..ede78db8c3 100644 --- a/python/freetoken/models/gemma4/moe.py +++ b/python/freetoken/models/gemma4/moe.py @@ -92,17 +92,25 @@ class Gemma4DenseMLP(BaseOP): (``mlp.* -> feed_forward.shared_mlp.*``, bare ``layer_scalar`` / ``post_feedforward_layernorm.``) resolves the dense checkpoint unchanged.""" - def __init__(self, config: ModelConfig, *, prefix: str = ""): - self.shared_mlp = GatedMLP(config, quant_config=config.quant, prefix=f"{prefix}.shared_mlp") + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + self.shared_mlp = GatedMLP( + config, + intermediate_size=config.intermediate_size_for_layer(layer_id), + quant_config=config.quant, + prefix=f"{prefix}.shared_mlp", + ) self.post_feedforward_layernorm = GemmaRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.layer_scalar = torch.empty(1) - def forward(self, pre_ff: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + def forward_unscaled(self, pre_ff: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """The sandwich before ``layer_scalar``; E2B / E4B add their per-layer input in between.""" h = self.shared_mlp.forward(pre_ff) - h = self.post_feedforward_layernorm.forward(h) - return (x + h) * self.layer_scalar + return x + self.post_feedforward_layernorm.forward(h) + + def forward(self, pre_ff: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + return self.forward_unscaled(pre_ff, x) * self.layer_scalar __all__ = [ diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index 7797c7842e..1b29a9e2e3 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -40,6 +40,9 @@ desc="Gemma4 NVFP4 experts", ) _LAYER_INDEX_PATTERN = re.compile(r"layers\.(\d+)\.") +# the attention leaves that change on a KV-shared layer; the vision tower's layers never match +_KV_SHARED_LEAF_PATTERN = re.compile(r"^model\.layers\.(?P\d+)\.self_attn\.(?P[qkv]_proj|[kv]_norm)\b") +_PER_LAYER_TABLE = "model.embed_tokens_per_layer.weight" _LAYER_FF_PREFIX_PATTERN = re.compile(r"^(model\.layers\.\d+)\.") _MERGE_RULES = { ".q_proj": MergeRule(".qkv_proj", "q", ("q", "k", "v")), @@ -180,6 +183,13 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: if isinstance(config.attention_group_for_layer(layer_id), FullAttentionGroupConfig) and config.attention_group_for_layer(layer_id).k_eq_v } + e_args = config.gemma4_e_args + + def kv_shared_leaf(name: str) -> str | None: + """The attention leaf of a KV-shared layer's key, or None for any other key.""" + m = _KV_SHARED_LEAF_PATTERN.match(name) if e_args is not None else None + return m["leaf"] if m is not None and e_args.is_kv_shared(int(m["layer"])) else None + merge_buf: dict[str, dict[str, torch.Tensor]] = {} gateup_buf: dict[str, dict[str, tuple]] = {} reader = ShardReader(model_path, device) @@ -229,8 +239,19 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: ) continue + # a KV-shared layer reuses another layer's k and v: the checkpoint still ships its own, which nothing reads + leaf = kv_shared_leaf(name) + if leaf not in (None, "q_proj"): + continue + + # the largest tensor in the file; read to host so --embed-weights host never touches VRAM for it + if name == _PER_LAYER_TABLE: + with safetensors.safe_open(file, framework="pt", device="cpu") as host_file: + yield name, host_file.get_tensor(raw_name) + continue + tensor = f.get_tensor(raw_name) - if is_vision or is_expert: + if is_vision or is_expert or leaf == "q_proj": yield name, tensor continue diff --git a/python/freetoken/models/gguf/tensors.py b/python/freetoken/models/gguf/tensors.py index f2cb0e0038..c1c2f2522f 100644 --- a/python/freetoken/models/gguf/tensors.py +++ b/python/freetoken/models/gguf/tensors.py @@ -29,6 +29,13 @@ def matrix(info: GgufTensorInfo) -> torch.Tensor: return packed_rows(info) if is_quantized(info.ggml_type) else float_tensor(info) +def table_matrix(info: GgufTensorInfo) -> torch.Tensor: + """``matrix`` for a table too large to copy first: float rows stay a read-only view of the mapped file, and the loader's copy to the table's final place is the only one.""" + if is_quantized(info.ggml_type): + return packed_rows(info) + return torch.from_numpy(info.data()).view(_FLOAT_DTYPES[info.ggml_type]).view(*reversed(info.ne)) + + def check_scales(model_path: str, *, nvfp4: bool = True, known: tuple[str, ...] = ()) -> None: """Refuse ``.scale`` / ``.input_scale`` tensors that nothing reads. @@ -78,7 +85,8 @@ def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k rows = [p.to(torch.bfloat16) for p in parts] if pad: rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=torch.bfloat16)) - yield f"{prefix}.weight", torch.cat(rows, dim=0) + # one part needs no copy: a BF16 per-layer table is 5.6 GB + yield f"{prefix}.weight", rows[0] if len(rows) == 1 else torch.cat(rows, dim=0) return segs = segments(tuple(types), tuple(p.shape[0] for p in parts)) for i, seg in enumerate(segs): @@ -97,4 +105,4 @@ def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k yield f"{prefix}.ggml_scale", scales -__all__ = ["check_scales", "float_tensor", "fused", "matrix", "packed_rows", "weight_scale"] +__all__ = ["check_scales", "float_tensor", "fused", "matrix", "packed_rows", "table_matrix", "weight_scale"] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 92722cec92..7b9db09ccb 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -2,7 +2,7 @@ import importlib from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable if TYPE_CHECKING: from .config import ModelConfig @@ -13,6 +13,8 @@ class EncoderSpec: kind: str # "vision" | "audio", the name --mm-disable takes config_key: str # the checkpoint config section that builds the tower modalities: tuple[str, ...] # inputs it serves, as implemented + # why a checkpoint's section builds no tower although the family registers one, or None; such a checkpoint serves the other modalities + unserved: Callable[[Any], str | None] | None = None @dataclass(frozen=True) @@ -84,7 +86,13 @@ class ModelSpec: _GEMMA4_PACKED = _DENSE_PACKED + _EXPERTS_PACKED _GEMMA4_PROCESSOR = "freetoken.mm.processors.gemma4:Gemma4MMProcessor" _GEMMA4_UNIFIED_PROCESSOR = "freetoken.mm.processors.gemma4:Gemma4UnifiedMMProcessor" -_GEMMA4_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),) + + +def _gemma4_vision_unserved(section: Any) -> str | None: + return "the tower does not apply clipped linears" if getattr(section, "use_clipped_linears", False) else None + + +_GEMMA4_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",), unserved=_gemma4_vision_unserved),) _MUSE_GLIMMER_PROCESSOR = "freetoken.mm.processors.muse_glimmer:MuseGlimmerMMProcessor" _MUSE_GLIMMER_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),) _MINIMAX_M3_PACKED = _DENSE_PACKED + ( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 3d75143e2e..3201b8de99 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -714,9 +714,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.embed_weights, choices=["gpu", "host"], help=( - "Where the input token embedding table lives. 'gpu' (default) keeps it in VRAM; 'host' keeps it in " - "pinned host RAM and each lookup reads its rows over PCIe, so its VRAM goes to the expert cache and KV. " - "Not for a table tied to the LM head." + "Where the input token embedding tables live. 'gpu' (default) keeps them in VRAM; 'host' keeps them in " + "pinned host RAM and each lookup reads its rows over PCIe, so their VRAM goes to the expert cache and KV. " + "A table tied to the LM head stays on the GPU; Gemma-4 E2B / E4B's per-layer table is the one that moves." ), ) diff --git a/tests/models/test_gemma4_e_series.py b/tests/models/test_gemma4_e_series.py new file mode 100644 index 0000000000..7156c54e07 --- /dev/null +++ b/tests/models/test_gemma4_e_series.py @@ -0,0 +1,319 @@ +"""Gemma-4 E2B / E4B: per-layer inputs, KV-shared layers and the double-wide MLP, against the shapes of the released checkpoints.""" +from __future__ import annotations + +import struct +from types import SimpleNamespace + +import numpy as np +import pytest +import safetensors.torch +import torch + +from freetoken.models.gemma4.config import parse_config +from freetoken.models.gemma4.gguf import parse_gguf_config, to_gguf + +# The layer types of the releases: every (period)th layer is full attention. +E2B = dict(layers=35, hidden=1536, ffn=6144, shared=20, double_wide=True, kv_heads=1, period=5) +E4B = dict(layers=42, hidden=2560, ffn=10240, shared=18, double_wide=False, kv_heads=2, period=6) + + +@pytest.fixture(autouse=True) +def _single_rank(monkeypatch): + from freetoken.distributed import info + + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=0, size=1)) + + +@pytest.fixture(autouse=True) +def _rope_tables(): + """Rope tables are cached per process and cannot live on the meta device these models build on.""" + from freetoken.layers import rotary + + device = rotary._ROPE_DEVICE + rotary.set_rope_device(torch.device("cuda" if torch.cuda.is_available() else "cpu")) + rotary.get_rope.cache_clear() + rotary._cos_sin_table.cache_clear() + yield + rotary.set_rope_device(device) + rotary.get_rope.cache_clear() + rotary._cos_sin_table.cache_clear() + + +def _is_full(release: dict, layer: int) -> bool: + return layer % release["period"] == release["period"] - 1 + + +def _hf_config(release: dict, **overrides): + n = release["layers"] + per_layer = { + "sliding_attention": SimpleNamespace(head_dim=256, num_key_value_heads=release["kv_heads"]), + "full_attention": SimpleNamespace(head_dim=512, num_key_value_heads=release["kv_heads"]), + } + fields = dict( + model_type="gemma4_text", + hidden_size=release["hidden"], + intermediate_size=release["ffn"], + num_hidden_layers=n, + num_attention_heads=8, + is_heterogeneous=True, + per_layer_config=per_layer, + vocab_size=262144, + rms_norm_eps=1e-6, + max_position_embeddings=131072, + sliding_window=512, + layer_types=["full_attention" if _is_full(release, i) else "sliding_attention" for i in range(n)], + rope_parameters={ + "sliding_attention": {"rope_theta": 10000.0}, + "full_attention": {"rope_theta": 1000000.0, "partial_rotary_factor": 0.25, "rope_type": "proportional"}, + }, + hidden_size_per_layer_input=256, + vocab_size_per_layer_input=262144, + num_kv_shared_layers=release["shared"], + use_double_wide_mlp=release["double_wide"], + tie_word_embeddings=True, + final_logit_softcapping=30.0, + ) + return SimpleNamespace(architectures=["Gemma4ForCausalLM"], **(fields | overrides)) + + +def _gguf_shim(tmp_path, release: dict, *, ffn): + from freetoken.models.gguf.config import GgufConfigShim + + n = release["layers"] + # no tensors and no KV of its own: the parser reads the shim's metadata + path = tmp_path / "metadata_only.gguf" + path.write_bytes(struct.pack("<4sIQQ", b"GGUF", 3, 0, 0)) + metadata = { + "gemma4.block_count": n, + "gemma4.embedding_length": release["hidden"], + "gemma4.embedding_length_per_layer_input": 256, + "gemma4.feed_forward_length": ffn, + "gemma4.context_length": 131072, + "gemma4.final_logit_softcapping": 30.0, + "gemma4.attention.head_count": 8, + "gemma4.attention.head_count_kv": release["kv_heads"], + "gemma4.attention.key_length": 512, + "gemma4.attention.key_length_swa": 256, + "gemma4.attention.value_length": 512, + "gemma4.attention.layer_norm_rms_epsilon": 1e-6, + "gemma4.attention.sliding_window": 512, + "gemma4.attention.shared_kv_layers": release["shared"], + "gemma4.attention.sliding_window_pattern": [not _is_full(release, i) for i in range(n)], + "gemma4.rope.freq_base": 1000000.0, + "gemma4.rope.freq_base_swa": 10000.0, + "gemma4.rope.dimension_count": 512, + "gemma4.rope.dimension_count_swa": 256, + } + return GgufConfigShim( + architectures=["Gemma4GGUFForCausalLM"], + model_path=str(path), + model_type="gemma4", + metadata=metadata, + vocab_size=262144, + tie_word_embeddings=True, + ) + + +def _meta_model(config): + from freetoken.models.gemma4.model import Gemma4ForCausalLM + from freetoken.utils.torch_utils import torch_dtype + + with torch.device("meta"), torch_dtype(torch.bfloat16): + return Gemma4ForCausalLM(config) + + +def test_e2b_config_doubles_the_shared_tail_only(): + config = parse_config(_hf_config(E2B)) + args = config.gemma4_e_args + + assert args.per_layer_hidden_size == 256 and args.first_kv_shared_layer == 15 + assert config.intermediate_size == 6144 + assert config.layer_intermediate_sizes == (6144,) * 15 + (12288,) * 20 + assert not args.is_kv_shared(14) and args.is_kv_shared(15) + + +def test_e4b_config_keeps_one_mlp_width(): + config = parse_config(_hf_config(E4B)) + + assert config.gemma4_e_args.first_kv_shared_layer == 24 + assert config.layer_intermediate_sizes == () + assert config.intermediate_size_for_layer(41) == 10240 + + +@pytest.mark.parametrize("fields", [ + {"hidden_size_per_layer_input": 0, "num_kv_shared_layers": 0, "use_double_wide_mlp": False}, # 26B-A4B, 31B + {"hidden_size_per_layer_input": None, "num_kv_shared_layers": None, "use_double_wide_mlp": None}, +]) +def test_other_releases_have_no_e_args(fields): + config = parse_config(_hf_config(E4B, **fields)) + + assert config.gemma4_e_args is None + assert config.layer_intermediate_sizes == () + + +def test_shared_layers_without_per_layer_inputs_are_refused(): + with pytest.raises(NotImplementedError, match="KV-shared"): + parse_config(_hf_config(E4B, hidden_size_per_layer_input=0)) + + +def test_kv_source_is_the_last_unshared_layer_of_each_attention_type(): + """E2B shares from layer 15: full attention hands over from 14, sliding from 13.""" + config = parse_config(_hf_config(E2B)) + full, swa = config.attention_groups + args = config.gemma4_e_args + + assert [i for i in full.layer_ids if args.is_kv_source(i, full.layer_ids)] == [14] + assert [i for i in swa.layer_ids if args.is_kv_source(i, swa.layer_ids)] == [13] + + +def test_gguf_reads_a_per_layer_mlp_width_list(tmp_path): + widths = [6144] * 15 + [12288] * 20 + config = parse_gguf_config(_gguf_shim(tmp_path, E2B, ffn=widths)) + + assert config.layer_intermediate_sizes == tuple(widths) + assert config.intermediate_size == 6144 + assert config.num_kv_heads == 1 + assert config.gemma4_e_args.first_kv_shared_layer == 15 + + +def test_gguf_reads_one_mlp_width_and_one_kv_head_count(tmp_path): + config = parse_gguf_config(_gguf_shim(tmp_path, E4B, ffn=10240)) + + assert config.layer_intermediate_sizes == () + assert config.intermediate_size == 10240 + assert config.num_kv_heads == 2 + assert config.gemma4_e_args.first_kv_shared_layer == 24 + + +def test_full_layers_store_v_on_e_series_gguf(tmp_path): + config = parse_gguf_config(_gguf_shim(tmp_path, E4B, ffn=10240)) + + assert config.attention_groups[0].k_eq_v is False + + +def test_to_gguf_names_the_per_layer_tensors(): + assert to_gguf("model.embed_tokens_per_layer", {}) == ("per_layer_token_embd",) + assert to_gguf("model.per_layer_model_projection", {}) == ("per_layer_model_proj",) + assert to_gguf("model.layers.3.per_layer_input_gate", {}) == ("blk.3.inp_gate",) + assert to_gguf("model.layers.3.per_layer_projection", {}) == ("blk.3.proj",) + + +def test_model_shapes_match_the_e2b_checkpoint(): + state = _meta_model(parse_config(_hf_config(E2B))).state_dict() + + assert state["model.embed_tokens_per_layer.weight"].shape == (262144, 35 * 256) + assert state["model.per_layer_model_projection.weight"].shape == (35 * 256, 1536) + assert state["model.layers.0.per_layer_input_gate.weight"].shape == (256, 1536) + assert state["model.layers.0.per_layer_projection.weight"].shape == (1536, 256) + assert state["model.layers.14.feed_forward.shared_mlp.gate_up_proj.weight"].shape == (2 * 6144, 1536) + assert state["model.layers.15.feed_forward.shared_mlp.gate_up_proj.weight"].shape == (2 * 12288, 1536) + assert state["model.layers.15.feed_forward.shared_mlp.down_proj.weight"].shape == (1536, 12288) + + +def test_shared_layers_hold_no_k_or_v_weights(): + state = _meta_model(parse_config(_hf_config(E2B))).state_dict() + + assert "model.layers.14.self_attn.qkv_proj.weight" in state + assert "model.layers.14.self_attn.k_norm.weight" in state + assert "model.layers.15.self_attn.qkv_proj.weight" not in state + assert "model.layers.15.self_attn.k_norm.weight" not in state + assert state["model.layers.15.self_attn.q_proj.weight"].shape == (8 * 256, 1536) + + +def test_other_releases_build_without_per_layer_weights(): + state = _meta_model(parse_config(_hf_config(E4B, hidden_size_per_layer_input=0, num_kv_shared_layers=0))).state_dict() + + assert not [k for k in state if "per_layer" in k] + assert "model.layers.41.self_attn.qkv_proj.weight" in state + + +def test_loader_drops_shared_k_v_but_not_the_vision_towers(tmp_path, monkeypatch): + """The checkpoint ships k, v and k_norm for shared layers too; the tower's layers reuse those names and indices.""" + from freetoken.models.gemma4 import weight + + release = dict(E2B, layers=3, shared=1, period=3) + prefix = "model.language_model.layers" + layers = { + f"{prefix}.{i}.self_attn.{leaf}.weight": torch.zeros(2, 2) + for i in range(3) + for leaf in ("q_proj", "k_proj", "v_proj", "k_norm") + } + tensors = { + **layers, + "model.language_model.embed_tokens_per_layer.weight": torch.zeros(2, 2), + "model.vision_tower.encoder.layers.2.self_attn.k_proj.linear.weight": torch.zeros(2, 2), + } + safetensors.torch.save_file(tensors, tmp_path / "model.safetensors") + monkeypatch.setattr(weight, "cached_load_hf_config", lambda _: _hf_config(release)) + + names = dict(weight.iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True)) + + assert "model.layers.2.self_attn.q_proj.weight" in names + assert not [k for k in names if k.startswith("model.layers.2.self_attn.") and "q_proj" not in k] + assert "model.layers.1.self_attn.qkv_proj.weight" in names + assert "vision_tower.encoder.layers.2.self_attn.k_proj.weight" in names + assert names["model.embed_tokens_per_layer.weight"].device.type == "cpu" + + +def test_embed_weights_host_moves_only_the_untied_table(): + from freetoken.engine.engine import _host_embedding_keys + + config = parse_config(_hf_config(E2B)) + model = _meta_model(config) + + assert _host_embedding_keys(model, SimpleNamespace(embed_weights="host", model_config=config)) == { + "model.embed_tokens_per_layer.weight" + } + assert _host_embedding_keys(model, SimpleNamespace(embed_weights="gpu", model_config=config)) == frozenset() + + +def test_embed_weights_host_is_refused_when_the_only_table_is_tied(): + from freetoken.engine.engine import _host_embedding_keys + + config = parse_config(_hf_config(E4B, hidden_size_per_layer_input=0, num_kv_shared_layers=0)) + + with pytest.raises(ValueError, match="tied to its LM head"): + _host_embedding_keys(_meta_model(config), SimpleNamespace(embed_weights="host", model_config=config)) + + +def test_a_float_table_is_read_without_a_copy(): + """A BF16 per-layer table is 5.6 GB: a copy before the loader's own would double the load's peak RAM.""" + from freetoken.models.gguf.tensors import fused, table_matrix + + data = np.arange(24, dtype=np.uint8) # 3 rows of 4 BF16 values, as ggml dims [4, 3] + info = SimpleNamespace(ggml_type=30, ne=(4, 3), data=lambda: data) + + table = table_matrix(info) + assert table.shape == (3, 4) and table.dtype == torch.bfloat16 + assert np.shares_memory(table.view(torch.uint8).numpy(), data) + ((name, weight),) = fused("t", [info], [table], 4) + assert name == "t.weight" and weight is table + + +def test_clipped_vision_towers_serve_text_only(): + """E2B / E4B's tower clips its linears, which the tower does not apply: text must still serve.""" + from freetoken.models.register import get_model_spec + + (encoder,) = get_model_spec("Gemma4ForConditionalGeneration").encoders + + assert "clipped linears" in encoder.unserved(SimpleNamespace(use_clipped_linears=True)) + assert encoder.unserved(SimpleNamespace(use_clipped_linears=False)) is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the rope kernels need CUDA") +def test_rotating_q_alone_matches_rotating_it_with_a_real_key(): + """A shared layer has no key to rotate, so it hands the kernel one head of zeros.""" + from freetoken.layers.rotary import get_rope + + with torch.device("cuda"): + rope = get_rope(head_dim=256, rotary_dim=256, max_position=512, base=10000.0, rope_scaling=None) + positions = torch.arange(16, device="cuda") + q = torch.randn(16, 8 * 256, device="cuda", dtype=torch.bfloat16) + with_key, alone = q.clone(), q.clone() + + rope.forward(positions, with_key, torch.randn(16, 256, device="cuda", dtype=torch.bfloat16)) + zeros = torch.zeros(16, 256, device="cuda", dtype=torch.bfloat16) + rope.forward(positions, alone, zeros) + + assert torch.equal(with_key, alone) + assert not zeros.any()