Bug fix: disable weight quantizer rotation after weight fold during vLLM fakequant export#1143
Bug fix: disable weight quantizer rotation after weight fold during vLLM fakequant export#1143kinjalpatel27 wants to merge 3 commits intomainfrom
Conversation
Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
📝 WalkthroughWalkthroughThe export flow was modified to preserve each weight quantizer's rotation configuration during the temporary disable/enable window used for snapshotting ModelOpt quantizer state. Rotation settings are now captured before disabling, restored to a disabled form while snapshotting, and then reset to their original values after completion. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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.
Actionable comments posted: 1
🧹 Nitpick comments (1)
modelopt/torch/export/plugins/vllm_fakequant_hf.py (1)
107-122: Please add arotate=Trueexport/reload regression test.This only fails when export, reload, and a later
fold_weight()interact, so it is easy to regress silently without an end-to-end check.As per coding guidelines, "Maintain minimum 70% code coverage on modelopt/* modules (enforced via coverage configuration)."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelopt/torch/export/plugins/vllm_fakequant_hf.py` around lines 107 - 122, Add a regression test that exercises export → reload → fold_weight to ensure the quantizer rotation state is preserved: create a small model using QuantModule with a TensorQuantizer that has rotate enabled, export it using the logic in vllm_fakequant_hf.py (so the code path that disables quantizer.rotate and stores orig_rotate in wqs_to_restore runs), reload the exported model, call fold_weight() on the reloaded model, and assert that weights are identical to the pre-export folded result and that the TensorQuantizer._rotate value is restored to its original boolean; target the test at the export/reload flow that manipulates QuantModule, TensorQuantizer, rotate/rotate_is_enabled, and fold_weight so regressions are caught.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@modelopt/torch/export/plugins/vllm_fakequant_hf.py`:
- Around line 118-122: The code disables quantizers and mutates
quantizer._rotate but does not guarantee restoration on error; wrap the section
that disables quantizers and the following processing (the region that builds
wqs_to_restore and does work up to where restores currently happen) in a
try/finally so restoration always runs: in the finally iterate wqs_to_restore
and call quantizer.enable() (or reverse the earlier disable call) and restore
quantizer._rotate to orig_rotate (using quantizer.rotate_is_enabled to decide if
it was changed), and swallow/log any exceptions during restore to avoid masking
the original exception; reference the quantizer object, wqs_to_restore list,
quantizer.disable/enable, quantizer._rotate and quantizer.rotate_is_enabled so
you restore state even if an exception occurs.
---
Nitpick comments:
In `@modelopt/torch/export/plugins/vllm_fakequant_hf.py`:
- Around line 107-122: Add a regression test that exercises export → reload →
fold_weight to ensure the quantizer rotation state is preserved: create a small
model using QuantModule with a TensorQuantizer that has rotate enabled, export
it using the logic in vllm_fakequant_hf.py (so the code path that disables
quantizer.rotate and stores orig_rotate in wqs_to_restore runs), reload the
exported model, call fold_weight() on the reloaded model, and assert that
weights are identical to the pre-export folded result and that the
TensorQuantizer._rotate value is restored to its original boolean; target the
test at the export/reload flow that manipulates QuantModule, TensorQuantizer,
rotate/rotate_is_enabled, and fold_weight so regressions are caught.
🪄 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: Pro
Run ID: 8f9e9ddd-4ce2-43db-b2d3-5afccc6c397a
📒 Files selected for processing (1)
modelopt/torch/export/plugins/vllm_fakequant_hf.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1143 +/- ##
==========================================
+ Coverage 70.18% 70.20% +0.02%
==========================================
Files 230 230
Lines 26080 26080
==========================================
+ Hits 18304 18310 +6
+ Misses 7776 7770 -6 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
LGTM. The bug is real — TensorQuantizer.forward() applies Hadamard rotation before the _disabled early-return, so when fold_weight is called on vLLM reload (which checks fake_quant but not is_enabled), the already-folded weight gets a second rotation. The fix correctly clears _rotate alongside disable() before saving the modelopt state, and restores both afterward to keep the export non-mutating.
One minor observation: _rotate can be a bool, dict, or RotateConfig — setting it to False works correctly for the serialized state (since rotate_is_enabled handles all three types), but preserving the original type (e.g. RotateConfig(enable=False)) would be slightly cleaner. Not a blocker.
Worth noting that the deeper root cause is QuantModule.fold_weight not guarding on is_enabled / _disabled before invoking the quantizer, but that's a broader change — this targeted fix is appropriate.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
modelopt/torch/export/plugins/vllm_fakequant_hf.py (1)
128-132: Logic is correct; consider optionaltry/finallyfor robustness.The sequence properly saves
orig_rotatebefore mutation and correctly usesrotate_is_enabledto avoid unnecessary changes. The fix addresses the double-rotation bug effectively.As noted in past review, wrapping lines 128-165 in a
try/finallywould guarantee restoration even if an exception occurs mid-export (relevant when callers catch exceptions and continue). This is optional given your deployment context but would make the function safer as reusable library code.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelopt/torch/export/plugins/vllm_fakequant_hf.py` around lines 128 - 132, Wrap the block that disables quantizers and records originals (the code that calls quantizer.disable(), saves orig_rotate from quantizer._rotate, checks quantizer.rotate_is_enabled and sets quantizer._rotate = disable_rotate(quantizer), and appends to wqs_to_restore) in a try/finally so that the finally always iterates wqs_to_restore and restores each quantizer._rotate and re-enables any state as needed; ensure you reference the same symbols (quantizer, orig_rotate, disable_rotate, wqs_to_restore) when restoring to guarantee restoration even if an exception occurs during export.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@modelopt/torch/export/plugins/vllm_fakequant_hf.py`:
- Around line 128-132: Wrap the block that disables quantizers and records
originals (the code that calls quantizer.disable(), saves orig_rotate from
quantizer._rotate, checks quantizer.rotate_is_enabled and sets quantizer._rotate
= disable_rotate(quantizer), and appends to wqs_to_restore) in a try/finally so
that the finally always iterates wqs_to_restore and restores each
quantizer._rotate and re-enables any state as needed; ensure you reference the
same symbols (quantizer, orig_rotate, disable_rotate, wqs_to_restore) when
restoring to guarantee restoration even if an exception occurs during export.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 6c0d7633-e5eb-4ca7-8003-e0d0f2fa9fe6
📒 Files selected for processing (1)
modelopt/torch/export/plugins/vllm_fakequant_hf.py
3 participants
What does this PR do?
Type of change: Bug fix
When
export_hf_vllm_fq_checkpointfolds fake-quantized weights into the HF checkpoint, it applies the weight quantizer's full forward pass — including any Hadamard rotation — to producefake_quant(rotate(W)). The weight quantizer is then disabled in the saved ModelOpt state so vLLM reload skips re-quantization.However, if
fold_weightis called after reload (e.g. infakequant_worker.py),QuantModule.fold_weightchecksfake_quantbut notis_enabled, so it calls the disabled weight quantizer. Withrotate=True, rotation runs before the_disabledearly-return inTensorQuantizer.forward, corrupting the already-folded weight with a second rotation.This fix clears
_rotateon weight quantizers alongsidedisable()before saving the ModelOpt state. Both are restored on the in-memory model after saving, so the export remains non-mutating.Usage
Testing
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Summary by CodeRabbit