Fix per-tensor FP8 weight dequantization#1962
Conversation
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1962 +/- ##
==========================================
- Coverage 77.81% 77.19% -0.62%
==========================================
Files 519 519
Lines 57960 57964 +4
==========================================
- Hits 45101 44747 -354
- Misses 12859 13217 +358
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:
|
Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
chadvoegele
force-pushed
the
cvoegele/mistral3-fp8-per-tensor-dequant
branch
from
7069c2b to
db3edfb
Compare
📝 WalkthroughWalkthroughChangesFP8 dequantization
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
| def _dequantize_weight(self, dtype: torch.dtype) -> Tensor: | ||
| weight, scale_inv = self._get_weight_and_scale_inv() | ||
| if self.block_size is None: | ||
| return weight.to(dtype) * scale_inv.to(dtype) | ||
| assert weight_dequant is not None, "Triton is not available" | ||
| return weight_dequant(weight, scale_inv, self.block_size, dtype=dtype) | ||
|
|
||
| def forward(self, input: Tensor) -> Tensor: | ||
| if self.weight.element_size() == 1: | ||
| with torch.cuda.device(self.weight.device): | ||
| weight, scale_inv = self._get_weight_and_scale_inv() | ||
| weight = weight_dequant(weight, scale_inv, self.block_size, dtype=input.dtype) | ||
| weight = self._dequantize_weight(input.dtype) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
torch.cuda.device() wrapping breaks the new scalar dequant path on CPU tensors.
forward and unpack_weight unconditionally wrap _dequantize_weight(...) in torch.cuda.device(self.weight.device). That context manager requires a CUDA device — _get_device_index raises ValueError: Expected a cuda device, but got: cpu for CPU tensors (confirmed via PyTorch source/behavior). The new block_size is None branch is a plain weight.to(dtype) * scale_inv.to(dtype) that works on any device and doesn't need this context at all; only the Triton weight_dequant call actually requires it. As written, calling forward() or unpack_weight() on a CPU-resident scalar-scale FP8 layer will crash — precisely the case this PR is trying to support independently of Triton/GPU. Note the new test_fp8_linear_per_tensor_dequant test calls _dequantize_weight directly, so it doesn't exercise this code path and misses the bug.
Move the CUDA context inside _dequantize_weight, scoped only to the Triton branch:
🐛 Proposed fix
def _dequantize_weight(self, dtype: torch.dtype) -> Tensor:
weight, scale_inv = self._get_weight_and_scale_inv()
if self.block_size is None:
return weight.to(dtype) * scale_inv.to(dtype)
assert weight_dequant is not None, "Triton is not available"
- return weight_dequant(weight, scale_inv, self.block_size, dtype=dtype)
+ with torch.cuda.device(weight.device):
+ return weight_dequant(weight, scale_inv, self.block_size, dtype=dtype)
def forward(self, input: Tensor) -> Tensor:
if self.weight.element_size() == 1:
- with torch.cuda.device(self.weight.device):
- weight = self._dequantize_weight(input.dtype)
+ weight = self._dequantize_weight(input.dtype)
else:
weight = self.weight
return linear(
self.input_quantizer(input),
self.weight_quantizer(weight),
self.bias,
)
def unpack_weight(self):
- with torch.cuda.device(self.weight.device):
- self.weight = nn.Parameter(
- self._dequantize_weight(torch.get_default_dtype()),
- requires_grad=False,
- )
+ self.weight = nn.Parameter(
+ self._dequantize_weight(torch.get_default_dtype()),
+ requires_grad=False,
+ )
if hasattr(self, "weight_scale_inv"):
del self.weight_scale_invAlso applies to: 1372-1377
🤖 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 `@modelopt/torch/quantization/plugins/huggingface.py` around lines 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.
cjluo-nv
left a comment
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.
2 participants
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
Bug Fixes
Tests