Repository navigation
Fix per-tensor FP8 weight dequantization - #1962
Conversation
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1962 +/- ##
==========================================
+ Coverage 76.50% 77.32% +0.81%
==========================================
Files 525 524 -1
Lines 59257 59251 -6
==========================================
+ Hits 45337 45814 +477
+ Misses 13920 13437 -483
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:
|
7069c2b to
db3edfb
Compare
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe change adds scalar-scale FP8 weight dequantization support, centralizes FP8 conversion paths, adds targeted unit coverage, and introduces a documented, loader-tested Mistral-Medium NVFP4 max-calibration recipe. ChangesFP8 dequantization
Mistral-Medium NVFP4 recipe
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/quantization/plugins/huggingface.py`:
- Around line 1353-1363: CPU scalar-scale dequantization fails because CUDA
device context wrapping occurs unconditionally in forward paths. Move the
torch.cuda.device(self.weight.device) context from forward and unpack_weight
into _dequantize_weight, applying it only around the Triton weight_dequant call;
leave the block_size is None multiplication path context-free so forward and
unpack_weight work with CPU tensors.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d0335345-d7f3-4221-b07b-bc382357a35c
📒 Files selected for processing (2)
modelopt/torch/quantization/plugins/huggingface.pytests/unit/torch/quantization/plugins/test_huggingface.py
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-4-8) — DM the bot to share feedback.
Clean, focused bug fix (+34/-12, 2 files) adding per-tensor (scalar weight_scale_inv) FP8 weight dequantization to _QuantFP8Linear.
Correctness verified:
_setupsetsblock_size = Nonefor 0D scale; 2D path preserves the existing block=128 assert/computation. Other ndims still fail thendim == 2assert as before.- New
_dequantize_weighthelper does directweight.to(dtype) * scale_inv.to(dtype)for the scalar path (correct per-tensor broadcast) and delegates to the unchanged Tritonweight_dequantkernel for the block path. Theassert weight_dequant is not Nonecorrectly moved into the block-only branch, so the scalar path no longer requires Triton. forwardandunpack_weightnow share the helper with behavior preserved for the block path (confirmed againstfp8_kernel.py, which is untouched).
Test: test_fp8_linear_per_tensor_dequant is CPU-runnable — it monkeypatches weight_dequant=None, asserts block_size is None, and checks the dequant result equals weight.float() * 2.0. Reasonable focused coverage; PR body documents manual full-model validation (Mistral-Medium-3.5, 616 modules, finite logits, unpack_weight() on q_proj).
No licensing changes (headers/LICENSE untouched), no local imports, no new subsystem/abstraction, small cohesive change. No prompt-injection attempts in the PR body/diff.
Minor (non-blocking, pre-existing): the scalar path in forward/unpack_weight still runs under torch.cuda.device(self.weight.device), and the test exercises _dequantize_weight directly rather than end-to-end forward()/unpack_weight() — both consistent with prior FP8 behavior and covered by the described manual GPU validation.
db3edfb to
9d6fdde
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unit/recipe/test_loader.py (1)
169-169: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd semantic coverage for the new recipe.
The current parameterized test only proves that the path parses and contains a
quantizesection. Add a focused assertion for this recipe’s expected max-calibration settings and representative NVFP4, FP8, and KV-cache overrides so incorrect wildcard ordering or missing imports cannot pass unnoticed.As per path instructions, tests should protect the expected behavior with focused pytest coverage.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/recipe/test_loader.py` at line 169, Extend the parameterized recipe test covering "huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib" with focused assertions for its max-calibration configuration, including representative NVFP4, FP8, and KV-cache overrides. Validate the resolved settings rather than only path parsing and the presence of quantize, ensuring wildcard ordering and required imports are exercised.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unit/recipe/test_loader.py`:
- Line 169: Extend the parameterized recipe test covering
"huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib"
with focused assertions for its max-calibration configuration, including
representative NVFP4, FP8, and KV-cache overrides. Validate the resolved
settings rather than only path parsing and the presence of quantize, ensuring
wildcard ordering and required imports are exercised.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 103c04e7-2eae-4ba4-9e29-36d9066f7217
📥 Commits
Reviewing files that changed from the base of the PR and between db3edfb and 9d6fdde481036efbdf17a2ce10d44832b3986a3c.
📒 Files selected for processing (4)
modelopt/torch/quantization/plugins/huggingface.pymodelopt_recipes/huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib.yamltests/unit/recipe/test_loader.pytests/unit/torch/quantization/plugins/test_huggingface.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/torch/quantization/plugins/test_huggingface.py
- modelopt/torch/quantization/plugins/huggingface.py
Add the published Mistral Medium 3.5 NVFP4 recipe. Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
9d6fdde to
3b1d157
Compare
Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
) ### 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 to `torch.get_default_dtype()`, which is fp32 unless the caller changed it. For a module whose format is `QUANTIZATION_NONE`, `_export_quant_linear` returns early, and nothing casts the weight back to the export dtype. The DeepSeek example scripts don't show the problem because they call `torch.set_default_dtype(torch.bfloat16)`. `examples/hf_ptq/hf_ptq.py` doesn'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 to `ctx.dtype` only when the module's quantization format is `QUANTIZATION_NONE`. Quantized layers are deliberately untouched. My first attempt passed the export dtype into `unpack_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 to `main`. ### Usage No API change. ```python # FP8 HF checkpoint, attention left unquantized by the recipe model = mtq.quantize(model, mtq.NVFP4_MLP_ONLY_CFG, forward_loop) export_hf_checkpoint(model, export_dir=out) # before: self_attn.*_proj.weight saved as F32; after: BF16 (matches config.torch_dtype) ``` ### Testing All of the following ran on an RTX 5060 (sm_120), Windows, torch 2.11+cu128, transformers 4.57.1, base `6a10687`. - **End-to-end repro via the public API:** I built a tiny `LlamaForCausalLM` (bf16) and replaced every projection with `transformers` `FP8Linear` (block 128×128). I quantized MLP only and exported with `export_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). - This PR: `self_attn.*` = **BF16**, everything else unchanged. - **Quantized path unchanged:** the 18 NVFP4 MLP tensors (weights + scales, 2 layers) are byte-identical between `main` and this PR (`torch.equal` on the uint8 views, 0 differing). - **New test** `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`). - Fails on `main` with `assert torch.float32 == torch.float16`. - Passes with this PR (2/2). - It lives under `tests/gpu/` because `unpack_weight()` runs under `torch.cuda.device`. - **Existing suites:** I ran `tests/unit/torch/export/test_unified_export_hf.py`, `test_quant_aware_conversion.py` and `tests/unit/torch/quantization/plugins/test_huggingface.py`. Both `main` and this PR give 70 passed, 50 skipped, and 4 failed. The same 4 fail on `main` without 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 --check` are clean on the new test. `ruff check` on `unified_export_hf.py` reports two `PLR0917` hits, but they're on untouched functions and also show up on `main` with 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 same `FP8Linear` class and the same `export_hf_checkpoint` entry point, just at toy size. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (unquantized FP8 layers change from fp32 to the export dtype; quantized layers are unchanged) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ (small fix; happy to add an entry if wanted) - Did you get Claude approval on this PR?: N/A (external contributor, can't trigger `/claude review`) ### Additional Information Related: #1962 introduced the shared `_dequantize_weight()` helper and kept `torch.get_default_dtype()` in `unpack_weight()`. This PR leaves that default as is and only fixes the export call site. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved export handling for FP8 linear layers: quantized layers retain packed weights, while unquantized weights are converted to the requested export dtype. * **Tests** * Added GPU coverage for exporting mixed quantized and unquantized FP8 layers in bfloat16 and float16. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: tritsystem <gbranaa4@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com>
Chad's Agent ### What does this PR do? Type of change: Bug fix. Native Transformers `FP8Experts` store fused FP8 weights with inverse scales. The generic wrapper exposed raw values to weight calibration and packing, while native FP8 dispatch bypassed ModelOpt's fake quantization. Dequantize one expert projection at a time for forward, calibration, and export, keeping scale multiplication in FP32 until the output cast. Export removes consumed inverse/activation scales. Native packed FP4 inputs are rejected. Explicit `unpack_weight()` materializes floating parameters before folding, matching FP8Linear's interface. Folding is inherited from the fused wrapper. Native forward still reads scale attributes after unpacking, so export removes them. Ordinary PTQ/export keeps per-expert dequantization without unpacking entire modules. Raw FP8 folding retains FP8Linear's existing unsupported behavior; shared folding code is unchanged. Independent of recipe PR #2817; analogous to #1962. FP8Linear is unchanged. ### Usage Use a recipe matching native expert names; M2.7 validation also needs #2817. Regenerate affected exports from the original FP8 checkpoint. Call `experts.unpack_weight()` explicitly before permanent or temporary folding. ### Testing - Current revision: 99 focused fused-expert/FP8Linear plugin tests pass on both Transformers 5.5.0 and 5.18.0. Gated/non-gated temporary-fold tests verify parameter identity, dtype, exact values, quantizer state, and forward results after normal and exceptional exits. The new tests fail on the previous head without explicit unpacking. - Six real native FP8Experts GPU cases pass across gated/non-gated layouts and FP32/FP16/BF16 inputs, with NVFP4 calibration/export and independent decoding. - Fresh full MiniMax M2.7 PTQ/export passes with #2817: 512 samples, batch 8, 64/64 calibration batches, 14 shards (134.52 GB versus 230.12 GB source payload). All 47,616 expert projections have correct packed/block-scale shapes. All 191,087 tensor headers and shard extents agree with the index; no stale inverse/activation scales or declaration mismatches; reverse naming succeeds. - 27 crops (layers 0/30/61, experts 0/127/255, gate/up/down) have 9.28–9.53% relative L2 error and norm ratios 0.997–1.003 against the **scaled FP8 source**. Independent E2M1 and ModelOpt NVFP4 decoders agree within FP32 rounding. Both generation previews complete. - Controlled environment: four GB300 GPUs, aarch64, `nvcr.io/nvidia/pytorch:26.08-py3`, Transformers 5.18.0, compressed-tensors 0.19.0, kernels 0.17.0, datasets 5.0.1; original ModelOpt 0.48.0rc0 source plus backend and recipe patches. Wrapper ASTs/export code match this PR. Model revision `d494266a4affc0d2995ba1fa35c8481cbd84294b`; CNN/DailyMail revision `96df5e686bee6baa90b8bee7c28b81fa3fa6223d`, config 3.0.0. The weight-quantizer guard remained active. - Fresh calibration: 26m10s; export: 150.1s; peak memory: 62.75–65.32 GB/GPU. Tokenizer/vocabulary/merges/chat-template files are byte-identical; generation settings differ only in the recorded Transformers version. Native config normalization and the loader-generated expert-parallel scale-sharding plan are accounted for. - Changed-file pre-commit hooks and `git diff --check` pass. Current repository CI: 32 successful checks, 1 skipped, 19 GPU/example jobs queued without assigned runners; no failures. Earlier revision: full-model smoke/full512 and 54 CI checks passed; the full512 results above are from the fresh rerun. ### Before your PR is "*Ready for review*" - Backward compatible: ✅ - Copied code/new dependencies: N/A - Necessary tests: ✅ - Changelog: ✅ - Claude approval: N/A ### Additional Information Full-model PTQ uses a four-GPU device map. FSDP2, offload, serving, and benchmark accuracy are untested. Excluded attention projections export as BF16, routers as FP32, and the LM head as BF16. --------- Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
Chad's Agent Type of change: Bug fix. Native Transformers `FP8Experts` store fused FP8 weights with inverse scales. The generic wrapper exposed raw values to weight calibration and packing, while native FP8 dispatch bypassed ModelOpt's fake quantization. Dequantize one expert projection at a time for forward, calibration, and export, keeping scale multiplication in FP32 until the output cast. Export removes consumed inverse/activation scales. Native packed FP4 inputs are rejected. Explicit `unpack_weight()` materializes floating parameters before folding, matching FP8Linear's interface. Folding is inherited from the fused wrapper. Native forward still reads scale attributes after unpacking, so export removes them. Ordinary PTQ/export keeps per-expert dequantization without unpacking entire modules. Raw FP8 folding retains FP8Linear's existing unsupported behavior; shared folding code is unchanged. Independent of recipe PR #2817; analogous to #1962. FP8Linear is unchanged. Use a recipe matching native expert names; M2.7 validation also needs Call `experts.unpack_weight()` explicitly before permanent or temporary folding. - Current revision: 99 focused fused-expert/FP8Linear plugin tests pass on both Transformers 5.5.0 and 5.18.0. Gated/non-gated temporary-fold tests verify parameter identity, dtype, exact values, quantizer state, and forward results after normal and exceptional exits. The new tests fail on the previous head without explicit unpacking. - Six real native FP8Experts GPU cases pass across gated/non-gated layouts and FP32/FP16/BF16 inputs, with NVFP4 calibration/export and independent decoding. - Fresh full MiniMax M2.7 PTQ/export passes with #2817: 512 samples, batch 8, 64/64 calibration batches, 14 shards (134.52 GB versus 230.12 GB source payload). All 47,616 expert projections have correct packed/block-scale shapes. All 191,087 tensor headers and shard extents agree with the index; no stale inverse/activation scales or declaration mismatches; reverse naming succeeds. - 27 crops (layers 0/30/61, experts 0/127/255, gate/up/down) have 9.28–9.53% relative L2 error and norm ratios 0.997–1.003 against the **scaled FP8 source**. Independent E2M1 and ModelOpt NVFP4 decoders agree within FP32 rounding. Both generation previews complete. - Controlled environment: four GB300 GPUs, aarch64, `nvcr.io/nvidia/pytorch:26.08-py3`, Transformers 5.18.0, compressed-tensors 0.19.0, kernels 0.17.0, datasets 5.0.1; original ModelOpt 0.48.0rc0 source plus backend and recipe patches. Wrapper ASTs/export code match this PR. Model revision `d494266a4affc0d2995ba1fa35c8481cbd84294b`; CNN/DailyMail revision `96df5e686bee6baa90b8bee7c28b81fa3fa6223d`, config 3.0.0. The weight-quantizer guard remained active. - Fresh calibration: 26m10s; export: 150.1s; peak memory: 62.75–65.32 GB/GPU. Tokenizer/vocabulary/merges/chat-template files are byte-identical; generation settings differ only in the recorded Transformers version. Native config normalization and the loader-generated expert-parallel scale-sharding plan are accounted for. - Changed-file pre-commit hooks and `git diff --check` pass. Current repository CI: 32 successful checks, 1 skipped, 19 GPU/example jobs queued without assigned runners; no failures. Earlier revision: full-model smoke/full512 and 54 CI checks passed; the full512 results above are from the fresh rerun. - Backward compatible: ✅ - Copied code/new dependencies: N/A - Necessary tests: ✅ - Changelog: ✅ - Claude approval: N/A Full-model PTQ uses a four-GPU device map. FSDP2, offload, serving, and benchmark accuracy are untested. Excluded attention projections export as BF16, routers as FP32, and the LM head as BF16. --------- Signed-off-by: Chad Voegele <cvoegele@nvidia.com> Signed-off-by: Rohan Joshi <rohjoshi@nvidia.com>
What does this PR do?
Type of change: Bug fix
Supports scalar
weight_scale_invvalues when converting Hugging Face FP8 linear layers. The scalar path dequantizes weights directly, while the existing block-scale Triton path remains unchanged. Forward execution andunpack_weight()share the same dequantization helper.Usage
No API changes. Hugging Face FP8 checkpoints using per-tensor weight scales convert without additional configuration.
Testing
test_fp8_linear_per_tensor_dequant— passed.mistralai/Mistral-Medium-3.5-128Bon AWS-PDX and converted all 616 scalar-scale FP8 linear modules.unpack_weight()passed onmodel.language_model.layers.0.self_attn.q_proj.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
This isolates the per-tensor FP8 dequantization fix from commit
2f913319d2ac2fbd47397ece871f18576f2ad009; unrelated recipe changes are excluded.Summary by CodeRabbit
Summary
block_sizeand using a dedicated dequantization path for more reliable results.