fix(megatron): restore untied output_layer quantization under Megatron-Bridge

modelopt/torch/distill/plugins/megatron.py

@@ -608,6 +608,46 @@ 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.
+        #
+        # 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.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} static-block weight quantizer(s) to "
+                    "StaticBlockScaleQuantizer after checkpoint load."
+                )
+            self._modelopt_nvfp4_promoted = True
         with torch.no_grad():
             self._teacher_model.eval()
             teacher_output = self._teacher_model(*args, **kwargs)

modelopt/torch/quantization/plugins/megatron.py

@@ -248,6 +248,52 @@ 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.
+    """
+    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.
 
@@ -260,6 +306,7 @@ 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):
         """Configure Attention module for KV cache quantization compatibility."""
@@ -287,11 +334,17 @@ 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 +360,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 +431,59 @@ 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
-            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
+            # 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.
+            # 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)
+                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:
+                        # 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", _wq._amax.detach().max().float().clone()
+                        )
+                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:

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