qwen3_5_moe: load mixed-precision compressed-tensors checkpoints (unsloth NVFP4+FP8) - #208
Open
chrisqianz wants to merge 3 commits into
Open
qwen3_5_moe: load mixed-precision compressed-tensors checkpoints (unsloth NVFP4+FP8)#208chrisqianz wants to merge 3 commits into
chrisqianz wants to merge 3 commits into
Conversation
…loth NVFP4+FP8) unsloth's dynamic per-module quant exports (e.g. unsloth/Qwen3.8-27B-NVFP4) store dense Qwen3.x checkpoints as per-module mixed precision: FP8 attention / GDN output linears (per-row scale), NVFP4 dense-MLP layers, FP8 dense-MLP layers, bf16 in_proj_b/a and norms. The loader crashed on them (fp8/bf16 promotion in the bf16 fusion, missing packed weights for natively-built linears) or materialized everything to bf16 (54 GB on a 27B -- no launch parameter fits a 32 GB card). Keep every dense linear in the storage the checkpoint actually uses, sniffed from model.safetensors.index.json: - config: _compressed_linear_storage() now reports per-module storage (attention nvfp4/fp8/none, per-layer dense-MLP overrides, fp8 lm_head); ModelConfig gains dense_mlp_storage (per-layer override map) and routes attn_quant/lm_head_quant to native fp8_pertensor when the export says so - moe: _SharedExpert builds per-layer native linears (NVFP4 W4A16 / FP8 W8A16 / bf16) from the override map; layer_id threaded through the dense MLP - model: lm_head built as native FP8 (Fp8PerTensorLinear) when the checkpoint stores it fp8 (halves the ~2.5 GB bf16 lm_head) - weight: the dense pass keeps fp8 parts native (q/k/v -> qkv_proj, in_proj_qkv/z -> in_proj_qkvz, dense gate/up -> gate_up_proj fusions; o_proj/out_proj/down_proj/lm_head singletons, per-row fp32 scales); a part buffered into the fp8 fusion never also enters the bf16 buffer (incomplete fusion assert); fp8 dequant (for the unscaled remainder) is a bf16 broadcast multiply (no fp32 copy); ShardReader gains has() for sibling-scale lookups Verified on unsloth/Qwen3.8-27B-NVFP4 (RTX 5090 D, 32 GB): weights resident at ~21.8 GB native (vs 54 GB bf16), CUDA graph capture, and live chat-completion requests all succeed; official dense-NVFP4 and routed-MoE layouts keep their native assumptions (sniffer fallbacks) and are unchanged. Tests: tests/models/test_qwen3_5_moe_config.py (8) + test_qwen3_5_moe_weight.py (13) -- storage sniffing, per-layer construction gates, fp8 native fusions, per-row/block scale dequant semantics.
…ported models unsloth's per-module mixed-precision dense exports (NVFP4 MLP + FP8 attention/GDN/lm_head + bf16 residual parts) load natively end-to-end; list the known-good checkpoint and document the layout.
No behavior change; the mixed-NVFP4+FP8 load path (unsloth/Qwen3.8-27B-NVFP4) stays byte-identical: - _pt_fp8_fuse: rename the 'scalar' param to 'scale' (it now accepts modelopt scalars AND unsloth's per-row [O, 1] scales) and fix the sloppy return annotation (bare 'list' -> list[tuple[str, torch.Tensor]] | None) - _per_row_scale: fail loud with a clear message if a non-scalar scale has the wrong element count instead of a cryptic reshape error - document the unscaled-fp8 fallthrough assumption in the .weight handler (an unscaled fp8 q/k/v or in_proj_qkv/z in the fp8-split layout would fail at load with a missing key, not silently dequant -- no real export has this) - weight tests: refresh the module docstring (native W8A16 + dequant, not dequant-only) and pin the local fp8-fusion-map extension (dense-MLP gate/up -> gate_up_proj native, per-row fp32 scales) with a unit test Verified: config tests 8/8, weight tests 14/14, parse_config on the real checkpoint, full-pass key/shape/dtype check (0 stray / 0 missing / 0 never-yielded) all green on the remote (RTX 5090 D, editable install). Co-Authored-By: GooeyPi <gpt-5.2@openai.com>
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
1 participant
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.
Problem
Mixed-precision compressed-tensors exports (e.g.
unsloth/Qwen3.8-27B-NVFP4) cannot load: the dense pass assumes either pure NVFP4 (weight_packed) or pure bf16, but unsloth's layout stores attention linears as weight-only FP8 (per-row scale), most dense-MLP layers as NVFP4, a few MLP layers as FP8, andlm_headas FP8. Raw e4m3 tensors die in the bf16 fusion (torch.catfp8/bf16 promotion), and bf16 dequant oflm_head/MLP doubles memory.Fix (per-module native storage — no new kernels)
Reuses the existing
Fp8PerTensor*kernels; no shared kernel / modelopt-pass changes:models/config.py:dense_mlp_storageper-layer override onModelConfigmodels/loader.py:ShardReader.has()for scale-sibling detectionqwen3_5_moe/config.py:_compressed_linear_storagesniffs the safetensors indexweight_map(order-independent per-layer classification:weight_packed→ nvfp4, scaled → fp8, plain → bf16) + FP8lm_headdetectionqwen3_5_moe/moe.py+model.py: overridden layers build shared-expert / dense-MLP /lm_headlinears as native W8A16qwen3_5_moe/weight.py: dense pass keeps scaled FP8 linears native (fused q/k/v →qkv_proj, in_proj_qkv/z →in_proj_qkvz, gate/up →gate_up_proj, singletons o_proj / GDN out_proj / down_proj / lm_head with per-row fp32 scales); unscaled fp8 dequantizes to bf16 (per-tensor / per-row / block scale). Official pure-NVFP4 checkpoints are unaffected (pass is a no-op there).docs/models.md: addunsloth/Qwen3.8-27B-NVFP4to the known-good listVerification (RTX 5090 D 32GB, driver 595.84, torch 2.11 cu130)
torch.catfp8/bf16 promotion crash on the mixed layoutft servee2e: 21.8 GB load, KV 16K, CUDA graph capture,/healthok, real chat generationUsage note: on a 32 GB card use
ft serve --model unsloth/Qwen3.8-27B-NVFP4 --num-tokens 32768 --max-prefill-length 1024(the large-vocab logits buffer OOMs at the default 8192-token prefill).