[Fix] Calibrate non-decoder modules during layerwise quantization - #2339
[Fix] Calibrate non-decoder modules during layerwise quantization#2339realAsma wants to merge 1 commit into
Conversation
Signed-off-by: realAsma <akuriparambi@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. |
|
Important Approval pendingCodeRabbit has no unresolved comments, but it has not reviewed the latest commit. Use the checkbox below to review the latest commit. CodeRabbit will approve the changes if it finds no blocking issues.
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2339 +/- ##
==========================================
+ Coverage 79.31% 79.32% +0.01%
==========================================
Files 527 527
Lines 61482 61522 +40
==========================================
+ Hits 48765 48804 +39
- Misses 12717 12718 +1
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.
| if has_enabled_outside_quantizer: | ||
| if any(device == "disk" for device in getattr(model, "hf_device_map", {}).values()): | ||
| warn_rank_0( | ||
| "Layerwise calibration found enabled quantizers outside transformer layers. " | ||
| "The required full-model calibration pass may be slow because disk-offloaded " | ||
| "decoder weights can be streamed for every batch." | ||
| ) | ||
|
|
||
| with _hide_modules_from_traversal(model, transformer_layers): |
There was a problem hiding this comment.
[SUGGESTION] has_enabled_outside_quantizer is a purely local decision, but the body it guards runs a full model forward plus max_calibrate's cross-rank amax syncs (sync_amax_across_distributed_group over DP/EP/TP groups). If two ranks disagree on the flag, one enters collectives the other never reaches and the job hangs rather than failing.
Today the common paths are safe: HF layerwise is effectively single-process, and get_mcore_layerwise_calibration_layers (plugins/megatron.py:997) appends model.output_layer to the layer list, so MCore's usual outside-quantizer candidate is already decoder-owned, and DP/TP replicas have identical module trees. The exposure is pipeline parallelism with a recipe that re-enables a stage-local module (first-stage embedding, mtp.*), where the flag differs by stage.
Cheap insurance: all-reduce the flag (logical OR / max) across the model's parallel groups before branching, so every rank takes the same path. A comment stating the single-process/replica-symmetry assumption would also be enough if you'd rather not add the collective.
| 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] Two small things on the proxy:
-
_ForwardOnlyLayerduplicates_SkipLayer's__getattr__delegation verbatim and reaches into_SkipLayer._PROXY_BLOCKLIST(line 108) — a private-attribute dependency between two sibling classes. Consider a shared private base that owns_PROXY_BLOCKLIST,__init__, and__getattr__, leaving each subclass with only itsforward. That keeps the accelerate blocklist a single source of truth if it ever grows. -
There's no
__setattr__override, so reads delegate to_originalbut writes don't. Anything that assigns onto the layer slot during the extra pass — a parent model caching per-layer state (layer.foo = ...), or acalib_funcattaching a hook/shared state to what it believes is the layer — lands on the throwaway proxy and is silently dropped whenfinallyswaps the originals back. Nothing in the current callers does this (SharedWeightGlobalAmaxState.attachmatches linear names inside the hidden subtree, so it finds nothing), but a docstring note like "writes to the proxy are discarded on exit; only forward behavior is preserved" would stop a future caller from being surprised.
| 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.
[SUGGESTION] The error message tells the user to drop export_dir, but most users hitting this actually want the other fix — keep progressive export and stop quantizing the outside module. Worth naming both options, and naming the offending quantizers so the user doesn't have to go find them:
raise ValueError(
"Layerwise export does not support enabled quantizers outside transformer layers "
f"(e.g. {sorted(outside_quantizer_names)[:5]}). Either disable them (e.g. add "
'{"quantizer_name": "*lm_head*", "enable": False} to quant_cfg), or calibrate '
"without export_dir and export the completed model separately."
)
That needs has_enabled_outside_quantizer to become a name list instead of a bool, which the any(...) above can be turned into cheaply since it already walks named_modules-equivalent state.
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.
1 participant
What does this PR do?
Type of change: Bug fix
Layerwise calibration previously calibrated only discovered transformer/decoder layers, so enabled quantizers outside those subtrees, such as a Hugging Face
lm_head, were left uncalibrated.This change adds one full-model calibration pass for enabled non-decoder quantizers after decoder calibration and checkpoint restoration. During that pass, decoder layers are replaced in their actual registration slots by forward-only proxies: their calibrated forwards and QDQ behavior still execute, while recursive module traversal cannot discover or recalibrate their quantizers. Aliased slots reuse the same proxy and every slot is restored in a
finallyblock.When
get_qdq_activations_from_prev_layer=False, decoder quantizers are temporarily disabled for the extra pass so downstream modules receive FP activations, matching the existing option semantics. Disk-offloaded models warn only when the extra full-model pass is required.Layerwise
export_dirnow fails early when enabled non-decoder quantizers are present. Progressive layer export leaves the root model unsuitable for this required forward pass, so rejecting the combination prevents a resume manifest from silently finalizing an uncalibrated tail.Usage
No new API is required. Existing layerwise calibration now includes enabled quantizers such as
lm_head:Testing
pytest_pwd tests/unit/torch/quantization/test_layerwise_calibrate.py -q— 44 passedpre-commit run --files modelopt/torch/quantization/model_calib.py modelopt/torch/quantization/utils/layerwise_calib.py tests/unit/torch/quantization/test_layerwise_calibrate.py— all hooks passedpython_pwd -m py_compile modelopt/torch/quantization/model_calib.py modelopt/torch/quantization/utils/layerwise_calib.py tests/unit/torch/quantization/test_layerwise_calibrate.py— passedBefore your PR is "Ready for review"
export_dircombination now raises an explicit error.CONTRIBUTING.md: N/AAdditional Information
The regression coverage includes real
lm_headmax calibration, QDQ versus FP propagation, traversal hiding and alias restoration on success and error, plain-list layer discovery, no-op behavior without outside quantizers, disk-offload warning gating, decoder calibration-state preservation, and export fail-closed behavior.