From e3a455d72109f6947d086605251167915d94cc43 Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 08:27:07 -0700 Subject: [PATCH 1/3] fix(megatron): restore untied output_layer quantization under Megatron-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 --- modelopt/torch/distill/plugins/megatron.py | 46 ++++++++ .../torch/quantization/plugins/megatron.py | 103 +++++++++++++++--- 2 files changed, 135 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 7e81c21a462..e02d65eeef0 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -608,6 +608,52 @@ def _set_input_tensor(self, input_tensors: list[Tensor]): # HACK: Concatenate output tensors when PP>1 so they can be passed between ranks. def _forward(self, *args, **kwargs): + # 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): + from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + StaticBlockScaleQuantizer, + TensorQuantizer, + ) + + n_promoted = n_skipped = 0 + 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 + if n_promoted or n_skipped: + logger.info( + f"Promoted {n_promoted} NVFP4 weight quantizer(s) to " + f"StaticBlockScaleQuantizer ({n_skipped} skipped)." + ) + self._modelopt_nvfp4_promoted = True with torch.no_grad(): self._teacher_model.eval() teacher_output = self._teacher_model(*args, **kwargs) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..2f742b87c5a 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -248,6 +248,22 @@ def _incompatible_method(self, *args, **kwargs): return _incompatible_method +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 + + def megatron_replace_quant_module_hook(model: torch.nn.Module): """Configure Megatron-Core model quantization support. @@ -261,6 +277,8 @@ def megatron_replace_quant_module_hook(model: torch.nn.Module): 3. For Attention modules, we configure them to use core_attention path for KV cache quantization. """ + untied = _resolve_output_layer_untied(model) + def _configure_attention_for_kv_cache_quant(module: Attention): """Configure Attention module for KV cache quantization compatibility.""" # Disable flash_decode if enabled - it bypasses core_attention (only called during inference) @@ -287,11 +305,19 @@ def _configure_attention_for_kv_cache_quant(module: Attention): def _register_extra_state_callbacks(model: torch.nn.Module): for name, module in model.named_modules(): if type(module) in QuantModuleRegistry: - # Skip output_layer w/o enabled weight_quantizer - if name.endswith("output_layer") and not getattr( - getattr(module, "weight_quantizer", None), "is_enabled", False - ): - continue + # 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( module, quant_module_get_extra_state, @@ -307,6 +333,10 @@ def _register_extra_state_callbacks(model: torch.nn.Module): if "vision_model" not in name: # We only enable hetereogenous_dist_checkpoint for language model, vision model is not quantized module.config.hetereogenous_dist_checkpoint = True + if untied is not None: + # Carried on the config so _MegatronParallelLinear.sharded_state_dict can read + # it without Megatron-LM global args (absent under Megatron-Bridge). + module.config.modelopt_output_layer_untied = untied _register_extra_state_callbacks(module) @@ -374,18 +404,63 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # output_layer.input_quantizer._amax but TP-only does not. This lead to # state_dict mismatch. 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) + # Materialize missing weight-quantizer scale buffers so their keys appear in the load + # plan -- the dist-checkpoint loader SILENTLY SKIPS any checkpoint key the model does + # not advertise, which leaves output_layer uncalibrated and exports it as BF16. + # ``_amax`` must be allocated FLAT ``[numel // block, 1]``: that is the in-memory + # layout every other block-quantized layer uses, and ``_process_quantizer_amax`` below + # exposes it to the checkpoint as a ``[out_features, blocks]`` VIEW sharing the same + # storage, so the loader writes straight through. Allocating the viewed shape instead + # loads fine but leaves the wrong in-memory shape, which breaks the export scale math. + _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), + ) + else: + # Leaving the buffers unallocated is the silent-drop failure this block exists + # to prevent, so say so rather than proceeding quietly. + warn_rank_0( + f"{prefix}weight_quantizer: cannot materialize scale buffers " + f"(block_size={_block}, in_features={self.weight.shape[-1]}); its " + "calibrated scales will not be restored from the checkpoint." + ) + quantizer_state_dict = {} for k, v in self.state_dict(prefix="", keep_vars=True).items(): if "_quantizer" in k and "_amax" in k: From c56408d4eca411bfabe21db6c0cf944d747a5e3f Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 09:02:54 -0700 Subject: [PATCH 2/3] chore: apply pre-commit formatting and document local quantization import - 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 --- modelopt/torch/distill/plugins/megatron.py | 3 +++ modelopt/torch/quantization/plugins/megatron.py | 5 +---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index e02d65eeef0..cea1a7e1ff8 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -621,6 +621,9 @@ def _forward(self, *args, **kwargs): # 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): + # Imported locally on purpose: this distillation plugin must stay usable without + # ``modelopt.torch.quantization`` installed/imported (a plain, non-quantized KD run + # never reaches this branch), so it must not take a module-scope dependency on it. from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( StaticBlockScaleQuantizer, TensorQuantizer, diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 2f742b87c5a..95d8ca1f3c8 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -276,7 +276,6 @@ def megatron_replace_quant_module_hook(model: torch.nn.Module): typing-matching the QuantModuleRegistry. 3. For Attention modules, we configure them to use core_attention path for KV cache quantization. """ - untied = _resolve_output_layer_untied(model) def _configure_attention_for_kv_cache_quant(module: Attention): @@ -313,9 +312,7 @@ def _register_extra_state_callbacks(model: torch.nn.Module): # 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 - ) + _skip = not getattr(_wq, "is_enabled", False) if _wq is not None else not untied if _skip: continue register_modelopt_extra_state_callbacks( From f46cc40933b51c28389fc8da83e8530d84c6ff56 Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 09:21:00 -0700 Subject: [PATCH 3/3] =?UTF-8?q?fix(megatron):=20address=20review=20?= =?UTF-8?q?=E2=80=94=20scope=20tiedness=20resolution,=20gate=20scale=20mat?= =?UTF-8?q?erialization,=20reuse=20promotion=20helper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- modelopt/torch/distill/plugins/megatron.py | 65 +++++++--------- .../torch/quantization/plugins/megatron.py | 76 +++++++++++++------ .../quantization/plugins/test_megatron.py | 51 +++++++++++++ 3 files changed, 130 insertions(+), 62 deletions(-) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index cea1a7e1ff8..a80201e20fd 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -613,48 +613,39 @@ def _forward(self, *args, **kwargs): # 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. + # In practice this converts exactly ONE module -- ``output_layer``; 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 restore path, before the model is wrapped in DDP; 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): # Imported locally on purpose: this distillation plugin must stay usable without # ``modelopt.torch.quantization`` installed/imported (a plain, non-quantized KD run # never reaches this branch), so it must not take a module-scope dependency on it. - from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( - StaticBlockScaleQuantizer, - TensorQuantizer, - ) - - n_promoted = n_skipped = 0 - 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 - if n_promoted or n_skipped: + from modelopt.torch.quantization.utils import promote_static_block_weight_quantizers + + # ``_global_amax`` restored from the checkpoint is a replicated scalar, but promotion + # recomputes it with ``reduce_amax`` over this rank's LOCAL weight shard -- which is + # rank-inconsistent for the column-parallel output_layer, whose ``_amax`` is sharded + # on dim 0. Keep whatever the checkpoint carried. + restored_global_amax = { + id(m): m._global_amax.detach().clone() + for m in self.modules() + if getattr(m, "_global_amax", None) is not None + } + # The teacher is a registered submodule; promote the student's quantizers only. + 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) + if n_promoted: logger.info( - f"Promoted {n_promoted} NVFP4 weight quantizer(s) to " - f"StaticBlockScaleQuantizer ({n_skipped} skipped)." + f"Promoted {n_promoted} static-block weight quantizer(s) to " + "StaticBlockScaleQuantizer after checkpoint load." ) self._modelopt_nvfp4_promoted = True with torch.no_grad(): diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 95d8ca1f3c8..7b75590a742 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -257,13 +257,43 @@ def _resolve_output_layer_untied(model: torch.nn.Module) -> bool | None: 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(model, "share_embeddings_and_output_weights", None) + if shared is not None: + return not bool(shared) + for name, module in model.named_modules(): + # Skip subtrees that do not own the language model's output_layer: the vision tower (never + # quantized here) and a distillation teacher, which may be tied differently from the + # student it is wrapped with. + if "vision_model" in name or "_teacher_model" in name: + continue shared = getattr(module, "share_embeddings_and_output_weights", None) if shared is not None: return not bool(shared) return None +def _output_layer_untied(config) -> bool: + """Whether ``output_layer`` is untied, for use from ``sharded_state_dict``. + + Precedence: the model-derived flag recorded by ``megatron_replace_quant_module_hook`` (the only + source available under Megatron-Bridge, which has no global args store), then Megatron-LM's + ``--untie-embeddings-and-output-weights``. The answer is cached back onto ``config`` so a model + carrying neither signal warns once instead of on every save and every load. + """ + untied = getattr(config, "modelopt_output_layer_untied", None) + if untied is not None: + return untied + 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 + config.modelopt_output_layer_untied = untied + return untied + + def megatron_replace_quant_module_hook(model: torch.nn.Module): """Configure Megatron-Core model quantization support. @@ -401,22 +431,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # output_layer.input_quantizer._amax but TP-only does not. This lead to # state_dict mismatch. if prefix.endswith("output_layer."): - # 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: + if not _output_layer_untied(self.config): return super().sharded_state_dict(prefix, sharded_offsets, metadata) # Materialize missing weight-quantizer scale buffers so their keys appear in the load @@ -427,27 +442,38 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # exposes it to the checkpoint as a ``[out_features, blocks]`` VIEW sharing the same # storage, so the loader writes straight through. Allocating the viewed shape instead # loads fine but leaves the wrong in-memory shape, which breaks the export scale math. + # Only STATIC block quant owns these buffers: a dynamic quantizer derives its scales + # per forward and deliberately never holds an ``_amax`` (its ``amax`` property asserts + # ``not self._dynamic``), so materializing one there would be actively wrong. _wq = getattr(self, "weight_quantizer", None) - if _wq is not None and getattr(_wq, "is_enabled", False): + if ( + _wq is not None + and getattr(_wq, "is_enabled", False) + and getattr(_wq, "is_static_block_quant", 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, + # Seed from the weights rather than zeros. For weight-only quantization + # this 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 (or NaN) at export and forward. + _wq.amax = ( + self.weight.detach() + .reshape(-1, int(_block)) + .abs() + .amax(dim=1, keepdim=True) + .float() ) # 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), + "_global_amax", _wq._amax.detach().max().float().clone() ) else: # Leaving the buffers unallocated is the silent-drop failure this block exists diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index d995c4b388c..8b457918545 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -1302,3 +1302,54 @@ def test_homogeneous_sharded_state_dict_te_spec(dist_workers, tmp_path): {"transformer_impl": "transformer_engine"}, ), ) + + +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 + + class _Flagged(torch.nn.Module): + def __init__(self, shared): + super().__init__() + self.share_embeddings_and_output_weights = shared + + # No signal anywhere -> unknown. + assert _resolve_output_layer_untied(torch.nn.Module()) is None + + # The root's own flag wins over any subtree. + root = _Flagged(False) + root.inner = _Flagged(True) + assert _resolve_output_layer_untied(root) is True + + # Otherwise fall back to a subtree scan. + root = torch.nn.Module() + root.language_model = _Flagged(True) + assert _resolve_output_layer_untied(root) is False + + # Subtrees that do not own the language model's output_layer are skipped: the vision tower + # and a distillation teacher, either of which may be tied differently from the student. + root = torch.nn.Module() + root.vision_model = _Flagged(True) + root._teacher_model = _Flagged(True) + root.language_model = _Flagged(False) + assert _resolve_output_layer_untied(root) is True + + +def test_output_layer_untied_precedence_and_caching(): + """The model-derived flag wins over Megatron-LM args, and the answer is cached.""" + from modelopt.torch.quantization.plugins.megatron import _output_layer_untied + + class _Config: + pass + + config = _Config() + config.modelopt_output_layer_untied = True + assert _output_layer_untied(config) is True + + # 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