Add reversible weight folding and support TE grouped weights - #2140
Add reversible weight folding and support TE grouped weights#2140mxinO wants to merge 17 commits into
Conversation
Signed-off-by: Meng Xin <mxin@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. |
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughWeight folding now supports temporary state tracking, shared-storage validation, LoRA-aware behavior, and Transformer Engine grouped linear layers. Temporary folding restores weights and quantizer state after normal or exceptional exit. Tests cover restoration, grouped quantizers, calibration attributes, and distributed weight access. ChangesWeight folding
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Model
participant FoldingContext
participant QuantModule
participant QuantizerState
Model->>FoldingContext: enter temporarily_fold_weights
FoldingContext->>QuantModule: fold eligible tensor weights
QuantModule->>QuantizerState: snapshot and update state
FoldingContext-->>Model: expose folded weights
Model->>FoldingContext: exit or raise exception
FoldingContext->>QuantizerState: restore state and weights
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2140 +/- ##
==========================================
+ Coverage 67.09% 67.14% +0.05%
==========================================
Files 522 522
Lines 60461 60509 +48
==========================================
+ Hits 40567 40631 +64
+ Misses 19894 19878 -16
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:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
The core change is correct and well-scoped: _fold_weight_quantizer now iterates the stages of a SequentialQuantizer (its disable()/disable_rotate() already broadcast via _QuantizerContainerBase), the base fold_weight accepts sequential weight quantizers, and the new _QuantTEGroupedLinear.fold_weight fixes a real bug — previously a GroupedQuantizer was silently skipped by the dir() scan, and a shared TensorQuantizer on TEGroupedLinear would even trip the assert hasattr(self, weight_name) since _setup deletes self.weight. Batching weights per quantizer before disable() is the right ordering, and the new CPU unit test verifies output invariance plus the keep_attrs matrix.
Three things worth addressing before merge:
-
The same bug is left unfixed one file over.
_QuantFusedExperts.fold_weight(plugins/huggingface.py) still gates onisinstance(q, TensorQuantizer), so per-expert quantizers thatset_quantizer_attributes_fullpromoted toSequentialQuantizer(a listcfgmatchesgate_up_proj_weight_quantizers.Nvia_normalize_fused_experts_quantizer_name, e.g. W4A8) are still silently skipped — folding becomes a no-op there. Since the helper you just generalized handles sequential fine, this is a one-token fix. -
Duplicated logic. The new TE
fold_weightis entirely generic overiter_weights_for_calibration(); it is a near-copy of_QuantFusedExperts.fold_weight, and vLLM's_QuantFusedMoEBase.fold_weightis a third variant. Consider hoisting the quantizer→weights batching loop intoQuantModule(e.g._fold_weights_from_calibration_iter) and having TE and fused-experts both call it — that removes the duplication and fixes (1) in the same stroke. -
Behavior change under
keep_attrs=Trueis undocumented. Setting_enable_pre_quant_scale = Falseis correct (the weight-quantizerpre_quant_scaleis applied before the_disabledearly-return inTensorQuantizer.forward, so a retained buffer would double-apply — and_apply_weight_pre_quant_scalewith_ENABLE_FOLDING_PQS_TO_WEIGHTS=Falsedoes put a live pqs on weight quantizers during auto_quantize). But it meansquantizer.pre_quant_scalenow returnsNoneafter folding while_pre_quant_scaleis still present, andpre_quant_scale's setter asserts on_enable_pre_quant_scale. Please state this in thefold_weightdocstring, and reconsider "Changelog: N/A" givenkeep_attrs=Trueis a public flag whose observable semantics changed.
Also flagging (nudge-level, no action strictly required): the only coverage for the new TE path is a tests/gpu_megatron test the author states was never executed locally, and it asserts only weights/_amax — not post-fold forward equality or the new _enable_pre_quant_scale flag.
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
mxinO
changed the title
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: 2
🧹 Nitpick comments (2)
tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py (1)
155-205: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the shared grouped-linear setup into a fixture.
Lines 157-166 duplicate lines 123-132 exactly. A small helper or fixture that returns
(model, calib_data, grouped_linear, weights, quantizers)for a givenshare_weight_quantizerwould keep both tests aligned when the grouped path changes.The coverage itself matches the earlier review request: the test now asserts folded output equality and pre-quant-scale inertness on the grouped and shared paths.
🤖 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/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py` around lines 155 - 205, Extract the duplicated grouped-linear setup from test_temporarily_fold_weight_grouped_linear and the earlier grouped-linear test into a shared fixture or helper. Have it accept share_weight_quantizer and return model, calib_data, grouped_linear, weights, and quantizers, then update both tests to reuse it while preserving their existing assertions.tests/unit/torch/quantization/test_tensor_quant_cpu.py (1)
417-422: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert that the first module was folded before the failure.
The test verifies restoration only. It does not prove that
firstwas folded beforesecondraised. If a future change makes the fold offirsta silent no-op, this test still passes. Add a probe that records the folded weight offirstat the moment the second fold fails, for example by wrappingsecond.fold_weightinstead of raising inside the backend.Attribution: the path instructions for
tests/**/*.pyrequire that tests exercise the behavior they claim to validate.🤖 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/torch/quantization/test_tensor_quant_cpu.py` around lines 417 - 422, Update the test around the first and second module folding so it records whether the first module’s weight changed before the second fold raises. Wrap or spy on second.fold_weight to capture first’s folded weight at failure, assert that this probe observes the folded state, then retain the existing assertions verifying weight and quantizer restoration.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.
Inline comments:
In `@modelopt/torch/quantization/nn/modules/quant_linear.py`:
- Around line 72-74: Ensure the standard PyTorch state_dict round trip preserves
the disabled SVDQuant LoRA state set by fold_weight(keep_attrs=True), so
reloaded modules do not re-enable retained factors or apply the residual twice.
Update the relevant quantized linear module serialization or reload behavior
around _enable_svdquant_lora and _svdquant_lora_a, then add a regression test
covering fold, state_dict save/load, and equivalent inference outputs.
In `@modelopt/torch/quantization/nn/modules/quant_module.py`:
- Around line 115-147: Update _shared_parameter_storages to track each
parameter’s actual memory range, accounting for shape, strides, storage offset,
and element size, and mark an allocation as shared only when parameter ranges
overlap; retain conservative handling for unsupported or ambiguous layouts.
Ensure _fold_weight_quantizer continues blocking overlapping aliases while
allowing folds for non-overlapping packed views, and add tests covering both
cases.
---
Nitpick comments:
In `@tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py`:
- Around line 155-205: Extract the duplicated grouped-linear setup from
test_temporarily_fold_weight_grouped_linear and the earlier grouped-linear test
into a shared fixture or helper. Have it accept share_weight_quantizer and
return model, calib_data, grouped_linear, weights, and quantizers, then update
both tests to reuse it while preserving their existing assertions.
In `@tests/unit/torch/quantization/test_tensor_quant_cpu.py`:
- Around line 417-422: Update the test around the first and second module
folding so it records whether the first module’s weight changed before the
second fold raises. Wrap or spy on second.fold_weight to capture first’s folded
weight at failure, assert that this probe observes the folded state, then retain
the existing assertions verifying weight and quantizer restoration.
🪄 Autofix
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: 6fa0afd9-94db-4484-a6a8-1c33c42a3a1f
📒 Files selected for processing (10)
CHANGELOG.rstmodelopt/torch/quantization/model_quant.pymodelopt/torch/quantization/nn/modules/quant_linear.pymodelopt/torch/quantization/nn/modules/quant_module.pymodelopt/torch/quantization/plugins/huggingface.pymodelopt/torch/quantization/plugins/transformer_engine.pytests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.pytests/unit/torch/quantization/plugins/test_huggingface.pytests/unit/torch/quantization/test_calib.pytests/unit/torch/quantization/test_tensor_quant_cpu.py
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/quantization/plugins/transformer_engine.py
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
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: 2
🤖 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/model_quant.py`:
- Around line 776-786: Update the restoration logic in the finally block to
catch and collect errors from each weight.copy_ and quantizer-state restoration,
continuing through all remaining weights and states. After both loops complete,
re-raise the first collected restoration error, while preserving the existing
no-grad context and missing-attribute handling.
In `@modelopt/torch/quantization/nn/modules/quant_module.py`:
- Around line 68-78: Update _tensor_storage_key to return None for empty or meta
tensors, and adjust _tied_parameter_storages to exclude None keys from its
counts. Preserve existing storage-key behavior for regular non-empty parameters
so only genuinely tied storages are reported.
🪄 Autofix
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: db6d62ca-3f03-43a2-854e-1a08e0e14176
📒 Files selected for processing (5)
CHANGELOG.rstmodelopt/torch/quantization/model_quant.pymodelopt/torch/quantization/nn/modules/quant_module.pytests/unit/torch/quantization/plugins/test_huggingface.pytests/unit/torch/quantization/test_tensor_quant_cpu.py
💤 Files with no reviewable changes (1)
- tests/unit/torch/quantization/plugins/test_huggingface.py
🚧 Files skipped from review as they are similar to previous changes (1)
- CHANGELOG.rst
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/quantization/nn/modules/quant_linear.py (1)
72-74: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftPrevent double application of the SVDQuant residual after permanent folding.
A non-temporary fold already adds the LoRA residual to
self.weight. These getters still return both retained buffers, soSVDQuantLinear.forwardcan add the same residual again. Keep the factors active only for temporary folds. Otherwise, remove them or persist an explicit inactive state after permanent folding. Add regression tests forfold_weight(keep_attrs=True)and a standardstate_dict()round trip. (raw.githubusercontent.com)Also applies to: 96-98
🤖 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/nn/modules/quant_linear.py` around lines 72 - 74, Update the SVDQuant LoRA-factor getters around _svdquant_lora_a and the corresponding factor getter so permanently folded residuals are not returned to SVDQuantLinear.forward; retain them only for temporary folds, or track an explicit inactive state after permanent folding. Ensure fold_weight(keep_attrs=True) and a standard state_dict() round trip do not reapply the residual, and add regression coverage for both cases.Source: MCP tools
🤖 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.
Outside diff comments:
In `@modelopt/torch/quantization/nn/modules/quant_linear.py`:
- Around line 72-74: Update the SVDQuant LoRA-factor getters around
_svdquant_lora_a and the corresponding factor getter so permanently folded
residuals are not returned to SVDQuantLinear.forward; retain them only for
temporary folds, or track an explicit inactive state after permanent folding.
Ensure fold_weight(keep_attrs=True) and a standard state_dict() round trip do
not reapply the residual, and add regression coverage for both cases.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 875b77a0-674a-4f86-b621-1b35699d6687
📒 Files selected for processing (6)
CHANGELOG.rstmodelopt/torch/quantization/model_quant.pymodelopt/torch/quantization/nn/modules/quant_linear.pymodelopt/torch/quantization/nn/modules/quant_module.pytests/unit/torch/quantization/test_calib.pytests/unit/torch/quantization/test_tensor_quant_cpu.py
💤 Files with no reviewable changes (1)
- tests/unit/torch/quantization/test_calib.py
🚧 Files skipped from review as they are similar to previous changes (3)
- modelopt/torch/quantization/model_quant.py
- tests/unit/torch/quantization/test_tensor_quant_cpu.py
- modelopt/torch/quantization/nn/modules/quant_module.py
Signed-off-by: Meng Xin <mxin@nvidia.com>
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 622-627: Update fold_weight to group weight views by each
TensorQuantizer identity before calling QuantModule._fold_weight_quantizer, so
every shared quantizer is folded once with all associated weights. Preserve the
existing fake_quant and keep_attrs conditions, and ensure duplicate quantizer
views are not processed through separate mutation calls.
🪄 Autofix
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: dde3c614-381c-4da2-b3cc-3fe8b3b7d692
📒 Files selected for processing (5)
CHANGELOG.rstmodelopt/torch/quantization/model_quant.pymodelopt/torch/quantization/nn/modules/quant_module.pymodelopt/torch/quantization/plugins/huggingface.pytests/unit/torch/quantization/test_tensor_quant_cpu.py
🚧 Files skipped from review as they are similar to previous changes (3)
- CHANGELOG.rst
- modelopt/torch/quantization/nn/modules/quant_module.py
- tests/unit/torch/quantization/test_tensor_quant_cpu.py
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
…izer-containers Signed-off-by: Meng Xin <mxin@nvidia.com> # Conflicts: # CHANGELOG.rst
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of PR #2140, now reshaped from "fix sequential folding" into a new public API: mtq.temporarily_fold_weights (+ TE grouped fold_weight). Size is fine (10 files, +325/-19) and the CPU tests are much better than the previous round (restore-after-exception, rollback when folding fails, SequentialQuantizer rejection, disabled-transform/weight=None, SVDQuant residual kept separate).
Design-review note (gate fired): the PR body doesn't discuss alternatives for the new context. The main design question is the discovery/mutation split: the context discovers what to snapshot via QuantModule.iter_weights_for_calibration() but delegates mutation to each plugin's fold_weight(). Those two sets are not guaranteed to agree, and where they don't the context silently fails to restore (see inline: vLLM _QuantFusedMoEBase). The repo also already has get_quantizer_state_dict/set_quantizer_state_dict and TensorQuantizer.get_modelopt_state() for capturing quantizer state, versus the four hand-picked private attrs here — please say in the PR body why those weren't reused, and why the snapshot is keyed off the calibration iterator rather than off what fold_weight actually touches.
Status of prior review points:
- ✅
keep_attrs=Truepre-quant-scale semantics are now documented in thefold_weightdocstring and CHANGELOG. - ✅ New CPU coverage for the temporary context is meaningful.
- ❌ TE GPU test: the reply says the tests "cover permanent and temporary folding ... verify retained PQS is inactive while folded, and verify weights/quantizer state are restored afterward", but the final diff only contains
test_fold_weight_grouped_linear(permanent fold; asserts forward equality and_amaxremoval). The temporary-fold grouped test and the_enable_pre_quant_scaleassertion aren't there. - 💬 Duplication (TE / HF fused-experts / vLLM fold loops) — author explained the deliberate choice to keep TE's batching module-specific; the third near-copy still stands, flagging for the owner rather than re-arguing.
- ❌
_QuantFusedExperts.fold_weight'sisinstance(q, TensorQuantizer)gate (silent no-op fold for W4A8 per-expertSequentialQuantizers) is unchanged, and the base_fold_weight_quantizer's sequential support from the earlier commits was reverted, somtq.fold_weightnow silently skips sequential weight quantizers while the new context raises for them. - Also: the earlier reply stated the context "raises before mutation" for weights/quantizers shared across modules and that a shared-storage test verifies rejection + rollback — that guard and test are not in the final diff; the limitation is now only prose in the docstring.
Nit: CHANGELOG mentions the pre-quant-scale change but not the SVDQuant fold_weight(keep_attrs=True) behavior change (residual is no longer baked into the weight), which is also public-facing.
No licensing concerns (standard NVIDIA/Apache headers, no vendored code).
Additional comments (outside the PR diff):
modelopt/torch/quantization/nn/modules/quant_module.py:175— > Bot comment.
With the sequential support from the earlier commits reverted, fold_weight now silently skips a SequentialQuantizer weight quantizer (the isinstance(attr, TensorQuantizer) gate), while temporarily_fold_weights raises NotImplementedError for the same case. That asymmetry means mtq.fold_weight(model, ...) on a W4A8 model reports success while folding nothing — the exact silent no-op flagged last round for _QuantFusedExperts.fold_weight (whose isinstance(q, TensorQuantizer) gate is also unchanged). If sequential folding is out of scope now, please make the permanent path warn or raise too, so the two entry points agree and the no-op isn't silent.
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
2 participants
What does this PR do?
Adds
mtq.temporarily_fold_weights(model, snapshot_device=None)for repeated inference over a frozen fake-quantized model. The context snapshots affected weights and quantizer runtime state, calls each module's nativefold_weight(keep_attrs=True), and restores state on exit, including after exceptions.It also:
GroupedLinearweights with grouped or sharedTensorQuantizers;w13/w2weights for weight calibration and temporary snapshots; andkeep_attrs=True.SequentialQuantizerweights and weights or quantizers shared acrossQuantModuleinstances are not supported by the temporary context.Related use case: NVIDIA-NeMo/RL#3441.
Testing
Focused CPU folding and restoration tests pass. Transformer Engine grouped folding is covered by GPU CI.