Fix MegatronFSDP root module hook dispatch

megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py

@@ -1487,7 +1487,7 @@ def forward(self, *inputs, **kwargs):
         self._replace_param_with_raw_if_needed()
         with torch.autograd.profiler.record_function("CustomFSDP.forward"):
             # Call the forward pass of the wrapped module.
-            output = self.module.forward(*inputs, **kwargs)
+            output = self.module(*inputs, **kwargs)
             return output
 
 

tests/unit_tests/distributed/mfsdp_v1/test_mfsdp_fully_shard.py

@@ -136,6 +136,23 @@ def forward(self, x, y):
         return x
 
 
+class RootParamModel(torch.nn.Module):
+    """Toy model with parameters owned directly by the root module."""
+
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.empty(DIM_SIZE, DIM_SIZE))
+        self.bias = torch.nn.Parameter(torch.empty(DIM_SIZE))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.bias)
+
+    def forward(self, x):
+        return torch.nn.functional.linear(x, self.weight, self.bias)
+
+
 class ToyTETransformer(torch.nn.Module):
     """Toy Transformer model for testing Megatron-FSDP with Transformer Engine."""
 
@@ -731,6 +748,33 @@ def test_fully_shard_ez(self, shard_strategy):
             optimizer.step()
             optimizer.zero_grad()
 
+    def test_root_module_forward_uses_gathered_parameters(self):
+        """
+        Test that root-owned parameters are gathered before the root forward.
+        """
+
+        model = RootParamModel().cuda()
+        with torch.no_grad():
+            model.weight.copy_(
+                torch.arange(DIM_SIZE * DIM_SIZE, dtype=torch.float32, device="cuda").view(
+                    DIM_SIZE, DIM_SIZE
+                )
+            )
+            model.bias.copy_(torch.arange(DIM_SIZE, dtype=torch.float32, device="cuda"))
+
+        model_input = torch.arange(DIM_SIZE * DIM_SIZE, dtype=torch.float32, device="cuda").view(
+            DIM_SIZE, DIM_SIZE
+        )
+        expected_output = model(model_input)
+
+        mfsdp_model = fully_shard_model(
+            module=model, fsdp_unit_modules=[RootParamModel], zero_dp_strategy=OPTIM_GRADS_PARAMS
+        )
+
+        output = mfsdp_model(model_input)
+
+        torch.testing.assert_close(output, expected_output)
+
     @pytest.mark.skipif(
         version.parse(torch.__version__) < version.parse('2.4.0'),
         reason="Megatron-FSDP requires PyTorch 2.4.0 or later.",