From f8d6a56b360ccadc7fc9d9cb6c9bd0752666a7ac Mon Sep 17 00:00:00 2001 From: "Stephen J. Thompson" Date: Tue, 8 Sep 2026 04:01:38 +0100 Subject: [PATCH] fix(models): support llm-compressor NVFP4 MoE export variants llm-compressor Qwen3.5-MoE NVFP4 exports come in layouts main could not convert or serve: per-expert vs stacked per-layer routed experts, a GDN left bf16 / fp8 by the ignore list or a ``gdn:fp8`` recipe, format-only configs with no ``config_groups``, native NVFP4 lm_head/embed_tokens, and single-file checkpoints without a shard index. Detect each from the on-disk layout / recipe, build the matching offload banks (per-expert and stacked loaders), and emit the state dict the model expects. Assisted-by: opencode --- python/freetoken/models/config.py | 11 + python/freetoken/models/nvfp4_banks.py | 142 +++- python/freetoken/models/qwen3_5_moe/config.py | 112 ++- python/freetoken/models/qwen3_5_moe/gdn.py | 8 +- python/freetoken/models/qwen3_5_moe/model.py | 1 + python/freetoken/models/qwen3_5_moe/weight.py | 209 ++++-- tests/models/test_qwen3_5_moe_ct_nvfp4.py | 696 ++++++++++++++++++ 7 files changed, 1111 insertions(+), 68 deletions(-) create mode 100644 tests/models/test_qwen3_5_moe_ct_nvfp4.py diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 1bce039cb..1e65138c0 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -69,6 +69,11 @@ def detect_compressed_tensors_nvfp4(hf_config: Any) -> bool: if str(get("quant_method") or "").lower() != "compressed-tensors": return False groups = get("config_groups") or {} + if not groups: + # No config_groups: some llm-compressor exports (e.g. the AEON / Kwaipilot Qwen3.6 + # MoE NVFP4 builds) carry only ``format: nvfp4-pack-quantized`` (+ a ``recipe`` + # string). Gate on the exact format string, like the config_groups branch below. + return str(get("format") or "").lower() == "nvfp4-pack-quantized" # Verdicts are collected across ALL groups before returning: an early return on # the first NVFP4 group would accept a mixed {nvfp4, mxfp4} checkpoint (and the # error would depend on the groups' key order). @@ -277,6 +282,12 @@ class ModelConfig: # scale and runs a W8A16 kernel (modelopt MIXED_PRECISION); "none" leaves them bf16 # (dequant-at-load for any other dense quant, e.g. NVFP4 shared_expert/lm_head). attn_quant: str = "none" + # Quantization of the GatedDeltaNet's ``out_proj`` only (qwen3_5_moe). Independent of + # ``attn_quant`` because llm-compressor checkpoints can quantize the full-attention + # projections while leaving the whole GDN bf16 (their ``ignore`` list names the + # ``linear_attn.*`` modules). "nvfp4" keeps ``out_proj`` packed (W4A16); "fp8_pertensor" + # keeps it per-tensor fp8 (modelopt MIXED_PRECISION); "none" -> bf16. + gdn_quant: str = "none" # Weight quantization of the *dense* NVFP4 MLP projections -- the shared expert, and dense # (non-MoE) MLP layers -- which NVFP4 checkpoints store as packed FP4 like the routed # experts. "nvfp4" keeps them packed and runs the W4A16 dense kernels (quartering their diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 6b933ff1d..097d3cd89 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -1,9 +1,11 @@ from __future__ import annotations import collections +import glob import json import os import re +import struct from dataclasses import dataclass from typing import Callable @@ -28,6 +30,10 @@ class Nvfp4ExpertSourceSpec: # The checkpoint stores the QUANT-side global scale (local fp8 scales were # multiplied by it before the cast); the banks keep its reciprocal. global_reciprocal: bool = False + # True when the experts are stacked per layer (``experts.gate_up_proj`` U8 + # [E*rows, cols] + a per-layer scalar global) instead of per expert; the stacked + # loader reshapes each bank tensor to [E, ...] and broadcasts the scalar global. + stacked: bool = False def _canon_kind(spec: "Nvfp4ExpertSourceSpec", kind: str) -> str: @@ -78,6 +84,24 @@ def _alloc_nvfp4_host_banks(num_layers: int, E: int, H: int, I: int): }, num_layers) +def _weight_map(folder: str) -> dict[str, str]: + """name -> shard (basename) from the index, or from each safetensors header when the + checkpoint ships a single shard without an index (llm-compressor single-file exports).""" + index = os.path.join(folder, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as f: + return json.load(f)["weight_map"] + weight_map: dict[str, str] = {} + for shard in sorted(os.path.basename(p) for p in glob.glob(os.path.join(folder, "*.safetensors"))): + with open(os.path.join(folder, shard), "rb") as fh: + n = struct.unpack(" int: } +def load_nvfp4_stacked_expert_sources( + model_path: str, + config, + spec: Nvfp4ExpertSourceSpec, + *, + drop_page_cache: DropPageCache, + primary: bool, + layer_sink=None, +) -> dict[str, list[torch.Tensor]]: + """Build the 6 native NVFP4 source banks for a STACKED (per-layer) expert layout. + + llm-compressor can store the routed experts as one packed tensor per layer instead of + per expert: ``...experts.gate_up_proj.weight_packed`` U8 [E*rows, cols] (rows are + expert-major, so the bank tensor just reshapes to [E, rows, cols]) plus ONE + layer-global ``weight_global_scale`` scalar (reciprocated at ingest). Placement and + the resulting 6-bank dict are identical to :func:`load_nvfp4_expert_source_banks` + (which is why the marlin/b12x repack and the offload cache never notice the + difference). ``layer_sink``: see :func:`load_nvfp4_expert_source_banks`.""" + folder = download_hf_weight(model_path) + weight_map = _weight_map(folder) + + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + num_layers = _num_moe_layers(config) + + weight_shards: dict[str, list[tuple[str, re.Match[str], int]]] = collections.defaultdict(list) + global_shards: dict[str, list[tuple[str, int, str]]] = collections.defaultdict(list) + for name, shard in weight_map.items(): + match = spec.key_pattern.match(name) + if match is None: + continue + bank_layer = _bank_layer(spec, int(match.group("layer")), config) + if bank_layer is None: + continue + kind = _canon_kind(spec, match.group("kind")) + if kind == "weight_scale_2": + global_shards[shard].append((name, bank_layer, match.group("proj"))) + elif kind in {"weight", "weight_scale"}: + weight_shards[shard].append((name, match, bank_layer)) + else: + raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") + + # Pass 1: the per-layer scalar globals (reciprocal at ingest). gate_up and down carry + # their own global, so key by (bank_layer, proj). + globals_map: dict[tuple[int, str], torch.Tensor] = {} + for shard in sorted(global_shards): + path = os.path.join(folder, shard) + drop_page_cache(path) + with safetensors.safe_open(path, framework="pt", device="cpu") as f: + for name, bank_layer, proj in global_shards[shard]: + globals_map[(bank_layer, proj)] = _ingest_global(spec, f.get_tensor(name)) + drop_page_cache(path) + + _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill + gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] + gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] + gate_up_global = [b.tensor for b in _hb["gate_up_global"]] + down_packed = [b.tensor for b in _hb["down_packed"]] + down_scale = [b.tensor for b in _hb["down_scale"]] + down_global = [b.tensor for b in _hb["down_global"]] + + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline + + def _load(sink) -> int: + tracker = LayerCompletionTracker(4, _hb, sink) # 2 gate_up + 2 down per layer + placed = 0 + for shard in tqdm(sorted(weight_shards), desc=f"Loading {spec.desc}", disable=not primary): + path = os.path.join(folder, shard) + with safetensors.safe_open(path, framework="pt", device="cpu") as f: + for name, match, bank_layer in weight_shards[shard]: + proj = match.group("proj") + kind = _canon_kind(spec, match.group("kind")) + tensor = f.get_tensor(name) + if kind == "weight": + if proj == "gate_up_proj": + gate_up_packed[bank_layer].copy_(tensor.view(E, 2 * I, H // 2)) + else: + down_packed[bank_layer].copy_(tensor.view(E, H, I // 2)) + else: + g = globals_map[(bank_layer, proj)] + if proj == "gate_up_proj": + gate_up_scale[bank_layer].copy_(tensor.view(E, 2 * I, H // 16)) + gate_up_global[bank_layer].fill_(g.item()) + else: + down_scale[bank_layer].copy_(tensor.view(E, H, I // 16)) + down_global[bank_layer].fill_(g.item()) + tracker.note(bank_layer) + placed += 1 + drop_page_cache(path) + return placed + + if layer_sink is not None: + placed = _load(layer_sink) + else: + with PinPipeline() as pins: + placed = _load(pins) + + expected = num_layers * 4 + assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" + return { + "gate_up_packed": gate_up_packed, + "gate_up_scale": gate_up_scale, + "gate_up_global": gate_up_global, + "down_packed": down_packed, + "down_scale": down_scale, + "down_global": down_global, + } + + __all__ = [ "Nvfp4ExpertSourceSpec", "load_nvfp4_expert_source_banks", "load_nvfp4_expert_source_banks_parallel", + "load_nvfp4_stacked_expert_sources", ] diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index 2ac4b607f..bb78aa85c 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/python/freetoken/models/qwen3_5_moe/config.py @@ -39,16 +39,110 @@ def _fp8_block_quant(hf_config: Any) -> tuple[str, tuple[int, int] | None]: return "none", None +def _ct_expert_groups_nvfp4(hf_config: Any) -> bool: + """compressed-tensors MoE checkpoint: are the *routed experts* NVFP4? Mirrors the + mixed-precision scan in ``models.config.detect_expert_quant``: groups whose + ``targets`` name the experts decide (a generic ``["Linear"]`` group falls back to + covering everything). Exports without ``config_groups`` (format-only, e.g. + doth4580/Kwaipilot-KAT-Coder-V2.5-Dev-NVFP4-MIXED) ride the top-level format.""" + get = _quant_accessor(hf_config) + if get is None: + return False + groups = get("config_groups") or {} + if not groups: + return str(get("format") or "").lower() == "nvfp4-pack-quantized" + groups = [g or {} for g in (groups.values() if isinstance(groups, dict) else [])] + expert_groups = [g for g in groups if any("experts" in str(t) for t in (g.get("targets") or []))] + for g in expert_groups or groups: + w = (g or {}).get("weights") or {} + if int(w.get("num_bits", 0) or 0) != 4 or str(w.get("type", "")).lower() != "float": + continue + if int(w.get("group_size", 0) or 0) == 16 and str(w.get("strategy", "")).lower() == "tensor_group": + return True + return False + + +def _ct_recipe(hf_config: Any) -> dict[str, str]: + """Parse llm-compressor's ``recipe`` string (``all,gdn:fp8,-router``) into a map of + module-key -> value. Tokens: ``name:quant`` (per-module override), ``-name`` (skip), + bare ``name`` (quantize). Unknown keys are ignored; absence means the default.""" + get = _quant_accessor(hf_config) + if get is None: + return {} + out: dict[str, str] = {} + for tok in str(get("recipe") or "").split(","): + tok = tok.strip() + if not tok: + continue + if tok.startswith("-"): + out.setdefault(tok[1:], "skip") + elif ":" in tok: + key, _, val = tok.partition(":") + out.setdefault(key, val.lower()) + else: + out.setdefault(tok, "quant") + return out + + +def _ct_ignored(hf_config: Any, probe: str) -> bool: + """Is the canonical module path ``probe`` excluded from quantization by the export's + ``ignore`` list (exact or ``re:`` entries) or a ``recipe`` skip token (``-lm_head``)?""" + get = _quant_accessor(hf_config) + if get is None: + return False + import re + + for entry in get("ignore") or []: + e = str(entry) + if e.startswith("re:"): + try: + if re.search(e[3:], probe): + return True + except re.error: + continue + elif e == probe: + return True + return _ct_recipe(hf_config).get(probe.split(".")[-1]) == "skip" + + +def _ct_gdn_nvfp4(hf_config: Any) -> bool: + """Is the GDN ``out_proj`` stored NVFP4 (as opposed to per-tensor fp8 / bf16)? + + The full-attention projections can be NVFP4 while the whole GDN is left bf16 (the + AEON Qwen3.6-35B-A3B ``ignore`` list) or per-tensor fp8 (a ``gdn:fp8`` recipe, e.g. + Kwaipilot-KAT-Coder-V2.5). Returns False in both cases; True when the GDN rides the + default NVFP4 format.""" + if _ct_ignored(hf_config, "model.language_model.layers.0.linear_attn.out_proj"): + return False + return _ct_recipe(hf_config).get("gdn", "nvfp4") in ("nvfp4", "fp4", "quant", "all") + + +def _ct_linear_attn_ignored(hf_config: Any) -> bool: + """Backward-compatible wrapper: the GDN out_proj is "ignored" exactly when it is not + NVFP4 (the ignore-list or ``recipe`` override left it fp8/bf16).""" + return not _ct_gdn_nvfp4(hf_config) + + def _expert_quant(hf_config: Any) -> str: """Quantization format of the *routed* experts (the only weights served from the offload cache). The nvidia/modelopt checkpoints are either plain NVFP4 (``quant_algo`` ``NVFP4``) or ``MIXED_PRECISION`` (per-layer ``quantized_layers`` map); in the mixed - case the routed experts carry their own ``W4A16_NVFP4``/``FP8`` algo. Dense quantized - weights (attention/shared-expert/lm_head) are handled separately by dequant-at-load.""" + case the routed experts carry their own ``W4A16_NVFP4``/``FP8`` algo. llm-compressor + (compressed-tensors) MoE checkpoints keep the routed experts NVFP4 in the offload + cache. Dense quantized weights (attention/shared-expert/lm_head) are handled + separately by dequant-at-load.""" get = _quant_accessor(hf_config) if get is None: return "none" algo = str(get("quant_algo") or get("quant_method") or "").lower() + if algo == "compressed-tensors": + # Dense exports (e.g. Qwen3.6-27B, no routed experts) stay "none" -- their + # weights ride the dense compressed-tensors reader. An MoE export (e.g. the + # Qwen3.6-35B-A3B NVFP4 builds) keeps its experts NVFP4 in the offload cache. + text = getattr(hf_config, "text_config", hf_config) + if int(getattr(text, "num_experts", 0) or 0) > 0 and _ct_expert_groups_nvfp4(hf_config): + return "nvfp4" + return "none" if "fp4" in algo: return "nvfp4" if "mixed" in algo: @@ -179,13 +273,18 @@ def parse_config(hf_config: Any) -> ModelConfig: dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config) lm_head_quant = _lm_head_quant(hf_config) - # compressed-tensors NVFP4 (dense Qwen3.6-27B): the attention (q/k/v/o, GDN out_proj) AND - # the dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms stay bf16. Wire the shared - # W4A16 kernels (attn_quant=="nvfp4" routes the attention/GDN linears through them too). + # compressed-tensors NVFP4: the attention (q/k/v/o) and dense MLP are W4A16 NVFP4. + # The GDN out_proj follows only when the export actually quantized it (its + # ``ignore``/``recipe`` may leave the GDN bf16 or fp8). lm_head is NVFP4 unless the + # export skipped it (dense Qwen3.6-27B and the AEON MoE builds put it in ``ignore``). if _compressed_tensors_nvfp4(hf_config): attn_quant = "nvfp4" dense_quant = "nvfp4" - lm_head_quant = "none" + lm_head_quant = "none" if _ct_ignored(hf_config, "lm_head") else "nvfp4" + # The GDN's out_proj follows attn_quant EXCEPT when the export left it non-NVFP4 + # (the AEON ``ignore`` list keeps the GDN bf16; a ``gdn:fp8`` recipe keeps it + # per-tensor fp8); the full-attention q/k/v/o are quantized either way. + gdn_quant = "none" if (attn_quant == "nvfp4" and not _ct_gdn_nvfp4(hf_config)) else attn_quant # Dense variants (e.g. Qwen3.6-27B) report num_experts==0: route the decoder MLP through # the dense Qwen3_5DenseMLP instead of the MoE block. @@ -255,6 +354,7 @@ def parse_config(hf_config: Any) -> ModelConfig: expert_quant=expert_quant, weight_block_size=weight_block_size, attn_quant=attn_quant, + gdn_quant=gdn_quant, dense_quant=dense_quant, lm_head_quant=lm_head_quant, ) diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..781c0c499 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -53,7 +53,7 @@ class Qwen3_5GatedDeltaNet(BaseOP): def __init__( self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size, rms_norm_eps, layer_id, expert_quant: str = "none", - attn_quant: str = "none", + attn_quant: str = "none", gdn_quant: str | None = None, ): self.layer_id = layer_id # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as @@ -100,9 +100,11 @@ def __init__( self.norm = _GatedRMSNorm(head_v_dim, eps=rms_norm_eps) # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors - # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4. + # NVFP4 checkpoint only makes out_proj native FP4 -- and only when the export actually + # quantized it (its ``ignore`` list may leave the whole GDN bf16, see config.gdn_quant). self.out_proj = make_replicated_quant( - expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False + expert_quant, attn_quant if gdn_quant is None else gdn_quant, + self.value_dim, hidden_size, has_bias=False, ) def _gate_params(self, a: torch.Tensor, b: torch.Tensor): diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24f..c138ee623 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -44,6 +44,7 @@ def __init__(self, config: ModelConfig, layer_id: int): layer_id=layer_id, expert_quant=config.expert_quant, attn_quant=config.attn_quant, + gdn_quant=config.gdn_quant, ) else: self.self_attn = Qwen3_5Attention(config, layer_id) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..05e3c2d70 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -47,6 +47,40 @@ layer_to_bank=lambda layer, config: layer, # every layer is MoE desc="Qwen3.5 NVFP4 experts", ) +# llm-compressor export (compressed-tensors): weight_packed | weight_scale | +# weight_global_scale (quant-side global -> reciprocal at ingest). ``input_global_scale`` +# (the calibrated W4A4 activation scale) deliberately does not match: our routed-expert +# paths are W4A16 and never quantize activations. +_NVFP4_CT_EXPERT_KEY_RE = re.compile( + r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight_packed|weight_global_scale|weight_scale)$" +) +_NVFP4_CT_SOURCE_SPEC = Nvfp4ExpertSourceSpec( + key_pattern=_NVFP4_CT_EXPERT_KEY_RE, + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc="Qwen3.5 NVFP4 experts (compressed-tensors)", + kind_map={"weight_packed": "weight", "weight_global_scale": "weight_scale_2"}, + global_reciprocal=True, +) +# Some llm-compressor exports stack the experts per layer instead of per expert: +# ``...experts.gate_up_proj.weight_packed`` U8 [E*rows, cols] with ONE layer-global +# ``weight_global_scale`` scalar (e.g. doth4580/Kwaipilot-KAT-Coder-V2.5-Dev-NVFP4-MIXED). +# The stacked rows are already expert-major, so each bank tensor just reshapes to [E, ...]. +_NVFP4_CT_STACKED_EXPERT_KEY_RE = re.compile( + r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\." + r"(?Pgate_up_proj|down_proj)\.(?Pweight_packed|weight_global_scale|weight_scale)$" +) +_NVFP4_CT_STACKED_SOURCE_SPEC = Nvfp4ExpertSourceSpec( + key_pattern=_NVFP4_CT_STACKED_EXPERT_KEY_RE, + proj_to_role={"gate_up_proj": "gate_up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc="Qwen3.5 NVFP4 experts (compressed-tensors, stacked)", + kind_map={"weight_packed": "weight", "weight_global_scale": "weight_scale_2"}, + global_reciprocal=True, + stacked=True, +) # Suffixes of the per-tensor modelopt quant scales; consumed alongside their ``.weight``, # never yielded on their own. _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") @@ -82,9 +116,14 @@ def _dequant_fp8_weight(weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor: - """Weight-only FP8 -> bf16 (per-tensor static scale). Activations stay bf16 (W8A16), - which is at least as precise as the checkpoint's intended W8A8.""" - return weight.to(torch.bfloat16) * weight_scale.to(torch.bfloat16) + """Weight-only FP8 -> bf16 (W8A16). ``weight_scale`` is either a scalar (per-tensor + FP8) or a per-output-row vector (modelopt / llm-compressor per-tensor fp8); broadcast + to the weight's rows in both cases. Activations stay bf16 (W8A16), which is at least + as precise as the checkpoint's intended W8A8.""" + if weight_scale.numel() == 1: + return weight.to(torch.bfloat16) * weight_scale.to(torch.bfloat16) + scale = weight_scale.reshape(-1, 1).to(torch.bfloat16) + return (weight.to(torch.bfloat16) * scale).contiguous() def _dequant_nvfp4_weight( @@ -186,6 +225,7 @@ def iter_weights( model_path, device, include_non_moe=include_non_moe, include_moe_experts=include_moe_experts, nvfp4=config.dense_quant == "nvfp4", + lmhead_nvfp4=config.lm_head_quant == "nvfp4", ) return if config.expert_quant == "fp8_block": @@ -544,6 +584,9 @@ def _iter_weights_attn_fp8( _CT_NVFP4_FUSE: dict[str, tuple[str, ...]] = { ".self_attn.qkv_proj": (".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj"), ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj"), + ".mlp.shared_expert.gate_up_proj": ( + ".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj", + ), } _CT_BF16_FUSE: dict[str, tuple[str, ...]] = { ".linear_attn.in_proj": ( @@ -564,20 +607,25 @@ def _ct_nvfp4_fuse(base: str, parts_tuple: tuple, buf: dict): def _iter_weights_compressed_tensors( model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool, - nvfp4: bool, + nvfp4: bool, lmhead_nvfp4: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: - """Dense pass for a compressed-tensors NVFP4 checkpoint (e.g. Qwen3.6-27B). - - Keeps the NVFP4 attention (q/k/v/o, GDN out_proj) and dense MLP (gate/up/down) native - (W4A16) -- ``.weight`` (uint8) + ``.weight_scale`` (fp8 block) + ``.weight_global`` (fp16 - per-row) -- when ``nvfp4``; otherwise dequantizes each to bf16. q/k/v -> ``qkv_proj``, dense gate/up -> ``gate_up_proj`` (output-dim concat). - GDN ``in_proj_{qkv,z,b,a}`` stay bf16 -> fused ``in_proj``; ``conv1d``/``A_log``/``dt_bias``/ - gated ``norm`` pass through (fp32 for A_log/dt_bias). Gemma (1+w) norms get +1. lm_head and - embeddings are bf16. The model is dense (no routed experts), so there is no experts pass.""" + """Dense pass for a compressed-tensors NVFP4 checkpoint. + + Keeps the NVFP4 attention (q/k/v/o), dense MLP (gate/up/down, shared_expert) and the + lm_head native (W4A16) -- ``.weight`` (uint8) + ``.weight_scale`` (fp8 block) + + ``.weight_global`` (fp16 per-row) -- when the corresponding quant flag is on; + otherwise dequantizes each to bf16. q/k/v -> ``qkv_proj``, dense gate/up -> + ``gate_up_proj`` (output-dim concat). GDN ``in_proj_{qkv,z,b,a}`` and ``out_proj`` + compute in bf16: NVFP4- or per-tensor-fp8-stored ones are dequantized (the export's + ``ignore``/``recipe`` may keep the whole GDN bf16). ``embed_tokens`` is always + dequantized (there is no NVFP4 embedding kernel). Routed experts never appear here -- + they ride the offload banks via ``load_nvfp4_expert_sources``. ``conv1d``/``A_log``/ + ``dt_bias``/gated ``norm`` pass through (fp32 for A_log/dt_bias). Gemma (1+w) norms + get +1.""" if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe weight loading currently supports TP=1 only") if not include_non_moe: - return # dense checkpoint: no routed experts to load + return # checkpoint's routed experts are not yielded by this pass tp_info = get_tp_info() nvfp4_buf: dict[str, dict[int, tuple]] = {} @@ -604,52 +652,69 @@ def _emit_bf16_weight(name: str, tensor: torch.Tensor): desc="Loading compressed-tensors weights", disable=not tp_info.is_primary(), ): - for raw_name in reader.names_in(file): - if raw_name.startswith(("mtp.", "model.visual.", "visual.")): - continue - if raw_name.endswith(_CT_SCALE_SUFFIXES): - continue # consumed with weight_packed (or unused W4A4 activation scales) + with safetensors.safe_open(file, framework="pt", device="cpu") as f: + keyset = set(f.keys()) + for raw_name in reader.names_in(file): + if raw_name.startswith(("mtp.", "model.visual.", "visual.")): + continue + # Routed experts go to the offload cache (load_nvfp4_expert_sources), + # not the dense pass: ``.mlp.experts..`` (per-expert) and + # ``.mlp.experts.gate_up_proj`` (stacked per-layer) keying. + if _NVFP4_EXPERT_RE.search(raw_name) or ".mlp.experts." in raw_name: + continue + if raw_name.endswith(_CT_SCALE_SUFFIXES): + continue # consumed with weight_packed (or unused W4A4 activation scales) - name = _rename(raw_name) - if name is None: - continue + name = _rename(raw_name) + if name is None: + continue - if raw_name.endswith(".weight_packed"): # NVFP4 projection - base = name[: -len(".weight_packed")] - raw_base = raw_name[: -len(".weight_packed")] - w, s, g = _nvfp4_parts_ct(reader, raw_base) - # GDN in_proj_* compute in bf16 (model contract) but some checkpoints - # (e.g. sakamakismile/Qwen3.6-27B-NVFP4) quantize them too: dequant to - # bf16 here and let the bf16 fusion assemble ``in_proj`` as usual. - if any(base.endswith(p) for ps in _CT_BF16_FUSE.values() for p in ps): + if raw_name.endswith(".weight_packed"): # NVFP4 projection + base = name[: -len(".weight_packed")] + raw_base = raw_name[: -len(".weight_packed")] + w, s, g = _nvfp4_parts_ct(reader, raw_base) + # GDN in_proj_* compute in bf16 (model contract) but some checkpoints + # (e.g. sakamakismile/Qwen3.6-27B-NVFP4) quantize them too: dequant to + # bf16 here and let the bf16 fusion assemble ``in_proj`` as usual. + if any(base.endswith(p) for ps in _CT_BF16_FUSE.values() for p in ps): + bf16 = _dequant_nvfp4_weight(w, s, g[:1]) + yield from _emit_bf16_weight(base + ".weight", bf16) + continue + if base == "model.embed_tokens": + # No NVFP4 embedding kernel: always dequant to bf16. + bf16 = _dequant_nvfp4_weight(w, s, g[:1]) + yield from _emit_bf16_weight(base + ".weight", bf16) + continue + if nvfp4 and (base != "lm_head" or lmhead_nvfp4): + # keep native (W4A16); lm_head only when the model wants it + emit = _ct_nvfp4_fuse(base, (w, s, g), nvfp4_buf) + if emit is not None: + yield from emit + else: # standalone: o_proj, linear_attn.out_proj, mlp.down_proj + yield base + ".weight", w + yield base + ".weight_scale", s + yield base + ".weight_global", g + continue + # bf16 A-B: dequant FP4 -> bf16, then merge q/k/v + gate/up as bf16. ``g`` is + # already the dequant global (1/weight_global_scale) per row; pass one element. bf16 = _dequant_nvfp4_weight(w, s, g[:1]) - yield from _emit_bf16_weight(base + ".weight", bf16) - continue - if nvfp4: # keep native (W4A16) - emit = _ct_nvfp4_fuse(base, (w, s, g), nvfp4_buf) + emit = _ct_bf16_fuse(base, bf16, bf16_buf, _CT_NVFP4_FUSE) if emit is not None: yield from emit - else: # standalone: o_proj, linear_attn.out_proj, mlp.down_proj - yield base + ".weight", w - yield base + ".weight_scale", s - yield base + ".weight_global", g + else: + yield base + ".weight", bf16 continue - # bf16 A-B: dequant FP4 -> bf16, then merge q/k/v + gate/up as bf16. ``g`` is - # already the dequant global (1/weight_global_scale) per row; pass one element. - bf16 = _dequant_nvfp4_weight(w, s, g[:1]) - emit = _ct_bf16_fuse(base, bf16, bf16_buf, _CT_NVFP4_FUSE) - if emit is not None: - yield from emit - else: - yield base + ".weight", bf16 - continue - if name.endswith(".weight"): - yield from _emit_bf16_weight(name, reader.get_tensor(raw_name)) - continue + if name.endswith(".weight"): + # bf16 weight, or a per-tensor fp8 projection (fp8-e4m3 weight + + # per-row ``weight_scale``, e.g. a ``gdn:fp8`` recipe) that the + # model computes in bf16 -- dequant, then feed the bf16 fusion. + tensor = _load_maybe_quantized(f, raw_name, keyset) + yield from _emit_bf16_weight(name, tensor) + continue - # A_log / dt_bias (kept fp32 by the model; the load downcast exempts them). - yield name, reader.get_tensor(raw_name) + # A_log / dt_bias (kept fp32 by the model; the load downcast exempts them). + yield name, reader.get_tensor(raw_name) finally: reader.close() @@ -1060,15 +1125,44 @@ def _load(sink) -> None: return ExpertBanks("bf16", banks, streamed=layer_sink is not None) +def _stacked_expert_keying(model_path: str) -> bool: + """Do the routed experts use per-layer stacked tensors (``...experts.gate_up_proj``) + instead of per-expert ones (``...experts.E.{gate,up,down}_proj``)?""" + from freetoken.models.nvfp4_banks import _weight_map + + folder = download_hf_weight(model_path) + for name in _weight_map(folder): + if ".mlp.experts." in name and name.endswith(".weight_packed"): + return re.search(r"\.mlp\.experts\.\d+\.", name) is None + return False + + +def _select_expert_source_spec(model_path: str) -> Nvfp4ExpertSourceSpec: + """Pick the NVFP4 expert bank spec from the checkpoint's ``quant_method``: modelopt + stores ``weight | weight_scale | weight_scale_2`` (dequant-side global); llm-compressor + (compressed-tensors) stores ``weight_packed | weight_scale | weight_global_scale`` + (quant-side global -> reciprocal at ingest), per-expert or stacked per-layer.""" + quant = getattr(cached_load_hf_config(model_path), "quantization_config", None) or {} + get = quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d)) + method = str(get("quant_method") or "").lower() + if method == "compressed-tensors": + return _NVFP4_CT_STACKED_SOURCE_SPEC if _stacked_expert_keying(model_path) else _NVFP4_CT_SOURCE_SPEC + return _NVFP4_SOURCE_SPEC + + def load_nvfp4_expert_sources( model_path: str, config, *, layer_sink=None ) -> dict[str, torch.Tensor]: """Build the CPU NVFP4 expert source banks for the offload cache (gate/up fused on the output-row axis, down separate; weight_scale_2 carried as the per-row global scale).""" - return load_nvfp4_expert_source_banks( + from freetoken.models.nvfp4_banks import load_nvfp4_stacked_expert_sources + + spec = _select_expert_source_spec(model_path) + loader = load_nvfp4_stacked_expert_sources if spec.stacked else load_nvfp4_expert_source_banks + return loader( model_path, config, - _NVFP4_SOURCE_SPEC, + spec, drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, @@ -1081,10 +1175,17 @@ def load_nvfp4_expert_sources_parallel( """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel + spec = _select_expert_source_spec(model_path) + if spec.stacked: + # stacked experts are a handful of large tensors per layer -- the serial read + # already saturates the disk, so the parallel reader has nothing to add. + raise NotImplementedError( + "parallel reader not implemented for stacked compressed-tensors experts" + ) return load_nvfp4_expert_source_banks_parallel( model_path, config, - _NVFP4_SOURCE_SPEC, + spec, drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), workers=workers, diff --git a/tests/models/test_qwen3_5_moe_ct_nvfp4.py b/tests/models/test_qwen3_5_moe_ct_nvfp4.py new file mode 100644 index 000000000..66d6aa4cc --- /dev/null +++ b/tests/models/test_qwen3_5_moe_ct_nvfp4.py @@ -0,0 +1,696 @@ +"""qwen3_5_moe compressed-tensors (llm-compressor) NVFP4 support. + +Qwen3.6-35B-A3B-class MoE checkpoints exported with ``quant_method: +compressed-tensors`` store their routed experts as ``weight_packed | weight_scale | +weight_global_scale`` (quant-side global), either per-expert +(``...experts.E.{gate,up,down}_proj``, AEON Qwen3.6-35B) or stacked per-layer +(``...experts.{gate_up,down}_proj`` [E*rows, cols], Kwaipilot-KAT-Coder). These tests +pin the fixes that let such checkpoints convert/serve with NVFP4 offload banks: config +detection (config_groups, format-only + ``recipe``), the CT expert-source specs, the +single-file no-index + stacked bank loaders, the shared-expert gate/up fusion, the fp8 +GDN (``gdn:fp8`` recipe) and NVFP4 embed_tokens dequantization, and lm_head kept native +NVFP4 when the export quantized it. +""" + +from __future__ import annotations + +import pytest + +from freetoken.models.qwen3_5_moe.config import parse_config + + +class _Cfg: + """Attribute-access shim over a dict (what cached_load_hf_config hands parse_config). + Only ``text_config`` is recursed; ``quantization_config`` stays a plain dict so + ``quant.get`` accessors fire.""" + + def __init__(self, data: dict): + for k, v in data.items(): + setattr(self, k, _Cfg(v) if k == "text_config" and isinstance(v, dict) else v) + + +def _text_config(num_layers: int, num_experts: int) -> dict: + return { + "hidden_size": 16, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 4, + "num_hidden_layers": num_layers, + "vocab_size": 32, + "hidden_act": "silu", + "rms_norm_eps": 1e-6, + "max_position_embeddings": 1000, + "tie_word_embeddings": False, + "num_experts": num_experts, + "num_experts_per_tok": 2, + "moe_intermediate_size": 8, + "shared_expert_intermediate_size": 8, + "layer_types": ["full_attention"] * num_layers, + "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"}, + "linear_num_key_heads": 2, + "linear_num_value_heads": 2, + "linear_key_head_dim": 4, + "linear_value_head_dim": 4, + "linear_conv_kernel_dim": 2, + } + + +def _ct_nvfp4_quant() -> dict: + """Shape of the Qwen3.6-35B-A3B-NVFP4 exports: one generic Linear group in the + NVFP4 geometry (no ``quant_algo``; ``quant_method`` carries it) plus an ``ignore`` + list that leaves the GDN (linear_attn.*), the routers, lm_head and vision bf16.""" + ignore = [ + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.in_proj_b", + "model.language_model.layers.0.linear_attn.in_proj_a", + "model.language_model.layers.1.linear_attn.out_proj", + "model.language_model.layers.0.mlp.gate", + "model.language_model.layers.0.mlp.shared_expert_gate", + "model.language_model.layers.1.mlp.gate", + "model.language_model.layers.1.mlp.shared_expert_gate", + "lm_head", + ] + return { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "ignore": ignore, + "config_groups": { + "group_0": { + "format": "nvfp4-pack-quantized", + "targets": ["Linear"], + "weights": { + "num_bits": 4, + "type": "float", + "group_size": 16, + "strategy": "tensor_group", + }, + } + }, + } + + +def _ct_mixed_quant() -> dict: + """A hypothetical mixed export: NVFP4 dense groups but fp8 routed experts. The + expert-targeted group decides -- experts are NOT nvfp4.""" + return { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": { + "group_0": { + "targets": ["re:.*self_attn.*_proj$"], + "weights": {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"}, + "format": "nvfp4-pack-quantized", + }, + "group_1": { + "targets": ["re:.*mlp\\.experts\\..*(gate|up|down)_proj$"], + "weights": {"num_bits": 8, "type": "float", "strategy": "block"}, + "format": "float-quantized", + }, + }, + } + + +def _ct_format_only_quant() -> dict: + """Shape of doth4580/Kwaipilot-KAT-Coder-V2.5-Dev-NVFP4-MIXED: no config_groups, no + ignore -- just ``format`` and a ``recipe`` that quantizes everything to NVFP4 except + the GDN (per-tensor fp8) and the routers (skipped).""" + return { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "recipe": "all,gdn:fp8,-router", + } + + +def _hf_config(num_layers: int = 2, num_experts: int = 4, quant: dict | None = None) -> _Cfg: + data = { + "architectures": ["Qwen3_5MoeForConditionalGeneration"], + "model_type": "qwen3_5_moe", + "text_config": _text_config(num_layers, num_experts), + } + if quant is not None: + data["quantization_config"] = quant + return _Cfg(data) + + +# ----------------------------------------------------------------------------------- +# config detection +# ----------------------------------------------------------------------------------- + + +def test_parse_config_ct_moe_nvfp4(): + """compressed-tensors MoE checkpoint: the routed experts resolve to nvfp4 (offload + banks), and the dense/attention weights keep native FP4 too -- but the GDN out_proj + stays bf16 because the export's ``ignore`` list skipped the whole linear_attn.""" + cfg = parse_config(_hf_config(quant=_ct_nvfp4_quant())) + assert cfg.expert_quant == "nvfp4" + assert cfg.dense_quant == "nvfp4" + assert cfg.attn_quant == "nvfp4" + assert cfg.gdn_quant == "none" + assert cfg.is_moe + + +def test_ct_linear_attn_ignored_detection(): + """The GDN-out_proj signal is read from the ignore list: exact per-module names + (AEON Qwen3.6-35B) and regexes both match; an export that only skipped in_proj_* + (dense Qwen3.6-27B) must NOT match.""" + from freetoken.models.qwen3_5_moe.config import _ct_linear_attn_ignored + + moe = _hf_config(quant=_ct_nvfp4_quant()) + assert _ct_linear_attn_ignored(moe) + # regex form of the same skip + regex_quant = dict(_ct_nvfp4_quant()) + regex_quant["ignore"] = ["re:.*linear_attn\\..*", "re:.*mlp\\.gate.*", "lm_head"] + assert _ct_linear_attn_ignored(_hf_config(quant=regex_quant)) + # only in_proj_* skipped -> out_proj stayed quantized (Qwen3.6-27B) + only_in = dict(_ct_nvfp4_quant()) + only_in["ignore"] = [x for x in only_in["ignore"] if "out_proj" not in x] + assert not _ct_linear_attn_ignored(_hf_config(quant=only_in)) + # no ignore list at all -> everything Linear quantized + no_ignore = dict(_ct_nvfp4_quant()) + no_ignore.pop("ignore", None) + assert not _ct_linear_attn_ignored(_hf_config(quant=no_ignore)) + + +def test_parse_config_ct_moe_gdn_quantized_when_not_ignored(): + """A CT export that quantizes the GDN (no linear_attn in its ignore list, e.g. the + dense Qwen3.6-27B) keeps out_proj native FP4 (gdn_quant == "nvfp4").""" + quant = dict(_ct_nvfp4_quant()) + quant["ignore"] = [x for x in quant["ignore"] if "linear_attn" not in x] + cfg = parse_config(_hf_config(quant=quant)) + assert cfg.attn_quant == "nvfp4" + assert cfg.gdn_quant == "nvfp4" + + +def test_parse_config_ct_format_only_moe(): + """Format-only export (Kwaipilot-KAT-Coder): no config_groups/ignore, just + ``format: nvfp4-pack-quantized`` + ``recipe: all,gdn:fp8,-router``. Everything is + NVFP4 except the GDN (fp8 -> gdn_quant none) and lm_head (kept NVFP4).""" + cfg = parse_config(_hf_config(quant=_ct_format_only_quant())) + assert cfg.expert_quant == "nvfp4" + assert cfg.attn_quant == "nvfp4" + assert cfg.dense_quant == "nvfp4" + assert cfg.gdn_quant == "none" + assert cfg.lm_head_quant == "nvfp4" + + +def test_ct_recipe_gdn_override(): + """The GDN-out_proj NVFP4 signal honors the ``recipe`` string: ``gdn:fp8`` leaves it + non-NVFP4; a gdn:nvfp4 token or a recipe that quantizes the GDN keeps it NVFP4.""" + from freetoken.models.qwen3_5_moe.config import _ct_gdn_nvfp4 + + base = dict(_ct_format_only_quant()) + assert not _ct_gdn_nvfp4(_hf_config(quant=base)) # gdn:fp8 + fp4 = dict(base) + fp4["recipe"] = "all,gdn:nvfp4,-router" + assert _ct_gdn_nvfp4(_hf_config(quant=fp4)) + no_recipe = dict(base) + no_recipe.pop("recipe", None) + assert _ct_gdn_nvfp4(_hf_config(quant=no_recipe)) # default: nvfp4 + + +def test_parse_config_ct_dense_keeps_experts_none(): + """Dense compressed-tensors export (e.g. Qwen3.6-27B, num_experts==0) must keep + expert_quant "none" -- the dense reader owns all of its weights.""" + cfg = parse_config(_hf_config(num_experts=0, quant=_ct_nvfp4_quant())) + assert cfg.expert_quant == "none" + assert cfg.dense_quant == "nvfp4" # dense MLP is still native FP4 + assert cfg.num_experts == 0 + assert not cfg.moe_enabled + + +def test_parse_config_ct_mixed_fp8_experts_keep_none(): + """A mixed export whose routed experts are fp8 (not nvfp4) must not route them + into the NVFP4 bank loader: the expert-targeted config group decides.""" + cfg = parse_config(_hf_config(quant=_ct_mixed_quant())) + assert cfg.expert_quant == "none" + + +# ----------------------------------------------------------------------------------- +# expert source spec selection +# ----------------------------------------------------------------------------------- + + +def test_ct_expert_source_spec(): + from freetoken.models.qwen3_5_moe.weight import _NVFP4_CT_SOURCE_SPEC, _NVFP4_SOURCE_SPEC + + ct = _NVFP4_CT_SOURCE_SPEC + m = ct.key_pattern.match( + "model.language_model.layers.5.mlp.experts.7.gate_proj.weight_packed" + ) + assert m and m.group("kind") == "weight_packed" + assert ct.kind_map["weight_packed"] == "weight" + assert ct.kind_map["weight_global_scale"] == "weight_scale_2" + assert ct.global_reciprocal + # W4A16 serving never consumes the calibrated activation scale. + assert ct.key_pattern.match( + "model.language_model.layers.5.mlp.experts.7.gate_proj.input_global_scale" + ) is None + assert _NVFP4_SOURCE_SPEC.kind_map is None + assert not _NVFP4_SOURCE_SPEC.global_reciprocal + + +def test_select_expert_source_spec(monkeypatch): + import freetoken.models.qwen3_5_moe.weight as w + from freetoken.models.qwen3_5_moe.weight import ( + _NVFP4_CT_SOURCE_SPEC, + _NVFP4_SOURCE_SPEC, + _select_expert_source_spec, + ) + + monkeypatch.setattr(w, "_stacked_expert_keying", lambda _p: False) + monkeypatch.setattr( + w, "cached_load_hf_config", + lambda _p: _Cfg({"quantization_config": {"quant_method": "compressed-tensors"}}), + ) + assert _select_expert_source_spec("x") is _NVFP4_CT_SOURCE_SPEC + monkeypatch.setattr( + w, "cached_load_hf_config", + lambda _p: _Cfg({"quantization_config": {"quant_algo": "NVFP4"}}), + ) + assert _select_expert_source_spec("x") is _NVFP4_SOURCE_SPEC + monkeypatch.setattr(w, "cached_load_hf_config", lambda _p: _Cfg({})) + assert _select_expert_source_spec("x") is _NVFP4_SOURCE_SPEC + + +def test_ct_stacked_expert_spec(): + """The stacked (per-layer) CT spec matches ``...experts.gate_up_proj.weight_packed`` + (one tensor per layer) but NOT the per-expert keying, and vice versa.""" + from freetoken.models.qwen3_5_moe.weight import ( + _NVFP4_CT_SOURCE_SPEC, + _NVFP4_CT_STACKED_SOURCE_SPEC, + ) + + stacked = _NVFP4_CT_STACKED_SOURCE_SPEC + m = stacked.key_pattern.match( + "model.language_model.layers.5.mlp.experts.gate_up_proj.weight_packed" + ) + assert m and m.group("proj") == "gate_up_proj" and m.group("kind") == "weight_packed" + assert stacked.stacked + assert stacked.kind_map["weight_global_scale"] == "weight_scale_2" + # per-expert keying must NOT match the stacked spec, and stacked not match per-expert + assert stacked.key_pattern.match( + "model.language_model.layers.5.mlp.experts.7.gate_proj.weight_packed" + ) is None + assert _NVFP4_CT_SOURCE_SPEC.key_pattern.match( + "model.language_model.layers.5.mlp.experts.gate_up_proj.weight_packed" + ) is None + + +def test_stacked_expert_keying_detection(tmp_path): + """``_stacked_expert_keying`` tells the per-expert from the stacked layout by probing + the weight map (down_proj key first must still resolve stacked).""" + import torch + + import freetoken.models.qwen3_5_moe.weight as w + + p = "model.language_model." + tensors = { + p + "layers.0.mlp.experts.down_proj.weight_packed": torch.zeros(4, 4, dtype=torch.uint8), + } + _write_single_file(tmp_path, tensors) + assert w._stacked_expert_keying(str(tmp_path)) + tensors2 = { + p + "layers.0.mlp.experts.0.down_proj.weight_packed": torch.zeros(4, 4, dtype=torch.uint8), + } + _write_single_file(tmp_path, tensors2) + assert not w._stacked_expert_keying(str(tmp_path)) + + +# ----------------------------------------------------------------------------------- +# single-file, no-index bank loading +# ----------------------------------------------------------------------------------- + + +def _ct_expert_tensors(p: str, layer: int, expert: int, H: int, I: int, + g_scale, up_scale, down_scale) -> dict: + import torch + + fp8 = torch.float8_e4m3fn + base = f"{p}layers.{layer}.mlp.experts.{expert}." + return { + base + "gate_proj.weight_packed": torch.randint(0, 256, (I, H // 2), dtype=torch.uint8), + base + "gate_proj.weight_scale": torch.randn(I, H // 16).abs().to(fp8), + base + "gate_proj.weight_global_scale": torch.tensor([g_scale]), + base + "gate_proj.input_global_scale": torch.tensor([1.0]), + base + "up_proj.weight_packed": torch.randint(0, 256, (I, H // 2), dtype=torch.uint8), + base + "up_proj.weight_scale": torch.randn(I, H // 16).abs().to(fp8), + base + "up_proj.weight_global_scale": torch.tensor([up_scale]), + base + "up_proj.input_global_scale": torch.tensor([1.0]), + base + "down_proj.weight_packed": torch.randint(0, 256, (H, I // 2), dtype=torch.uint8), + base + "down_proj.weight_scale": torch.randn(H, I // 16).abs().to(fp8), + base + "down_proj.weight_global_scale": torch.tensor([down_scale]), + base + "down_proj.input_global_scale": torch.tensor([1.0]), + } + + +def _write_single_file(tmp_path, tensors: dict) -> None: + """Write one ``model.safetensors`` with NO index -- the llm-compressor single-file + layout (the checkpoints that broke the old unconditional index read).""" + import safetensors.torch + + safetensors.torch.save_file(tensors, str(tmp_path / "model.safetensors")) + + +def test_load_nvfp4_expert_source_banks_ct_single_file(tmp_path): + """The CT spec + the no-index weight-map fallback must place every (layer, expert) + into the six native banks with the reciprocal globals. Fails on the old code + (unconditional ``model.safetensors.index.json`` open -> FileNotFoundError).""" + import torch + from types import SimpleNamespace + + from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks + from freetoken.models.qwen3_5_moe.weight import _NVFP4_CT_SOURCE_SPEC + + L, E, H, I = 2, 2, 32, 32 + p = "model.language_model." + tensors = {} + for li in range(L): + for ei in range(E): + tensors |= _ct_expert_tensors(p, li, ei, H, I, 4.0, 2.0, 3.0) + _write_single_file(tmp_path, tensors) + + cfg = SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, + num_layers=L, num_moe_layers=L, first_k_dense_replace=0) + collected: dict[int, dict[str, torch.Tensor]] = {} + + class _Sink: + def __call__(self, layer_id: int, banks: dict) -> None: + collected[layer_id] = {k: v.tensor.clone() for k, v in banks.items()} + + banks = load_nvfp4_expert_source_banks( + str(tmp_path), cfg, _NVFP4_CT_SOURCE_SPEC, + drop_page_cache=lambda _path: None, primary=True, layer_sink=_Sink(), + ) + assert banks.keys() == { + "gate_up_packed", "gate_up_scale", "gate_up_global", + "down_packed", "down_scale", "down_global", + } + assert set(collected) == {0, 1} + gu = collected[0] + assert gu["gate_up_packed"].shape == (E, 2 * I, H // 2) + assert gu["gate_up_scale"].shape == (E, 2 * I, H // 16) + assert gu["gate_up_global"].shape == (E, 2 * I) + assert gu["down_packed"].shape == (E, H, I // 2) + assert gu["down_scale"].shape == (E, H, I // 16) + assert gu["down_global"].shape == (E, H) + + # gate fills the fused rows [0:I), up the rows [I:2I) -- exact placement. + src = tensors + e0 = f"{p}layers.0.mlp.experts.0." + assert torch.equal(gu["gate_up_packed"][0, :I], src[e0 + "gate_proj.weight_packed"]) + assert torch.equal(gu["gate_up_packed"][0, I:], src[e0 + "up_proj.weight_packed"]) + assert torch.equal(gu["down_packed"][0], src[e0 + "down_proj.weight_packed"]) + # globals are the reciprocal of the stored quant-side scales, per output row. + assert gu["gate_up_global"][0, 0].item() == pytest.approx(0.25) # 1 / 4.0 + assert gu["gate_up_global"][0, I].item() == pytest.approx(0.5) # 1 / 2.0 + assert gu["down_global"][0, 0].item() == pytest.approx( + torch.tensor(1 / 3).to(torch.float16).item() # fp16-rounded reciprocal + ) + + +def _ct_stacked_expert_tensors(p: str, layer: int, E: int, H: int, I: int, + g_scale: float, down_scale: float | None = None) -> dict: + """One layer's stacked (per-layer) expert tensors, real llm-compressor layout.""" + import torch + + fp8 = torch.float8_e4m3fn + if down_scale is None: + down_scale = g_scale + base = f"{p}layers.{layer}.mlp.experts." + return { + base + "gate_up_proj.weight_packed": torch.randint(0, 256, (E * 2 * I, H // 2), dtype=torch.uint8), + base + "gate_up_proj.weight_scale": torch.randn(E * 2 * I, H // 16).abs().to(fp8), + base + "gate_up_proj.weight_global_scale": torch.tensor([g_scale]), + base + "down_proj.weight_packed": torch.randint(0, 256, (E * H, I // 2), dtype=torch.uint8), + base + "down_proj.weight_scale": torch.randn(E * H, I // 16).abs().to(fp8), + base + "down_proj.weight_global_scale": torch.tensor([down_scale]), + } + + +def test_load_nvfp4_stacked_expert_sources(tmp_path): + """The stacked loader reshapes each [E*rows, cols] tensor into [E, rows, cols] and + broadcasts the layer-global (reciprocated) scale into the *_global banks.""" + import torch + from types import SimpleNamespace + + from freetoken.models.nvfp4_banks import load_nvfp4_stacked_expert_sources + from freetoken.models.qwen3_5_moe.weight import _NVFP4_CT_STACKED_SOURCE_SPEC + + L, E, H, I = 2, 2, 32, 32 + p = "model.language_model." + tensors = {} + for li in range(L): + tensors |= _ct_stacked_expert_tensors(p, li, E, H, I, 4.0, down_scale=8.0) + _write_single_file(tmp_path, tensors) + + cfg = SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, + num_layers=L, num_moe_layers=L, first_k_dense_replace=0) + collected: dict[int, dict[str, torch.Tensor]] = {} + + class _Sink: + def __call__(self, layer_id: int, banks: dict) -> None: + collected[layer_id] = {k: v.tensor.clone() for k, v in banks.items()} + + banks = load_nvfp4_stacked_expert_sources( + str(tmp_path), cfg, _NVFP4_CT_STACKED_SOURCE_SPEC, + drop_page_cache=lambda _path: None, primary=True, layer_sink=_Sink(), + ) + assert banks.keys() == { + "gate_up_packed", "gate_up_scale", "gate_up_global", + "down_packed", "down_scale", "down_global", + } + assert set(collected) == {0, 1} + gu = collected[0] + assert gu["gate_up_packed"].shape == (E, 2 * I, H // 2) + assert gu["gate_up_scale"].shape == (E, 2 * I, H // 16) + assert gu["gate_up_global"].shape == (E, 2 * I) + assert gu["down_packed"].shape == (E, H, I // 2) + assert gu["down_scale"].shape == (E, H, I // 16) + assert gu["down_global"].shape == (E, H) + + src = tensors + l0 = f"{p}layers.0.mlp.experts." + assert torch.equal( + gu["gate_up_packed"], + src[l0 + "gate_up_proj.weight_packed"].view(E, 2 * I, H // 2), + ) + assert torch.equal( + gu["down_packed"], + src[l0 + "down_proj.weight_packed"].view(E, H, I // 2), + ) + # the layer-globals are reciprocated and broadcast to every expert row, per proj. + assert gu["gate_up_global"][0, 0].item() == pytest.approx(0.25) # 1 / 4.0 + assert gu["gate_up_global"].unique().numel() == 1 + assert gu["down_global"][0, 0].item() == pytest.approx( + torch.tensor(1 / 8).to(torch.float16).item() # 1 / 8.0, fp16-rounded + ) + assert gu["down_global"].unique().numel() == 1 + + +def test_dequant_fp8_weight_per_row(): + """Per-tensor fp8 (per-output-row scale, e.g. a ``gdn:fp8`` recipe) must broadcast + the row scale, not fail on the row count vs 1.""" + import torch + + from freetoken.models.qwen3_5_moe.weight import _dequant_fp8_weight + + w8 = torch.randint(-128, 128, (6, 4), dtype=torch.int8).to(torch.float8_e4m3fn) + rows = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + out = _dequant_fp8_weight(w8, rows) + assert out.shape == (6, 4) + assert out.dtype == torch.bfloat16 + assert torch.allclose(out[2].float(), (w8[2].float() * 3.0), atol=0.5) + # a scalar scale still works + out2 = _dequant_fp8_weight(w8, torch.tensor(2.0)) + assert out2.shape == (6, 4) + assert torch.allclose(out2[0].float(), (w8[0].float() * 2.0), atol=0.5) + + +# ----------------------------------------------------------------------------------- +# dense pass: expert exclusion + shared-expert gate/up fusion +# ----------------------------------------------------------------------------------- + + +def test_iter_weights_ct_moe_dense_pass(tmp_path, monkeypatch): + """compressed-tensors MoE dense pass: routed experts must NOT leak into the dense + weights, and the shared expert's gate/up must fuse into the native ``gate_up_proj`` + (W4A16, per-part globals) the model's Nvfp4DenseColMerged expects. Fails on the old + code: experts were emitted as dense weights and shared_expert gate/up stayed split.""" + import torch + + import freetoken.models.qwen3_5_moe.weight as w + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.models.qwen3_5_moe.weight import iter_weights + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + L, E, H, I, Q, KV, SHARED, VOCAB = 2, 4, 32, 16, 32, 16, 8, 32 + p = "model.language_model." + fp8 = torch.float8_e4m3fn + tensors = {} + + def nvfp4(base: str, o: int, i: int) -> None: + tensors[base + ".weight_packed"] = torch.randint(0, 256, (o, i // 2), dtype=torch.uint8) + tensors[base + ".weight_scale"] = torch.randn(o, i // 16).abs().to(fp8) + tensors[base + ".weight_global_scale"] = torch.tensor([2.0]) + tensors[base + ".input_global_scale"] = torch.tensor([1.0]) + + for li in range(L): + lp = f"{p}layers.{li}." + nvfp4(lp + "self_attn.q_proj", Q, H) + nvfp4(lp + "self_attn.k_proj", KV, H) + nvfp4(lp + "self_attn.v_proj", KV, H) + nvfp4(lp + "self_attn.o_proj", H, Q) + nvfp4(lp + "mlp.shared_expert.gate_proj", SHARED, H) + nvfp4(lp + "mlp.shared_expert.up_proj", SHARED, H) + nvfp4(lp + "mlp.shared_expert.down_proj", H, SHARED) + tensors[lp + "mlp.gate.weight"] = torch.randn(E, H, dtype=torch.bfloat16) + tensors[lp + "mlp.shared_expert_gate.weight"] = torch.randn(1, H, dtype=torch.bfloat16) + tensors[lp + "input_layernorm.weight"] = torch.randn(H, dtype=torch.bfloat16) + tensors[lp + "post_attention_layernorm.weight"] = torch.randn(H, dtype=torch.bfloat16) + for ei in range(E): + tensors |= _ct_expert_tensors(p, li, ei, H, I, 2.0, 2.0, 2.0) + tensors[p + "embed_tokens.weight"] = torch.randn(VOCAB, H, dtype=torch.bfloat16) + tensors[p + "norm.weight"] = torch.randn(H, dtype=torch.bfloat16) + tensors["lm_head.weight"] = torch.randn(VOCAB, H, dtype=torch.bfloat16) + _write_single_file(tmp_path, tensors) + + monkeypatch.setattr(w, "cached_load_hf_config", lambda _p: _hf_config(num_layers=L, num_experts=E, quant=_ct_nvfp4_quant())) + + loaded = dict( + iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True) + ) + # routed experts never reach the dense pass. + assert not any(".mlp.experts." in k for k in loaded) + + gu = "model.layers.0.mlp.shared_expert.gate_up_proj" + lp = f"{p}layers.0." + assert loaded[gu + ".weight"].shape == (2 * SHARED, H // 2) + assert loaded[gu + ".weight_scale"].shape == (2 * SHARED, H // 16) + assert loaded[gu + ".weight_global"].shape == (2 * SHARED,) + assert torch.equal(loaded[gu + ".weight"][:SHARED], + tensors[lp + "mlp.shared_expert.gate_proj.weight_packed"]) + assert torch.equal(loaded[gu + ".weight"][SHARED:], + tensors[lp + "mlp.shared_expert.up_proj.weight_packed"]) + assert loaded[gu + ".weight_global"][0].item() == pytest.approx(0.5) # 1 / 2.0 + + dp = "model.layers.0.mlp.shared_expert.down_proj" + assert torch.equal(loaded[dp + ".weight"], + tensors[lp + "mlp.shared_expert.down_proj.weight_packed"]) + + qkv = "model.layers.0.self_attn.qkv_proj" + assert loaded[qkv + ".weight"].shape == (Q + 2 * KV, H // 2) + + +def test_iter_weights_ct_format_only_mixed(tmp_path, monkeypatch): + """Kwaipilot-style export (format-only + recipe): the GDN is per-tensor fp8 + (dequantized to bf16, fused in_proj), stacked experts never reach the dense pass, + embed_tokens dequantizes to bf16, lm_head stays native NVFP4 (the model wants it), + and the full-attn/shared-expert projections stay native.""" + import torch + + import freetoken.models.qwen3_5_moe.weight as w + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.models.qwen3_5_moe.weight import iter_weights + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + # _dequant_nvfp4_weight runs on CUDA; fake it (CPU bf16 of the right shape). + monkeypatch.setattr( + w, "_dequant_nvfp4_weight", + lambda weight, scale, g: torch.zeros( + weight.shape[0], weight.shape[1] * 2, dtype=torch.bfloat16 + ), + ) + + L, E, H, I, Q, KV, SHARED, VOCAB = 2, 4, 32, 16, 32, 16, 8, 64 + CONV, VAL = 24, 8 # conv_dim = 2*num_k*key_dim + value_dim; value_dim = num_v*value_head_dim + p = "model.language_model." + fp8 = torch.float8_e4m3fn + tensors = {} + + def nvfp4(base: str, o: int, i: int) -> None: + tensors[base + ".weight_packed"] = torch.randint(0, 256, (o, i // 2), dtype=torch.uint8) + tensors[base + ".weight_scale"] = torch.randn(o, i // 16).abs().to(fp8) + tensors[base + ".weight_global_scale"] = torch.tensor([2.0]) + tensors[base + ".input_global_scale"] = torch.tensor([1.0]) + + def pt_fp8(base: str, o: int, i: int) -> None: + # per-tensor fp8: fp8 weight + per-output-row F32 scale + tensors[base + ".weight"] = torch.randint(-128, 128, (o, i), dtype=torch.int8).to(fp8) + tensors[base + ".weight_scale"] = torch.rand(o).abs() + 0.5 + + for li in range(L): + lp = f"{p}layers.{li}." + # GDN is fp8 (recipe gdn:fp8): in_proj_* + out_proj. + pt_fp8(lp + "linear_attn.in_proj_qkv", CONV, H) + pt_fp8(lp + "linear_attn.in_proj_z", VAL, H) + pt_fp8(lp + "linear_attn.in_proj_b", 2, H) + pt_fp8(lp + "linear_attn.in_proj_a", 2, H) + pt_fp8(lp + "linear_attn.out_proj", VAL, H) + tensors[lp + "linear_attn.conv1d.weight"] = torch.randn(CONV, 1, 2, dtype=torch.bfloat16) + tensors[lp + "linear_attn.A_log"] = torch.randn(2, dtype=torch.float32) + tensors[lp + "linear_attn.dt_bias"] = torch.randn(2, dtype=torch.float32) + tensors[lp + "linear_attn.norm.weight"] = torch.randn(4, dtype=torch.bfloat16) + # full-attn projections stay NVFP4. + nvfp4(lp + "self_attn.q_proj", Q, H) + nvfp4(lp + "self_attn.k_proj", KV, H) + nvfp4(lp + "self_attn.v_proj", KV, H) + nvfp4(lp + "self_attn.o_proj", H, Q) + tensors[lp + "self_attn.q_norm.weight"] = torch.randn(4, dtype=torch.bfloat16) + tensors[lp + "self_attn.k_norm.weight"] = torch.randn(4, dtype=torch.bfloat16) + # shared expert NVFP4. + nvfp4(lp + "mlp.shared_expert.gate_proj", SHARED, H) + nvfp4(lp + "mlp.shared_expert.up_proj", SHARED, H) + nvfp4(lp + "mlp.shared_expert.down_proj", H, SHARED) + tensors[lp + "mlp.gate.weight"] = torch.randn(E, H, dtype=torch.bfloat16) + tensors[lp + "mlp.shared_expert_gate.weight"] = torch.randn(1, H, dtype=torch.bfloat16) + tensors[lp + "input_layernorm.weight"] = torch.randn(H, dtype=torch.bfloat16) + tensors[lp + "post_attention_layernorm.weight"] = torch.randn(H, dtype=torch.bfloat16) + # stacked (per-layer) experts -- must be skipped by the dense pass. + tensors |= _ct_stacked_expert_tensors(p, li, E, H, I, 2.0) + # embed_tokens NVFP4 (dequantized), lm_head NVFP4 (kept native), norm bf16. + nvfp4(p + "embed_tokens", VOCAB, H) + nvfp4("lm_head", VOCAB, H) + tensors[p + "norm.weight"] = torch.randn(H, dtype=torch.bfloat16) + _write_single_file(tmp_path, tensors) + + monkeypatch.setattr( + w, "cached_load_hf_config", + lambda _p: _hf_config(num_layers=L, num_experts=E, quant=_ct_format_only_quant()), + ) + + loaded = dict( + iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True) + ) + # stacked experts never reach the dense pass. + assert not any(".mlp.experts." in k for k in loaded) + + lp = f"{p}layers.0." + # GDN: bf16 fused in_proj + bf16 out_proj (dequantized from fp8). + in_proj = "model.layers.0.linear_attn.in_proj" + assert loaded[in_proj + ".weight"].dtype == torch.bfloat16 + assert loaded[in_proj + ".weight"].shape == (CONV + VAL + 4, H) + assert loaded["model.layers.0.linear_attn.out_proj.weight"].dtype == torch.bfloat16 + assert loaded["model.layers.0.linear_attn.out_proj.weight"].shape == (VAL, H) + + # embed_tokens dequantized to bf16. + assert loaded["model.embed_tokens.weight"].dtype == torch.bfloat16 + assert loaded["model.embed_tokens.weight"].shape == (VOCAB, H) + + # lm_head kept native NVFP4 (lm_head_quant == "nvfp4"). + lh = "lm_head" + assert loaded[lh + ".weight"].dtype == torch.uint8 + assert loaded[lh + ".weight"].shape == (VOCAB, H // 2) + assert loaded[lh + ".weight_scale"].shape == (VOCAB, H // 16) + assert loaded[lh + ".weight_global"].shape == (VOCAB,) + + # full-attn qkv + shared expert native. + qkv = "model.layers.0.self_attn.qkv_proj" + assert loaded[qkv + ".weight"].dtype == torch.uint8 + gu = "model.layers.0.mlp.shared_expert.gate_up_proj" + assert loaded[gu + ".weight"].dtype == torch.uint8 + assert loaded[gu + ".weight"].shape == (2 * SHARED, H // 2) \ No newline at end of file