fix(qwen3_5_moe): detect shared-expert quant for modelopt mixed checkpoints

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

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

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"