Skip to content

qwen4_exp: load modelopt MIXED_PRECISION (NVFP4 experts + block-FP8 dense) checkpoints - #320

Open
gberasmus87 wants to merge 1 commit into
FlashML-org:mainfrom
gberasmus87:support-mixed-precision-qwen4-checkpoints
Open

qwen4_exp: load modelopt MIXED_PRECISION (NVFP4 experts + block-FP8 dense) checkpoints#320
gberasmus87 wants to merge 1 commit into
FlashML-org:mainfrom
gberasmus87:support-mixed-precision-qwen4-checkpoints

Conversation

@gberasmus87

Copy link
Copy Markdown

What

Adds support for loading Qwen3.8-Flash-Next checkpoints exported with modelopt MIXED_PRECISION quantization - NVFP4 routed experts + 128x128 block-FP8 dense (attention / GDN) weights. The common community build of this shape is lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8, the ~124GB quant that fits the model on a 24GB GPU via the offload backend.

Before this, such a checkpoint fails to load in two places:

  1. Config detection - parse_config only recognises a top-level quant_algo of fp8 or *fp4*. MIXED_PRECISION declares the per-module algo in quantized_layers instead, so expert_quant fell through to "none" and the offload cache tried to load the NVFP4 experts as bf16 banks -> ValueError: Missing MoE expert source layers.
  2. Dense weight loading - the dense attn / GDN projections are stored as 128x128 block-FP8 (.weight_scale_inv), but iter_weights assumed all non-expert weights are bf16 and passed them straight to _try_fuse -> RuntimeError: Promotion for Float8 Types is not supported, attempted to promote Float8_e4m3fn and BFloat16 on the GDN in_proj fuse.

Changes

  • config.py: detect quant_algo == "mixed_precision", read quantized_layers to set expert_quant = "nvfp4" (experts stay native for the offload cache), and keep every dense module bf16.
  • weight.py: _load_maybe_block_fp8 dequantizes a dense .weight to bf16 when a sibling .weight_scale_inv is present (reusing the existing dequant_block_fp8 kernel), and .weight_scale_inv is added to _SCALE_SUFFIXES so the scale tensor is not yielded as a weight. Only block-FP8 is added - it is the only dense quant these checkpoints use.

Testing

Loads and serves lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 (~124GB) on an RTX PRO 4000 Blackwell (24GB VRAM) backed by 172GB system RAM, with --moe-backend hybrid - the offload backend holds the ~124GB checkpoint resident in host RAM and streams experts to the GPU per token. Output is coherent and correct (reasoning / math / coding verified) and tool-calling works, ~22 tok/s single-stream. The new branches are gated on mixed_precision / a present .weight_scale_inv, so no other checkpoint path changes.

@gdevenyi

gdevenyi commented Sep 4, 2026

Copy link
Copy Markdown

Tested with lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8 (123 GiB: 192 NVFP4 expert shards, 4 block-FP8 dense shards with weight_scale_inv, 10 FP8 PLE shards) on 2 x RTX 6000 Ada at TP=2 (#385), --moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --max-running-requests 16. The patch did not apply onto my TP-modified weight.py, so I merged the mixed_precision config branch and the weight_scale_inv dequant by hand; same logic.

It loads and serves. Because the FP8 dense projections are dequantized to bf16 at load, speed is the same as the modelopt NVFP4 build under the same flags:

checkpoint single-stream tok/s 8 concurrent tok/s expert residency TTFT (1k) 8-question probe
RadixArk/Qwen3.8-Flash-Next-NVFP4 90.8 326.0 94.5% 0.81 s 6/8
lovedheart NVFP4-FP8 via this PR 90.2 325.6 94.6% 0.82 s 7/8

Greedy 256-token continuations of three prompts differ from the modelopt build after 7 to 29 words (different dense weights, so expected); the probe difference is within what eight questions can resolve. The obvious follow-up on top of this is keeping the dense weights in FP8 (fp8_block linears) instead of dequantizing: the decode step here is bound by the 8 GiB bf16 dense read per token at TP=1.

@gberasmus87

Copy link
Copy Markdown
Author

@gdevenyi Thanks for testing it on hardware and a checkpoint I don't have. That's a more useful signal than anything I could produce on my own.

On the patch not applying: #320 is mergeable/clean against current main (2 files, +39/-2, both under models/qwen4_exp/). I think the conflict is #385 rather than this branch - it's open, unmerged, and changes models/qwen4_exp/weight.py too (+64/-7). Whichever of the two lands second will need the merge; happy for that to be this one if #385 is closer to ready.

Agreed the divergence is expected - different dense weights, and greedy decoding amplifies the first differing logit. I wouldn't read much into 7/8 vs 6/8 either way at that sample size.

On keeping the dense weights in FP8 rather than dequantizing at load: agreed, and I'll open it as a separate PR rather than growing this one, since it changes the compute path rather than the loading path and deserves its own review and bisect point.

A second data point on the sizing, since my checkpoint is the modelopt MIXED_PRECISION build rather than lovedheart's: there the block-FP8 dense is 156 tensors, all F8_E4M3, 2.67G elements - 2.49 GiB stored, 4.98 GiB once expanded to bf16 at load. Keeping them quantized would halve that read and hand back ~2.5 GiB, which on a 24 GB card is worth having by itself. I'm at TP=1 on a single RTX PRO 4000 Blackwell with --moe-backend hybrid, so the dense read isn't split the way it is in your TP=2 run. I'll benchmark it there and post numbers with the PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants