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.

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.

1 participant