From 7f5095fb95a2c81acc456d4b28380c7e7bb42b9c Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 29 Jul 2026 11:22:37 -0700 Subject: [PATCH] fix: refresh TE fused MLP DDP hooks dynamically Signed-off-by: Nitin Vegesna --- .../core/extensions/transformer_engine.py | 100 ++++++++++-------- .../transformer/test_te_fused_mlp_hooks.py | 74 +++++++++++++ 2 files changed, 128 insertions(+), 46 deletions(-) create mode 100644 tests/unit_tests/transformer/test_te_fused_mlp_hooks.py diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 8f0117c68ec..6eb9062c793 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2902,42 +2902,52 @@ def _register_hooks_on_fused_impl(self, fused_impl: torch.nn.Module) -> None: """ - # Get submodule hooks - forward_pre_hooks = [] - forward_post_hooks = [] - backward_pre_hooks = [] - backward_post_hooks = [] - for submodule in self.modules(): - for hook_id, hook in submodule._forward_pre_hooks.items(): - with_kwargs = hook_id in submodule._forward_pre_hooks_with_kwargs - forward_pre_hooks.append((submodule, hook, with_kwargs)) - for hook_id, hook in submodule._forward_hooks.items(): - with_kwargs = hook_id in submodule._forward_hooks_with_kwargs - forward_post_hooks.append((submodule, hook, with_kwargs)) - for hook in submodule._backward_pre_hooks.values(): - backward_pre_hooks.append((submodule, hook)) - for hook in submodule._backward_hooks.values(): - backward_post_hooks.append((submodule, hook)) - - # Pre-forward hooks - # Note: DDP pre-forward hooks are safe since they do not - # interact with input tensor. - if forward_pre_hooks: - from megatron.core.distributed import distributed_data_parallel - - if any( - inspect.getmodule(hook) != distributed_data_parallel - for _, hook, _ in forward_pre_hooks - ): - warnings.warn( - "TEFusedMLP module has a submodule with a pre-forward hook. " - "TEFusedMLP module does not expose intermediate tensors, " - "so the hook may have incorrect behavior if it attempts to " - "access the input tensor." + # Hooks on TEFusedMLP itself are executed by its normal Module.__call__. + # Cache only the descendants whose calls the fused implementation skips. + skipped_submodules = tuple( + submodule for submodule in self.modules() if submodule is not self + ) + for submodule in skipped_submodules: + if submodule._backward_pre_hooks: + raise RuntimeError( + "TEFusedMLP module does not support submodules with " "pre-backward hooks" + ) + if submodule._backward_hooks: + raise RuntimeError( + "TEFusedMLP module does not support submodules with " "post-backward hooks" ) - def forward_pre_hook(module, *_) -> None: - for submodule, hook, with_kwargs in forward_pre_hooks: + distributed_data_parallel = None + warned_non_ddp_hooks = set() + + def forward_pre_hook(module, *_) -> None: + # DDP may disable its parameter-gather hooks for the first iteration and + # re-enable them after this fused implementation has been cached. Read the + # current hook registries on every invocation instead of capturing a stale + # construction-time snapshot. + nonlocal distributed_data_parallel + for submodule in skipped_submodules: + hooks_with_kwargs = submodule._forward_pre_hooks_with_kwargs + for hook_id, hook in list(submodule._forward_pre_hooks.items()): + if distributed_data_parallel is None: + from megatron.core.distributed import ( + distributed_data_parallel as ddp_module, + ) + + distributed_data_parallel = ddp_module + hook_key = (id(submodule), hook_id) + if ( + inspect.getmodule(hook) != distributed_data_parallel + and hook_key not in warned_non_ddp_hooks + ): + warnings.warn( + "TEFusedMLP module has a submodule with a pre-forward hook. " + "TEFusedMLP module does not expose intermediate tensors, " + "so the hook may have incorrect behavior if it attempts to " + "access the input tensor." + ) + warned_non_ddp_hooks.add(hook_key) + with_kwargs = hook_id in hooks_with_kwargs if with_kwargs: ret = hook(submodule, (), {}) else: @@ -2948,9 +2958,17 @@ def forward_pre_hook(module, *_) -> None: "submodule has pre-forward hook that modifies input tensor." ) - fused_impl.register_forward_pre_hook(forward_pre_hook) + # Install the pre-hook forwarder even when source hook registries are currently + # empty. DDP changes their contents throughout the training lifecycle. + fused_impl.register_forward_pre_hook(forward_pre_hook) - # Post-forward hooks + # Preserve the existing post-hook behavior without adding an unconditional + # forward-hook dispatch to models that do not use post hooks. + forward_post_hooks = [] + for submodule in skipped_submodules: + hooks_with_kwargs = submodule._forward_hooks_with_kwargs + for hook_id, hook in submodule._forward_hooks.items(): + forward_post_hooks.append((submodule, hook, hook_id in hooks_with_kwargs)) if forward_post_hooks: warnings.warn( "TEFusedMLP module has a submodule with a post-forward hook. " @@ -2973,16 +2991,6 @@ def forward_post_hook(module, *_) -> None: fused_impl.register_forward_hook(forward_post_hook) - # Backward hooks - if backward_pre_hooks: - raise RuntimeError( - "TEFusedMLP module does not support submodules with pre-backward hooks" - ) - if backward_post_hooks: - raise RuntimeError( - "TEFusedMLP module does not support submodules with post-backward hooks" - ) - def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Optional[Tensor]]: """Forward.""" diff --git a/tests/unit_tests/transformer/test_te_fused_mlp_hooks.py b/tests/unit_tests/transformer/test_te_fused_mlp_hooks.py new file mode 100644 index 00000000000..d8b4d6f2f2a --- /dev/null +++ b/tests/unit_tests/transformer/test_te_fused_mlp_hooks.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import warnings + +import pytest +import torch + +from megatron.core.extensions.transformer_engine import TEFusedMLP + +pytestmark = pytest.mark.skipif(TEFusedMLP is None, reason="Transformer Engine is unavailable") + + +def _make_fused_mlp_shell(): + module = TEFusedMLP.__new__(TEFusedMLP) + torch.nn.Module.__init__(module) + module.linear_fc1 = torch.nn.Linear(2, 2) + module.linear_fc2 = torch.nn.Linear(2, 2) + return module + + +def test_fused_mlp_forwards_current_submodule_pre_hooks(): + module = _make_fused_mlp_shell() + fused_impl = torch.nn.Identity() + + # The real training lifecycle lazily constructs the fused implementation while + # DDP parameter-gather hooks are disabled for the first iteration. + module._register_hooks_on_fused_impl(fused_impl) + + events = [] + + def old_hook(submodule, _inputs): + events.append(("old", submodule, tuple(submodule.parameters(recurse=False)))) + + old_handles = [ + module.linear_fc1.register_forward_pre_hook(old_hook), + module.linear_fc2.register_forward_pre_hook(old_hook), + ] + module.register_forward_pre_hook(old_hook) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + fused_impl(torch.ones(1, 2)) + + assert [event[0] for event in events] == ["old", "old"] + assert [event[1] for event in events] == [module.linear_fc1, module.linear_fc2] + assert events[0][2][0] is module.linear_fc1.weight + assert events[1][2][0] is module.linear_fc2.weight + + for handle in old_handles: + handle.remove() + events.clear() + + def replacement_hook(submodule, _inputs, kwargs): + events.append(("replacement", submodule, kwargs)) + + module.linear_fc2.register_forward_pre_hook(replacement_hook, with_kwargs=True) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + fused_impl(torch.ones(1, 2)) + + assert events == [("replacement", module.linear_fc2, {})] + + +def test_fused_mlp_rejects_input_modifying_submodule_hook_added_after_construction(): + module = _make_fused_mlp_shell() + fused_impl = torch.nn.Identity() + module._register_hooks_on_fused_impl(fused_impl) + + module.linear_fc1.register_forward_pre_hook(lambda _module, inputs: inputs) + + with pytest.warns(UserWarning, match="pre-forward hook"): + with pytest.raises(RuntimeError, match="modifies input tensor"): + fused_impl(torch.ones(1, 2))