Fix FSDP activation recompute prefetch

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

@@ -36,6 +36,9 @@ class FsdpContext:
     # unnecessary because it can be detected when ``model_weight``, after syncing
     # from ``main_weight``, has placements different from ``Placements.optimizer``.
     is_last_microbatch: bool
+    # True from the root pre-backward hook until autograd completes. Forward
+    # hooks use this to identify activation recomputation inside backward.
+    is_backward: bool
     root_module: "FsdpModule"
     # Static orders used to drive all-gather prefetch. We may want to switch to
     # capturing runtime order if static module order proves too fragile. Each
@@ -52,6 +55,7 @@ def __init__(self, device: torch.device, root_module: "FsdpModule") -> None:
         """
         self.root_module = root_module
         self.is_last_microbatch = True
+        self.is_backward = False
         self.forward_order = IndexedOrder()
         self.backward_order = IndexedOrder()
         with torch.cuda.device(device):
@@ -73,6 +77,7 @@ def register_post_backward_final_callback(self) -> None:
 
         def post_backward_final_callback() -> None:
             self.current_stream().wait_stream(self.reduce_scatter_stream)
+            self.is_backward = False
 
         torch.autograd.Variable._execution_engine.queue_callback(post_backward_final_callback)
 
@@ -244,9 +249,13 @@ def pre_forward(self) -> None:
         # issued afterwards, so it is free to run concurrently with this FsdpModule).
         current_stream.wait_event(self._unshard_event)
 
-        next_module = context.forward_order.next_item(self)
-        if next_module is not None:
-            next_module._unshard_parameter_groups()
+        # Activation recomputation runs forward hooks inside backward. Do not
+        # prefetch the next module in forward order: its backward may already
+        # be complete, so no later backward hook would reshard it.
+        if not context.is_backward:
+            next_module = context.forward_order.next_item(self)
+            if next_module is not None:
+                next_module._unshard_parameter_groups()
 
     def _unshard_parameter_groups(self) -> None:
         """Unshard this FsdpModule's parameter groups on the all-gather stream.
@@ -267,7 +276,11 @@ def _unshard_parameter_groups(self) -> None:
 
     def post_forward(self) -> None:
         """Return parameters to their sharded resting state after forward compute."""
-        self._reshard_parameter_groups()
+        # Recomputed parameters are consumed immediately by this module's
+        # backward. Keep them materialized to avoid an unnecessary all-gather;
+        # post_backward() will reshard them after gradient reduction.
+        if not self.context.is_backward:
+            self._reshard_parameter_groups()
         torch.cuda.nvtx.range_pop()
 
     def _reshard_parameter_groups(self) -> None:
@@ -294,6 +307,7 @@ def pre_backward(self) -> None:
         context = self.context
         current_stream = context.current_stream()
         if self.is_root():
+            context.is_backward = True
             context.register_post_backward_final_callback()
             # Fork the reduce-scatter stream from the current stream once, at the
             # start of backward, so every module's post-backward reduce-scatter is

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

@@ -11,6 +11,7 @@
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.tensor import DTensor
 from torch.profiler import ProfilerActivity, profile
+from torch.utils.checkpoint import checkpoint
 
 from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import (
     Flat,
@@ -44,6 +45,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.fc2(self.relu(self.fc1(x)))
 
 
+class CheckpointedTinyModel(TinyModel):
+    """Tiny model that activation-checkpoints each shardable module."""
+
+    def __init__(self, use_reentrant: bool) -> None:
+        super().__init__()
+        self.use_reentrant = use_reentrant
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Run each linear layer through activation checkpointing."""
+        x = checkpoint(self.fc1, x, use_reentrant=self.use_reentrant)
+        return checkpoint(self.fc2, self.relu(x), use_reentrant=self.use_reentrant)
+
+
 class NestedModel(nn.Module):
     """Model with direct and child-owned parameters."""
 
@@ -189,6 +203,28 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
     )
 
 
+@pytest.mark.parametrize("use_reentrant", [True, False], ids=["reentrant", "non_reentrant"])
+def test_fully_shard_activation_recompute_reshards_parameters(distributed_setup, use_reentrant):
+    """Activation recomputation should leave every FSDP module resharded."""
+    world_size = distributed_setup.world_size
+    device = distributed_setup.device
+    if world_size < 2:
+        pytest.skip("This test requires at least 2 ranks.")
+
+    mesh = init_device_mesh(device.type, (world_size,))
+    model = CheckpointedTinyModel(use_reentrant=use_reentrant).to(device)
+    fully_shard(model.fc1, mesh=mesh, placements=_flat_placements())
+    fully_shard(model.fc2, mesh=mesh, placements=_flat_placements())
+    fully_shard(model, mesh=mesh, placements=_flat_placements())
+
+    x = torch.randn(2, 8, device=device, requires_grad=True)
+    model(x).sum().backward()
+
+    assert isinstance(model.fc1.weight, DTensor)
+    assert isinstance(model.fc2.weight, DTensor)
+    assert not model.context.is_backward
+
+
 @pytest.mark.parametrize("set_to_none", [True, False])
 @pytest.mark.parametrize("num_microbatches", [1, 3])
 def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches, set_to_none):