qwen4_exp: load modelopt MIXED_PRECISION (NVFP4 experts + block-FP8 dense) checkpoints - #320
Conversation
…ense) checkpoints
|
Tested with 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:
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 ( |
|
@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. |
What
Adds support for loading Qwen3.8-Flash-Next checkpoints exported with modelopt
MIXED_PRECISIONquantization - NVFP4 routed experts + 128x128 block-FP8 dense (attention / GDN) weights. The common community build of this shape islovedheart/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:
parse_configonly recognises a top-levelquant_algooffp8or*fp4*.MIXED_PRECISIONdeclares the per-module algo inquantized_layersinstead, soexpert_quantfell through to"none"and the offload cache tried to load the NVFP4 experts as bf16 banks ->ValueError: Missing MoE expert source layers..weight_scale_inv), butiter_weightsassumed 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 BFloat16on the GDNin_projfuse.Changes
config.py: detectquant_algo == "mixed_precision", readquantized_layersto setexpert_quant = "nvfp4"(experts stay native for the offload cache), and keep every dense module bf16.weight.py:_load_maybe_block_fp8dequantizes a dense.weightto bf16 when a sibling.weight_scale_invis present (reusing the existingdequant_block_fp8kernel), and.weight_scale_invis added to_SCALE_SUFFIXESso 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 onmixed_precision/ a present.weight_scale_inv, so no other checkpoint path changes.