feat(qwen3_5_moe): keep a natively-fp8 lm_head native instead of failing to load - #416
feat(qwen3_5_moe): keep a natively-fp8 lm_head native instead of failing to load#416salekseev wants to merge 2 commits into
Conversation
…loader
_per_row_scale unconditionally did scale.reshape(1), which is only valid for a
per-tensor scale. compressed-tensors also emits strategy: "channel", one scalar
per output row stored as [rows, 1], and on such a checkpoint the reshape raises
RuntimeError: shape '[1]' is invalid for input of size 248320
Reshape to [-1] and branch on the element count instead: 1 broadcasts as before,
rows passes through, anything else raises rather than being broadcast -- a
mis-shaped scale that loaded would apply one row's factor to every output row and
serve fluent, wrong tokens.
Fp8PerTensorLinear.weight_scale is already declared per-output-row, so nothing
downstream changes and no kernel work follows.
…ing to load
A compressed-tensors mixed-precision checkpoint can put lm_head in its fp8
group (unsloth/Qwen3.6-35B-A3B-NVFP4-Fast). Today that checkpoint cannot load:
_iter_weights_attn_fp8's per-tensor-fp8 branch emits any fp8 .weight natively,
including lm_head, but the model always builds the bf16 ParallelLMHead for a
non-NVFP4 head, which has no weight_scale buffer:
Unexpected keys in state_dict: ['lm_head.weight_scale']
Add the missing head class rather than dequantizing. Nvfp4LMHead is FP4, and
glm_moe_dsa/glm5_next's fp8 heads subclass the bf16 ParallelLMHead and quantize
at load, so no existing class could consume a head that is already fp8 on disk.
Fp8LMHead is Fp8PerTensorLinear plus ParallelLMHead's prefill slice -- the
weight buffers, the input_scale handling and the uniform-scale/segment
precompute are all inherited.
_lm_head_quant gains the matching "fp8" verdict, gated on the same weights
geometry _attn_quant uses (type float, 8 bits, no group_size), and the loader
only keeps the head fp8 when the model actually built an Fp8LMHead.
On [248320, 2048] this is 0.474 GiB native versus 0.947 GiB dequantized. On a
16 GiB card that is not only decode traffic: the dequantized head left 118 MiB
free and the GDN chunked-prefill workspace then OOM'd an 8k prefill.
|
Some independent evidence for this one turned up while I was testing a different checkpoint for #390.
That's 14.3 tok/s, and the head accounts for most of it: 0.947 GiB read per decoded token versus 0.474. The only other difference between the two is the shared expert (fp8 vs NVFP4, worth roughly +52 MiB/token), so the head is about 90% of the ~525 MiB traffic delta. ΔTPOT is ~1.10 ms, which against ~525 MiB implies ~480 GB/s effective — the right order for this card, so the traffic explanation at least hangs together. Caveat it properly: two checkpoints differing in more than one variable, and primitive-ai is a single run. It is not a controlled experiment. But it's an independent measurement of the thing this PR is for, on a checkpoint I didn't pick for the purpose. Also updating one number in the description: the decode figure there was a single 120.05 run. Four runs give a mean of 121.26. Same reason I corrected it on #390 — the original could be read as the fp8 path costing decode, and it doesn't; against For what it's worth, I've promoted this checkpoint to my serving default on the strength of the tool-calling score and the checkpoint-calibrated KV scales, so the |
1 participant
Note
Stacked on #415 — that's the first commit here. Review the second commit, or merge #415 first and this becomes a single commit.
A compressed-tensors mixed-precision checkpoint can put
lm_headin its FP8 group (unsloth/Qwen3.6-35B-A3B-NVFP4-Fastdoes). Right now that combination can't load at all, and the reason is a mismatch between two places that each look locally correct.The FP8 branch in
_iter_weights_attn_fp8emits any FP8.weightnatively —lm_headincluded. But the model always builds the bf16ParallelLMHeadfor a head that isn't NVFP4, and that class has noweight_scalebuffer. So:The easy fix would be to dequantize the head at load and move on. I went the other way and added the missing layer class, because nothing existing can consume a head that's already FP8 on disk —
Nvfp4LMHeadis FP4, and the FP8 heads inglm_moe_dsa/glm5_nextsubclass the bf16ParallelLMHeadand quantize during load.Fp8LMHeadis justFp8PerTensorLinearplusParallelLMHead's prefill slice; the weight buffers, theinput_scalehandling and the uniform-scale/segment precompute all come from the parent, so the subclass is the slice and its docstring._lm_head_quantgains the matching"fp8"verdict, gated on the same weights geometry_attn_quantalready uses, and the loader only keeps the head native when the model actually built anFp8LMHead— otherwise it falls through and dequantizes as before.Why this is more than a nice-to-have on small cards
On
[248320, 2048]it's 0.474 GiB native vs 0.947 GiB dequantized. On a 16 GiB card that isn't just decode bandwidth: with the bf16 head there was 118 MiB free, and the GDN chunked-prefill workspace then OOM-killed the backend on an 8192-token prefill. Keeping it native is what makes that prefill fit.One thing I'd flag for review
_lm_head_quantalso skips a head that the checkpoint lists inignore. compressed-tensors givesignoreprecedence over a group'stargets, and the common llm-compressor shape is a broad target plus anignorelist carving modules back out — such a head is bf16 on disk, so claiming it as FP8 would build anFp8LMHeadfor a bf16 weight and die on the dtype check.Worth knowing: the match there is a substring test, not
fnmatch.models/qwen4_exp/config.pyconsultsignoreviafnmatch, which suits ModelOpt's plain module names, but compressed-tensors writes regex entries likere:.*lm_headandfnmatch("lm_head", "re:.*lm_head")isFalse. So reusing that helper here would silently not match. Happy to change the approach if you'd rather it were shared.Reachability
Same caveat as #415: on
mainyou can't get here, because the only checkpoints reaching_iter_weights_attn_fp8are ModelOpt ones and theirlm_headis NVFP4, not FP8. It becomes reachable once compressed-tensors per-channel detection lands (#390 plus accepting"channel").Tests
tests/models/test_qwen3_5_moe_fp8_lm_head.py— detection for both FP8 strategies, plus the cases that must not be claimed: a head outside the FP8 group, a 4-bit group, a grouped/block scale, and an ignored head. Also asserts the layer's buffers match the per-output-row contract the loader feeds it.Testing done
unsloth/Qwen3.6-35B-A3B-NVFP4-Fastloads and serves on an RTX 4080 SUPER (16 GiB, sm_89, TP=1): 121.26 tok/s decode at 1200 tokens of context (mean of 4 runs, 120.05-123.3), 3311-3418 tok/s prefill at 8192 (8/8), 95.00 on BFCL parallel tool calls, 9/9 exact-match needle recall.nvidia/Qwen3.6-35B-A3B-NVFP4is unaffected — its four quant verdicts are unchanged. Note that build also carries an unrelated FP8 KV-cache change, so don't read those throughput numbers as coming from a clean tree. More detail in #252.