Skip to content

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

Open
JIAQI13 wants to merge 1 commit into
FlashML-org:mainfrom
JIAQI13:fix/qwen35-moe-shared-expert-quant
Open

fix(qwen3_5_moe): detect shared-expert quant for modelopt mixed checkpoints#294
JIAQI13 wants to merge 1 commit into
FlashML-org:mainfrom
JIAQI13:fix/qwen35-moe-shared-expert-quant

Conversation

@JIAQI13

@JIAQI13 JIAQI13 commented Aug 30, 2026

Copy link
Copy Markdown

Fixes #164
Fixes #183
Related: #274 (same shared-expert crash family on another mixed-quant uncensored checkpoint; the dense side routes through #296, but that exact checkpoint is untested)

Problem

apodex/Apodex-1.1-mini-NVFP4 (Qwen3.5 MoE A3B, 256 experts) is an nvidia modelopt MIXED_PRECISION checkpoint: the routed experts are NVFP4, but the shared expert and attn/GDN projections are per-tensor FP8 (quant_algo: "FP8" in quantized_layers, fp8-e4m3 weight + scalar weight_scale + input_scale, no weight_scale_2).

parse_config assumed the shared expert always matched the routed experts:

dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config)

So dense_quant was forced to "nvfp4", the shared expert was built as Nvfp4DenseColMerged / Nvfp4DenseLinear, and weight loading crashed with

KeyError: 'model.layers.0.mlp.shared_expert.gate_up_proj.weight'
  at freetoken/kernel/triton/nvfp4_linear.py:866

because an FP8 shared expert has no weight_scale_2 / weight_global buffers. vLLM/SGLang serve the same checkpoint fine, so this is purely the FreeToken mixed-precision path.

Changes

models/qwen3_5_moe/config.py

  • Add _shared_expert_quant(): probe the modelopt quantized_layers map for the shared-expert quant_algo. Returns "nvfp4" when packed FP4 (kept native W4A16) or "none" when FP8/bf16 (dequantized to bf16 at load). Pure-NVFP4 checkpoints have no per-layer map and keep the native FP4 default; dense (num_experts == 0) checkpoints are unaffected.

models/qwen3_5_moe/weight.py

  • In _iter_weights_attn_fp8, the native fp8 (W8A16) path now only applies to attn/GDN projections (.self_attn. / .linear_attn.), which actually have fp8 linears. Any other per-tensor FP8 weight — the MoE shared expert on mixed checkpoints — falls through to the existing bf16 dequant path, where the already-present shared-expert gate/up fusion (shared_expert.{gate,up}_proj -> gate_up_proj) builds the bf16 state dict.

tests/models/test_qwen3_5_moe_shared_expert_quant.py

  • Unit tests for _shared_expert_quant: no quant config, NVFP4 without layer map, mixed with FP8 shared expert, mixed with NVFP4 shared expert, and a map without a shared-expert entry. The FP8 case fails on main (the detection is absent and the checkpoint crashes at load) and passes with this change.

Pure NVFP4, block-FP8, dense NVFP4, and compressed-tensors NVFP4 checkpoints take unchanged code paths.

Tested on (real hardware)

  • GPU: NVIDIA GeForce RTX 3080 Laptop, 16 GB VRAM (sm_86), driver 610.57.04
  • CPU: Intel Core i7-10870H (8 cores), 32 GB RAM + 31 GB swap
  • OS: CachyOS (Linux, x86_64)
  • FreeToken 0.1.2, torch 2.11.0+cu130 (CUDA 13.0)
  • Tested commit: e6e2b26 on top of main (4b94bdc)
  • Checkpoint: apodex/Apodex-1.1-mini-NVFP4
  • Command: ft serve --model-path <Apodex-1.1-mini-NVFP4> --moe-backend offload --moe-cache-auto --port 8899 --host 127.0.0.1
  • main: crashes at weight load with the KeyError above.
  • This branch: model loads, API server is ready, /v1/chat/completions returns correct generations; decode throughput 44–60 tok/s in the engine logs (gen throughput (token/s) lines); GPU utilization shows the expected MoE-offload bursts (~87–100% during active decode, power up to ~114 W vs ~34 W idle); ~14.6/16 GB VRAM used.
  • The new unit tests were run as standalone pytest-style functions (no weights / GPU required).

Notes

  • The FP8 shared expert is dequantized to bf16 at load (W8A16 precision). It is small relative to the 256 routed experts, but a native fp8 shared-expert module would preserve the weight-bandwidth benefit — that is a separate enhancement, not needed to fix the crash.
  • This covers the modelopt MIXED_PRECISION format. A compressed-tensors mixed-precision checkpoint (128x128 block-FP8 shared expert + NVFP4 experts, e.g. some Qwen3.6-35B-A3B releases) uses config_groups rather than quantized_layers and is not handled here.

…points

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 FlashML-org#183
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant