fix: refresh TE fused MLP DDP hooks dynamically

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."""
 

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))