diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index ba990c007..f0ff81460 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -6,12 +6,15 @@ import safetensors import torch from freetoken.distributed import get_tp_info +from freetoken.layers.quantization import QuantKind, get_quant_config from freetoken.models.config import FullAttentionGroupConfig from freetoken.models.loader import ( MergeRule, ShardReader, + ct_nvfp4_fuse, drop_page_cache, iter_weight_files, + nvfp4_parts_ct, ) from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, @@ -39,6 +42,15 @@ layer_to_bank=lambda layer, config: layer, # every layer is MoE (no dense prefix) desc="Gemma4 NVFP4 experts", ) +_NVFP4_EXPERT_KEY_TEMPLATE = ( + r"^model\.language_model\.layers\.(?P\d+)\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\.(?P{kinds})$" +) +_NVFP4_BANK_KINDS = { + "weight": "weight", + "weight_scale": "weight_scale", + "weight_global": "weight_scale_2", +} _LAYER_INDEX_PATTERN = re.compile(r"layers\.(\d+)\.") _LAYER_FF_PREFIX_PATTERN = re.compile(r"^(model\.layers\.\d+)\.") _MERGE_RULES = { @@ -63,6 +75,17 @@ # The scales are consumed with their .weight. _NVFP4_DENSE_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") _NVFP4_DENSE_MLP_RE = re.compile(r"\.mlp\.(gate_proj|up_proj|down_proj)\.weight$") +_CT_NVFP4_FUSIONS = { + ".feed_forward.shared_mlp.gate_up_proj": ( + ".feed_forward.shared_mlp.gate_proj", + ".feed_forward.shared_mlp.up_proj", + ), +} +_CT_NVFP4_SCALE_SUFFIXES = ( + ".weight_scale", + ".weight_global_scale", + ".input_global_scale", +) def _nvfp4_dense_parts(reader: ShardReader, raw_base: str): @@ -156,6 +179,10 @@ def iter_weights( include_vision: bool = True, ) -> Iterator[tuple[str, torch.Tensor]]: def rename_key(raw_name: str) -> str | None: + # Static KV quantizer metadata is not part of FreeToken's runtime KV + # cache state (same rule used by the generic compressed-tensors reader). + if raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): + return None prefix = "model.language_model." if raw_name.startswith(prefix): return _rename_language_key(raw_name) @@ -165,8 +192,11 @@ def rename_key(raw_name: str) -> str | None: def merge_info(key: str) -> tuple[str, MergeRule] | None: for suffix, rule in _MERGE_RULES.items(): - if key.endswith(suffix + ".weight") or key.endswith(suffix): - return key.replace(suffix, rule.fused_suffix), rule + # Quantized tensors carry roles after the projection name, e.g. + # ``q_proj.weight_scale``. Every role must follow the same fusion as + # its weight or the model asks for a fused tensor that was never emitted. + if key.endswith(suffix) or suffix + "." in key: + return key.replace(suffix, rule.fused_suffix, 1), rule return None config = parse_config(cached_load_hf_config(model_path)) @@ -200,8 +230,19 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: if _NVFP4_EXPERT_RE.search(raw_name): continue - # NVFP4 dense-MLP scales are consumed with their .weight (below), never yielded. - if raw_name.endswith(_NVFP4_DENSE_SCALE_SUFFIXES): + # NVFP4 scales are consumed with their packed weight. Do not + # blanket-drop ``weight_scale``: FP8 attention projections use + # that same suffix and must be fused q/k/v-wise. + ct_scale_base = next( + (raw_name[: -len(s)] for s in _CT_NVFP4_SCALE_SUFFIXES if raw_name.endswith(s)), + None, + ) + if ct_scale_base is not None and reader.has(ct_scale_base + ".weight_packed"): + continue + if ( + config.dense_quant == "nvfp4" + and raw_name.endswith(_NVFP4_DENSE_SCALE_SUFFIXES) + ): continue is_vision = name.startswith(("vision_tower.", "embed_vision.", "vision_embedder.")) @@ -215,6 +256,25 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: if not is_expert and not include_non_moe: continue + # llm-compressor / compressed-tensors NVFP4. The stored names + # differ from ModelOpt but map to the same FreeToken buffers. + if raw_name.endswith(".weight_packed"): + raw_base = raw_name[: -len(".weight_packed")] + base = name[: -len(".weight_packed")] + parts = nvfp4_parts_ct(reader, raw_base) + emitted = ct_nvfp4_fuse(base, parts, gateup_buf, _CT_NVFP4_FUSIONS) + if emitted is None: + w, s, g, a = parts + emitted = [ + (base + ".weight", w), + (base + ".weight_scale", s), + (base + ".weight_global", g), + ] + if a is not None: + emitted.append((base + ".input_scale", a)) + yield from emitted + continue + # Native W4A16 NVFP4 dense MLP: the .weight is FP4-packed and carries block + per-tensor scales. # The weight_scale_2 sibling guard is defense-in-depth beyond config.dense_quant -- the sibling MoE checkpoint's bf16 shared_mlp has no such sibling, so it falls through to the bf16 path. if ( @@ -230,6 +290,15 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: continue tensor = f.get_tensor(raw_name) + # compressed-tensors stores FP8 per-channel scales as + # ``[out_features, 1]`` (often bf16); FreeToken's FP8 linear + # buffer is a flat fp32 vector. + if ( + raw_name.endswith(".weight_scale") + and tensor.ndim == 2 + and tensor.shape[1] == 1 + ): + tensor = tensor.reshape(-1).to(torch.float32) if is_vision or is_expert: yield name, tensor continue @@ -296,7 +365,23 @@ def _expert_name(raw_name: str) -> str | None: def nvfp4_expert_spec(model_path: str, config): - return _NVFP4_SOURCE_SPEC + quant = get_quant_config() + stored = quant.stored_tensors(QuantKind.NVFP4) + kind_map = {stored[role].name: kind for role, kind in _NVFP4_BANK_KINDS.items()} + # Preserve the original ModelOpt descriptor exactly; compressed-tensors uses + # weight_packed / weight_global_scale and reciprocal global scales. + if set(kind_map) == {"weight", "weight_scale", "weight_scale_2"}: + return _NVFP4_SOURCE_SPEC + return Nvfp4ExpertSourceSpec( + key_pattern=re.compile( + _NVFP4_EXPERT_KEY_TEMPLATE.format(kinds="|".join(map(re.escape, kind_map))) + ), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, + desc=f"Gemma4 NVFP4 experts ({quant.dialect})", + kind_map=kind_map, + global_reciprocal=stored["weight_global"].reciprocal, + ) __all__ = [ diff --git a/tests/models/test_gemma4_compressed_tensors.py b/tests/models/test_gemma4_compressed_tensors.py new file mode 100644 index 000000000..c7949ea12 --- /dev/null +++ b/tests/models/test_gemma4_compressed_tensors.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.layers.quantization import set_quant_config +from freetoken.layers.quantization.configs.compressed_tensors import CompressedTensorsConfig +from freetoken.layers.quantization.configs.modelopt import ModelOptConfig +from freetoken.models.gemma4.weight import iter_weights, nvfp4_expert_spec + + +@pytest.fixture(scope="session", autouse=True) +def _tp_info(): + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +def _nvfp4(base: str, rows: int, cols: int, global_scale: float = 2.0): + return { + base + ".weight_packed": torch.randint( + 0, 256, (rows, cols // 2), dtype=torch.uint8 + ), + base + ".weight_scale": torch.ones( + rows, cols // 16, dtype=torch.float8_e4m3fn + ), + base + ".weight_global_scale": torch.tensor([global_scale]), + base + ".input_global_scale": torch.tensor([2.0]), + } + + +def test_mixed_compressed_tensors_dense_weights(tmp_path, monkeypatch): + """Gemma 4 Unsloth exports mix FP8 attention with compressed-tensors + NVFP4 shared MLPs; both roles must land in the buffers the model builds.""" + h, q, kv, intermediate = 16, 8, 4, 32 + layer = "model.language_model.layers.0" + attn = layer + ".self_attn" + mlp = layer + ".mlp" + raw = { + attn + ".q_proj.weight": torch.ones(q, h, dtype=torch.float8_e4m3fn), + attn + ".q_proj.weight_scale": torch.full((q, 1), 2.0, dtype=torch.bfloat16), + attn + ".k_proj.weight": torch.ones(kv, h, dtype=torch.float8_e4m3fn), + attn + ".k_proj.weight_scale": torch.full((kv, 1), 3.0, dtype=torch.bfloat16), + attn + ".v_proj.weight": torch.ones(kv, h, dtype=torch.float8_e4m3fn), + attn + ".v_proj.weight_scale": torch.full((kv, 1), 4.0, dtype=torch.bfloat16), + attn + ".o_proj.weight": torch.ones(h, q, dtype=torch.float8_e4m3fn), + attn + ".o_proj.weight_scale": torch.full((h, 1), 5.0, dtype=torch.bfloat16), + attn + ".k_scale": torch.tensor(1.0), + attn + ".v_scale": torch.tensor(1.0), + } + raw |= _nvfp4(mlp + ".gate_proj", intermediate, h) + raw |= _nvfp4(mlp + ".up_proj", intermediate, h) + raw |= _nvfp4(mlp + ".down_proj", h, intermediate, global_scale=4.0) + save_file(raw, tmp_path / "model.safetensors") + + import freetoken.models.gemma4.weight as weight + + config = SimpleNamespace( + num_layers=1, + dense_quant="none", + attention_group_for_layer=lambda _layer: None, + ) + monkeypatch.setattr(weight, "cached_load_hf_config", lambda _path: object()) + monkeypatch.setattr(weight, "parse_config", lambda _hf: config) + + loaded = dict( + iter_weights( + str(tmp_path), + torch.device("cpu"), + include_moe_experts=False, + include_non_moe=True, + include_vision=False, + ) + ) + + qkv = "model.layers.0.self_attn.qkv_proj" + assert loaded[qkv + ".weight"].shape == (q + kv + kv, h) + assert loaded[qkv + ".weight_scale"].shape == (q + kv + kv,) + assert loaded[qkv + ".weight_scale"].dtype is torch.float32 + assert torch.equal( + loaded[qkv + ".weight_scale"], + torch.tensor([2.0] * q + [3.0] * kv + [4.0] * kv), + ) + assert loaded["model.layers.0.self_attn.o_proj.weight_scale"].shape == (h,) + assert not any(name.endswith((".k_scale", ".v_scale")) for name in loaded) + + gate_up = "model.layers.0.feed_forward.shared_mlp.gate_up_proj" + assert loaded[gate_up + ".weight"].shape == (2 * intermediate, h // 2) + assert loaded[gate_up + ".weight_scale"].shape == (2 * intermediate, h // 16) + assert loaded[gate_up + ".weight_global"].shape == (2 * intermediate,) + assert loaded[gate_up + ".weight_global"][0].item() == pytest.approx(0.5) + assert loaded[gate_up + ".input_scale"].item() == pytest.approx(0.5) + + down = "model.layers.0.feed_forward.shared_mlp.down_proj" + assert loaded[down + ".weight"].shape == (h, intermediate // 2) + assert loaded[down + ".weight_scale"].shape == (h, intermediate // 16) + assert loaded[down + ".weight_global"][0].item() == pytest.approx(0.25) + assert not any("weight_packed" in name or "global_scale" in name for name in loaded) + + +def test_nvfp4_expert_spec_uses_checkpoint_dialect(): + ct = CompressedTensorsConfig( + { + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": { + "num_bits": 4, + "type": "float", + "strategy": "tensor_group", + "group_size": 16, + }, + } + } + } + ) + set_quant_config(ct) + spec = nvfp4_expert_spec("unused", object()) + match = spec.key_pattern.match( + "model.language_model.layers.2.experts.7.gate_proj.weight_packed" + ) + assert match and match.group("kind") == "weight_packed" + assert spec.kind_map == { + "weight_packed": "weight", + "weight_scale": "weight_scale", + "weight_global_scale": "weight_scale_2", + } + assert spec.global_reciprocal + assert spec.key_pattern.match( + "model.language_model.layers.2.experts.7.gate_proj.input_global_scale" + ) is None + + set_quant_config(ModelOptConfig({"quant_algo": "NVFP4"})) + modelopt = nvfp4_expert_spec("unused", object()) + assert modelopt.kind_map is None + assert not modelopt.global_reciprocal + assert modelopt.key_pattern.match( + "model.language_model.layers.2.experts.7.gate_proj.weight" + )