diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 7e81c21a462..a80201e20fd 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/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) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..7b75590a742 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/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: 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