feat(qwen4_exp): load block-fp8 dense projections natively - #428
feat(qwen4_exp): load block-fp8 dense projections natively#428jason-fxz wants to merge 1 commit into
Conversation
|
Happy to see this land properly — One behaviour question, not an objection.
if scheme is None:
if is_fp8:
raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized")So a checkpoint that stores block-fp8 dense weights without declaring them per module now fails to load. #320 dequantized that case to bf16 at load and served it, which is how I first got the lovedheart build running — my early notes on it say the block-fp8 attention/GDN tensors were not flagged in Raising it only because the failure mode changes from "loads, dequantized" to "raises", and community requants are the population most likely to store quantized weights without describing them. If that's a deliberate call — an explicit error beats a silent divergence between the buffers the model built and what the reader emits — then it's the right one and worth a line in the docstring so the next person doesn't re-add a fallback. If it isn't, a On the run you asked for: yes. My box is the regime this feature is actually for — a single 24 GB RTX PRO 4000 Blackwell on a PCIe gen 3 board, offload backend, where the model is nowhere near resident and expert-cache residency is the binding constraint. That's the opposite end from your H100 80GB and @gdevenyi's 2 x 48 GB Ada, where the freed dense bytes have much less to buy. On the earlier form of this reader I measured 27.4 -> 34.2 tok/s there, and the gain tracked Also flagging, since #392's reader half is superseded by this: the |
b16877b to
540b2ce
Compare
|
Deliberate. Since #418 a module the config does not list is built bf16, so an fp8 tensor there means the checkpoint's declaration and its tensors disagree. We treat that as a broken checkpoint, not a layout to support; the error names the module to fix. Workaround: add the module to quantization_config.quantized_layers ({"quant_algo": "FP8_PB_WO", "group_size": 128}) in config.json, and in hf_quant_config.json if present. Thanks for the run and for splitting out #429. |
|
Ran it. The reader works — I can't give you decode numbers, though, because this tree runs 40-60x slower than the build I'm coming from on this configuration, and I think that matters more than the numbers I owed you. What I measuredSame box, same checkpoint, same flags, back-to-back:
End-to-end, counting every streamed delta (this model emits most tokens as What it looks like while it is slow
What I can't tell youI cannot attribute this to #428, and I do not think it is #428. My baseline is I also could not get a stack. Why this configuration is probably not coveredYou tested on an H100 80 GB and @gdevenyi is on 2 x RTX 6000 Ada 48 GB. On both, this model is close to resident and the offload expert path barely engages. Here it carries essentially the whole model — 123 GiB of weights against 24 GB of VRAM on a gen-3 board, which is the regime the offload backend exists for, and the one where the dense-side savings in this PR are worth the most. If the slowdown is in that path, it would be invisible on either of your boxes and fatal on mine. Happy to bisect |
Loads modelopt MIXED_PRECISION Qwen3.8-Flash-Next checkpoints whose attention / GDN projections are 128x128 block-fp8 (
FP8_PB_WO), e.g.lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8. The model already builds the fp8 buffers (#418); the reader still fused GDNin_proj_{qkv,z,b,a}into one bf16 tensor and died intorch.catwhenin_proj_qkvwas fp8._DenseFuserreplaces_FUSIONS/_try_fuse: parts come frompacked_modules_mapping,.weightand.weight_scale_invfuse per kind, GDN splits intoin_proj_qkvz(fp8) +in_proj_ba(bf16) when the QuantConfig says so.Tests: synthetic bf16, NVFP4 and block-fp8 checkpoints in
tests/models/qwen4_exp/test_weight.py(emitted keys == model state dict, fp8 slice-back, four rejection cases); scheme resolution intest_config.py.Tested on H100 80GB / Xeon 8480+, driver 580.95.05, torch 2.11.0+cu130, TP=1, offload +
--moe-cache-auto:lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8passestests/e2e/test_aime.py, all 926 emitted tensors match the model's buffers, 18 fused fp8 tensors are bit-exact against the raw shards; RadixArk and nvidia NVFP4 unchanged.FREETOKEN_TEST_MODEL=<dir> FREETOKEN_TEST_MOE_CACHE_AUTO=1 pytest -x -s tests/e2e/test_aime.py.Depends on #427. Supersedes #392 / #320 by @gberasmus87, which reached the same reader design before #418. @gberasmus87 @gdevenyi: a run on your cards would be a useful check.