fix(megatron): restore untied output_layer quantization under Megatron-Bridge - #2112
fix(megatron): restore untied output_layer quantization under Megatron-Bridge#2112yueshen2016 wants to merge 3 commits into
Conversation
…n-Bridge A quantized output_layer (lm_head) is silently exported as BF16 when the model is built through Megatron-Bridge: no error, no warning, just an unquantized output layer in the exported checkpoint. Two independent causes, both in quantization/plugins/megatron.py: 1. _register_extra_state_callbacks skips output_layer unless its weight_quantizer is already enabled. That hook also runs BEFORE QuantModule replacement (e.g. on restore), when weight_quantizer does not exist yet, so the check always skipped and output_layer never received the ModelOpt extra-state callbacks. Its quantizer state (promotion to StaticBlockScaleQuantizer, _amax, _global_amax) was therefore never saved or restored. 2. sharded_state_dict decides whether the output layer is quantizable by reading Megatron-LM's get_args().untie_embeddings_and_output_weights. Megatron-Bridge has no such global args store, so the import raises, the except path treats the layer as tied, and quantization is dropped. Both are fixed by resolving tiedness from the model itself: * _resolve_output_layer_untied() walks named_modules() for the framework-agnostic share_embeddings_and_output_weights flag and records the result on the model config as modelopt_output_layer_untied. * sharded_state_dict prefers that flag and falls back to get_args() only when it is absent, so Megatron-LM behaviour is unchanged. * Callback registration falls back to the tiedness flag when weight_quantizer does not exist yet, so an untied output_layer is registered. * When the load plan is built, the per-block NVFP4 scale buffers are materialized so the loader has somewhere to write: a Megatron dist-checkpoint load silently skips any key the model does not advertise. Divisibility is checked against weight.shape[-1] (the axis _process_quantizer_amax later views over) and a warning is emitted rather than leaving the buffers unallocated. The distillation plugin gains a matching one-shot promotion on the first forward for the single module the restore path does not promote (measured on Nemotron-Nano-3: already_sbsq=460, converted=1, disabled=0 -- exactly output_layer). Verified on Nemotron-Nano-3 W4A16 NVFP4 four_over_six: the exported checkpoint carries lm_head.weight (U8 [131072, 1344]) plus weight_scale and weight_scale_2, matching a known-good reference export (52 shards, 18487 keys, 72 exclusions), and serves correctly on vLLM 0.26.0. Signed-off-by: James Shen <yueshen@nvidia.com>
📝 WalkthroughWalkthroughMegatron quantization now resolves output-layer tying from the model, preserves quantizer state for sharded checkpoints, materializes block-scale buffers, and performs one-time NVFP4 quantizer promotion before teacher and student inference. ChangesMegatron quantization and distillation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MegatronModel
participant QuantizationHook
participant OutputLayerCallback
participant ShardedState
MegatronModel->>QuantizationHook: resolve output-layer tying state
QuantizationHook->>OutputLayerCallback: register output-layer quantizer callback
OutputLayerCallback->>ShardedState: store untied flag and materialize scale buffers
sequenceDiagram
participant Forward
participant NVFP4Quantizers
participant Teacher
participant Student
Forward->>NVFP4Quantizers: promote eligible quantizers using loaded amax values
NVFP4Quantizers-->>Forward: return promotion statistics
Forward->>Teacher: run inference
Forward->>Student: run inference
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/distill/plugins/megatron.py`:
- Around line 624-627: Move the StaticBlockScaleQuantizer and TensorQuantizer
import from the local scope to module scope in megatron.py. Only retain it
locally if the dependency is circular, optional, or unusually heavy, and
document that reason beside the import.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2717a12c-a49e-4acf-8aee-997ebfcbea38
📒 Files selected for processing (2)
modelopt/torch/distill/plugins/megatron.pymodelopt/torch/quantization/plugins/megatron.py
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2112 +/- ##
==========================================
- Coverage 78.60% 78.53% -0.08%
==========================================
Files 522 522
Lines 60167 60212 +45
==========================================
- Hits 47294 47285 -9
- Misses 12873 12927 +54
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[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| _wq = getattr(self, "weight_quantizer", None) | ||
| if _wq is not None and getattr(_wq, "is_enabled", False): | ||
| _block_sizes = getattr(_wq, "_block_sizes", None) or {} | ||
| _block = _block_sizes.get(-1) or _block_sizes.get(1) | ||
| # `_process_quantizer_amax` later does `v.view(weight.shape[0], -1)`, which | ||
| # requires in_features (not just numel) to divide evenly by the block size. | ||
| if _block and self.weight.shape[-1] % int(_block) == 0: | ||
| if getattr(_wq, "_amax", None) is None: | ||
| _wq.amax = torch.zeros( | ||
| self.weight.numel() // int(_block), | ||
| 1, | ||
| dtype=torch.float32, | ||
| device=self.weight.device, | ||
| ) | ||
| # register_buffer directly: the ``global_amax`` property lives on | ||
| # StaticBlockScaleQuantizer, and on restore this is still a plain | ||
| # TensorQuantizer (promotion happens later), so the setter is unavailable. | ||
| if getattr(_wq, "_global_amax", None) is None: | ||
| _wq.register_buffer( | ||
| "_global_amax", | ||
| torch.zeros((), dtype=torch.float32, device=self.weight.device), | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Algorithm] This materialization is not gated on the quantizer being static block quant, and it runs on the save path as well as the load path.
sharded_state_dict() is called both when writing a checkpoint and when building the load plan. The only guards here are is_enabled and block-size divisibility, so:
-
Dynamic NVFP4 (the default) gets a bogus
_amax.modelopt_recipes/configs/numerics/nvfp4.yaml— used byw4_nvfp4(W4A16) andw4a4_nvfp4_nvfp4— sets the weight quantizer toblock_sizes: {-1: 16, type: dynamic, scale_bits: e4m3}. A dynamic quantizer intentionally never has_amax(TensorQuantizer.forwardskips calibration whenself._dynamic), andTensorQuantizer.amaxassertsnot self._dynamic. With an untied, enabledoutput_layerunder dynamic NVFP4, this block registers a zero_amaxon a dynamic quantizer, which then (a) trips that assertion on any later.amaxaccess, and (b) is picked up by the loop below ("_quantizer" in k and "_amax" in k) and written into the checkpoint as an all-zeroweight_quantizer._amaxfor a layer where no other rank/layer has one. The author's verified config (nvfp4_four_over_six) happens to betype: static, so this path wasn't exercised. -
Zeros are not a safe default even in the static case. If the checkpoint being loaded doesn't carry these keys (e.g. a checkpoint saved before this fix, or a save-path invocation), the buffers stay all-zero, and a zero amax/global_amax yields zero or NaN scales at export/forward — worse than the BF16 fallback this PR is fixing, and equally silent.
Suggested fix: gate on static block quant, and only materialize when actually building a load plan (or at minimum verify post-load that the values are non-zero):
_wq = getattr(self, "weight_quantizer", None)
if (
_wq is not None
and getattr(_wq, "is_enabled", False)
and getattr(_wq, "is_static_block_quant", False)
):
...and consider initializing _amax from the current weight (reduce_amax over blocks) rather than zeros, so a checkpoint that lacks the key degrades to "recalibrated from weights" instead of "scale = 0".
There was a problem hiding this comment.
Accepted, both halves — fixed in f46cc40.
- The materialization is now gated on
is_static_block_quant, so a dynamic NVFP4 weight quantizer is left alone. You are right that it would otherwise have registered an_amaxon a quantizer whoseamaxproperty assertsnot self._dynamic, and written an all-zero amax into the checkpoint. Verifiednvfp4.yamlusestype: dynamicwhile the config this PR was developed against (nvfp4_four_over_six) istype: static— so the bad path was real and simply never exercised. - The buffers are now seeded from the weights rather than zeros:
For weight-only quantization that is exactly what max calibration produces, so a checkpoint that turns out not to carry these keys degrades to "recalibrated from weights" instead of to scale=0/NaN — which, as you note, would be worse than the BF16 fallback this PR is fixing.
_wq.amax = self.weight.detach().reshape(-1, int(_block)).abs().amax(dim=1, keepdim=True).float() ... _wq.register_buffer("_global_amax", _wq._amax.detach().max().float().clone())
| for name, module in self.named_modules(): | ||
| if not isinstance(module, TensorQuantizer) or isinstance( | ||
| module, StaticBlockScaleQuantizer | ||
| ): | ||
| continue | ||
| block_sizes = getattr(module, "_block_sizes", None) | ||
| is_nvfp4 = getattr(module, "_num_bits", None) == (2, 1) and ( | ||
| isinstance(block_sizes, dict) and block_sizes.get("scale_bits") == (4, 3) | ||
| ) | ||
| if not is_nvfp4: | ||
| continue | ||
| amax = getattr(module, "_amax", None) | ||
| if amax is None: | ||
| # Uncalibrated: leave it alone rather than silently changing precision. | ||
| logger.warning(f"NVFP4 weight quantizer {name} has no _amax; not promoted.") | ||
| n_skipped += 1 | ||
| continue | ||
| StaticBlockScaleQuantizer.from_tensor_quantizer( | ||
| module, global_amax=amax.detach().float().abs().max() | ||
| ) | ||
| n_promoted += 1 |
There was a problem hiding this comment.
[CRITICAL Algorithm] This hand-rolled promotion loop diverges from the canonical one in three ways that change numerics, and it should just call the existing helper.
modelopt.torch.quantization.utils.promote_static_block_weight_quantizers() (utils/core_utils.py:1075) does exactly this job. Compared to it, this copy:
-
Promotes activation and KV quantizers, not just weight quantizers. The canonical version iterates
module.iter_weights_for_calibration(); this one iterates everyTensorQuantizerin the model. Under W4A4 NVFP4 (w4a4_nvfp4_nvfp4) the*input_quantizeris alsonum_bits: e2m1+scale_bits: e4m3, andkv_nvfp4gives*[kv]_bmm_quantizerthe same attributes — so any of those that carries an_amaxgets converted toStaticBlockScaleQuantizer, silently changing the activation/KV quant path during QAD. The comment claims "the student's weight quantizers", but nothing here restricts to weights. -
Missing the
is_static_block_quant/is_enabledguards. Theis_nvfp4predicate here checks only_num_bits == (2, 1)andscale_bits == (4, 3); it does not checkblock_sizes["type"] != "dynamic"or_fake_quant, whichis_nvfp4_static(tensor_quantizer.py:573) does. A dynamic NVFP4 quantizer that happens to hold a stale_amaxwould be promoted to a static one. -
self.named_modules()includes_teacher_model. The teacher is a registered submodule (distillation_model.py:102), so a quantized teacher's quantizers are promoted too — and the count in the comment (460/1/0) wouldn't distinguish that.
Replacing the block with the shared helper fixes all three and satisfies the DRY guidance in CONTRIBUTING:
if not getattr(self, "_modelopt_nvfp4_promoted", False):
# Local import: quantization is an optional dependency of distillation.
from modelopt.torch.quantization.utils import promote_static_block_weight_quantizers
n = promote_static_block_weight_quantizers(self)
if n:
logger.info(f"Promoted {n} static-block weight quantizer(s) after checkpoint load.")
self._modelopt_nvfp4_promoted = True(If the teacher must be excluded, call it on the student submodule rather than self.)
There was a problem hiding this comment.
Accepted — replaced with the canonical helper in f46cc40.
Confirmed all three divergences: promote_static_block_weight_quantizers() (utils/core_utils.py:965) iterates iter_weights_for_calibration() rather than every TensorQuantizer, checks is_enabled / is_static_block_quant / is_nvfp4_static, and additionally ties grouped siblings to a shared global_amax — none of which the copy did.
The teacher is excluded via the existing hide_teacher_model() context manager rather than by calling the helper on a submodule, so the helper still sees the whole student root and can discover shared-state groups that span children.
| amax = getattr(module, "_amax", None) | ||
| if amax is None: | ||
| # Uncalibrated: leave it alone rather than silently changing precision. | ||
| logger.warning(f"NVFP4 weight quantizer {name} has no _amax; not promoted.") | ||
| n_skipped += 1 | ||
| continue | ||
| StaticBlockScaleQuantizer.from_tensor_quantizer( | ||
| module, global_amax=amax.detach().float().abs().max() | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Algorithm] Passing global_amax= unconditionally overwrites the _global_amax this PR's other half just restored from the checkpoint, with a per-TP-rank recomputation.
from_tensor_quantizer → _preserve_and_set_global_amax does tq.global_amax = global_amax whenever the argument is not None, and the global_amax setter copies into the existing buffer (tensor_quantizer.py:1588-1601). So for output_layer — the one module this block exists for, and the module whose _global_amax the sharded_state_dict change now materializes and loads — the checkpoint value is discarded here and replaced by amax.abs().max().
That substitution is not value-preserving under TP: output_layer is column-parallel and its per-block _amax is sharded (_MegatronColumnParallelLinear._get_shard_axis_dict shards _amax on dim 0 and explicitly skips _global_amax as "a replicated scalar"). Taking max() over the local shard therefore produces a different global scale on every TP rank, whereas the checkpointed scalar is replicated. The result is rank-inconsistent weight_scale_2 during QAD, and a mismatch against the calibrated value the checkpoint carried.
Fix: only supply global_amax when the buffer is genuinely absent, so a restored value wins.
already_has_global = getattr(module, "_global_amax", None) is not None
StaticBlockScaleQuantizer.from_tensor_quantizer(
module,
global_amax=None if already_has_global else reduce_amax(amax.clone().detach(), axis=None),
)Using the shared promote_static_block_weight_quantizers() helper (see the other comment) does not by itself solve this — it also passes a locally-reduced global_amax — so if you switch to the helper, this needs handling there or a guard here.
There was a problem hiding this comment.
Accepted — fixed in f46cc40, and thank you for flagging that switching to the helper does not fix it by itself. It does not: core_utils.py:1028 computes global_amax = reduce_amax(amax.clone().detach(), axis=None) and _preserve_and_set_global_amax assigns it whenever it is not None.
Rather than special-casing the argument, the restored values are snapshotted and written back around the promotion, so the checkpoint's replicated scalar wins over any rank-local recomputation:
restored_global_amax = {
id(m): m._global_amax.detach().clone()
for m in self.modules()
if getattr(m, "_global_amax", None) is not None
}
with self.hide_teacher_model():
n_promoted = promote_static_block_weight_quantizers(self)
for m in self.modules():
saved = restored_global_amax.get(id(m))
if saved is not None:
m._global_amax.copy_(saved)This also keeps working if the helper's own global_amax handling changes later.
| # Skip output_layer w/o enabled weight_quantizer. This hook also runs BEFORE | ||
| # QuantModule replacement (e.g. on restore), when ``weight_quantizer`` does not | ||
| # exist yet -- the old check then always skipped, so output_layer never received | ||
| # ModelOpt extra-state callbacks and its quantizer state (promotion to | ||
| # StaticBlockScaleQuantizer, ``_amax``, ``_global_amax``) was never restored. | ||
| # Fall back to the tying flag: an untied output_layer is quantizable. | ||
| if name.endswith("output_layer"): | ||
| _wq = getattr(module, "weight_quantizer", None) | ||
| _skip = ( | ||
| not getattr(_wq, "is_enabled", False) if _wq is not None else not untied | ||
| ) | ||
| if _skip: | ||
| continue | ||
| register_modelopt_extra_state_callbacks( |
There was a problem hiding this comment.
[IMPORTANT Compatibility] The fallback registers extra-state callbacks for every untied output_layer, including the (very common) case where its quantizer is disabled — which changes the set of keys existing checkpoints advertise.
Since this hook runs before QuantModule replacement, _wq is always None in that pass, so _skip = not untied is the only thing deciding registration. modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml disables *output_layer* and *lm_head* by default, so for the mainstream recipes on any untied model (Llama-style, most HF configs), output_layer now gets get_extra_state/set_extra_state patched where it previously did not.
_modelopt_get_extra_state returns a non-empty dict as soon as any callback fires (quant_module_get_extra_state always emits modelopt_quantizer_state, plus real-quant keys), so output_layer._extra_state appears in the sharded state dict. That is a state-dict schema change in both directions:
- new code loading a pre-PR checkpoint: the model advertises
output_layer._extra_state, the checkpoint has no such key — Megatron's dist loader errors on keys the model requires but the checkpoint lacks (the silent-skip asymmetry you documented is the other direction); - pre-PR code loading a new checkpoint: the extra key is silently dropped, which is benign but leaves the two formats non-identical.
Suggestion: keep the fallback narrow so it only fires when output_layer will actually be quantized. Two options — either consult the resolved quant config for the layer instead of just tiedness, or (simpler) restrict the _wq is None fallback to the restore path, e.g. register only when the module already carries ModelOpt quantizer state. Whichever you choose, please note the checkpoint-format implication in the PR description so it isn't discovered on a resume.
There was a problem hiding this comment.
Half accepted: I will document this in the PR description, but I do not think the code can be narrowed the way you suggest.
You are right about the mechanism — at this point in the hook weight_quantizer does not exist yet, so untied is the only available signal, and an untied output_layer under a recipe that disables it will now get extra-state callbacks it did not get before.
The problem with both proposed narrowings is that they need information this hook cannot have:
- "consult the resolved quant config" — the quant config is not passed to
CUSTOM_MODEL_PLUGINScallbacks; the hook only receives the model. - "register only when the module already carries ModelOpt quantizer state" — that is false on the save side of a fresh
mtq.quantize, where the state is created by the very replacement this hook runs before. Gating on it would reinstate exactly the silent-BF16 bug this PR fixes, just on a different path.
Scoping note on the blast radius: this hook only runs inside ModelOpt quantize/restore flows, so a plain Megatron checkpoint is unaffected; the change is confined to ModelOpt checkpoints where output_layer is untied.
If a maintainer prefers, I am happy to thread the resolved quant config into the hook in a follow-up so registration can key off the layer's actual config rather than tiedness — that is the only version of this that is correct in both directions. Flagging it in the PR description meanwhile.
| def _resolve_output_layer_untied(model: torch.nn.Module) -> bool | None: | ||
| """Whether ``output_layer`` weights are untied from the input embeddings, or None if unknown. | ||
|
|
||
| Megatron-Core models carry ``share_embeddings_and_output_weights`` (Megatron-Bridge sets it | ||
| from the HF config, Megatron-LM from ``--untie-embeddings-and-output-weights``), so reading | ||
| it off the model works under both frameworks. ``megatron.training.get_args()`` does not: | ||
| Bridge has no global args store, and defaulting to "tied" there silently drops the | ||
| ``output_layer`` weight-quantizer state from the sharded checkpoint. | ||
| """ | ||
| for _, module in model.named_modules(): | ||
| shared = getattr(module, "share_embeddings_and_output_weights", None) | ||
| if shared is not None: | ||
| return not bool(shared) | ||
| return None |
There was a problem hiding this comment.
[IMPORTANT Compatibility] "First module in named_modules() that has the attribute" is fragile for multi-model graphs, and the resolved value is then broadcast onto all non-vision configs.
Two concrete cases:
- Distillation / QAD: this hook is invoked via
CUSTOM_MODEL_PLUGINSon whatever root it's given. For aDistillationModel,named_modules()walks both_teacher_modeland the student; whichever appears first decides tiedness for both, and the flag is then written onto everyMegatronModule.configbelow (line 339). If teacher and student differ in tying, one of them gets the wrong answer. - VLM: the flag is looked up across the whole tree including
vision_model, but only written to non-vision_modelconfigs._resolve_output_layer_untiedshould be scoped to the language model it's describing.
Since output_layer is what the flag is about, resolving from the module that actually owns it is both narrower and unambiguous — walk to the enclosing model of each output_layer, or at minimum prefer the root's own attribute before falling back to a subtree scan:
def _resolve_output_layer_untied(model):
shared = getattr(model, "share_embeddings_and_output_weights", None)
if shared is not None:
return not bool(shared)
for name, module in model.named_modules():
if "vision_model" in name:
continue
shared = getattr(module, "share_embeddings_and_output_weights", None)
if shared is not None:
return not bool(shared)
return NoneNote the existing consumers of this flag elsewhere in the codebase all read it off the specific model (unified_export_megatron.py:480, megatron_importer.py:825), not off an arbitrary first match.
There was a problem hiding this comment.
Accepted — fixed in f46cc40.
_resolve_output_layer_untied now prefers the root's own share_embeddings_and_output_weights and skips vision_model / _teacher_model subtrees in the fallback scan. The DistillationModel case was the real one: named_modules() walks both teacher and student, so whichever came first decided tiedness for both.
Added unit coverage for all four cases (no signal, root wins over subtree, subtree fallback, vision/teacher skipped).
| # Static-block NVFP4: promote the student's weight quantizers once, after the | ||
| # checkpoint amax/scales have been loaded, so the training forward takes the | ||
| # StaticBlockScaleQuantizer path rather than the generic FP8 (E4M3) path. Promotion | ||
| # cannot happen at build time because the scales only exist after the load. | ||
| # | ||
| # NOTE: in practice this converts exactly ONE module -- ``output_layer``. A measured run | ||
| # reports ``already promoted 460, converted 1, skipped 0``: every other quantizer is | ||
| # already a StaticBlockScaleQuantizer by the time training starts. So this is a workaround | ||
| # for output_layer being the one module the restore path does not promote (the same | ||
| # asymmetry behind its weight-quantizer scales not being restored). The better fix is to | ||
| # promote it on the normal path; until then, without this block the output projection | ||
| # would train through the generic FP8 path instead of static-block NVFP4. | ||
| if not getattr(self, "_modelopt_nvfp4_promoted", False): |
There was a problem hiding this comment.
[SUGGESTION] Two smaller points on placing this in _forward:
-
The comment documents the workaround as a workaround. It says outright that "the better fix is to promote it on the normal path" — and the other half of this PR is touching that normal path. Since promotion after
set_extra_stateis whatquant_module_set_extra_state→maybe_promote_nvfp4_static_quantizeralready does for every other module, it would be worth a follow-up issue reference here so this doesn't become permanent. -
Promotion mutates module classes and can register a new
_global_amaxbuffer at first forward, i.e. after the model has been wrapped in DDP / the distributed optimizer and afterparam_and_grad_bufferbucketing. Newly registered buffers aren't broadcast by DDP, so ranks rely on each computing the same value locally — which is the concern raised in theglobal_amax=comment above. If promotion instead happens at the end of restore (before wrapping), both problems disappear.
Also: the per-run numbers in the comment (already promoted 460, converted 1, skipped 0) will read as stale the first time someone runs a different model/recipe. Consider dropping the counts and keeping the "exactly output_layer in practice" statement.
There was a problem hiding this comment.
Accepted the concrete parts — f46cc40.
- Dropped the run-specific counts from the comment; kept the "exactly
output_layerin practice" statement. - Reworded the follow-up note to name the destination explicitly: promotion belongs on the normal restore path, before the model is wrapped in DDP. That is the same fix that removes your second concern (a buffer registered at first forward is not broadcast), so I have kept them as one item rather than two.
I have deliberately not moved the promotion in this PR — doing it at restore time touches the shared quant_module_set_extra_state path for every module, which is a larger change than the output_layer fix this PR is scoped to. Happy to open it as a follow-up issue.
| if prefix.endswith("output_layer."): | ||
| try: | ||
| from megatron.training import get_args as _mlm_get_args | ||
|
|
||
| _untied = bool( | ||
| getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) | ||
| ) | ||
| except Exception as e: | ||
| warn_rank_0(f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}") | ||
| _untied = False | ||
| # Prefer the model-derived flag (set by megatron_replace_quant_module_hook); it is the | ||
| # only source available under Megatron-Bridge, which has no global args store. | ||
| _untied = getattr(self.config, "modelopt_output_layer_untied", None) | ||
| if _untied is None: | ||
| try: | ||
| from megatron.training import get_args as _mlm_get_args | ||
|
|
||
| _untied = bool( | ||
| getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) | ||
| ) | ||
| except Exception as e: | ||
| warn_rank_0( | ||
| f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}" | ||
| ) | ||
| _untied = False | ||
| if not _untied: | ||
| return super().sharded_state_dict(prefix, sharded_offsets, metadata) |
There was a problem hiding this comment.
[SUGGESTION] The tiedness resolution now exists in two places with a precedence rule between them (self.config.modelopt_output_layer_untied first, get_args() second). That logic — "model flag wins, Megatron-LM global args as fallback" — would be clearer as a single helper next to _resolve_output_layer_untied, e.g. _output_layer_untied(config), so the fallback can't drift out of sync if a third source ever appears.
Minor related point: when _resolve_output_layer_untied returns None and get_args() raises, warn_rank_0 fires once per output_layer per sharded_state_dict() call. Under Megatron-Bridge with a model that carries neither signal, that's a warning on every save and every load. Caching the resolved value on the config (as the hook already does for the success case) would keep it to one.
There was a problem hiding this comment.
Accepted — fixed in f46cc40.
Both points are addressed by a single _output_layer_untied(config) helper that owns the precedence rule and caches the answer back onto the config, so the get_args()-raises path warns once rather than on every save and every load. sharded_state_dict is now a one-liner:
if not _output_layer_untied(self.config):
return super().sharded_state_dict(prefix, sharded_offsets, metadata)Unit-tested for both precedence and the caching behaviour.
There was a problem hiding this comment.
Claude review — 3 CRITICAL, 2 IMPORTANT, 2 SUGGESTION
Scope: full review. Both changed files (modelopt/torch/quantization/plugins/megatron.py, modelopt/torch/distill/plugins/megatron.py) reviewed in full; traced the new modelopt_output_layer_untied flag and _resolve_output_layer_untied to their consumers, and compared the new promotion loop against the canonical promote_static_block_weight_quantizers in utils/core_utils.py.
The root-cause analysis in the description is right and the fix direction is correct — resolving tiedness from the model rather than from megatron.training.get_args() is the right call, and the callback-registration bug is real. The findings below are about the blast radius of the two mechanisms, not the diagnosis.
Most impactful
-
Zero-initialized scale buffers on a dynamic weight quantizer, on both save and load paths (
quantization/plugins/megatron.py:433-454). The materialization block is guarded only byis_enabled+ divisibility, not byis_static_block_quant. The default NVFP4 numerics (configs/numerics/nvfp4.yaml, used byw4_nvfp4andw4a4_nvfp4_nvfp4) set the weight quantizer totype: dynamic, where_amaxintentionally never exists andTensorQuantizer.amaxassertsnot self._dynamic. Your verified recipe (nvfp4_four_over_six) istype: static, so this path wasn't exercised. Separately,sharded_state_dict()also runs on save, so an all-zero_amaxcan be written into a checkpoint — and a zero scale is a worse failure than the BF16 fallback being fixed. -
The promotion loop promotes activation and KV quantizers, not just weights (
distill/plugins/megatron.py:630-650). It iterates everyTensorQuantizerand matches on_num_bits/scale_bitsonly. Underw4a4_nvfp4_nvfp4the*input_quantizermatches; underkv_nvfp4the*[kv]_bmm_quantizermatches. It also walks_teacher_model, which is a registered submodule. The canonical helperpromote_static_block_weight_quantizers()already restricts toiter_weights_for_calibration()and checksis_static_block_quant/is_enabled— reusing it fixes all of these. -
Promotion clobbers the restored
_global_amaxwith a per-TP-rank value (distill/plugins/megatron.py:641-649).from_tensor_quantizer(..., global_amax=...)copies into the existing buffer, so the value thesharded_state_dicthalf of this PR just loaded foroutput_layeris overwritten byamax.abs().max()over the local shard._MegatronColumnParallelLinear._get_shard_axis_dictshards_amaxon dim 0 and explicitly treats_global_amaxas a replicated scalar — so the two halves of this PR fight each other, and ranks end up disagreeing onweight_scale_2.
Also flagged as IMPORTANT: the registration fallback fires for every untied output_layer even when its quantizer is disabled (which is the default via default_disabled_quantizers.yaml), adding an output_layer._extra_state key that pre-PR checkpoints don't carry — a two-way state-dict format change worth calling out explicitly; and _resolve_output_layer_untied's "first matching module wins" scan is ambiguous for DistillationModel (teacher + student) and unscoped w.r.t. vision_model.
Risk assessment
Moderate-to-high. The change is small and the target case is verified end-to-end, but the guards are tuned to one recipe (nvfp4_four_over_six, static) and one framework path; the default dynamic-NVFP4 recipes take a code path this PR newly makes reachable. Given the PR has no unit tests, a focused test that (a) builds an untied output_layer under dynamic NVFP4 and round-trips sharded_state_dict, and (b) asserts a restored _global_amax survives promotion, would cover the two riskiest findings cheaply.
…port - ruff-format: drop the blank line after the docstring and collapse the wrapped _skip ternary in megatron_replace_quant_module_hook. - Explain why StaticBlockScaleQuantizer/TensorQuantizer are imported inside _forward: the distillation plugin must stay importable without modelopt.torch.quantization, so it cannot take a module-scope dependency on it (addresses CodeRabbit review comment). Signed-off-by: James Shen <yueshen@nvidia.com>
… materialization, reuse promotion helper quantization/plugins/megatron.py: - `_resolve_output_layer_untied` now prefers the root's own `share_embeddings_and_output_weights` and skips subtrees that do not own the language model's output_layer (`vision_model`, `_teacher_model`). "First module in named_modules() with the attribute" gave the wrong answer for a DistillationModel, whose teacher and student are both walked. - New `_output_layer_untied(config)` holds the single precedence rule (model-derived flag, then Megatron-LM args) and caches the result on the config, so a model carrying neither signal warns once instead of on every save and every load. - Scale-buffer materialization is now gated on `is_static_block_quant`. A dynamic NVFP4 weight quantizer (the default in configs/numerics/nvfp4.yaml) deliberately never holds an `_amax` -- its `amax` property asserts `not self._dynamic` -- so materializing one there registered a bogus buffer and wrote an all-zero amax to the checkpoint. - The buffers are seeded from the weights instead of zeros. For weight-only quantization that is exactly what max calibration produces, so a checkpoint that turns out not to carry these keys degrades to "recalibrated from weights" rather than to scale=0/NaN at export and forward. distill/plugins/megatron.py: - Replace the hand-rolled promotion loop with the canonical `promote_static_block_weight_quantizers()`. The copy iterated every TensorQuantizer rather than weight quantizers only (so W4A4 activation and KV quantizers were eligible), skipped the `is_static_block_quant` / `is_enabled` guards, and walked the teacher. - Preserve the `_global_amax` restored from the checkpoint across promotion. It is a replicated scalar, but promotion recomputes it with `reduce_amax` over the rank-local weight shard, which is rank-inconsistent for the column-parallel output_layer -- and discarded the value the other half of this PR had just restored. - Drop the run-specific promotion counts from the comment. tests: cover both resolvers (root-flag precedence, subtree scan, vision/teacher skipping, args fallback and caching). Signed-off-by: James Shen <yueshen@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py`:
- Line 1309: Move the plugin imports for _resolve_output_layer_untied at
tests/gpu_megatron/torch/quantization/plugins/test_megatron.py lines 1309-1309
and _output_layer_untied at lines 1340-1340 into the module-level import
section, removing the imports from the test methods; no direct changes are
required elsewhere.
- Around line 1349-1355: Update the test for _output_layer_untied to patch
megatron.training.get_args and cover both True and False fallback results.
Assert each resolved value, verify it is cached on
config.modelopt_output_layer_untied, and assert the second invocation returns
the cached value without calling get_args again.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c02ef6fc-b388-488a-b819-b9c57f0ef274
📒 Files selected for processing (3)
modelopt/torch/distill/plugins/megatron.pymodelopt/torch/quantization/plugins/megatron.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.py
🚧 Files skipped from review as they are similar to previous changes (2)
- modelopt/torch/quantization/plugins/megatron.py
- modelopt/torch/distill/plugins/megatron.py
|
|
||
| def test_resolve_output_layer_untied(): | ||
| """The tiedness signal is read off the model, not from Megatron-LM global args.""" | ||
| from modelopt.torch.quantization.plugins.megatron import _resolve_output_layer_untied |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Move the plugin imports to module scope.
These imports have no circular-dependency, optional-dependency, or deferred-heavy-import justification. Import errors should fail during test collection.
tests/gpu_megatron/torch/quantization/plugins/test_megatron.py#L1309-L1309: move_resolve_output_layer_untiedto the module import section.tests/gpu_megatron/torch/quantization/plugins/test_megatron.py#L1340-L1340: move_output_layer_untiedto the module import section.
As per path instructions, imports inside test methods require an explicit justification and otherwise belong at the top of the file.
📍 Affects 1 file
tests/gpu_megatron/torch/quantization/plugins/test_megatron.py#L1309-L1309(this comment)tests/gpu_megatron/torch/quantization/plugins/test_megatron.py#L1340-L1340
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py` at line 1309,
Move the plugin imports for _resolve_output_layer_untied at
tests/gpu_megatron/torch/quantization/plugins/test_megatron.py lines 1309-1309
and _output_layer_untied at lines 1340-1340 into the module-level import
section, removing the imports from the test methods; no direct changes are
required elsewhere.
Sources: Coding guidelines, Path instructions
| # With no model-derived flag the args fallback runs, and its answer is cached back onto the | ||
| # config so a model carrying neither signal does not warn on every save and every load. | ||
| config = _Config() | ||
| resolved = _output_layer_untied(config) | ||
| assert isinstance(resolved, bool) | ||
| assert config.modelopt_output_layer_untied is resolved | ||
| assert _output_layer_untied(config) is resolved |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Test the Megatron-LM fallback result.
The test passes when megatron.training is unavailable because _output_layer_untied catches the import failure and returns False. It also passes if the fallback ignores a true untie_embeddings_and_output_weights value.
Patch megatron.training.get_args to return both True and False values. Assert the resolved value and assert that the second call uses the cached value without another get_args call. As per coding guidelines, tests must exercise the behavior they claim to validate.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py` around lines
1349 - 1355, Update the test for _output_layer_untied to patch
megatron.training.get_args and cover both True and False fallback results.
Assert each resolved value, verify it is cached on
config.modelopt_output_layer_untied, and assert the second invocation returns
the cached value without calling get_args again.
Source: Coding guidelines
1 participant
What does this PR do ?
Type of change: Bug fix
Overview: A quantized
output_layer(lm_head) is silently exported as BF16 when themodel is built through Megatron-Bridge — no error, no warning, just an unquantized output layer
in the exported checkpoint. This restores it.
Root cause — two independent failures, both still on
main1. The extra-state callbacks are never registered.
This hook also runs before
QuantModulereplacement (e.g. on restore), whenweight_quantizerdoes not exist yet — so the check always skipped, andoutput_layerneverreceived the ModelOpt extra-state callbacks. Its quantizer state (promotion to
StaticBlockScaleQuantizer,_amax,_global_amax) was therefore never saved or restored.2. Tiedness is read from a Megatron-LM global that Megatron-Bridge does not have.
Megatron-Bridge has no
megatron.trainingglobal args store, so the import raises, theexceptpath treats the layer as tied, and quantization is dropped.The two combine into a silent failure: the save path writes nothing for the output layer, and
the load path has nothing to restore.
The fix
Resolve tiedness from the model instead of from a framework global:
_resolve_output_layer_untied()walksnamed_modules()for the framework-agnosticshare_embeddings_and_output_weightsflag and records it on the model config asmodelopt_output_layer_untied.sharded_state_dictprefers that flag and falls back toget_args()only when it is absent— Megatron-LM behaviour is unchanged.
weight_quantizerdoes not existyet, so an untied
output_layeris registered.has somewhere to write. A Megatron dist-checkpoint load silently skips any key the model
does not advertise, which is why the missing buffers produced no error. Divisibility is
checked against
weight.shape[-1](the axis_process_quantizer_amaxlater views over), anda warning is emitted rather than leaving the buffers unallocated.
modelopt/torch/distill/plugins/megatron.pygains a matching one-shot promotion on the firstforward, for the single module the restore path does not promote. Measured on Nemotron-Nano-3:
already_sbsq=460 converted=1 disabled=0— exactlyoutput_layer.Testing
Verified end-to-end on Nemotron-Nano-3, W4A16 NVFP4
four_over_six, 4x GB200:lm_head.weight(U8[131072, 1344]) plusweight_scaleandweight_scale_2disabled=0(wasdisabled=1before the fix) independently confirms the outputlayer stays quantized through distillation
Without the fix the same pipeline produces a BF16
lm_headwith no diagnostic of any kind.Before your PR is "Ready for review"
get_args()path; the new model-derived flag is only consulted first, and only affects cases that are currently broken._resolve_output_layer_untied,_output_layer_untied), plus end-to-end verification on a real model as above.Checkpoint-format implication
megatron_replace_quant_module_hookruns beforeQuantModulereplacement, so at that pointweight_quantizerdoes not exist yet and tiedness is the only signal available. As a result, an untiedoutput_layernow receives ModelOpt extra-state callbacks even under a recipe that disables its quantizer (e.g.configs/ptq/units/default_disabled_quantizers.yaml), which adds anoutput_layer._extra_statekey to the sharded state dict that was previously absent.Scope: this hook only runs inside ModelOpt quantize/restore flows, so plain Megatron checkpoints are unaffected — but a ModelOpt checkpoint saved before this change and loaded after it (or vice versa) will differ by that key for untied models. Narrowing the registration to layers that are actually quantized needs the resolved quant config threaded into the hook, which is left as a follow-up; gating on the quantizer's presence here would reinstate the silent-BF16 bug this PR fixes.
🤖 Generated with Claude Code
Summary by CodeRabbit