Autoquant and GPTQ in support in Megatron-Core [OMNIML-3151] - #1562
Merged
Conversation
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
…izers
The branch previously short-circuited mse_calibrate's Step 2 with an early
`continue` that skipped any quantizer whose config didn't match the NVFP4
static pattern (num_bits=(2,1) + scale_bits=(4,3)). This broke main's
contract that:
- fp8_scale_sweep=True + registered backend -> backend factory called
- any enabled quantizer -> calibrator replaced with
MseCalibrator (default)
Tests TestRegisterFP8SweepCalibrator::{
test_mse_calibrate_dispatches_to_registered_factory,
test_unregistered_backend_uses_default_mse_calibrator,
} regressed on this branch because they use INT8 quantizers which were
silently skipped.
Restructure so:
1. NVFP4-static promotion runs only when applicable (gated on
module.is_nvfp4_static)
2. Backend factory dispatch runs for any backend with fp8_scale_sweep=True
3. NVFP4MSECalibrator runs only for NVFP4-static + fp8_scale_sweep
4. MseCalibrator default fallback runs for everything else (INT8, FP8,
non-sweep NVFP4)
Also drops the misleading 'skipped non-NVFP4' warning (it implied we skip,
but we now always set a calibrator).
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jenny Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
jenchen13
requested review from
ajrasane,
realAsma and
sugunav14
and removed request for
a team
May 28, 2026 21:30
jenchen13
changed the title
Contributor
|
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:
📝 WalkthroughWalkthroughAutoQuantize now synchronizes score, cost, and final format selection across expert model parallel groups, derives weight budgets from candidate statistics, and adds Megatron-specific auto-quant and calibration hooks. ChangesAutoQuantize Expert Model Parallelism and Megatron Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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
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/utils/calib_utils.py (1)
60-61:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate the docstring note to reflect the new behavior.
The note states that "input must be non-empty" and "a zero-sized input causes division by zero", but the new guard clause at lines 66-67 now handles
batch_size == 0gracefully. Update the docstring to reflect that empty inputs are now supported.📝 Proposed docstring update
- Note: input must be non-empty (batch_size > 0); a zero-sized input causes division by zero. + Note: Empty inputs (batch_size == 0) are handled gracefully and return unchanged hessian/n_samples. + This can occur in MoE models when some experts receive no tokens.🤖 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/utils/calib_utils.py` around lines 60 - 61, Update the docstring Note to reflect that empty inputs are now supported: replace "input must be non-empty (batch_size > 0); a zero-sized input causes division by zero" with a sentence stating that the function now handles batch_size == 0 via the guard clause (which returns early when batch_size == 0) and will not raise a division-by-zero error; mention that non-empty inputs are still processed normally. Target the docstring for the function that contains the guard checking batch_size == 0 (the docstring immediately above that guard) and keep the wording brief and clear.
🧹 Nitpick comments (2)
modelopt/torch/quantization/plugins/megatron.py (1)
810-837: ⚡ Quick winDocument and export the newly added public APIs.
register_megatron_autoquant_supportandget_mcore_decoder_layersare public (non-underscore) but only one has a docstring, and neither is reflected in__all__.As per coding guidelines, "Document public APIs with docstrings, including examples when useful" and "Define the public API with
__all__at the top of each module".🤖 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/megatron.py` around lines 810 - 837, Add a docstring to the newly public function get_mcore_decoder_layers describing purpose, parameters, return type and an example, and ensure register_megatron_autoquant_support also has appropriate public-docstring coverage if needed; then export both symbols by adding "register_megatron_autoquant_support" and "get_mcore_decoder_layers" to the module's __all__ list at the top of the file so they are part of the public API surface.modelopt/torch/quantization/model_quant.py (1)
510-515: ⚡ Quick winDon’t silently swallow plugin import failures.
Line 514 currently suppresses all
ImportErrors, which can hide real regressions and make Megatron auto-quant support silently disappear. Emit a warning (or gate the exception type more narrowly) so failures are diagnosable.Proposed change
try: from .plugins.megatron import register_megatron_autoquant_support register_megatron_autoquant_support() - except ImportError: - pass + except ImportError as exc: + warnings.warn( + f"Skipping Megatron auto-quant support registration due to import error: {exc}", + RuntimeWarning, + stacklevel=2, + )🤖 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/model_quant.py` around lines 510 - 515, The current try/except around importing and calling register_megatron_autoquant_support silently swallows ImportError; update the block to either catch a more specific exception (e.g., ModuleNotFoundError for the plugin import) or log a warning when import/call fails so failures are visible; specifically wrap the import and call to register_megatron_autoquant_support() and on failure call the module's logger or warnings.warn/processLogger.warning with a clear message including the exception text and that Megatron auto-quant support is disabled.
🤖 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/megatron.py`:
- Around line 830-831: get_mcore_decoder_layers is mutating model.decoder.layers
by appending model.output_layer which causes duplicated entries on repeated
calls; instead return a new nn.ModuleList (e.g., copy model.decoder.layers into
a fresh list/ModuleList) and append the output_layer to that new collection or
check for existence before appending so augmentation is idempotent; update
get_mcore_decoder_layers (and calls from
LayerActivationCollector.get_decoder_layers /
LayerActivationCollector._patch_all_layers) to use the non-mutating copy so
_cleanup_layers need not undo permanent changes.
---
Outside diff comments:
In `@modelopt/torch/quantization/utils/calib_utils.py`:
- Around line 60-61: Update the docstring Note to reflect that empty inputs are
now supported: replace "input must be non-empty (batch_size > 0); a zero-sized
input causes division by zero" with a sentence stating that the function now
handles batch_size == 0 via the guard clause (which returns early when
batch_size == 0) and will not raise a division-by-zero error; mention that
non-empty inputs are still processed normally. Target the docstring for the
function that contains the guard checking batch_size == 0 (the docstring
immediately above that guard) and keep the wording brief and clear.
---
Nitpick comments:
In `@modelopt/torch/quantization/model_quant.py`:
- Around line 510-515: The current try/except around importing and calling
register_megatron_autoquant_support silently swallows ImportError; update the
block to either catch a more specific exception (e.g., ModuleNotFoundError for
the plugin import) or log a warning when import/call fails so failures are
visible; specifically wrap the import and call to
register_megatron_autoquant_support() and on failure call the module's logger or
warnings.warn/processLogger.warning with a clear message including the exception
text and that Megatron auto-quant support is disabled.
In `@modelopt/torch/quantization/plugins/megatron.py`:
- Around line 810-837: Add a docstring to the newly public function
get_mcore_decoder_layers describing purpose, parameters, return type and an
example, and ensure register_megatron_autoquant_support also has appropriate
public-docstring coverage if needed; then export both symbols by adding
"register_megatron_autoquant_support" and "get_mcore_decoder_layers" to the
module's __all__ list at the top of the file so they are part of the public API
surface.
🪄 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: 30c2390a-c99c-4b41-8c0c-0be68734dc77
📒 Files selected for processing (7)
modelopt/torch/quantization/algorithms.pymodelopt/torch/quantization/model_quant.pymodelopt/torch/quantization/nn/modules/tensor_quantizer.pymodelopt/torch/quantization/plugins/megatron.pymodelopt/torch/quantization/utils/calib_utils.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.pytests/unit/torch/quantization/test_autoquant.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1562 +/- ##
===========================================
+ Coverage 61.17% 76.84% +15.67%
===========================================
Files 515 515
Lines 57207 57684 +477
===========================================
+ Hits 34994 44328 +9334
+ Misses 22213 13356 -8857
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: Jennifer Chen <jennifchen@nvidia.com>
sugunav14
reviewed
Jun 2, 2026
|
|
||
|
|
||
| # GPTQ layerwise calibration support | ||
| def get_mcore_decoder_layers(model: torch.nn.Module) -> torch.nn.ModuleList | None: |
Contributor
There was a problem hiding this comment.
since we are returning both decoder layers and output layer could we rename this to better reflect that?
Contributor
Author
There was a problem hiding this comment.
i can rename to get_mcore_model_layers, but it's still being called by LayerActivationCollector.register_decoder_layer_support
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
sugunav14
reviewed
Jun 3, 2026
sugunav14
left a comment
sugunav14
left a comment
Contributor
There was a problem hiding this comment.
Reviewed the GPTQ support! LGTM
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
jenchen13
force-pushed
the
jennifchen/mcore_autoquant_gptq
branch
from
July 1, 2026 16:33
a196424 to
c976a23
Compare
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
jenchen13
force-pushed
the
jennifchen/mcore_autoquant_gptq
branch
from
July 1, 2026 17:47
77b8ed3 to
8c307de
Compare
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
jenchen13
force-pushed
the
jennifchen/mcore_autoquant_gptq
branch
from
July 1, 2026 18:15
d8b85bc to
2e5100c
Compare
kevalmorabia97
approved these changes
Jul 1, 2026
kevalmorabia97
left a comment
kevalmorabia97
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. Please also wait for 2-gpu tests to pass: https://github.com/NVIDIA/Model-Optimizer/actions/runs/28538480157
Contributor
Author
|
/ok to test 2e5100c |
realAsma
reviewed
Jul 1, 2026
realAsma
reviewed
Jul 1, 2026
| total_weight_size = self._cost_model.total_weight_size( | ||
| self.model.named_modules(), self._is_auto_quantize_module, self.config["cost"] | ||
| ) | ||
| if self.candidate_stats: |
Contributor
There was a problem hiding this comment.
@jenchen13 why do we need the else branch here? self.candidate_stats is always set in before_search which gets called before run_search
Contributor
There was a problem hiding this comment.
The else branch is never used.
Contributor
Author
There was a problem hiding this comment.
hmm I think i kept it for backward compatibility
realAsma
reviewed
Jul 1, 2026
Comment on lines
+919
to
+926
| if self.candidate_stats: | ||
| total_weight_size = self._get_total_weight_size_from_candidate_stats( | ||
| self.candidate_stats | ||
| ) | ||
| else: | ||
| total_weight_size = self._cost_model.total_weight_size( | ||
| self.model.named_modules(), self._is_auto_quantize_module, self.config["cost"] | ||
| ) |
Contributor
There was a problem hiding this comment.
Suggested change
| if self.candidate_stats: | |
| total_weight_size = self._get_total_weight_size_from_candidate_stats( | |
| self.candidate_stats | |
| ) | |
| else: | |
| total_weight_size = self._cost_model.total_weight_size( | |
| self.model.named_modules(), self._is_auto_quantize_module, self.config["cost"] | |
| ) | |
| total_weight_size = self._get_total_weight_size_from_candidate_stats( | |
| self.candidate_stats | |
| ) |
kevalmorabia97
disabled auto-merge
July 1, 2026 22:01
run_search's `else` branch recomputed the effective-bits denominator via _cost_model.total_weight_size(live model), but candidate_stats is always populated by before_search() (restored from checkpoint -> early return, or built by initialize_candidate_stats), so the branch was dead. Replace it with an assertion documenting the invariant and always derive the denominator from candidate_stats -- the source the solver's per-layer costs already come from, and which is DP/TP/EP-synced and merge-group aware (unlike the live-model recompute). Remove the now prod-dead AutoQuantizeCostModel.total_weight_size (its only caller was that else) and rewrite its two unit tests to assert the exclusion rule and MoE active-expert scaling directly against the production-used primitives (module_cost_weight + _get_module_weight_numel). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Contributor
Author
|
/ok to test 054e353 |
jenchen13
commented
Jul 1, 2026
jenchen13
commented
Jul 1, 2026
talorabr
pushed a commit
that referenced
this pull request
Jul 7, 2026
### What does this PR do? Type of change: New Feature Autoquant and GPTQ in support in Megatron-Core - Add EP support to AutoQuantize - Register MCore support in AutoQuantize - Add decoder `output_layer` (lm head) to layerwise hook so that GPTQ can register all decoder layers & lm head - Split dataloader helper function out of megatron calibration utils so that AutoQuantize in Megatron-LM can reuse the same dataloader ### Usage See NVIDIA/Megatron-LM#4821 for Autoquant usage in Megatron ```python # For GPTQ pick a recipe that uses gptq algorithm and run mtq.quantize # e.g. general/ptq/nvfp4_default-kv_none-gptq ``` ### Testing Tested AutoQuant on Nemotron Nano and Ultra. Tested GPTQ on Nano 3. Added unit tests for both AutoQuant and GPTQ ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A <!--- Mandatory --> - Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> - Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run `/claude review`. NVIDIA org members can self-trigger for complex changes; orthogonal to CodeRabbit. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Added lazy Megatron-Core AutoQuant integration with Megatron-specific quantization hooks and better decoder-layer discovery for layerwise calibration. * Improved AutoQuantize for expert-parallel (EP) models, including consistent per-layer recipe selection across DP/TP/EP. * Extended quant-layer grouping for NemotronH MCore fused “local_experts” linear layers. * **Bug Fixes** * Prevented division-by-zero when calibration inputs are empty during Hessian updates. * **Tests** * Added/extended unit and GPU coverage for EP AutoQuant, decoder-layer calibration discovery behavior, and zero-token Hessian no-op. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: Jenny Chen <jennifchen@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: talora <talora@nvidia.com>
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
6 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: New Feature
Autoquant and GPTQ in support in Megatron-Core
output_layer(lm head) to layerwise hook so that GPTQ can register all decoder layers & lm headUsage
See NVIDIA/Megatron-LM#4821 for Autoquant usage in Megatron
Testing
Tested AutoQuant on Nemotron Nano and Ultra.
Tested GPTQ on Nano 3.
Added unit tests for both AutoQuant and GPTQ
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
Summary by CodeRabbit
New Features
Bug Fixes
Tests