Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions python/freetoken/models/gemma4/gguf.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,14 @@ def g(key: str):
full_kv = int(kv_per_layer[full_layer_ids[0]]) if full_layer_ids else int(kv_per_layer[0])

max_pos = int(g("context_length"))
# Dense Gemma-4 checkpoints (gemma-4-12B-it, gemma-4-31B-it) carry no expert fields
# in their GGUF metadata; only MoE members (e.g. Gemma-4-26B-A4B) do. Mirror the HF
# side (gemma4.config.parse_config): default the expert fields to 0 and derive
# moe_enabled from them instead of assuming every GGUF is MoE.
num_experts = int(m.get("gemma4.expert_count", 0) or 0)
num_experts_per_tok = int(m.get("gemma4.expert_used_count", 0) or 0)
moe_intermediate_size = int(m.get("gemma4.expert_feed_forward_length", 0) or 0)
moe_enabled = num_experts > 0
full_rotary = RotaryConfig(
head_dim=full_head_dim,
rotary_dim=_full_rotary_dim(shim, full_head_dim),
Expand Down Expand Up @@ -105,15 +113,16 @@ def g(key: str):
rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")),
tie_word_embeddings=bool(shim.tie_word_embeddings),
rotary_config=full_rotary,
num_experts=int(g("expert_count")),
num_experts_per_tok=int(g("expert_used_count")),
moe_intermediate_size=int(g("expert_feed_forward_length")),
num_experts=num_experts,
num_experts_per_tok=num_experts_per_tok,
moe_intermediate_size=moe_intermediate_size,
norm_topk_prob=True,
model_type="gemma4",
architectures=list(shim.architectures),
moe_enabled=True,
expert_quant="q4_0",
moe_weight_format="q4_0",
moe_enabled=moe_enabled,
# Native-Q4_0 offload-cache path for the routed experts (MoE checkpoints only).
expert_quant="q4_0" if moe_enabled else "none",
moe_weight_format="q4_0" if moe_enabled else "none",
use_qk_norm=True,
attn_sm_scale=1.0,
final_logit_softcapping=float(g("final_logit_softcapping")),
Expand Down
89 changes: 89 additions & 0 deletions tests/models/test_gemma4_gguf_config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
"""parse_gguf_config must handle dense Gemma-4 GGUFs (issue #357).

Dense checkpoints (gemma-4-12B-it, gemma-4-31B-it) carry no gemma4.expert_* keys
in their GGUF metadata. The old parser unconditionally required expert_count and
hardcoded moe_enabled=True, so dense GGUFs could not load at all.
"""

import struct

import pytest

from freetoken.models.gguf.config import GgufConfigShim
from freetoken.models.gemma4.gguf import parse_gguf_config


def make_shim(metadata_overrides: dict, tmp_path) -> GgufConfigShim:
"""A metadata-only shim over an empty file; _full_rotary_dim falls back to head_dim//4."""
empty = tmp_path / "meta_only.gguf"
# Minimal GGUF: magic + version 3 + kv_count=0 + tensor_count=0. GGUFReader opens
# it fine with no fields/tensors, so _full_rotary_dim takes its metadata-only
# fallback (head_dim//4) without needing a real checkpoint.
empty.write_bytes(b"GGUF" + struct.pack("<IQQ", 3, 0, 0))
base = {
"gemma4.block_count": 4,
"gemma4.embedding_length": 256,
"gemma4.attention.head_count": 4,
"gemma4.attention.head_count_kv": [2, 2, 2, 2],
# One SWA layer + three full layers.
"gemma4.attention.sliding_window_pattern": [True, False, False, False],
"gemma4.attention.key_length_swa": 128,
"gemma4.attention.key_length": 256,
"gemma4.attention.sliding_window": 512,
"gemma4.context_length": 1024,
"gemma4.rope.freq_base": 1000000.0,
"gemma4.rope.freq_base_swa": 1000000.0,
"gemma4.rope.dimension_count_swa": 128,
"gemma4.feed_forward_length": 512,
"gemma4.attention.layer_norm_rms_epsilon": 1e-6,
"gemma4.final_logit_softcapping": 30.0,
}
base.update(metadata_overrides)
return GgufConfigShim(
architectures=["Gemma4GGUFForCausalLM"],
model_path=str(empty),
model_type="gemma4",
metadata=base,
vocab_size=262144,
tie_word_embeddings=True,
)


def test_moe_gguf_keeps_moe_path(tmp_path):
"""A MoE GGUF (has expert_* keys) keeps moe_enabled + q4_0 expert quant."""
shim = make_shim({
"gemma4.expert_count": 128,
"gemma4.expert_used_count": 8,
"gemma4.expert_feed_forward_length": 1024,
}, tmp_path)
cfg = parse_gguf_config(shim)
assert cfg.moe_enabled is True
assert cfg.num_experts == 128
assert cfg.num_experts_per_tok == 8
assert cfg.moe_intermediate_size == 1024
assert cfg.expert_quant == "q4_0"


def test_dense_gguf_loads_without_expert_keys(tmp_path):
"""A dense GGUF (no expert_* keys) must parse and route to the dense path."""
cfg = parse_gguf_config(make_shim({}, tmp_path))
assert cfg.moe_enabled is False
assert cfg.num_experts == 0
assert cfg.num_experts_per_tok == 0
assert cfg.moe_intermediate_size == 0
assert cfg.expert_quant == "none"
# Sanity: the rest of the geometry still parses.
assert cfg.num_layers == 4
assert cfg.model_type == "gemma4"


def test_explicit_zero_experts_is_dense(tmp_path):
"""expert_count=0 in metadata must behave like an absent key."""
shim = make_shim({
"gemma4.expert_count": 0,
"gemma4.expert_used_count": 0,
"gemma4.expert_feed_forward_length": 0,
}, tmp_path)
cfg = parse_gguf_config(shim)
assert cfg.moe_enabled is False
assert cfg.expert_quant == "none"