[Fix] Calibrate non-decoder modules during layerwise quantization - #2339
[Fix] Calibrate non-decoder modules during layerwise quantization#2339realAsma wants to merge 6 commits into
Conversation
|
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. |
|
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 (4)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. 📝 WalkthroughWalkthroughLayerwise calibration now handles enabled quantizers outside transformer decoder layers. It hides decoder subtrees during full-model calibration, integrates the process with export and restoration paths, and adds coverage for calibration, offload, cleanup, and export behavior. ChangesOutside Quantizer Calibration
Priority: ⬇️ Low Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant LayerwiseCalibration
participant OutsideQuantizerCalibrator
participant DecoderLayers
participant OutsideQuantizers
LayerwiseCalibration->>OutsideQuantizerCalibrator: identify enabled outside quantizers
OutsideQuantizerCalibrator->>DecoderLayers: hide decoder subtrees
OutsideQuantizerCalibrator->>OutsideQuantizers: run full-model calibration
OutsideQuantizerCalibrator->>DecoderLayers: restore decoder subtrees
LayerwiseCalibration->>OutsideQuantizers: retain calibrated state
Suggested reviewers: Merge Risk: ⚪ Minimal · up to No merge-blocking calibration risk was identified. 🚥 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 #2339 +/- ##
==========================================
+ Coverage 75.52% 78.65% +3.12%
==========================================
Files 542 542
Lines 63778 63838 +60
==========================================
+ Hits 48167 50209 +2042
+ Misses 15611 13629 -1982
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:
|
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| decoder_owned_ids = {id(module) for layer in transformer_layers for module in layer.modules()} | ||
| has_enabled_outside_quantizer = any( | ||
| isinstance(module, TensorQuantizer) | ||
| and module.is_enabled | ||
| and id(module) not in decoder_owned_ids | ||
| for module in model.modules() | ||
| ) | ||
|
|
||
| if export_dir is not None and has_enabled_outside_quantizer: | ||
| raise ValueError( | ||
| "Layerwise export does not support enabled quantizers outside transformer layers. " | ||
| "Calibrate without export_dir, then export the completed model separately." | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Performance] The gate is module.is_enabled, not "this quantizer still needs data-driven calibration". Any enabled non-decoder quantizer — including a weight-only quantizer, a top-level type: dynamic quantizer, an MX (MXFP4/MXFP8) quantizer, or one pinned via constant_amax — flips has_enabled_outside_quantizer to True.
Why it matters: for a weight-only recipe that enables lm_head (e.g. INT8_WEIGHT_ONLY_CFG, INT4_BLOCKWISE_WEIGHT_ONLY_CFG, W4A16 AWQ), the only outside quantizer is lm_head.weight_quantizer, which max_calibrate calibrates directly on the weight tensor via weight_only_quantize() — the forward_loop contributes nothing. But the block at line 2232 unconditionally runs calib_func(model, forward_loop, ...), i.e. a full extra pass of the entire calibration dataset through the whole model. On the exact models layerwise calibration exists for (large, accelerate/disk-offloaded), that is the single most expensive thing in the run, and for these recipes it is pure waste. The same applies to a fully-dynamic activation quantizer outside the decoder, which needs no amax at all.
This file already has the precise predicate for the forward question — _needs_activation_forward_for_max_calib() (line 268) — and max_calibrate already accepts skip_forward_without_activation_calib.
Suggested shape: keep the calib_func invocation gated on "some outside quantizer needs any calibration" (weight amax counts), but gate the forward on whether an outside activation quantizer needs data — e.g. compute the flag over the hidden-decoder view and pass skip_forward_without_activation_calib=True for this extra pass when calib_func supports it:
if has_enabled_outside_quantizer:
...
with _hide_modules_from_traversal(model, transformer_layers):
extra_kwargs = dict(calib_kwargs)
if calib_func is max_calibrate:
# Outside quantizers may be weight-only / dynamic / MX; let max_calibrate
# skip the (full-model, full-dataset) forward when no activation stats are needed.
extra_kwargs.setdefault("skip_forward_without_activation_calib", True)
...At minimum, please make the export_dir rejection at line 2097 use the narrower predicate too, so recipes whose outside quantizers need nothing data-driven don't lose layerwise export for no reason.
There was a problem hiding this comment.
🤖 Bot comment.
Published in ed7b4d7. The outside calibration pass now runs only when needed, preserves omitted versus explicit skip-forward behavior, and blocks export whenever that pass would execute the progressively packed model. Focused pre-commit and all 50 layerwise calibration tests pass.
There was a problem hiding this comment.
Claude review — layerwise non-decoder calibration
Scope: full review (trigger comment had no scoping instructions). 3 files changed (274+/2-); reviewed all three: modelopt/torch/quantization/model_calib.py, modelopt/torch/quantization/utils/layerwise_calib.py, tests/unit/torch/quantization/test_layerwise_calibrate.py. Also read plugins/megatron.py / huggingface.py decoder-layer registration and modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml to judge blast radius.
Findings — CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 3
| # | Severity | Where | Issue |
|---|---|---|---|
| 1 | IMPORTANT Performance | model_calib.py:2089-2101 |
Gate is is_enabled, not "needs data-driven calibration" — weight-only / dynamic / MX / constant-amax outside quantizers trigger a full extra pass of the whole calibration dataset through the whole model |
| 2 | SUGGESTION | model_calib.py:2224-2232 |
has_enabled_outside_quantizer is a rank-local decision guarding collectives; PP stage asymmetry could hang |
| 3 | SUGGESTION | layerwise_calib.py:105-124 |
_ForwardOnlyLayer duplicates _SkipLayer delegation and reaches into its private blocklist; proxy silently discards attribute writes |
| 4 | SUGGESTION | model_calib.py:2097-2101 |
Export error message omits the fix most users want (disable the outside quantizer) and doesn't name the offenders |
Most impactful
Finding 1 is the only one I would hold the PR for. The bug being fixed is real and the mechanism is sound, but the trigger condition is broader than the need. For a weight-only recipe that enables lm_head, the only outside quantizer is lm_head.weight_quantizer, which max_calibrate calibrates directly on the weight tensor in weight_only_quantize() — the forward_loop contributes nothing, yet the new block runs a complete extra dataset pass over the full (often disk-offloaded) model. This file already has the exact predicate (_needs_activation_forward_for_max_calib, line 268) and max_calibrate already has the skip_forward_without_activation_calib knob. The same over-broad predicate makes the new export_dir ValueError fire for configs where nothing was actually miscalibrated.
What I verified as correct
- Traversal hiding actually hides.
_originalis set viaobject.__setattr__, so it never lands in_modules;enable_stats_collection/finish_stats_collection/_needs_activation_forward_for_max_calib/SharedWeightGlobalAmaxState.attachall walknamed_modules()and therefore cannot re-touch decoder quantizers. Decoder_amaxsurvives the extra pass — and the test asserts exactly that. - Forward fidelity.
forwardcallsself._original(...), i.e.type(original).__call__, so HFGradientCheckpointingLayer.__call__overrides and accelerate's_hf_hook-wrapped forward still run._hf_hook/_old_forwardstaying on_PROXY_BLOCKLISTcorrectly keeps accelerate from trying to manage the parameter-free proxy. - Slot restoration. Parents come from
model.modules(), aliased slots share one proxy keyed byid, proxies are built before thetry, and reassigning an existing_moduleskey preserves insertion order — sonn.ModuleListindexing and ordering are intact on restore. Covered on both the success and exception paths bytest_hide_modules_from_traversal_restores_aliases. get_qdq_activations_from_prev_layersemantics match the per-layer loop in both directions (QDQ propagation when True, all decoder quantizers disabled viaExitStackso the tail sees FP when False), and the test pins the tail input to 1.0 / 2.0 accordingly.- Offloaded weight writeback is not a gap. The extra pass is not wrapped in
persistent_materialization, but the weight-mutating algorithms (gptq,awq_lite,smoothquant,svdquant) each useenable_weight_access_and_writebackinternally, solm_headupdates are not dropped on accelerate-offloaded models. - Compat blast radius is small.
default_disabled_quantizers.yamlalready disables*lm_head*,*output_layer*, embeddings, routers, and the vision branch, and MCore's discoverer foldsoutput_layerinto the layer list — so the newexport_direrror only fires for custom recipes that deliberately enable an outside quantizer. Ordering the check beforeLayerwiseExporterconstruction (fail in seconds, not hours) matches the existing convention in this function. - Test coverage is genuinely targeted, not incidental: alias restoration on success and error, traversal-hiding assertions, QDQ-vs-FP tail input, decoder amax preservation, plain-list discovery, the no-op path, disk-warning gating across four device maps, and
export_dirfail-closed with an assertion that no directory was created.
Risk: low-to-moderate. The mechanism is well-contained and well-tested, and the state-composition story checks out — nothing here can corrupt already-calibrated decoder quantizers. The residual risk is cost, not correctness: as written, a class of common weight-only recipes pays a full extra calibration pass, and loses layerwise export, for calibration work that needs no forward at all. Narrowing the predicate addresses both.
Nothing here duplicates CodeRabbit's pattern gate — no security anti-patterns, style, or typo findings.
|
/claude review |
There was a problem hiding this comment.
Claude review — layerwise non-decoder calibration (re-review at ed7b4d7d)
Scope: full review (trigger comment carried no scoping instructions). 4 files changed (453+/8-); reviewed all four: modelopt/torch/quantization/mode.py, model_calib.py, utils/layerwise_calib.py, tests/unit/torch/quantization/test_layerwise_calibrate.py. Because the diff changes a calibration config field's effective default, I also read quantization/config.py (MaxCalibConfig, _coerce_layerwise_input, _gptq_qdq_default), mode.py:get_modelike_from_algo_cfg, model_quant.py:quantize, export/layerwise_export.py (LayerwiseExporter.__init__ / finalize), and the Megatron/HF decoder-layer discoverers to judge blast radius.
Findings this round — CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 1
| # | Severity | Where | Issue |
|---|---|---|---|
| 1 | IMPORTANT Compatibility | mode.py:229-230 |
model_fields_set is not a reliable "the user set this" signal here: get_modelike_from_algo_cfg dumps a config object without exclude_unset=True, so algorithm=mtq.MaxCalibConfig(layerwise={...}) marks every field set, the pop is skipped, and the same logical config raises ValueError under export_dir where the dict form works |
| 2 | SUGGESTION | model_calib.py:2105 |
The skip_forward_without_activation_calib=True default is applied only to the non-decoder pass; the per-layer loop still replays every batch through every layer for a weight-only recipe — num_layers x the cost this commit just removed |
Prior round: what is resolved
- Resolved — the previous IMPORTANT (over-broad
is_enabledgate forcing a full extra dataset pass for weight-only / dynamic / MX / constant-amax outside quantizers) is exactly whated7b4d7dfixes. The predicate is now_needs_activation_forward_for_max_calibevaluated under_hide_modules_from_traversal, reusing the existingmax_calibrateknob rather than inventing a new one, andtest_layerwise_export_allows_weight_only_outside_quantizer/test_layerwise_max_offload_warning_matches_outside_forwardpin both the forward count and the warning. It also correctly narrows theexport_dirrejection so a weight-onlylm_headno longer loses layerwise export. - Still open, non-blocking — the three prior SUGGESTIONs are unaddressed: rank-local
has_enabled_outside_quantizerguardingmax_calibrate's collectives (PP-stage asymmetry; low practical risk, since MCore foldsoutput_layerinto the discovered layer list sohas_enabled_outside_quantizerisFalsethere);_ForwardOnlyLayerduplicating_SkipLayer's__getattr__and reaching into its_PROXY_BLOCKLIST; and the export error message not naming the offending quantizers. Finding 1 gives that last one extra weight — with an object-form config the message nameslm_headwhen the actual trigger is a flag the user never wrote.
What I re-verified as correct
- The
export_dirrejection is well-founded, not over-strict. I initially read the gate onoutside_calib_runs_forward(rather thanoutside_calib_needs_forward) as too conservative, butexport_layerprogressively replaces quantizer modules and leaves the decoder in export form, so any full-model forward after the loop would run over packed weights. Rejecting whenever a forward will run — including the explicitskip_forward_without_activation_calib=Falsecase — is the right predicate, and the check sits beforeLayerwiseExporter.__init__so itsmkdirnever fires (asserted by the tests). - The newly-allowed export path produces a consistent checkpoint.
finalize()walks non-decoder modules through_dispatch_export_handler, solm_headis packed from the amax the new pass just wrote, and_calibrate_outside_quantizers()runs beforefinalize()in both the normal and resume-complete branches.self._quant_configis snapshotted in__init__before that amax exists, butget_quantization_formatkeys offis_enabled/num_bits/block_sizesand never_amax, soconfig.jsonand the tensors agree — no excluded-BF16-read-as-quantized mismatch. outside_calib_needs_forwardhas a real use path (it feedsoutside_calib_runs_forward), and evaluating the predicate at the top rather than at pass time is safe: it is structural, and the state that changes in between (_amax, SharedQuantState modules added inside layers) is either irrelevant to it or hidden.- Traversal hiding still holds.
_originalset viaobject.__setattr__keeps it out of_modules, soenable_stats_collection/weight_only_quantize/finish_stats_collection/_needs_activation_forward_for_max_calib/SharedWeightGlobalAmaxState.attach— allnamed_modules()walkers — cannot re-touch decoder quantizers. Slots come from every parent's_modules, aliases share one proxy keyed byid, and both the Megatron (list(decoder.layers) + [output_layer]) and HF discoverers return modules that are registered children, so no target is silently missed. get_qdq_activations_from_prev_layersemantics match the per-layer loop in both directions, andset_quantizer_by_cfg_context'spreserve_quantizer_attributes_contextrestores decoder enable state on exit — the test asserts both the tail input value (1.0 vs 2.0) anddecoder_quantizer.is_enabledafterwards.- The
mode.pypop is harmless outside the layerwise path (whole-modelmax_calibratealready defaults toFalse) and safe for configs lacking the field (pop(..., None));MseCalibConfig/LocalHessianCalibConfigdo not inherit it. modelopt_stateis untouched — no schema, mode-registration, or public-API change, so no migration concern.
Risk: low. The mechanism is well-contained and the regression coverage is genuinely targeted rather than incidental. Finding 1 is a config-plumbing footgun with a real user-visible failure (a hard error under export_dir), not a numerical or state-corruption bug, and the fix is local — move the differing default into MaxCalibConfig instead of inferring intent from model_fields_set.
Nothing here duplicates CodeRabbit's gate — no security anti-patterns, style, or typo findings.
|
/claude review |
There was a problem hiding this comment.
Claude review — layerwise non-decoder calibration (re-review at 5ff8457e)
Scope: full review (the trigger comment carried no scoping instructions). 6 files changed (472+/17-); reviewed all six. Because the diff changes the effective default of a public calibration config field, I also read wrapped_calib_func / get_modelike_from_algo_cfg in mode.py, max_calibrate + _needs_activation_forward_for_max_calib in model_calib.py, _SkipLayer in layerwise_calib.py, the per-layer qdq_from_prev capture block, the config-signature handling in algorithms.py, and CHANGELOG.rst to judge blast radius. Note: the local checkout has origin/main ahead of the base of this PR, so I used gh pr diff rather than a two-dot diff — a two-dot diff here shows unrelated reversals.
Findings this round — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3
| # | Severity | Where | Issue |
|---|---|---|---|
| 1 | SUGGESTION | model_calib.py:2100-2109 |
The None sentinel is resolved across three files, and this setdefault yields the inverse of the documented behavior if None survives into calib_kwargs. max_calibrate already treats None as falsy, so the mode.py pop is unnecessary — resolve the sentinel where it is consumed |
| 2 | SUGGESTION | layerwise_calib.py:129-136 |
_hide_modules_from_traversal silently no-ops for targets that are not registered _modules children; that no-op degenerates the non-decoder-only pass into a full-model calibration forward and overwrites decoder amax, with no error or warning |
| 3 | SUGGESTION | CHANGELOG.rst |
Not in the diff, so noted here rather than inline (see below) |
Finding 3 detail. The PR marks the changelog N/A. That is defensible for the bug fix itself: the shipped recipes already disable *lm_head* / *output_layer* via default_disabled_quantizers.yaml, so the miscalibration only bit custom recipes that deliberately enabled an outside quantizer, which is a fair reading of "critical or known bugs from a previous release". What still seems worth documenting is the config change. MaxCalibConfig.skip_forward_without_activation_calib was introduced in 0.46.0 (released 2026-08-18) and that entry says "opt-in, default False", which this PR makes stale. Rather than editing the released entry, add one line under 0.48.0: the field now accepts None (the new default), which behaves as False for whole-model and per-layer calibration and as True for the generated layerwise non-decoder pass. The public type also widened from bool to bool | None, which is visible to anyone asserting on the value.
Prior rounds: what is resolved
- Resolved — the previous IMPORTANT finding, that
model_fields_setis not a reliable "the user set this" signal here, becauseget_modelike_from_algo_cfgdumps a config object withoutexclude_unset=True, soalgorithm=mtq.MaxCalibConfig(layerwise={...})raisedValueErrorunderexport_dirwhere the dict form worked.5ff8457ereplaces that intent inference with an explicitNonesentinel checked by value, so an explicitNonein a dict is handled identically to an omitted key.test_layerwise_export_allows_weight_only_outside_quantizer[algorithm_as_config=True]pins exactly the object-form config that used to fail — a good regression choice. - Resolved earlier — the over-broad
is_enabledgate that forced a full extra dataset pass for weight-only / dynamic / MX / constant-amax outside quantizers. - Still open, non-blocking — the earlier SUGGESTIONs stand and I did not re-post them inline: rank-local
has_enabled_outside_quantizerguarding the collectives inmax_calibrate(a PP stage with no enabled outside quantizer skips_calibrate_outside_quantizers()entirely while another rank entersforward_loopand the amax all-reduce — a hang, not a slowdown; practical risk stays low because MCore foldsoutput_layerinto the discovered layer list, so this needs something like an enabled embedding quantizer on stage 0 only);_ForwardOnlyLayerduplicating the__getattr__of_SkipLayerand reaching into its_PROXY_BLOCKLIST; the export error message not naming the offending quantizers; and the per-layer loop still replaying every cached batch through every layer for a weight-only recipe (num_layers xthe cost this PR removes from the outside pass) becauseNoneresolves toFalsethere by design.
What I verified as correct this round
- The sentinel is genuinely well-contained in-tree.
skip_forward_without_activation_calibis declared only onMaxCalibConfig, which has no subclasses, so it can only ever reachmax_calibrate;_calib_func = max_calibrateatmode.py:435is a plain function attribute, socalib_func is max_calibrateatmodel_calib.py:2101isTrue— not apartialorstaticmethodwrapper that would silently disable the whole optimization. Themodel_dump()atalgorithms.py:1194feeds a signature/stored config, not calib kwargs. Finding 1 is future-proofing, not a live bug. - No
modelopt_statebreak. Old states carryingfalsere-validate againstbool | Nonefine; the field is a calibration knob rather than restore-affecting state, and no mode registration, schema key, or public export changed. - The
qdq_from_prev=Falsesemantics match the per-layer loop exactly. The extra pass usesset_quantizer_by_cfg_context(layer, [{"quantizer_name": "*", "enable": False}])— the same pattern, weight quantizers included, thatmodel_calib.py:2218-2223uses when capturing next-layer inputs. Disabling every decoder layer is the right analogue of disabling the one current layer in a full-model forward, andset_quantizer_by_cfg_contextoperates on the originallayerobjects, so traversal hiding does not interfere. The test pins the tail input to 1.0 / 2.0 and assertsdecoder_quantizer.is_enabledafterwards. - Ordering is right in both export branches. The
ValueErrorat 2111 fires beforeLayerwiseExporter(model, export_dir)at 2149, so no directory is created (asserted)._calibrate_outside_quantizers()runs beforefinalize()in both the resume-complete early return and the normal tail, and always outside the hiding context — sofinalize()walks the real modules and packslm_headfrom the amax the pass just wrote.test_layerwise_export_completed_resume_calibrates_weight_only_tailassertsforward_calls == 0together with a non-None tail amax, which is the right assertion pair for that branch. - Traversal hiding still holds, and forward fidelity is preserved.
_originalset viaobject.__setattr__never enters_modules, soenable_stats_collection,weight_only_quantize,finish_stats_collection,_needs_activation_forward_for_max_calib, andSharedWeightGlobalAmaxState.attach— allnamed_modules()walkers — cannot re-touch decoder quantizers, and the test asserts decoder_amaxis unchanged across the extra pass.forwardcallsself._original(...), i.e.type(original).__call__, so the_hf_hook-wrapped forward from accelerate and the HFGradientCheckpointingLayer.__call__override both still run; keeping_hf_hook/_old_forwardon the blocklist is correct here for the same reason as in_SkipLayer— those hooks live on the original, and the parameter-free proxy must not be managed. - Alias and slot restoration. Slots come from the
_modulesof every parent, aliased slots share one proxy keyed byid, proxies are built before thetry, and reassigning an existing_moduleskey preserves insertion order, sonn.ModuleListindexing survives. Covered on both the success and exception paths.
Risk: low. The mechanism is well-contained, the state-composition story checks out (nothing here can corrupt already-calibrated decoder quantizers), and the regression coverage is targeted rather than incidental. Both code findings this round harden invariants that hold today rather than fixing live defects — approving.
Nothing here duplicates the CodeRabbit gate: no security anti-patterns, style, or typo findings.
🤖 Generated with Claude Code
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>
5ff8457 to
a776897
Compare
|
/claude review |
| decoder_owned_ids = { | ||
| id(module) for layer in transformer_layers for module in layer.modules() | ||
| } | ||
| self.enabled = any( | ||
| isinstance(module, TensorQuantizer) | ||
| and module.is_enabled | ||
| and id(module) not in decoder_owned_ids | ||
| for module in model.modules() | ||
| ) | ||
| self.runs_forward = self.enabled | ||
| if self.enabled and activation_forward_predicate is not None: | ||
| with _hide_modules_from_traversal(model, transformer_layers): | ||
| needs_forward = activation_forward_predicate(model) | ||
| self.calib_kwargs = dict(calib_kwargs) | ||
| self.calib_kwargs.setdefault("skip_forward_without_activation_calib", True) | ||
| self.runs_forward = ( | ||
| needs_forward or not self.calib_kwargs["skip_forward_without_activation_calib"] | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Algorithm] The detection treats every enabled non-decoder quantizer as a downstream tail, but the calibration pass is only correct for quantizers that sit after the decoder stack.
Consider a recipe that also enables quantizers upstream of the decoder — a quantized embed_tokens, a multimodal projector / vision tower, or an input projection. Two things go wrong:
- Ordering. Those quantizers are live during the entire per-layer calibration loop, but they have no
_amaxyet (enable_stats_collectionis only ever called on the currentlayer). With_amaxunregistered,TensorQuantizer.forwardfalls through to the dynamic path (cf.extra_repr:if not hasattr(self, "_amax"): return "dynamic"), so every decoder layer is calibrated against dynamic-amax QDQ activations, and then the upstream quantizer is switched to a static amax at the very end. The decoder amaxes no longer match the activations inference will produce. get_qdq_activations_from_prev_layer=Falsesemantics. The stated contract is "downstream modules receive FP activations".calibrate()honours that for the tail pass by disabling decoder quantizers, but nothing disables upstream outside quantizers during the decoder loop — so decoder inputs already carry QDQ error while the tail sees FP. The two directions are handled inconsistently.
Since lm_head is the motivating case, the cheapest correct fix is to scope this explicitly rather than to "all non-decoder quantizers": either (a) disable all non-decoder quantizers for the duration of the decoder loop, so decoder calibration is against a clean FP prefix and the tail pass is the only place they are active, or (b) partition outside quantizers into pre-/post-decoder by module order and only defer the post-decoder ones, raising/warning for pre-decoder ones. Either way, please state the assumption in the class docstring.
| if export_dir is not None and outside_calibrator.runs_forward: | ||
| raise ValueError( | ||
| "Layerwise export does not support enabled quantizers outside transformer layers. " | ||
| "Calibrate without export_dir, then export the completed model separately." | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This rejects on runs_forward, not on "enabled quantizers outside transformer layers", and the two diverge in a way that breaks a previously working configuration.
runs_forward = needs_forward or not skip_forward_without_activation_calib. So for a weight-only outside quantizer (needs_forward=False), a user who explicitly sets skip_forward_without_activation_calib=False — the documented thing to do when the forward_loop has side effects the caller relies on, e.g. materializing ZeRO-3 shards — now gets a hard ValueError for a export_dir run that worked before this PR. test_layerwise_export_rejects_explicit_outside_forward codifies exactly that case, while test_layerwise_export_allows_weight_only_outside_quantizer allows the identical model with the default. Nothing about the export path actually got worse because the user asked not to skip an unnecessary forward.
Two things to fix:
- Gate on the real blocker — an outside quantizer that genuinely needs activation data — rather than on the
skip_*flag. Something likeoutside_calibrator.needs_activation_forward(i.e. the predicate result, ignoring the skip flag) keeps the fail-closed guarantee for the case the PR description is about (uncalibrated activation quantizers in the tail) without regressing weight-only + export. - Make the message name the condition it actually checks. As written it claims outside quantizers are unsupported, which the passing test above contradicts — and it will send users hunting for a quantizer to disable when the fix is to drop
skip_forward_without_activation_calib=False.
Also worth reconciling the PR description ("Layerwise export_dir now fails early when enabled non-decoder quantizers are present") with the code, which deliberately permits the weight-only case.
There was a problem hiding this comment.
🤖 Bot comment.
Addressed in dfde796 with the conservative export contract: layerwise export now rejects whenever any enabled outside quantizer requires the post-layer calibration step, including weight-only cases. The tests and PR description now consistently reflect that condition.
| if self.runs_forward and has_accelerate_offload(self.model): | ||
| warn_rank_0( | ||
| "Layerwise calibration found enabled quantizers outside transformer layers. " | ||
| "The required full-model calibration pass may be slow because CPU- or " | ||
| "disk-offloaded decoder weights can be transferred for every batch." | ||
| ) | ||
|
|
||
| with _hide_modules_from_traversal(self.model, self.transformer_layers): | ||
| if self.qdq_from_prev: | ||
| self.calib_func(self.model, self.forward_loop, **self.calib_kwargs) | ||
| return | ||
|
|
||
| # Inline import breaks conversion -> utils -> layerwise_calib import cycle. | ||
| from ..conversion import set_quantizer_by_cfg_context | ||
|
|
||
| with ExitStack() as stack: | ||
| for layer in self.transformer_layers: | ||
| stack.enter_context( | ||
| set_quantizer_by_cfg_context( | ||
| layer, [{"quantizer_name": "*", "enable": False}] | ||
| ) | ||
| ) | ||
| self.calib_func(self.model, self.forward_loop, **self.calib_kwargs) |
There was a problem hiding this comment.
[IMPORTANT Performance] The warning ("may be slow") understates what happens for every calib_func other than max_calibrate, and it undercuts the reason layerwise calibration exists ("performant layerwise calibration for large models that don't fit on GPU", per the CHANGELOG).
self.calib_func(self.model, self.forward_loop, ...) re-runs the entire algorithm on the whole model, not just an amax collection:
gptq(model, forward_loop, ...)→max_calibrate(model, forward_loop)(full-model forward Update README.md #1) → thenforward_loop(model)again underset_quantizer_by_cfg_context(full-model forward [E] Uncaught exception detected: Unable to open library: libnvinfer_plugin.so.9 due to libnvinfer_plugin.so.9: cannot open shared object file #2) → Hessian accumulation + weight updates.awq_lite/svdquant→awq(model, ...)alpha search (several full-model forwards), pluscreate_and_replace_svdquant_linear_on_the_fly(model=model)which mutates module structure, and a finalmax_calibrate(model, forward_loop).
Peak memory is the bigger issue than wall-clock. The per-layer loop keeps roughly two decoder layers resident (_SkipLayer frees the rest, and the meta-device placeholders were added specifically for this). This pass needs every decoder layer live and executing simultaneously, so a model that only fits under layerwise now has to fit — or thrash offload — for one extra full pass. For DeepSeek-R1 / Kimi-K2 class models that is the difference between "runs" and "does not".
Suggestions, in order of preference:
- Restrict the outside pass to the amax-only calib funcs (
max, andmse/local_hessianif they are cheap enough), and forgptq/awq_lite/svdquanteither run plainmax_calibrateon the outside quantizers instead of the full algorithm, or raise with a pointer to calibrating the tail separately.lm_headalmost never wants GPTQ/AWQ treatment anyway. - If the full algorithm is intentional, say so in the warning — name the number of extra full-model forwards and the fact that all decoder weights must be resident — and gate the warning on more than
has_accelerate_offload, since a fully resident model can OOM here too.
| @@ -938,10 +938,10 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): | |||
| "dynamic, or MX (MXFP4/MXFP8) quantization. Weight calibration still runs on the " | |||
| "weight tensors directly, so the quantized weights are unchanged; only the wasted " | |||
| "forward is avoided. " | |||
| "Opt-in (default False) because the provided ``forward_loop`` can carry side " | |||
| "effects the caller relies on — most notably materializing sharded parameters under " | |||
| "DeepSpeed ZeRO-3 — so enable it per-recipe when the calibration data is known to be " | |||
| "unnecessary." | |||
| "Set False when the provided ``forward_loop`` carries side effects the caller relies " | |||
| "on — most notably materializing sharded parameters under DeepSpeed ZeRO-3. The " | |||
| "default None behaves as False for whole-model and per-layer calibration, and as True " | |||
| "for ModelOpt's generated layerwise non-decoder pass." | |||
| ), | |||
| ) | |||
There was a problem hiding this comment.
[IMPORTANT Compatibility] skip_forward_without_activation_calib is a released public config field (CHANGELOG.rst:129, in a shipped section, documented as "opt-in, default False"). Widening it to bool | None and moving the default to None is an observable public-API change with no changelog entry, and the PR checklist marks CHANGELOG as N/A.
Two concrete consequences:
mtq.MaxCalibConfig().skip_forward_without_activation_calibnow returnsNonewhere it returnedFalse. Truthiness checks survive; anything that round-trips the value through abool-typed field, ais Falsecomparison, or an external schema does not.modelopt_statewritten by this version storesnullfor the key. An older ModelOpt (where the annotation isbool) will fail pydantic validation onrestore()instead of loading. This repo treats that class of change as noteworthy — see thelayerwiselegacy-key entry atCHANGELOG.rst:78, which explicitly documents a restore-validation break.
Per CLAUDE.md, backward-breaking changes get a CHANGELOG.rst entry; please add a one-or-two-sentence one under the unreleased Backward Breaking Changes section stating the new tri-state default and that checkpoints saved with it do not restore on older releases.
Separately on the mechanism: the None → effective-default resolution is now spread across three files — the None default here, the kwargs.pop() in mode.py:229, and self.calib_kwargs.setdefault(..., True) in layerwise_calib.py:186. Popping the key in wrapped_calib_func means the sentinel is resolved by absence, so max_calibrate's own signature default silently becomes the source of truth and any future caller that forwards config.model_dump() without the pop gets skip_forward_without_activation_calib=None passed straight through. Resolving None in one place (a validator here, or at the top of max_calibrate) would keep the tri-state without that coupling.
There was a problem hiding this comment.
🤖 Bot comment.
Addressed in dfde796: the released bool field and False default are restored. model_fields_set distinguishes omission from an explicit False internally, preserving the serialized schema and direct-call behavior.
| class _ForwardOnlyLayer(nn.Module): | ||
| """Hide a layer from module traversal while preserving its forward execution.""" | ||
|
|
||
| _PROXY_BLOCKLIST = _SkipLayer._PROXY_BLOCKLIST | ||
|
|
||
| def __init__(self, original: nn.Module): | ||
| super().__init__() | ||
| object.__setattr__(self, "_original", original) | ||
|
|
||
| def __getattr__(self, name: str): | ||
| try: | ||
| return super().__getattr__(name) | ||
| except AttributeError: | ||
| if name in self._PROXY_BLOCKLIST: | ||
| raise | ||
| return getattr(object.__getattribute__(self, "_original"), name) | ||
|
|
||
| def forward(self, *args, **kwargs): | ||
| return self._original(*args, **kwargs) | ||
|
|
||
|
|
There was a problem hiding this comment.
[SUGGESTION] The docstring says "hide a layer from module traversal", which undersells the blast radius. Because _original is set with object.__setattr__, the proxy has an empty _modules/_parameters, so for the duration of the context the model also loses the decoder layers from named_parameters(), parameters(), buffers(), state_dict(), and therefore from .to(), .cuda(), FSDP/DTensor parameter collection, and any next(model.parameters()).device idiom.
That is fine for the max_calibrate path this PR exercises, but the outside pass hands the model to arbitrary calib_funcs and to the user's forward_loop. A user forward_loop that does inputs.to(next(model.parameters()).device) will now pick the device of whatever non-decoder parameter happens to come first, and an FSDP2-wrapped model will see a parameter set that does not match its flat-param groups.
Please widen the docstring to say explicitly that parameters, buffers and state_dict are hidden too — not only modules() — so the next person reasoning about a new calib_func knows what they are working with. Worth noting the same for _OutsideQuantizerCalibrator.calibrate, since that is where the model is exposed to caller code.
There was a problem hiding this comment.
🤖 Bot comment.
Addressed in dfde796: the traversal helper and outside-calibrator docstrings now state that modules, parameters, buffers, and state-dict traversal are hidden during the outside pass.
| self.runs_forward = self.enabled | ||
| if self.enabled and activation_forward_predicate is not None: | ||
| with _hide_modules_from_traversal(model, transformer_layers): | ||
| needs_forward = activation_forward_predicate(model) | ||
| self.calib_kwargs = dict(calib_kwargs) | ||
| self.calib_kwargs.setdefault("skip_forward_without_activation_calib", True) | ||
| self.runs_forward = ( | ||
| needs_forward or not self.calib_kwargs["skip_forward_without_activation_calib"] | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] enabled / runs_forward are derived purely from the local module tree, so they can diverge across ranks whose structure is not homogeneous. Under pipeline parallelism only the last stage owns lm_head / output_layer, so only that rank would take the calibrate() branch — and then call forward_loop(model) and max_calibrate(..., distributed_sync=True) (which all-reduces amax) while the other ranks have already returned. That is a hang, not a wrong number.
Layerwise calibration may already be single-stage-only in practice (the per-layer loop drives forward_loop a different number of times per rank when layer counts differ), in which case this is pre-existing and only needs a note. If layerwise is meant to work under PP, an dist.allreduce/any-across-ranks on enabled and runs_forward before acting on them would make the decision collective and keep the pass symmetric.
There was a problem hiding this comment.
Claude review — Layerwise non-decoder calibration
Scope: trigger comment was a bare /claude review, so this is a full review. All 6 changed files (508+/12-) were opened: modelopt/torch/quantization/{config.py, mode.py, model_calib.py, utils/layerwise_calib.py} plus both test files. Surrounding context read for max_calibrate / _needs_activation_forward_for_max_calib, TensorQuantizer.forward, _SkipLayer / LayerActivationCollector patch-unpatch, _CheckpointState.full_restore, gptq, and wrapped_calib_func.
Findings
CRITICAL: 0 · IMPORTANT: 3 · SUGGESTION: 2
The core mechanism is sound. I specifically checked and could not find problems in the places most likely to break:
_hide_modules_from_traversalalias handling — slots are collected per(parent, child_name)and one proxy is shared per child id, so multi-slot aliases restore correctly, and thefinallycovers the raising path.- Decoder amax preservation —
enable_stats_collection/finish_stats_collection/ the MoE amax sync /_check_nvfp4_static_tp_supportedall traversenamed_modules(), so hidden decoder quantizers keep their calibrated_amax. - Accelerate offload —
self._original(...)goes through the original module's__call__, so its_hf_hookstill fires and offloaded weights are still fetched;_PROXY_BLOCKLISTcorrectly keeps accelerate off the parameter-free proxy.has_accelerate_offloadis called before hiding, which is the only ordering that works. - Object identity across the run —
_cleanup_layersrestores the original layer objects andfull_restoremutates in place, soself.transformer_layersis never stale by the timecalibrate()hides slots. - No recursion risk:
gptqdoes not re-readlayerwise.enableoff the model. - Ordering of the
_reconcile_export_with_resumeearly return correctly calibrates the tail before returning, so a completed-resumefinalize()no longer writes an uncalibrated tail shard.
Most impactful
-
Upstream outside quantizers are mis-scoped (
layerwise_calib.py:172-189). The detection catches every enabled non-decoder quantizer, but the deferred pass is only correct for ones downstream of the decoder. A quantizedembed_tokens/ projector runs with_amaxunregistered — i.e. on the dynamic path — through the whole per-layer loop, so decoder amaxes are calibrated against activations that differ from inference. It also breaks theget_qdq_activations_from_prev_layer=Falsecontract asymmetrically: decoder quantizers are disabled for the tail pass, but upstream outside quantizers are never disabled for the decoder loop. -
export_dirgate keys off the wrong condition (model_calib.py:2097-2101). It rejects onruns_forward, soskip_forward_without_activation_calib=False+ a weight-onlylm_headnow raises even though that combination is fine — and the message blames "enabled quantizers outside transformer layers", which the passingtest_layerwise_export_allows_weight_only_outside_quantizercontradicts. This is a regression for the documented ZeRO-3 use of that flag. -
The outside pass re-runs the whole algorithm, not just amax collection (
layerwise_calib.py:196-218). Forgptqthat is two extra full-model forwards plus Hessian accumulation; forawq_lite/svdquantan alpha search plus module replacement. Peak memory is the real cost: the per-layer loop keeps ~2 decoder layers resident by design (_SkipLayer+ meta placeholders), and this pass needs all of them live at once. For the DeepSeek-R1 / Kimi-K2 class of model that layerwise exists to serve, that can turn "runs" into "OOMs", and the warning only says "may be slow".
Plus two SUGGESTIONs: the _ForwardOnlyLayer docstring understates that parameters/buffers/state_dict are hidden as well as modules() (matters because arbitrary calib_funcs and the user's forward_loop see that model), and enabled/runs_forward are rank-local so they diverge under pipeline parallelism where only the last stage owns lm_head.
Risk
Moderate. The change is well contained and the test coverage is genuinely good — alias restore on success and error, QDQ-vs-FP propagation, decoder amax preservation, offload-warning gating, and the resume/export matrix are all exercised. The residual risk is concentrated in the "all non-decoder quantizers" scoping: for the lm_head case this PR is aimed at, it does the right thing, but the same code path is what a quantized embedding or a multimodal projector will hit, and there it is silently wrong rather than loudly unsupported. Finding 2 is a small, mechanical fix; findings 1 and 3 are mostly about narrowing the contract and saying so.
🤖 Generated with Claude Code
| decoder_owned_ids = { | ||
| id(module) for layer in transformer_layers for module in layer.modules() | ||
| } | ||
| self.enabled = any( |
There was a problem hiding this comment.
can we move this self.enabled inside _hide_modules_from_traversal ? this way way we can simplify this code further and avoid iteration over decoder layer quantizers. Further, can we cache the name to module map during init, then _hide_modules_from_traversal can simply look a decoder modules parent name fro the name to module map. this way multiple _hide_modules_from_traversal does not iterate over the whole model.
There was a problem hiding this comment.
🤖 Bot comment.
Implemented in dfde796: the calibrator now caches unique physical decoder registration slots once, computes enabled while those slots are hidden, and preserves direct plus aliased-parent registrations with shared-proxy restoration. Focused success/error coverage was added.
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
/claude review |
| So lets wrap them to be compatible with the ModelOpt convert entrypoint. | ||
| """ | ||
| kwargs = config.model_dump() | ||
| if "skip_forward_without_activation_calib" not in config.model_fields_set: |
There was a problem hiding this comment.
why is this needed?
There was a problem hiding this comment.
🤖 Bot comment.
Removed this in 76106e2. With both the config field and max_calibrate defaulting to False, the conditional was redundant; I also removed the corresponding model_fields_set test assertions.
| if "skip_forward_without_activation_calib" not in config.model_fields_set: | ||
| kwargs.pop("skip_forward_without_activation_calib", None) |
There was a problem hiding this comment.
[SUGGESTION] This model_fields_set pop has no behavioral effect — it looks like a leftover from the earlier None-sentinel revision.
skip_forward_without_activation_calib is declared only on MaxCalibConfig (config.py:931, no subclasses), whose _calib_func is max_calibrate, and max_calibrate's own default for that parameter is already False (model_calib.py:321). So:
- dict form (
{"method": "max", ...}) → field not inmodel_fields_set→ popped →max_calibrateuses its defaultFalse - object form (
mtq.MaxCalibConfig(...)) → every field is inmodel_fields_set→ not popped →max_calibratereceivesFalseexplicitly
Both paths end at False, and nothing downstream inspects calib_kwargs keys (layerwise_calibrate forwards the rest verbatim, and _OutsideQuantizerCalibrator.calibrate passes **self.calib_kwargs unchanged). Your own new test agrees: test_layerwise_max_outside_calibration_uses_configured_forward_behavior expects None and False to produce the same 2 forward calls.
Why it matters: it reads as load-bearing intent inference, and model_fields_set is an unreliable "the user set this" signal here anyway (the object form marks all fields set), so a future change that does branch on presence will silently behave differently for object vs. dict configs. The paired test_calib.py assertions then pin pydantic bookkeeping rather than any behavior.
Suggestion: drop these two lines (and the two model_fields_set assertions in test_calib.py). If a different default for the generated non-decoder pass is wanted later, set it explicitly at the consumption site in _OutsideQuantizerCalibrator.calibrate() rather than inferring it from field-set state.
There was a problem hiding this comment.
🤖 Bot comment.
Addressed in 76106e237: removed the redundant field-set branch and the bookkeeping-only assertions. Both config paths now rely on the existing False default.
| if exporter is not None and _reconcile_export_with_resume( | ||
| exporter, checkpoint_dir, start_layer, num_layers | ||
| ): | ||
| outside_calibrator.calibrate() |
There was a problem hiding this comment.
[SUGGESTION] This call is unreachable-as-effective and can be dropped.
The branch is guarded by exporter is not None, and exporter is only assigned inside if export_dir is not None: (line 2106). But line 2096 now raises whenever export_dir is not None and outside_calibrator.enabled. So reaching line 2132 implies export_dir is not None implies outside_calibrator.enabled is False, and calibrate() returns immediately on if not self.enabled.
It was meaningful in the previous revision, where layerwise export allowed a weight-only outside quantizer and the completed-resume path had to calibrate the tail before finalize(). With the simplified unconditional rejection, it is dead.
Why it matters: it implies the resume-complete export path still calibrates something, which invites a future reader to assume outside quantizers are handled under export_dir — the opposite of what line 2096 enforces. Either remove the call, or keep it and narrow the line-2096 gate back to "the outside pass would actually run a forward" so the call becomes live again for the weight-only case.
There was a problem hiding this comment.
Claude review — layerwise non-decoder calibration (re-review at dfde7960)
Scope: the trigger comment was a bare /claude review, so this is a full review. All 5 changed files (419+/4-) were opened: modelopt/torch/quantization/{mode.py, model_calib.py, utils/layerwise_calib.py} plus tests/unit/torch/quantization/{test_calib.py, test_layerwise_calibrate.py}. For surrounding context I read max_calibrate / _needs_activation_forward_for_max_calib, MaxCalibConfig and the sibling *CalibConfig classes in config.py, wrapped_calib_func and MaxCalibrateModeDescriptor, _SkipLayer / get_input_activations / get_first_layer_inputs in layerwise_calib.py, set_quantizer_by_cfg_context + preserve_quantizer_attributes_context, and utils/core_utils.py (to check the new module-level import).
Method note: this branch is a single commit on top of 7f7c46d8 and the two-dot diff against origin/main matches gh pr diff --name-only exactly, so the two-dot diff is authoritative this round. I was not able to run the test suite here (command approval denied), so test assertions below are reviewed by reading, not by execution.
Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3
| # | Severity | Where | Issue |
|---|---|---|---|
| 1 | SUGGESTION | mode.py:229-230 |
The model_fields_set pop is a behavioral no-op left over from the reverted None-sentinel revision; the paired test_calib.py assertions pin pydantic bookkeeping rather than behavior |
| 2 | SUGGESTION | model_calib.py:2132 |
outside_calibrator.calibrate() in the completed-resume export branch is unreachable-as-effective: that branch implies export_dir is not None, which line 2096 guarantees implies not enabled |
| 3 | SUGGESTION | layerwise_calib.py:107-136 |
Not posted inline (overlaps a prior round). Beyond modules() / named_parameters() / state_dict(), hiding also blocks mutation propagation: model.eval() / model.train() / model.to() / model.apply() called by an arbitrary calib_func or by the user forward_loop during the pass silently skip the decoder layers, and attribute writes land on a throwaway proxy. A forward_loop that toggles model.eval() internally leaves decoder layers in whatever mode they were in — dropout-active if the model was in train mode. Worth a sentence in the _ForwardOnlyLayer docstring, and arguably an eval-mode assert before the pass |
Both inline findings are the same shape: vestigial code from the earlier revision that now reads as load-bearing. Neither changes behavior today, which is why they are SUGGESTIONs.
What I verified as correct this round
- The new import is cycle-free.
utils/core_utils.pyimports onlyquantization.configandmodelopt.torch.utils, so the new module-levelfrom .core_utils import has_accelerate_offloaddoes not reintroduce thenn -> qtensor -> utils -> layerwise_calibcycle the inline imports exist to break;..nnand..conversionare correctly kept function-local. - Slot discovery and restoration. Slots are collected as
(parent, child_name, child)over every parent_modulesdict, aliased slots share one proxy keyed byid, proxies are built before thetry, and reassigning an existing_moduleskey preserves insertion order — sonn.ModuleListindexing survives. Restoration is infinally, covered on both the success and raising paths bytest_outside_calibrator_hides_and_restores_layer_aliases. - Traversal hiding actually hides.
_originalis set viaobject.__setattr__so it never lands in_modules, and the proxy own_modulesis empty. Everynamed_modules()walker in the deferred pass —enable_stats_collection,weight_only_quantize,finish_stats_collection, the MoE amax sync,_check_nvfp4_static_tp_supported,SharedWeightGlobalAmaxState.attach, and thedistributed_syncamax reduction — therefore cannot re-touch decoder quantizers, and the test asserts decoder_amaxis unchanged across the extra pass. - Forward fidelity.
forwardcallsself._original(...), i.e.type(original).__call__, so accelerate_hf_hook-wrapped forwards and the HFGradientCheckpointingLayer.__call__override both still run; keeping_hf_hook/_old_forwardon_PROXY_BLOCKLISTcorrectly keeps accelerate from managing the parameter-free proxy.has_accelerate_offload(self.model)is evaluated before hiding, which is the only ordering that works, and the warning fires lazily inside the wrapper so it is not emitted whenmax_calibrateskips the forward — exactly what the parametrized warning tests pin. qdq_from_prevsemantics mirror the per-layer loop. TheFalsepath disables every decoder quantizer viaset_quantizer_by_cfg_context(layer, [{"quantizer_name": "*", "enable": False}])— the same pattern, weight quantizers included, used atmodel_calib.py:2218-2223when capturing next-layer inputs — so the tail sees FP activations; theTruepath leaves them enabled so the tail sees QDQ error. That context operates on the originallayerobjects, so hiding does not interfere, andpreserve_quantizer_attributes_contextrestores enable state on exit (asserted).- Downstream ordering is sound for the
lm_headcase this PR targets.get_input_activations/cache_outputs_for_next_layer_calibearly-stop via_EarlyStopForwardErrorat the target decoder layer, solm_headis never reached during the per-layer loop — it only ever runs in the deferred full-model pass, with the decoder already calibrated. That pass runs after_unpatch_all_layers()and afterckpt.full_restore(...), and always outside the hiding context, so nothing sees a patched or proxied model afterwards. - Config plumbing is contained.
skip_forward_without_activation_calibis declared only onMaxCalibConfig, which has no subclasses (NVFP4ActHeadroomCalibConfig,MseCalibConfig,LocalHessianCalibConfigall derive fromQuantizeAlgorithmConfigdirectly), so themode.pypop can neither reach a calib func that wouldTypeErroron the kwarg nor change any default. - No mode/state or public-API change.
modelopt_stateschema, mode registration, and__init__exports are untouched, so there is no migration concern; the changelog being N/A is defensible since the shipped recipes already disable the affected quantizers.
Residual risk and prior rounds — these are unchanged in substance, so I did not re-post them:
- Non-decoder scoping is still uniform for upstream and downstream quantizers. A quantized embedding or VLM projector is upstream of the decoder, so it runs uncalibrated (the dynamic amax path, since
_amaxis unregistered) throughout the per-layer loop and only gets a static amax afterwards — decoder amaxes are then calibrated against activations that differ slightly from inference. Second-order numerically, and not a regression relative tomain(where those quantizers were never calibrated at all), but the contract is worth stating in the docstring. - The deferred pass needs every decoder layer resident at once, the opposite of the ~2-layers-resident invariant layerwise exists to provide; the warning says "may be slow" rather than "may OOM".
enabledis rank-local, so a PP stage with no enabled non-decoder quantizer skips the pass entirely while another stage enters forward_loop and the amax reduction. Practical risk stays low because MCore foldsoutput_layerinto the discovered layer list.- Blast radius of the new
export_dirrejection. With the gate back to plainis_enabled,export_dirplus any enabled non-decoder quantizer is now a hardValueError, including the weight-onlylm_headcase the previous revision allowed and which needs no forward at all. Shipped recipes are unaffected (default_disabled_quantizers.yamldisables lm_head, output_layer, embeddings, routers and the vision branch), so this only bites custom configs — e.g. a plainFP8_DEFAULT_CFGon a VLM, where projector / vision-tower Linears are enabled by default and layerwise export would now fail where it previously ran with a silently uncalibrated quantizer. The PR body states fail-early as the intent and the message is actionable, so I read this as a deliberate trade-off rather than a defect — noting it so the choice is explicit.
Risk: low. The mechanism is well contained, the state-composition story checks out (nothing here can corrupt already-calibrated decoder quantizers), and the regression coverage is targeted rather than incidental — alias restore on success and error, traversal-hiding assertions, QDQ-vs-FP tail input, decoder amax preservation, offload-warning gating, the skip-forward count matrix, and export fail-closed with an assertion that no directory was created. Both code findings are cleanups of vestigial code rather than live defects — approving.
Nothing here duplicates the CodeRabbit gate: no security anti-patterns, style, or typo findings.
🤖 Generated with Claude Code
Signed-off-by: realAsma <akuriparambi@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Logic looks correct and CPU test coverage is genuinely thorough, but a few things need an owner's call before merge — most notably an unreachable calibrate() call and no offload-path test for the new full-model pass.
Needs action:
- Remove or comment the
outside_calibrator.calibrate()added in the export-resume early-return inmodel_calib.py: that branch is only reachable whenexport_dir is not None, which the newValueErroralready guaranteesenabled is False, so the call can never do work. - Add a GPU test in
tests/gpu/torch/quantization/plugins/test_accelerate_gpu.pyexercising the new pass on a CPU/disk-offloaded model —_hide_modules_from_traversalswapsparent._modulesunder live accelerate hooks and only thehas_accelerate_offloadflag is currently monkeypatched. - Confirm the new hard
ValueErrorfor layerwise export + enabled outside quantizers is acceptable as a silent break for existing configs, and add a CHANGELOG line if so. - Confirm the extra full-model forward should run by default even when the outside quantizers are weight-only (see
test_layerwise_max_outside_calibration_uses_configured_forward_behavior,expected_forward_calls == 2). - Revert the unrelated cosmetic edit in
tests/unit/torch/quantization/test_calib.py.
| if exporter is not None and _reconcile_export_with_resume( | ||
| exporter, checkpoint_dir, start_layer, num_layers | ||
| ): | ||
| outside_calibrator.calibrate() |
There was a problem hiding this comment.
Bot comment.
This call is unreachable in effect: exporter is non-None only when export_dir is not None, and the new check above raises whenever export_dir is not None and outside_calibrator.enabled. So enabled is always False here and calibrate() returns immediately. Drop the line, or add a comment saying why it's kept defensively.
| self.enabled = any( | ||
| isinstance(module, TensorQuantizer) and module.is_enabled | ||
| for module in model.modules() | ||
| ) |
There was a problem hiding this comment.
Bot comment.
The new pass replaces entries in parent._modules while accelerate/FSDP hooks are live on the real layers. The CPU tests only monkeypatch has_accelerate_offload, so the proxy is never exercised against an actual _hf_hook/AlignDevicesHook layout. tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py already has make_cpu_offloaded_model and disk-offload helpers — please add a case there with an enabled lm_head quantizer.
| def test_max_calib_config_skip_is_opt_in(): | ||
| """The flag is opt-in (default False) so it does not change behavior for direct callers.""" | ||
| assert MaxCalibConfig().skip_forward_without_activation_calib is False | ||
| default_config = MaxCalibConfig() |
There was a problem hiding this comment.
Bot comment.
This is unrelated churn — hoisting MaxCalibConfig() into default_config doesn't change the assertion or the test. Please revert to keep the diff focused.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Is this feature mainly targeting |
| ) | ||
|
|
||
|
|
||
| class _ForwardOnlyLayer(nn.Module): |
There was a problem hiding this comment.
Do you think it's better to reuse _SkipLayer here? The proxy machinery can be reused and we need a new forward function for this feature
There was a problem hiding this comment.
BB: This is a good point, can you give a minimal design for this?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 I’m taking a look and will follow up with a minimal design.
There was a problem hiding this comment.
🤖 Bot comment.
Minimal design: extract only the shared proxy mechanics into a small _LayerProxy(nn.Module) base. Its __init__ stores _original without registering it, and its __getattr__ applies the existing hook blocklist before forwarding to _original. Then keep _SkipLayer(_LayerProxy) responsible for skip state plus the zero/meta forward, and _ForwardOnlyLayer(_LayerProxy) responsible only for delegating forward to _original. That reuses the machinery without making the forward-preserving case inherit skip semantics. I’d cover the base behavior through the existing alias/blocklist tests rather than add another abstraction or callback strategy.
There was a problem hiding this comment.
Can we have a simpler design -> Just add the forward with hide module and without any meta data manipulation support as well to skip layer ? then for outside we use forward with hide module and without any meta data
There was a problem hiding this comment.
🤖 Bot comment.
I read this as folding the two wrappers into _SkipLayer: keep the current metadata-backed zero-output mode for calibrated layers, and add a forward-through-original mode that hides registered modules without creating _LayerCalibState. _hide_modules_from_traversal would use the latter, so _ForwardOnlyLayer disappears. Please confirm that is the intended shape and I’ll implement it.
3 participants
What does this PR do?
Type of change: Bug fix
Layerwise calibration now calibrates enabled quantizers outside transformer layers, such as
lm_head, while hiding decoder layers from the additional calibration traversal. It also fails early when these quantizers are combined with progressive layerwise export, whose in-place conversion makes the required model calibration unsafe.Usage
No API changes.
Testing
pytest_pwd tests/unit/torch/quantization/test_calib.py tests/unit/torch/quantization/test_layerwise_calibrate.py -q— 68 passedgit diff --check— passedBefore your PR is "Ready for review"
CONTRIBUTING.md: N/ASummary by CodeRabbit
New Features
Bug Fixes
Behavior Changes