diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 2ce7462b8..4a2a5b615 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -2,7 +2,7 @@ Three separate paths, because the checkpoint's three weight classes live in different places: -* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_FUSIONS``. +* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseFuser``. * :func:`load_ple_table` -- the 47.7 GiB FP8 n-gram table, 128 checkpoint shards concatenated into one pinned :class:`HostBank`. * :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named, for the offload cache's expert reader. @@ -25,8 +25,10 @@ from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) +from freetoken.layers.quantization import get_quant_config +from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into -from freetoken.utils import download_hf_weight +from freetoken.utils import cached_load_hf_config, download_hf_weight from freetoken.utils.progress import byte_bar from tqdm import tqdm @@ -70,32 +72,13 @@ ".self_attn.indexer.k_layernorm.weight", ) -# Fused projections: concat the checkpoint parts along dim 0 in this exact order. A nonzero pad -# rounds the merged row count up; the model splits the result back with the same sizes. -_FUSIONS: dict[str, tuple[tuple[str, ...], int]] = { - # q carries the output gate, so its half is twice the attention width: [2*qo | kv | kv]. - ".self_attn.qkv_proj.weight": (( - ".self_attn.q_proj.weight", ".self_attn.k_proj.weight", ".self_attn.v_proj.weight", - ), 0), - ".linear_attn.in_proj.weight": (( - ".linear_attn.in_proj_qkv.weight", ".linear_attn.in_proj_z.weight", - ".linear_attn.in_proj_b.weight", ".linear_attn.in_proj_a.weight", - ), 0), - ".mlp.shared_expert.gate_up_proj.weight": (( - ".mlp.shared_expert.gate_proj.weight", ".mlp.shared_expert.up_proj.weight", - ), 0), - # HC mix reads the low-rank down projection and the injection logits from one GEMM; vLLM - # pads the merged output to a multiple of 16 rows for cuBLAS (hyperconnection.py pad_size). - # The top-level hyper_connection_mixer has no injection and so never fuses. - ".attn_hyper_connection.input_mix_weight_down_block_inject.weight": (( - ".attn_hyper_connection.input_mix_weight_down.weight", - ".attn_hyper_connection.block_inject_weight.weight", - ), 16), - ".mlp_hyper_connection.input_mix_weight_down_block_inject.weight": (( - ".mlp_hyper_connection.input_mix_weight_down.weight", - ".mlp_hyper_connection.block_inject_weight.weight", - ), 16), -} +# The per-layer HC mix reads the low-rank down projection and the injection logits from one GEMM; vLLM pads the merged rows to a multiple of 16 for cuBLAS (hyperconnection.py pad_size). +# The top-level hyper_connection_mixer has no injection and never fuses. +_PAD_TO = {"input_mix_weight_down_block_inject": 16} +_HC_WITH_INJECT = (".attn_hyper_connection", ".mlp_hyper_connection") +_KIND_SUFFIXES = (".weight_scale_inv", ".weight") +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) +_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn} def _rename(raw_name: str) -> str | None: @@ -115,26 +98,95 @@ def _rename(raw_name: str) -> str | None: return raw_name -def _try_fuse( - name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]] -) -> tuple[str, torch.Tensor] | tuple[()] | None: - """Buffer a fusion part; return the merged ``(name, tensor)`` once all parts arrive, ``()`` while incomplete, ``None`` if ``name`` is not a fusion part.""" - for fused_suffix, (parts, pad_to) in _FUSIONS.items(): - for idx, part in enumerate(parts): - if not name.endswith(part): - continue - key = name[: -len(part)] + fused_suffix - slots = buf.setdefault(key, {}) - slots[idx] = tensor - if len(slots) < len(parts): - return () - del buf[key] - rows = [slots[i] for i in range(len(parts))] - pad = (-sum(t.shape[0] for t in rows)) % pad_to if pad_to else 0 - if pad: - rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=rows[0].dtype, device=rows[0].device)) - return key, torch.cat(rows, dim=0) - return None +def _split_kind(name: str) -> tuple[str, str]: + """``name`` -> ``(module, kind)``; kind is "" for tensors that are neither a weight nor a block scale.""" + for suffix in _KIND_SUFFIXES: + if name.endswith(suffix): + return name[: -len(suffix)], suffix + return name, "" + + +class _DenseFuser: + """Concatenates checkpoint projection parts into the model's merged buffers, per kind (weight / block scale). + + The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and validates each part against the scheme the model built its buffer from. + """ + + def __init__(self, quant, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None: + self.quant = quant + self.groups = {fused: parts for fused, parts in packed if fused != "experts"} # experts: bank reader + self.by_part: dict[str, list[tuple[str, int]]] = {} + for fused, parts in self.groups.items(): + for idx, part in enumerate(parts): + self.by_part.setdefault(part, []).append((fused, idx)) + self.buf: dict[tuple[str, str], dict[int, torch.Tensor]] = {} + + def scheme(self, module: str): + return None if self.quant is None else self.quant.scheme_for(module) + + def _target(self, parent: str, leaf: str) -> tuple[str, int] | None: + candidates = self.by_part.get(leaf) + if not candidates: + return None + if len(candidates) > 1: + # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as gdn.py + split = self.scheme(f"{parent}.in_proj_qkvz") is not None + keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"} + candidates = [c for c in candidates if c[0] in keep] + if not candidates: + raise ValueError(f"{parent}.{leaf}: no merged projection for the {'split' if split else 'fused'} GDN layout") + fused, idx = candidates[0] + if fused in _PAD_TO and not parent.endswith(_HC_WITH_INJECT): + return None + return f"{parent}.{fused}", idx + + def check(self, module: str, name: str, tensor: torch.Tensor) -> None: + """``tensor`` (checkpoint key ``name``) must match the scheme the model built ``module`` from.""" + scheme = self.scheme(module) + if name.endswith(".weight_scale_inv"): + if scheme is None or not scheme.has("weight_scale_inv"): + raise ValueError(f"{name}: {module} has no block scale in the checkpoint's quant config ({scheme})") + return + is_fp8 = tensor.dtype in _FP8_DTYPES + if scheme is None: + if is_fp8: + raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized") + return + expected = _ELEM_DTYPES.get(scheme.weight.elem) + if expected is not None and tensor.dtype is not expected: + raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} {scheme}") + rows, cols = (scheme.weight.group or (1, 1)) + if rows > 1 and tensor.shape[0] % rows or cols > 1 and tensor.shape[1] % cols: + raise ValueError(f"{name}: {tuple(tensor.shape)} is not a multiple of the {rows}x{cols} scale block of {module}") + + def check_unfused(self, name: str, tensor: torch.Tensor) -> None: + module, kind = _split_kind(name) + if kind == ".weight_scale_inv" or (kind == ".weight" and tensor.dtype in _FP8_DTYPES): + self.check(module, name, tensor) + + def fuse(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None: + """Buffer a part; return the merged ``[(name, tensor)]`` once its kind is complete, ``[]`` while incomplete, ``None`` if ``name`` is not a part.""" + module, kind = _split_kind(name) + if not kind: + return None + parent, _, leaf = module.rpartition(".") + hit = self._target(parent, leaf) + if hit is None: + return None + fused, idx = hit + self.check(fused, name, tensor) + slots = self.buf.setdefault((fused, kind), {}) + slots[idx] = tensor + parts = self.groups[fused.rpartition(".")[2]] + if len(slots) < len(parts): + return [] + del self.buf[(fused, kind)] + rows = [slots[i] for i in range(len(parts))] + pad_to = _PAD_TO.get(fused.rpartition(".")[2], 0) if kind == ".weight" else 0 + pad = (-sum(t.shape[0] for t in rows)) % pad_to if pad_to else 0 + if pad: + rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=rows[0].dtype, device=rows[0].device)) + return [(fused + kind, torch.cat(rows, dim=0))] def iter_weights( @@ -146,24 +198,19 @@ def iter_weights( ) -> Iterator[tuple[str, torch.Tensor]]: """Yield the dense (non-expert) weights, prefix-stripped and fused to the model's buffers. - Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the - model's state dict minus the routed experts. Nothing here is quantized: every release's skip - list (modelopt ``ignore``, fp8 ``modules_to_not_convert``) covers everything except those experts, - so attention, GDN, HC, PLE, the shared expert and lm_head are all plain bf16 (the n-gram hash - constants stay int64). Fusions: - attention q|k|v -> ``qkv_proj``, GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, shared-expert - gate|up -> ``gate_up_proj``, and each per-layer HC's ``input_mix_weight_down`` | - ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``. - - ``include_moe_experts`` is accepted for the loader contract but never yields anything: the - routed experts are NVFP4 and always come from the offload cache's expert reader. + Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the model's state dict minus the routed experts. + A dense projection is bf16 or 128x128 block-fp8 (``.weight`` e4m3 + ``.weight_scale_inv``) as the checkpoint's QuantConfig says: the official releases skip everything but the routed experts, the community NVFP4-FP8 requants quantize the attention / GDN projections. + Fusions, per kind: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + bf16 ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``. + ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from the offload cache's expert reader. """ if get_tp_info().size > 1: raise NotImplementedError("qwen4_exp weight loading supports TP=1 only") if not include_non_moe: return - fuse_buf: dict[str, dict[int, torch.Tensor]] = {} + hf_config = cached_load_hf_config(model_path) + spec = get_model_spec(hf_config.architectures[0]) + fuser = _DenseFuser(get_quant_config(), spec.packed_modules_mapping) for file in tqdm( iter_weight_files(model_path), desc="Loading weights", @@ -175,14 +222,14 @@ def iter_weights( if name is None: continue tensor = f.get_tensor(raw_name) - fused = _try_fuse(name, tensor, fuse_buf) - if fused is not None: - if fused != (): # () means buffered, not yet complete - yield fused - continue - yield name, tensor - - assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}" + fused = fuser.fuse(name, tensor) + if fused is None: + fuser.check_unfused(name, tensor) + yield name, tensor + else: + yield from fused + + assert not fuser.buf, f"Incomplete projection fusions: {sorted(k[0] + k[1] for k in fuser.buf)}" # ====================================================================================== diff --git a/tests/models/qwen4_exp/common.py b/tests/models/qwen4_exp/common.py index 1f9c117bd..aa6f23995 100644 --- a/tests/models/qwen4_exp/common.py +++ b/tests/models/qwen4_exp/common.py @@ -255,3 +255,118 @@ def spy(self, index, md, slot): monkeypatch.setattr(QSASparseAttnBackend, "_select", spy) return seen + + +LM = "model.language_model" + +# quantization_config of each released checkpoint, trimmed to the entries parse_config and the reader look at + +# RadixArk/Qwen3.8-Flash-Next-NVFP4: modelopt NVFP4 everywhere except the ignore list +RADIXARK_NVFP4 = { + "quant_algo": "NVFP4", + "quant_method": "modelopt", + "ignore": [ + "model.embed_tokens", + "mtp.*", + "model.mtp.*", + "*.self_attn.*", + "*.linear_attn.*", + "*.mlp.gate*", + "*.mlp.shared_expert.*", + "*.mlp.shared_expert_gate*", + "*hyper_connection*", + "*.ple.*", + "model.visual.*", + "model.language_model.embed_tokens", + "lm_head", + ], +} + +# nvidia/Qwen3.8-Flash-Next-NVFP4: modelopt MIXED_PRECISION, the per-module algo sits in quantized_layers +NVIDIA_NVFP4 = { + "quant_algo": "MIXED_PRECISION", + "quant_method": "modelopt", + "quantized_layers": { + **{f"model.language_model.layers.{i}.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16} for i in range(48)}, + "model.language_model.layers.1.ple.ple_embedding.ngram_embedding": {"quant_algo": "FP8"}, + "mtp.layers.0.mlp.experts": {"quant_algo": "FP8_PB_WO", "group_size": 128}, + }, + "ignore": ["lm_head", "model.language_model.embed_tokens", "model.language_model.layers.0.mlp.shared_expert*", "model.visual*"], +} + +# Qwen/Qwen3.8-Flash-Next-FP8: 128x128 block-fp8 experts, everything else listed in modules_to_not_convert +QWEN_FP8 = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_per_tensor": False, + "act_per_tensor": False, + "weight_block_size": [128, 128], + "modules_to_not_convert": [ + "lm_head", + "model.language_model.embed_tokens", + "model.language_model.hyper_connection_mixer.input_mix_weight_up", + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.mlp.gate", + "model.language_model.layers.3.mlp.shared_expert.gate_proj", + ], + "modules_to_convert": ["ple.ple_embedding.ngram_embedding"], +} + + +def mixed_precision_quant(gdn_layers, attn_layers, moe_layers) -> dict: + """modelopt MIXED_PRECISION with NVFP4 routed experts and FP8_PB_WO attention / GDN projections, ignore list as in lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8.""" + return { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + **{f"{LM}.layers.{i}.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16} for i in moe_layers}, + **{f"{LM}.layers.{i}.linear_attn.{p}": {"quant_algo": "FP8_PB_WO", "group_size": 128} + for i in gdn_layers for p in ("in_proj_qkv", "in_proj_z", "out_proj")}, + **{f"{LM}.layers.{i}.self_attn.{p}_proj": {"quant_algo": "FP8_PB_WO", "group_size": 128} + for i in attn_layers for p in "qkvo"}, + }, + "ignore": [ + "model.embed_tokens", "mtp.*", "model.mtp.*", "*.mlp.gate*", "*.mlp.shared_expert.*", + "*.mlp.shared_expert_gate*", "*hyper_connection*", "*.ple.*", "model.visual.*", + "model.language_model.embed_tokens", "lm_head", "*.self_attn.indexer*", + ], + } + + +# lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8, trimmed to layers 0 (GDN) and 3 (attention) +LOVEDHEART_NVFP4_FP8 = mixed_precision_quant(gdn_layers=(0,), attn_layers=(3,), moe_layers=(0, 3)) + + +def install_quant_config(model_path: str) -> None: + """Install ``model_path``'s QuantConfig process-wide, as EngineConfig does before the reader runs.""" + from freetoken.layers.quantization import set_quant_config + from freetoken.models.register import checkpoint_quant_config, get_model_spec + from freetoken.utils import cached_load_hf_config + + hf = cached_load_hf_config(model_path) + set_quant_config(checkpoint_quant_config(model_path, hf, get_model_spec(hf.architectures[0]))) + + +def meta_state_dict(model_path: str) -> dict[str, torch.Tensor]: + """State dict of the model the engine builds for ``model_path`` (experts offloaded), on the meta device.""" + from freetoken.engine.config import EngineConfig + from freetoken.engine.engine import _decode_target + from freetoken.layers import rotary + from freetoken.models import create_model + from freetoken.utils.torch_utils import torch_dtype + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + config = EngineConfig(model_path=model_path, tp_info=try_get_tp_info(), dtype=torch.bfloat16, moe_strategy="offload") + object.__setattr__(config.model_config, "moe_strategy", "offload") + object.__setattr__(config.model_config, "decode_target", _decode_target(config)) + saved = rotary._ROPE_DEVICE + rotary.set_rope_device(torch.device("cpu")) # get_rope refuses to build on meta + rotary.get_rope.cache_clear() + try: + with torch.device("meta"), torch_dtype(torch.bfloat16): + return create_model(config.model_config).state_dict() + finally: + rotary.set_rope_device(saved) + rotary.get_rope.cache_clear() # the cpu rope must not leak into the GPU tests' cache diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 771356304..b81e02f4d 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -8,6 +8,8 @@ from freetoken.models.config import FullAttentionGroupConfig, LinearGatedDeltaGroupConfig from freetoken.models.qwen4_exp.config import parse_config +from .common import LOVEDHEART_NVFP4_FP8, NVIDIA_NVFP4, QWEN_FP8, RADIXARK_NVFP4 + def _text_config(): return SimpleNamespace( @@ -61,61 +63,6 @@ def _text_config(): ) -# quantization_config of each released checkpoint, trimmed to the entries parse_config looks at - -# RadixArk/Qwen3.8-Flash-Next-NVFP4: modelopt NVFP4 everywhere except the ignore list -RADIXARK_NVFP4 = { - "quant_algo": "NVFP4", - "quant_method": "modelopt", - "ignore": [ - "model.embed_tokens", - "mtp.*", - "model.mtp.*", - "*.self_attn.*", - "*.linear_attn.*", - "*.mlp.gate*", - "*.mlp.shared_expert.*", - "*.mlp.shared_expert_gate*", - "*hyper_connection*", - "*.ple.*", - "model.visual.*", - "model.language_model.embed_tokens", - "lm_head", - ], -} - -# nvidia/Qwen3.8-Flash-Next-NVFP4: modelopt MIXED_PRECISION, the per-module algo sits in quantized_layers -NVIDIA_NVFP4 = { - "quant_algo": "MIXED_PRECISION", - "quant_method": "modelopt", - "quantized_layers": { - **{f"model.language_model.layers.{i}.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16} for i in range(48)}, - "model.language_model.layers.1.ple.ple_embedding.ngram_embedding": {"quant_algo": "FP8"}, - "mtp.layers.0.mlp.experts": {"quant_algo": "FP8_PB_WO", "group_size": 128}, - }, - "ignore": ["lm_head", "model.language_model.embed_tokens", "model.language_model.layers.0.mlp.shared_expert*", "model.visual*"], -} - -# Qwen/Qwen3.8-Flash-Next-FP8: 128x128 block-fp8 experts, everything else listed in modules_to_not_convert -QWEN_FP8 = { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "weight_per_tensor": False, - "act_per_tensor": False, - "weight_block_size": [128, 128], - "modules_to_not_convert": [ - "lm_head", - "model.language_model.embed_tokens", - "model.language_model.hyper_connection_mixer.input_mix_weight_up", - "model.language_model.layers.0.linear_attn.in_proj_qkv", - "model.language_model.layers.3.self_attn.q_proj", - "model.language_model.layers.3.mlp.gate", - "model.language_model.layers.3.mlp.shared_expert.gate_proj", - ], - "modules_to_convert": ["ple.ple_embedding.ngram_embedding"], -} - - def _hf_config(quantization_config=RADIXARK_NVFP4): return SimpleNamespace( model_type="qwen4_exp", @@ -208,3 +155,38 @@ def test_eos_token_id_list_uses_the_first_entry(): hf = _hf_config() hf.text_config.eos_token_id = [base, base + 1] assert parse_config(hf).qwen4_args.ngram_boundary_token_id == base + + +# the merged-projection prefixes the model asks the QuantConfig about (attention.py / gdn.py) +DENSE_PREFIXES = ( + "model.layers.3.self_attn.qkv_proj", "model.layers.3.self_attn.o_proj", + "model.layers.0.linear_attn.in_proj_qkvz", "model.layers.0.linear_attn.out_proj", +) +BF16_PREFIXES = ( + "model.layers.0.linear_attn.in_proj_ba", "model.layers.0.mlp.shared_expert.gate_up_proj", + "model.layers.3.self_attn.indexer.index_qk_proj", "lm_head", +) + + +def _quant(hf, tmp_path): + from freetoken.models.register import checkpoint_quant_config, get_model_spec + + return checkpoint_quant_config(str(tmp_path), hf, get_model_spec(hf.architectures[0])) + + +def test_block_fp8_dense_schemes(tmp_path): + quant = _quant(_hf_config(LOVEDHEART_NVFP4_FP8), tmp_path) + for prefix in DENSE_PREFIXES: + scheme = quant.scheme_for(prefix) + assert str(scheme.kind) == "fp8_block" and scheme.has("weight_scale_inv"), prefix + for prefix in BF16_PREFIXES: + assert quant.scheme_for(prefix) is None, prefix + assert str(quant.scheme_for("model.layers.0.mlp.experts").kind) == "nvfp4" + assert parse_config(_hf_config(LOVEDHEART_NVFP4_FP8)).expert_quant == "nvfp4" + + +@pytest.mark.parametrize("quantization_config", [RADIXARK_NVFP4, NVIDIA_NVFP4, None], ids=["RadixArk", "nvidia", "bf16"]) +def test_released_checkpoints_keep_the_dense_projections_bf16(quantization_config, tmp_path): + quant = _quant(_hf_config(quantization_config), tmp_path) + for prefix in DENSE_PREFIXES + BF16_PREFIXES: + assert quant.scheme_for(prefix) is None, prefix diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index b3f1f851f..1ddbf4250 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -1,11 +1,12 @@ -"""qwen4_exp weight loading against a synthetic checkpoint shaped like the RadixArk NVFP4 one. +"""qwen4_exp weight loading against synthetic checkpoints shaped like the released ones. The tensors are tiny but the key names, dtypes and the fusion geometry that matters -(hc_lowrank=320 + hc_count=4 -> a 12-row zero pad) are the real ones. +(hc_lowrank=320 + hc_count=4 -> a 12-row zero pad; 128-row block scales) are the real ones. """ from __future__ import annotations +import json import random from types import SimpleNamespace @@ -17,18 +18,23 @@ from freetoken.kernel.aot_models import SUPPORTED_MODELS, expert_bank_row_bytes from freetoken.models.qwen4_exp.weight import ( _ZERO_CENTERED_NORM_SUFFIXES, + _DenseFuser, iter_weights, load_ple_table, ) +from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into -H = 32 # hidden_size +from .common import LM, RADIXARK_NVFP4, hf_config, install_quant_config, meta_state_dict, mixed_precision_quant + +H = 128 # hidden_size; every block-fp8 projection needs in/out multiples of 128 HC = 4 # hc_count LR = 320 # hc_lowrank; kept real so the merged HC pad is the real (-(320+4)) % 16 = 12 HCH = HC * H # hyper-connection stream width -KH, VH, HD = 2, 6, 8 # GDN key / value heads, head dim -QH, KVH, AHD = 4, 2, 16 # QSA q / kv heads, head dim -IHD = 8 # indexer head dim +KH, VH, HD = 2, 4, 32 # GDN key / value heads, head dim: qkv rows 256, z rows 128 +QH, KVH, AHD = 4, 2, 64 # QSA q / kv heads, head dim: q rows 512, k / v rows 128 +IHD = 64 # indexer head dim +BLOCK = 128 E, I = 3, 6 # routed experts, moe_intermediate_size NGRAM_DIM, NGRAM_ROWS, NGRAM_SHARDS = 4, 7, 4 @@ -54,8 +60,15 @@ def _hc_weights(prefix: str, inject: bool) -> dict[str, torch.Tensor]: return w -def _raw_checkpoint() -> dict[str, torch.Tensor]: - """Layer 0 = GDN + PLE, layer 1 = QSA; plus the mtp / visual / routed-expert noise.""" +def _fp8_scale(weight: torch.Tensor) -> torch.Tensor: + return torch.rand(weight.shape[0] // BLOCK, weight.shape[1] // BLOCK) + 0.5 + + +def _raw_checkpoint(dense_fp8: bool = False) -> dict[str, torch.Tensor]: + """Layer 0 = GDN + PLE, layer 1 = QSA; plus the mtp / visual / routed-expert noise. + + ``dense_fp8`` stores the attention and GDN qkv|z / out projections as 128x128 block-fp8 (e4m3 ``.weight`` + fp32 ``.weight_scale_inv``) like the community NVFP4-FP8 requants. + """ lm = "model.language_model" raw: dict[str, torch.Tensor] = { f"{lm}.embed_tokens.weight": _bf16(11, H), @@ -127,9 +140,31 @@ def _raw_checkpoint() -> dict[str, torch.Tensor]: "model.visual.blocks.0.attn.qkv.weight": _bf16(3 * H, H), "model.visual.merger.norm.weight": _bf16(H), }) + if dense_fp8: + for module in (f"{gdn}.in_proj_qkv", f"{gdn}.in_proj_z", f"{gdn}.out_proj", + *(f"{attn}.{p}_proj" for p in "qkvo")): + weight = raw[f"{module}.weight"] + raw[f"{module}.weight"] = weight.to(torch.float8_e4m3fn) + raw[f"{module}.weight_scale_inv"] = _fp8_scale(weight) return raw +FP8_DENSE_QUANT = mixed_precision_quant(gdn_layers=(0,), attn_layers=(1,), moe_layers=(0, 1)) + + +def _config_json(quantization_config) -> dict: + cfg = hf_config( + num_layers=2, head_dim=AHD, num_q=QH, num_kv=KVH, index_head_dim=IHD, index_heads=2, + budget=16, hidden=H, max_position=4096, rope_theta=10000.0, + layer_types=["linear_attention", "full_attention"], + linear_num_key_heads=KH, linear_num_value_heads=VH, + linear_key_head_dim=HD, linear_value_head_dim=HD, + hc_lowrank=LR, ple_layer_ids=[1], + num_experts=E, moe_intermediate_size=I, shared_expert_intermediate_size=I, + ) + return {**vars(cfg), "text_config": vars(cfg.text_config), "quantization_config": quantization_config} + + def _ngram_table() -> tuple[dict[str, torch.Tensor], torch.Tensor]: prefix = "model.language_model.layers.0.ple.ple_embedding.ngram_embedding" shards = { @@ -144,11 +179,7 @@ def _ngram_table() -> tuple[dict[str, torch.Tensor], torch.Tensor]: return shards, scale -@pytest.fixture(scope="module") -def checkpoint(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]: - torch.manual_seed(0) - folder = tmp_path_factory.mktemp("qwen4_exp_ckpt") - raw = _raw_checkpoint() +def _write_checkpoint(folder, raw: dict[str, torch.Tensor], quantization_config) -> tuple[str, dict[str, torch.Tensor]]: table, _scale = _ngram_table() # Spread the dense tensors over two shards so the fusion buffer has to survive a file # boundary, and put the n-gram table in its own shards like the real checkpoint does. @@ -158,12 +189,12 @@ def checkpoint(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]: shard_names = sorted(table) save_file({n: table[n] for n in shard_names[:2]}, str(folder / "model-plefp8-00000.safetensors")) save_file({n: table[n] for n in shard_names[2:]}, str(folder / "model-plefp8-00001.safetensors")) + (folder / "config.json").write_text(json.dumps(_config_json(quantization_config))) return str(folder), {**raw, **table} -@pytest.fixture(scope="module") -def loaded(checkpoint) -> dict[str, torch.Tensor]: - folder, _raw = checkpoint +def _load(folder: str) -> dict[str, torch.Tensor]: + install_quant_config(folder) return { name: tensor.clone() for name, tensor in iter_weights( @@ -172,6 +203,30 @@ def loaded(checkpoint) -> dict[str, torch.Tensor]: } +@pytest.fixture(scope="module") +def checkpoint(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]: + torch.manual_seed(0) + return _write_checkpoint(tmp_path_factory.mktemp("qwen4_exp_ckpt"), _raw_checkpoint(), None) + + +@pytest.fixture(scope="module") +def loaded(checkpoint) -> dict[str, torch.Tensor]: + return _load(checkpoint[0]) + + +@pytest.fixture(scope="module") +def checkpoint_fp8(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]: + torch.manual_seed(1) + return _write_checkpoint( + tmp_path_factory.mktemp("qwen4_exp_fp8_ckpt"), _raw_checkpoint(dense_fp8=True), FP8_DENSE_QUANT + ) + + +@pytest.fixture(scope="module") +def loaded_fp8(checkpoint_fp8) -> dict[str, torch.Tensor]: + return _load(checkpoint_fp8[0]) + + def _expected_names() -> set[str]: names = {"model.embed_tokens.weight", "lm_head.weight"} names |= {f"model.hyper_connection_mixer.{leaf}" for leaf in @@ -206,7 +261,7 @@ def test_mtp_visual_experts_and_table_never_loaded(loaded): assert not name.startswith(("mtp.", "model.visual.")) assert ".mlp.experts." not in name assert "ngram_embedding" not in name - assert not name.endswith((".weight_scale", ".weight_scale_2", ".input_scale")) + assert not name.endswith((".weight_scale", ".weight_scale_2", ".input_scale", ".weight_scale_inv")) def test_hc_merge_is_down_then_inject_then_zero_pad(loaded, checkpoint): @@ -399,12 +454,92 @@ def test_every_registry_architecture_is_claimed_by_an_aot_entry(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs cuda") def test_fusion_pad_rides_the_tensor_device(): """safetensors loads straight to cuda; a cpu-allocated pad row would break torch.cat.""" - from freetoken.models.qwen4_exp.weight import _try_fuse - - buf = {} + fuser = _DenseFuser(None, get_model_spec("Qwen4ExpForConditionalGeneration").packed_modules_mapping) down = torch.randn(320, 64, device="cuda", dtype=torch.bfloat16) inject = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) - assert _try_fuse("model.layers.0.attn_hyper_connection.input_mix_weight_down.weight", down, buf) == () - key, fused = _try_fuse("model.layers.0.attn_hyper_connection.block_inject_weight.weight", inject, buf) + assert fuser.fuse("model.layers.0.attn_hyper_connection.input_mix_weight_down.weight", down) == [] + [(key, fused)] = fuser.fuse("model.layers.0.attn_hyper_connection.block_inject_weight.weight", inject) + assert key == "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight" assert fused.device.type == "cuda" and fused.shape[0] == 336 assert torch.equal(fused[324:], torch.zeros(12, 64, device="cuda", dtype=torch.bfloat16)) + + +# ====================================================================================== +# the reader against the model the engine builds, for each released quant layout +# ====================================================================================== + + +@pytest.fixture(scope="module") +def checkpoint_nvfp4(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]: + """bf16 dense tensors under a real ModelOptConfig whose ignore list covers them (RadixArk).""" + torch.manual_seed(2) + return _write_checkpoint(tmp_path_factory.mktemp("qwen4_exp_nvfp4_ckpt"), _raw_checkpoint(), RADIXARK_NVFP4) + + +FP8_MODULES = ( + "model.layers.1.self_attn.qkv_proj", "model.layers.1.self_attn.o_proj", + "model.layers.0.linear_attn.in_proj_qkvz", "model.layers.0.linear_attn.out_proj", +) + + +@pytest.mark.parametrize("fixture", ["checkpoint", "checkpoint_nvfp4", "checkpoint_fp8"]) +def test_emitted_keys_are_the_model_state_dict(fixture, request): + """The reader fills exactly the buffers the engine builds from the same config, block-fp8 ones with the stored dtypes.""" + folder, _raw = request.getfixturevalue(fixture) + loaded, state = _load(folder), meta_state_dict(folder) + assert set(loaded) == set(state) + if fixture != "checkpoint_fp8": + assert loaded["model.layers.0.linear_attn.in_proj.weight"].dtype is torch.bfloat16 + return + for module in FP8_MODULES: + for kind in (".weight", ".weight_scale_inv"): + assert loaded[module + kind].shape == state[module + kind].shape, module + kind + assert loaded[module + ".weight"].dtype is state[module + ".weight"].dtype is torch.float8_e4m3fn + assert loaded[module + ".weight_scale_inv"].dtype is torch.float32 # the engine casts it to the bf16 buffer at load + + +def _assert_fused_per_kind(loaded, raw, fused: str, parts: list[str]) -> None: + for kind in (".weight", ".weight_scale_inv"): + sources = [raw[f"{p}{kind}"].view(torch.uint8) for p in parts] + merged = loaded[fused + kind] + assert merged.dtype is raw[f"{parts[0]}{kind}"].dtype + for source, back in zip(sources, torch.split(merged.view(torch.uint8), [s.shape[0] for s in sources], dim=0)): + assert torch.equal(source, back) + + +def test_fp8_projections_fuse_per_kind(loaded_fp8, checkpoint_fp8): + _folder, raw = checkpoint_fp8 + attn, gdn = f"{LM}.layers.1.self_attn", f"{LM}.layers.0.linear_attn" + _assert_fused_per_kind(loaded_fp8, raw, "model.layers.1.self_attn.qkv_proj", [f"{attn}.{p}_proj" for p in "qkv"]) + _assert_fused_per_kind(loaded_fp8, raw, "model.layers.0.linear_attn.in_proj_qkvz", [f"{gdn}.in_proj_qkv", f"{gdn}.in_proj_z"]) + assert torch.equal(loaded_fp8["model.layers.0.linear_attn.out_proj.weight_scale_inv"], raw[f"{gdn}.out_proj.weight_scale_inv"]) + assert torch.equal( + loaded_fp8["model.layers.0.linear_attn.in_proj_ba.weight"], + torch.cat([raw[f"{gdn}.in_proj_b.weight"], raw[f"{gdn}.in_proj_a.weight"]], dim=0), + ) + for name in ("model.layers.0.linear_attn.in_proj_ba.weight", "model.layers.1.mlp.shared_expert.gate_up_proj.weight", + "model.layers.1.self_attn.indexer.index_qk_proj.weight", "lm_head.weight", + "model.hyper_connection_mixer.input_mix_weight_down.weight", + "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight"): + assert loaded_fp8[name].dtype is torch.bfloat16 + + +ATTN = f"{LM}.layers.1.self_attn" +REJECTED = [ + pytest.param(None, lambda w: {f"{ATTN}.q_proj.weight": w, f"{ATTN}.q_proj.weight_scale_inv": _fp8_scale(w)}, + r"q_proj\.weight_scale_inv", id="scale the config does not declare"), + pytest.param(None, lambda w: {f"{ATTN}.o_proj.weight": w.to(torch.float8_e4m3fn)}, + r"o_proj\.weight is torch\.float8", id="fp8 weight the config declares bf16"), + pytest.param(FP8_DENSE_QUANT, lambda w: {f"{ATTN}.{p}_proj.weight": w.clone() for p in "qkv"}, + r"[qkv]_proj\.weight is torch\.bfloat16", id="bf16 weight the config declares fp8"), + pytest.param(FP8_DENSE_QUANT, lambda w: {f"{ATTN}.q_proj.weight": w[:-64].to(torch.float8_e4m3fn), f"{ATTN}.q_proj.weight_scale_inv": _fp8_scale(w)}, + "128x128", id="part that is not a 128-row multiple"), +] + + +@pytest.mark.parametrize("quantization_config, tensors, match", REJECTED) +def test_checkpoint_disagreeing_with_its_quant_config_is_rejected(tmp_path, quantization_config, tensors, match): + save_file(tensors(_bf16(2 * QH * AHD, H)), str(tmp_path / "model.safetensors")) + (tmp_path / "config.json").write_text(json.dumps(_config_json(quantization_config))) + with pytest.raises(ValueError, match=match): + _load(str(tmp_path)) diff --git a/tests/models/qwen4_exp/test_weight_ckpt.py b/tests/models/qwen4_exp/test_weight_ckpt.py index 82f9a01d2..605c40f18 100644 --- a/tests/models/qwen4_exp/test_weight_ckpt.py +++ b/tests/models/qwen4_exp/test_weight_ckpt.py @@ -31,6 +31,8 @@ from freetoken.moe.host_banks import HostResidency from freetoken.utils import cached_load_hf_config +from .common import install_quant_config, meta_state_dict + MODEL_PATH = os.environ.get("FREETOKEN_QWEN4EXP_MODEL") pytestmark = [ pytest.mark.needs_weights, @@ -166,6 +168,7 @@ def dense_pass() -> tuple[list[str], dict[str, torch.Tensor]]: wanted = {name for name, _raw, _mode in SAMPLES} names: list[str] = [] sampled: dict[str, torch.Tensor] = {} + install_quant_config(MODEL_PATH) for name, tensor in iter_weights( MODEL_PATH, torch.device("cpu"), include_moe_experts=True, include_non_moe=True ): @@ -190,27 +193,12 @@ def test_emitted_names_are_unique_and_complete(dense_pass): @pytest.fixture(scope="module") def model_state_dict_keys() -> set[str]: """Keys ``Qwen4ExpForCausalLM`` declares -- the authoritative target the loader must fill.""" - from freetoken.layers import rotary - from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM - - config = parse_config(cached_load_hf_config(MODEL_PATH)) - saved = rotary._ROPE_DEVICE - rotary.set_rope_device(torch.device("cpu")) # get_rope refuses to build on meta - rotary.get_rope.cache_clear() - try: - with torch.device("meta"): - return set(Qwen4ExpForCausalLM(config).state_dict()) - finally: - rotary.set_rope_device(saved) - rotary.get_rope.cache_clear() + return set(meta_state_dict(MODEL_PATH)) def test_emitted_names_are_the_model_state_dict(dense_pass, model_state_dict_keys): names, _sampled = dense_pass - # The routed NVFP4 experts come from the offload source banks, never from the dense pass. - expected = {k for k in model_state_dict_keys - if not k.endswith((".mlp.experts.gate_up_proj", ".mlp.experts.down_proj"))} - assert set(names) == expected + assert set(names) == model_state_dict_keys def test_every_zero_centered_norm_is_present_and_raw(dense_pass, reader):