Repository navigation
Export unquantized FP8Linear layers in the export dtype, not fp32 - #2651
Merged
Merged
Conversation
_dispatch_export_handler() unpacks every HF FP8Linear via unpack_weight(), which dequantizes to torch.get_default_dtype() (fp32). For a layer the recipe leaves unquantized (QUANTIZATION_NONE, e.g. nvfp4_mlp_only on an FP8 checkpoint), the export handler returns early, so the fp32 weight is written as-is into an otherwise bf16/fp16 checkpoint. Cast those layers to the export dtype after unpacking. Quantized layers are unchanged: they still dequantize in fp32 before scale computation, and their packed output is byte-identical to before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: tritsystem <gbranaa4@gmail.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe export handler now casts unpacked, unquantized FP8 weights to the requested export dtype. CUDA tests cover models with NVFP4-quantized MLP weights and FP8 attention weights. ChangesFP8 export weight handling
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to Unquantized FP8 weights are converted to the requested export dtype before the exporter collects its state, while the quantized path remains unchanged. No actionable merge risk is evident, though the reported CUDA tests were not independently run here. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
shengliangxu
approved these changes
Oct 6, 2026
shengliangxu
left a comment
shengliangxu
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. Thanks for working on the fix!
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2651 +/- ##
==========================================
+ Coverage 69.19% 78.44% +9.25%
==========================================
Files 620 618 -2
Lines 69588 69588
==========================================
+ Hits 48150 54588 +6438
+ Misses 21438 15000 -6438
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Collaborator
|
/ok to test 10fffda |
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.
What does this PR do?
Type of change: Bug fix
When an HF FP8 checkpoint is exported with a recipe that leaves some FP8 layers unquantized (for example
--qformat nvfp4_mlp_only), those layers are written to the checkpoint as fp32, while the rest of the model is in the export dtype (bf16/fp16).Mechanism:
_dispatch_export_handler()(modelopt/torch/export/unified_export_hf.py) calls_QuantFP8Linear.unpack_weight()on every FP8 linear.unpack_weight()dequantizes totorch.get_default_dtype(), which is fp32 unless the caller changed it. For a module whose format isQUANTIZATION_NONE,_export_quant_linearreturns early, and nothing casts the weight back to the export dtype. The DeepSeek example scripts don't show the problem because they calltorch.set_default_dtype(torch.bfloat16).examples/hf_ptq/hf_ptq.pydoesn't, so the main HF path does show it.Effects: those tensors take 2× the bytes, their dtype doesn't match
config.torch_dtype, and the in-memory model can't run forward after export (expected mat1 and mat2 to have the same dtype ... BFloat16 != float).Fix: after
unpack_weight(), cast the weight toctx.dtypeonly when the module's quantization format isQUANTIZATION_NONE. Quantized layers are deliberately untouched. My first attempt passed the export dtype intounpack_weight()itself, but that changed the NVFP4 packed bytes (12 of 18 MLP weight/scale tensors differed), because the per-block scales were then computed from a bf16 dequant instead of fp32. The version here keeps the fp32 dequant for quantized layers, and their output is byte-identical tomain.Usage
No API change.
Testing
All of the following ran on an RTX 5060 (sm_120), Windows, torch 2.11+cu128, transformers 4.57.1, base
6a10687.LlamaForCausalLM(bf16) and replaced every projection withtransformersFP8Linear(block 128×128). I quantized MLP only and exported withexport_hf_checkpoint, then read the safetensors dtypes.main:self_attn.{q,k,v,o}_proj.weight= F32, layernorms = BF16, MLP = U8 (NVFP4) / I8 (INT8 weight-only variant).self_attn.*= BF16, everything else unchanged.mainand this PR (torch.equalon the uint8 views, 0 differing).tests/gpu/torch/export/test_export_fp8_linear.py, parametrized over bf16 and fp16. It checks that the unquantized FP8 layer's exported weight has the export dtype and equals the fp32 dequant cast to that dtype exactly (torch.equal).mainwithassert torch.float32 == torch.float16.tests/gpu/becauseunpack_weight()runs undertorch.cuda.device.tests/unit/torch/export/test_unified_export_hf.py,test_quant_aware_conversion.pyandtests/unit/torch/quantization/plugins/test_huggingface.py. Bothmainand this PR give 70 passed, 50 skipped, and 4 failed. The same 4 fail onmainwithout this change (test_autoquantize_huggingface[*-gradient],test_quantized_transformers_save_restore[*]), so I believe they come from my Windows environment rather than this change.ruff check/ruff format --checkare clean on the new test.ruff checkonunified_export_hf.pyreports twoPLR0917hits, but they're on untouched functions and also show up onmainwith my local ruff version.Not tested: a real large FP8 checkpoint (e.g. Qwen3-FP8) through
hf_ptq.py, the streaming/offload export path, or Linux CI. The repro uses the sameFP8Linearclass and the sameexport_hf_checkpointentry point, just at toy size.Before your PR is "Ready for review"
CONTRIBUTING.md: N/A/claude review)Additional Information
Related: #1962 introduced the shared
_dequantize_weight()helper and kepttorch.get_default_dtype()inunpack_weight(). This PR leaves that default as is and only fixes the export call site.🤖 Generated with Claude Code
Summary by CodeRabbit