fix(qwen3_5): load quantized lm_head, not just BF16 (#164) - #169
Merged
Conversation
Reported by rohitpaul: both public Qwen3.6-27B NVFP4 checkpoints die at load
with "dense loader: expected BF16 for lm_head.weight". Confirmed in-tree and,
more usefully, confirmed WHY.
Our benchmarks are not wrong, and the checkpoint is not wrong: the repo was
re-quantized under us. Reading the safetensors headers of both snapshots we
hold on the lab box:
unsloth/Qwen3.6-27B-NVFP4 @890bdef7 lm_head.weight BF16 [248320, 5120]
unsloth/Qwen3.6-27B-NVFP4 @ccdaab7e lm_head.weight F8_E4M3 [248320, 5120]
lm_head.weight_scale BF16 [248320, 1]
@890bdef7 is the snapshot every recorded 27B-NVFP4 number ran on, which is why
an unconditional BF16 head survived this long: the body has had real NVFP4
support all along, the head was simply never quantized under us. The reporter
used @ccdaab7e. nvidia/Qwen3.6-27B-NVFP4 ships a third form, a ModelOpt NVFP4
head. So the support claim was silently revision-locked.
LoadLmHeadAnyDtype now dispatches on the stored dtype:
BF16 weight [V,H] -> transpose (unchanged)
F8_E4M3 weight [V,H] + weight_scale [V,1] or scalar -> per-row/per-tensor
U8 weight [V,H/2] + weight_scale F8 [V,H/16]
+ weight_scale_2 | weight_global_scale -> nvfp4
Note the scale is PER-OUTPUT-CHANNEL in @ccdaab7e, not per-tensor as the issue
guessed, so the existing scalar-scale FP8 loader would have been quietly wrong
rather than loud. All three forms land on the same bf16 [in, out] Matmul-B
operand the logits GEMM already consumes, so the forward is untouched and a
BF16 head takes the identical call as before -- byte-exact, no dequant.
ModelOpt spells the global scale weight_scale_2 and stores the scale itself;
compressed-tensors spells it weight_global_scale and stores the reciprocal.
Both are accepted, converted to the CT divisor convention so the shared
DequantCtNvfp4WeightToF32 computes each exactly.
An unsupported dtype now names what it saw and what is accepted, instead of
claiming BF16 was expected.
Keeping the head QUANTIZED end to end would save ~2.3 GiB over dequantizing to
bf16; that needs an lm_head_fp4-style field plus a forward branch and is a
follow-up, not this fix. This restores loadability at the memory profile we
already benchmark.
Test: tests/vllm/models/test_qwen3_5_lm_head_dtypes.cpp, 5 cases / 32
assertions, synthetic tensors only (no checkpoint, no GPU): BF16 unchanged,
FP8 per-output-channel, FP8 per-tensor, unsupported dtype names itself, and
FP8 missing its scale fails loudly.
FOLLOWING_AGENTS_PROTOCOL
Assisted-by: Claude Opus 5 (1M context)
mudler
force-pushed
the
row/MODEL-QWEN35-LMHEAD-DTYPES
branch
from
August 8, 2026 22:24
619aba4 to
6f3a3d0
Compare
mudler
added a commit
that referenced
this pull request
Aug 8, 2026
…llow-up) `nvidia/Qwen3.6-27B-NVFP4` and `unsloth/Qwen3.6-27B-NVFP4` @ccdaab7e both failed to load. #169 fixed the reported symptom, the BF16-only `lm_head`; running it against the real checkpoints then exposed two more layers behind it. All three are the same root cause: the dense loader assumed BF16 wherever the compressed-tensors NVFP4 probe did not match. Layer 2, FP8 tower tensors. Both publishers quantize parts of the tower to FP8 while leaving the rest BF16, and they disagree on which parts and on the scale layout: nvidia ships FP8 `linear_attn` in_proj_qkv/in_proj_z/out_proj with a per-tensor F32 scale, unsloth @ccdaab7e went FP8 across the whole tower with a per-output-channel BF16 scale. Rather than teach each loader its own rules, one `MaterializeBf16Source` now sits under LoadBf16Direct, LoadBf16Transposed and LoadMergedBf16RawNK: BF16 stays zero-copy on the mmap (unchanged), FP8 is dequantized to BF16. Reading a per-channel scale as per-tensor would be silently WRONG rather than loud, so the element count decides and anything else is rejected. Layer 3, ModelOpt NVFP4 naming. Our probe was `has(<proj>.weight_packed)`, which is compressed-tensors only. nvidia's checkpoint is ModelOpt: `<proj>.weight` U8 + `.weight_scale` F8 + `.weight_scale_2` F32. Every probe missed, so a genuinely NVFP4 tower fell through to the BF16 path. `IsNvfp4Projection` now accepts both spellings and `LoadNvfp4AnyNaming` reads either, converting ModelOpt's direct scale to the CT divisor convention the shared dequant expects (the same conversion #169 made for lm_head). W4A16 is the default for ModelOpt. These checkpoints carry a per-tensor `input_scale`; consuming it sets `alpha`, flips `IsTrueW4A4()` and routes to the fp4-activation GEMM, which produced incoherent output on nvidia's checkpoint. Leaving `alpha` at 0 takes the weight-only dispatcher and generates correctly. `VT_MODELOPT_W4A4=1` restores the other arm for A/B. Verified on dgx GB10 (sm_121) against the real 21 GB `nvidia/Qwen3.6-27B-NVFP4`: before, "expected BF16 for ...in_proj_qkv.weight"; after, greedy "The capital of France is" -> " Paris." Clean CUDA build; clean CPU build. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude:claude-opus-5 [ClaudeCode]
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
2 participants
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, reported by @rohitpaul. Their diagnosis was right; this adds the part they could not see from outside.
Our benchmarks were not measured on a broken checkpoint. The repo was re-quantized under us.
Reading the safetensors headers of both snapshots we hold on the lab box:
lm_head.weightunsloth/Qwen3.6-27B-NVFP4@890bdef7[248320, 5120]unsloth/Qwen3.6-27B-NVFP4@ccdaab7e[248320, 5120]weight_scaleBF16[248320, 1]@
890bdef7is the snapshot every recorded 27B-NVFP4 number ran on, which is why an unconditional BF16 head survived this long: the transformer body has had real compressed-tensors NVFP4 support all along, the head simply was never quantized under us. The reporter used @ccdaab7e.nvidia/Qwen3.6-27B-NVFP4ships a third form (ModelOpt NVFP4). So the support claim was silently revision-locked — that is the actual defect.The fix
LoadLmHeadAnyDtypedispatches on the stored dtype:BF16F8_E4M3weight_scaleper-output-channel[V,1]or per-tensorU8NVFP4weight_scaleF8[V,H/16]+weight_scale_2(ModelOpt) orweight_global_scale(CT)Two details worth flagging:
ccdaab7eis per-output-channel, not per-tensor as the issue guessed. Reusing the existing scalar-scale FP8 loader would have been quietly wrong rather than loud, which is why this gets its own path and its own test.DequantCtNvfp4WeightToF32reciprocates internally. Both spellings are accepted and normalized to the CT divisor convention so the one shared dequant computes each exactly.All three land on the same
bf16 [in, out]Matmul-B operand the logits GEMM already consumes, so the forward is untouched. An unsupported dtype now names what it saw instead of claiming BF16 was expected.Scope note
Keeping the head quantized end-to-end would save ~2.3 GiB over dequantizing to bf16, but needs an
lm_head_fp4-style field plus a forward branch. That is a follow-up. This restores loadability at the memory profile we already benchmark.Verification
tests/vllm/models/test_qwen3_5_lm_head_dtypes.cpp: 5 cases / 32 assertions, synthetic tensors only (no checkpoint, no GPU, runs in CI) — BF16 unchanged, FP8 per-output-channel, FP8 per-tensor, unsupported dtype names itself, FP8 missing its scale fails loudly.scripts/check-*.pybattery green.Docs:
BENCHMARKS.mdnow pins the revision the 27B NVFP4 grid actually ran on,USAGE.mdgains a table of accepted head forms,FEATURES.md/STATUS.mdupdated.