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 |
|
Closing this -- the quantization refactor in #418/#427 replaced the mechanism this PR is built on, and did it better. Every piece of scaffolding this hangs off is gone from
And the behaviour it was chasing now falls out of the new config layer on its own. Probing
Nice work; this is a better design than bolting another head class on beside One thing that is still open, and it's the reason I'm not just closing this quietly. I'd rather spend the effort there than rebase this. I have #415 I'll leave open for now, since it fixes a live bug in the legacy |
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.