From e6e2b261f63ee4b6bbc8fcd0e94421b8366894df Mon Sep 17 00:00:00 2001 From: JIAQI13 <35429081+JIAQI13@users.noreply.github.com> Date: Sun, 30 Aug 2026 10:41:45 -0700 Subject: [PATCH] fix(qwen3_5_moe): detect shared-expert quant for modelopt mixed checkpoints MoE checkpoints quantize routed experts and the shared expert independently. modelopt MIXED_PRECISION checkpoints such as Apodex-1.1-mini-NVFP4 ship NVFP4 experts with a per-tensor FP8 shared expert, but parse_config forced dense_quant=nvfp4 whenever the experts were NVFP4. That built Nvfp4DenseColMerged for the shared expert and crashed at weight load with KeyError on model.layers.0.mlp.shared_expert.gate_up_proj.weight (an FP8 shared expert has no weight_scale_2 / weight_global buffers). - add _shared_expert_quant(): probe quantized_layers for the shared expert quant_algo; pure NVFP4 checkpoints without a per-layer map keep the native W4A16 default - _iter_weights_attn_fp8: keep native fp8 (W8A16) only for attn/GDN projections; FP8 shared-expert weights fall through to bf16 dequant plus the existing gate/up fusion - tests: shared-expert detection for no-map / FP8 / NVFP4 mixed cases Fixes #183 --- python/freetoken/models/qwen3_5_moe/config.py | 38 +++++++- python/freetoken/models/qwen3_5_moe/weight.py | 9 +- .../test_qwen3_5_moe_shared_expert_quant.py | 88 +++++++++++++++++++ 3 files changed, 133 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_qwen3_5_moe_shared_expert_quant.py diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index 2ac4b607f..12d3b510a 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/python/freetoken/models/qwen3_5_moe/config.py @@ -104,6 +104,27 @@ def _dense_mlp_quant(hf_config: Any) -> str: return "none" +def _shared_expert_quant(hf_config: Any) -> str: + """Quant format of the MoE shared expert (``.mlp.shared_expert.gate_proj``). + + Returns ``"nvfp4"`` when it is packed FP4 (kept native W4A16) or ``"none"`` when it + is per-tensor FP8 / bf16 (dequantized to bf16 at load). Pure-NVFP4 checkpoints have + no per-layer ``quantized_layers`` map, so the shared expert is native FP4; modelopt + MIXED_PRECISION checkpoints tag it explicitly (FP8 on Apodex-1.1-mini even though the + routed experts are NVFP4).""" + get = _quant_accessor(hf_config) + if get is None: + return "nvfp4" + layers = get("quantized_layers") or {} + if not isinstance(layers, dict): + return "nvfp4" + for name, spec in layers.items(): + if name.endswith(".mlp.shared_expert.gate_proj"): + algo = str((spec or {}).get("quant_algo", "")).lower() + return "nvfp4" if "fp4" in algo else "none" + return "nvfp4" + + def _attn_quant(hf_config: Any) -> str: """Per-tensor FP8 on the *dense* attention/GDN projections. The modelopt ``MIXED_PRECISION`` checkpoints tag ``self_attn.{q,k,v,o}_proj`` and @@ -176,7 +197,22 @@ def parse_config(hf_config: Any) -> ModelConfig: # NVFP4. The lm_head is detected separately (only the mixed checkpoint quantizes it). # MoE-NVFP4 keeps the shared_expert dense MLP native FP4 (expert_quant=="nvfp4"); a dense # (non-MoE) modelopt checkpoint instead tags the bare .mlp.{gate,up,down}_proj as NVFP4. - dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config) + if expert_quant == "nvfp4": + # In a MoE checkpoint the routed experts and the shared expert are quantized + # independently. Most NVFP4 MoE checkpoints keep the shared expert packed FP4 + # (native W4A16), but modelopt MIXED_PRECISION variants (e.g. Apodex-1.1-mini: + # NVFP4 experts + per-tensor FP8 shared expert) do not -- forcing W4A16 there + # builds Nvfp4DenseColMerged and then fails looking for weight_scale_2 / + # weight_global buffers that do not exist on an FP8 shared expert. Dense + # (num_experts == 0) NVFP4 checkpoints keep the .mlp projections packed FP4 + # like the experts and stay on the native path. + dense_quant = ( + _shared_expert_quant(hf_config) + if int(getattr(text, "num_experts", 0) or 0) > 0 + else "nvfp4" + ) + else: + dense_quant = _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 diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..83c2dd8dc 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -474,7 +474,14 @@ def _iter_weights_attn_fp8( raw_base = raw_name[: -len(".weight")] has_s2 = raw_base + ".weight_scale_2" in keyset has_s = raw_base + ".weight_scale" in keyset - if has_s and not has_s2: # per-tensor FP8 dense projection + # Native FP8 (W8A16) only for the attn/GDN projections: those have + # fp8 linears in the model. Any other per-tensor FP8 .weight (the + # MoE shared_expert on a mixed checkpoint whose experts are NVFP4, + # e.g. Apodex-1.1-mini) has no fp8 module -- let it fall through to + # bf16 dequant + the shared-expert gate/up fusion below. + if has_s and not has_s2 and ( + ".self_attn." in base or ".linear_attn." in base + ): w = f.get_tensor(raw_name) # fp8-e4m3, kept verbatim sc = f.get_tensor(raw_base + ".weight_scale") # modelopt's calibrated activation scale: kept (not dropped with the diff --git a/tests/models/test_qwen3_5_moe_shared_expert_quant.py b/tests/models/test_qwen3_5_moe_shared_expert_quant.py new file mode 100644 index 000000000..54345b3a2 --- /dev/null +++ b/tests/models/test_qwen3_5_moe_shared_expert_quant.py @@ -0,0 +1,88 @@ +"""Shared-expert quant detection for MoE checkpoints (issue #183). + +modelopt MIXED_PRECISION checkpoints can quantize the routed experts and the +shared expert independently: Apodex-1.1-mini-NVFP4 ships NVFP4 experts with a +per-tensor FP8 shared expert. The old code assumed the shared expert matched +the experts (``dense_quant = "nvfp4"`` whenever experts are NVFP4), built an +``Nvfp4DenseColMerged`` module, and crashed at load with +``KeyError: 'model.layers.0.mlp.shared_expert.gate_up_proj.weight'`` because an +FP8 shared expert has no ``weight_scale_2`` / ``weight_global`` buffers. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from freetoken.models.qwen3_5_moe.config import _shared_expert_quant + + +def _hf(quantization_config): + return SimpleNamespace(quantization_config=quantization_config) + + +def test_no_quant_config_defaults_to_nvfp4(): + # Pure NVFP4 checkpoints ship without a quant_config accessor; the shared + # expert is packed FP4 like the experts and stays native W4A16. + assert _shared_expert_quant(SimpleNamespace()) == "nvfp4" + + +def test_modelopt_nvfp4_without_layer_map_defaults_to_nvfp4(): + cfg = _hf({"quant_algo": "NVFP4"}) + assert _shared_expert_quant(cfg) == "nvfp4" + + +def test_modelopt_mixed_fp8_shared_expert_is_dequantized(): + # Apodex-1.1-mini-NVFP4 layout: NVFP4 routed experts, per-tensor FP8 + # shared expert tagged explicitly in quantized_layers. Must NOT take the + # native W4A16 path (that module needs weight_scale_2 / weight_global). + cfg = _hf( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.language_model.layers.0.mlp.experts": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + "model.language_model.layers.0.mlp.shared_expert.gate_proj": { + "quant_algo": "FP8" + }, + "model.language_model.layers.0.mlp.shared_expert.up_proj": { + "quant_algo": "FP8" + }, + "model.language_model.layers.0.mlp.shared_expert.down_proj": { + "quant_algo": "FP8" + }, + }, + } + ) + assert _shared_expert_quant(cfg) == "none" + + +def test_modelopt_mixed_nvfp4_shared_expert_stays_native(): + cfg = _hf( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.language_model.layers.0.mlp.shared_expert.gate_proj": { + "quant_algo": "W4A16_NVFP4" + }, + }, + } + ) + assert _shared_expert_quant(cfg) == "nvfp4" + + +def test_mixed_map_without_shared_expert_entry_defaults_to_nvfp4(): + # Per-layer map present but no shared_expert tag: keep the family default + # (MoE NVFP4 checkpoints keep the shared expert native FP4). + cfg = _hf( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.language_model.layers.0.self_attn.q_proj": { + "quant_algo": "FP8" + }, + }, + } + ) + assert _shared_expert_quant(cfg) == "nvfp4"