fix(qwen3_5_moe): detect shared-expert quant for modelopt mixed checkpoints - #294
Open
JIAQI13 wants to merge 1 commit into
Open
fix(qwen3_5_moe): detect shared-expert quant for modelopt mixed checkpoints#294JIAQI13 wants to merge 1 commit into
JIAQI13 wants to merge 1 commit into
Conversation
…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
This was referenced Aug 30, 2026
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 modeloptMIXED_PRECISIONcheckpoint: the routed experts are NVFP4, but the shared expert and attn/GDN projections are per-tensor FP8 (quant_algo: "FP8"inquantized_layers, fp8-e4m3 weight + scalarweight_scale+input_scale, noweight_scale_2).parse_configassumed the shared expert always matched the routed experts:So
dense_quantwas forced to"nvfp4", the shared expert was built asNvfp4DenseColMerged/Nvfp4DenseLinear, and weight loading crashed withbecause an FP8 shared expert has no
weight_scale_2/weight_globalbuffers. vLLM/SGLang serve the same checkpoint fine, so this is purely the FreeToken mixed-precision path.Changes
models/qwen3_5_moe/config.py_shared_expert_quant(): probe the modeloptquantized_layersmap for the shared-expertquant_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_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_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 onmain(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)
e6e2b26on top ofmain(4b94bdc)apodex/Apodex-1.1-mini-NVFP4ft serve --model-path <Apodex-1.1-mini-NVFP4> --moe-backend offload --moe-cache-auto --port 8899 --host 127.0.0.1main: crashes at weight load with theKeyErrorabove.API server is ready,/v1/chat/completionsreturns 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.Notes
MIXED_PRECISIONformat. A compressed-tensorsmixed-precisioncheckpoint (128x128 block-FP8 shared expert + NVFP4 experts, e.g. some Qwen3.6-35B-A3B releases) usesconfig_groupsrather thanquantized_layersand is not handled here.