Qwen3.5 MoE modelopt_mixed loader assumes shared experts use routed-expert NVFP4 quantization

[peterhon168]

Qwen3.5 MoE modelopt_mixed loader assumes shared experts use routed-expert NVFP4 quantization

Summary

FreeToken 0.1.2 fails to load apodex/Apodex-1.1-mini-NVFP4, a Qwen3.5-MoE checkpoint whose routed experts are NVFP4 while its shared-expert projections are per-tensor FP8.

The Qwen3.5-MoE config currently derives dense_quant="nvfp4" from expert_quant=="nvfp4". That allocates NVFP4 shared-expert modules even though the checkpoint's quantized_layers map declares FP8 for *.mlp.shared_expert.{gate,up,down}_proj. Loading then looks for a fused NVFP4 key that cannot exist.

Environment

• FreeToken: 0.1.2
• Python: 3.12.3
• PyTorch: 2.11.0+cu130
• Transformers: 5.15.1
• CUDA runtime/toolkit: 13.0
• GPU: NVIDIA GeForce RTX 3080 Ti 12 GB (Ampere)
• Host: Linux x86_64, dual Xeon E5-2682 v4, 110 GiB RAM
• Model: apodex/Apodex-1.1-mini-NVFP4
• Architecture: Qwen3_5MoeForConditionalGeneration
• Quantization: modelopt_mixed; routed experts NVFP4, shared expert/attention FP8 per-tensor, lm_head BF16

Minimal reproduction

ft serve \
  --model /path/to/Apodex-1.1-mini-NVFP4 \
  --served-model-name Apodex-1.1-mini-NVFP4 \
  --host 127.0.0.1 --port 1919 \
  --max-running-requests 1 \
  --graph 1 \
  --max-seq-len-override 32768 \
  --num-tokens 32768 \
  --kv-reserve-tokens 32768 \
  --max-prefill-length 2048 \
  --memory-ratio 0.80 \
  --cache-type radix \
  --expert-load auto \
  --moe-backend hybrid \
  --moe-cache-auto \
  --moe-prefill-hit-d2d \
  --moe-cpu-threads 32 \
  --nvfp4-backend auto

Actual behavior

All three mixed-FP8 shards are read, then startup fails in the NVFP4 dense loader:

KeyError: 'model.layers.0.mlp.shared_expert.gate_up_proj.weight'

freetoken/kernel/triton/nvfp4_linear.py:866

The checkpoint contains separate FP8 tensors such as:

model.layers.0.mlp.shared_expert.gate_proj.weight
model.layers.0.mlp.shared_expert.up_proj.weight
model.layers.0.mlp.shared_expert.down_proj.weight

with their FP8 scale/input-scale tensors; it does not contain an NVFP4 shared_expert.gate_up_proj.weight.

Expected behavior

For modelopt_mixed, shared-expert quantization should be detected independently from routed-expert quantization. This checkpoint should instantiate FP8 per-tensor shared-expert modules and fuse the FP8 gate/up pair using the existing per-tensor FP8 fusion path.

Tested workaround

A local three-part compatibility patch made the model load and serve correctly:

1. Inspect quantized_layers for *.mlp.shared_expert.* and derive a separate shared-MLP quant mode.
2. Instantiate Fp8PerTensorColMerged for shared_expert.gate_up_proj and Fp8PerTensorLinear for shared_expert.down_proj when that mode is FP8 per-tensor.
3. Add this mapping to the existing FP8 gate/up fusion table:

".mlp.shared_expert.gate_up_proj": (
    ".mlp.shared_expert.gate_proj",
    ".mlp.shared_expert.up_proj",
),

After the patch:

• 32K startup succeeded and completed 20/20 benchmark requests without service errors.
• A 196,608-token KV configuration started successfully.
• Nine retrieval probes from 8K through 188,423 actual prompt tokens returned exact answers.
• The model ran for more than an hour without loader/runtime errors.

Suggested fix/tests

• Do not derive shared-expert quantization solely from routed expert_quant for modelopt_mixed checkpoints.
• Add a loader fixture with NVFP4 routed experts + FP8 shared expert + BF16 lm_head.
• Assert that fused FP8 gate/up weights, per-row scales, and input scales all map correctly.
• Preserve the current pure-NVFP4 fallback for checkpoints without a per-layer quantization map.

[JIAQI13]

This root-cause analysis matches what I found while fixing it. PR #294 implements the fix:

#294

• parse_config now derives the shared-expert quant from the quantized_layers map (_shared_expert_quant()) instead of forcing dense_quant="nvfp4" from expert_quant. A per-tensor FP8 shared expert now resolves to bf16-at-load, so the shared expert is built as the bf16 merged linears instead of Nvfp4DenseColMerged.
• The modelopt mixed-precision dense pass keeps native FP8 (W8A16) only for attn/GDN projections; FP8 shared-expert weights fall through to the existing bf16 dequant + gate/up fusion path.

Verified end-to-end on apodex/Apodex-1.1-mini-NVFP4 (RTX 3080 Laptop, FreeToken 0.1.2): main fails to load with the shared-expert key error; with the PR it converts/loads and serves /v1/chat/completions at ~50–60 tok/s decode, with the expected MoE-offload GPU utilization bursts. Includes detection unit tests (fail on main / pass with the fix).

Note for the sibling format: the same crash family appears with compressed-tensors mixed checkpoints (128×128 block-FP8 shared expert, e.g. the llm-compressor uncensored/abliterated Qwen3.6-35B-A3B exports in #263/#274) — that needs a different detection/routing path and is handled separately in #296.