Validate MFSDP module phase transitions

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

@@ -144,8 +144,10 @@ class Phase(enum.Enum):
     # ``None`` lets pre_forward enqueue an all-gather unless an earlier FsdpModule
     # already prefetched this module.
     _unshard_event: torch.cuda.Event | None
-    # Backward-pre hook sets this to BACKWARD before activation recomputation
-    # can run. Forward and backward hooks own all other transitions.
+    # ``phase`` is FORWARD between pre_forward() and post_forward(), BACKWARD
+    # between pre_backward() and post_backward(), and RESTING otherwise. The only
+    # exception is non-reentrant activation recomputation: it runs between pre_backward()
+    # and post_backward(), preserving BACKWARD through its nested forward hooks.
     _phase: Phase
 
     def __init__(
@@ -191,6 +193,24 @@ def context(self) -> FsdpContext:
         """Return the FSDP context."""
         return self._context
 
+    @property
+    def phase(self) -> Phase:
+        """Return this module's lifecycle phase."""
+        return self._phase
+
+    @phase.setter
+    def phase(self, phase: Phase) -> None:
+        """Transition this module between its valid lifecycle phases."""
+        allowed_transitions = {
+            (FsdpModule.Phase.RESTING, FsdpModule.Phase.FORWARD),
+            (FsdpModule.Phase.FORWARD, FsdpModule.Phase.RESTING),
+            (FsdpModule.Phase.RESTING, FsdpModule.Phase.BACKWARD),
+            (FsdpModule.Phase.BACKWARD, FsdpModule.Phase.RESTING),
+        }
+        if (self._phase, phase) not in allowed_transitions:
+            raise RuntimeError(f"Invalid FSDP module phase transition: {self._phase} -> {phase}.")
+        self._phase = phase
+
     @property
     def name(self) -> str:
         """Return this FsdpModule's name."""
@@ -254,17 +274,14 @@ def pre_forward(self) -> None:
         on the comm stream, so ``AG_{i+1}`` is launched before ``F_i`` finishes.
         """
         context = self.context
+        # This is the first MFSDP hook to run, so finalize the context here once
+        # before any module begins communication.
         context.ensure_finalized()
-        # post_forward() resets the phase after a non-recomputed forward, so a
-        # FORWARD phase here means this forward-pre hook ran while the previous
-        # forward was still in progress.
-        assert self._phase is not FsdpModule.Phase.FORWARD
         # A reentrant checkpoint recomputes before the child module's backward-pre
-        # hook can set its phase. Its forward still runs inside the active autograd
-        # GraphTask, which is the signal PyTorch FSDP2 uses as well.
-        is_recomputing = self._phase is FsdpModule.Phase.BACKWARD or _is_in_backward()
-        if not is_recomputing:
-            self._phase = FsdpModule.Phase.FORWARD
+        # hook runs. The active autograd GraphTask identifies that recomputation.
+        is_recomputing = self.phase is FsdpModule.Phase.BACKWARD or _is_in_backward()
+        if self.phase is not FsdpModule.Phase.BACKWARD:
+            self.phase = FsdpModule.Phase.FORWARD
         torch.cuda.nvtx.range_push(self._nvtx_label("forward"))
         self._num_ready_grad_parameters = 0
         allgather_stream = context.allgather_stream
@@ -309,10 +326,11 @@ def post_forward(self) -> None:
         # 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.
-        is_recomputing = self._phase is FsdpModule.Phase.BACKWARD or _is_in_backward()
+        is_recomputing = self.phase is FsdpModule.Phase.BACKWARD or _is_in_backward()
         if not is_recomputing:
             self._reshard_parameter_groups()
-            self._phase = FsdpModule.Phase.RESTING
+        if self.phase is FsdpModule.Phase.FORWARD:
+            self.phase = FsdpModule.Phase.RESTING
         torch.cuda.nvtx.range_pop()
 
     def _reshard_parameter_groups(self) -> None:
@@ -335,7 +353,7 @@ def _reshard_parameter_groups(self) -> None:
 
     def pre_backward(self) -> None:
         """Prepare full parameters and prefetch the next FsdpModule in backward order."""
-        self._phase = FsdpModule.Phase.BACKWARD
+        self.phase = FsdpModule.Phase.BACKWARD
         torch.cuda.nvtx.range_push(self._nvtx_label("backward"))
         context = self.context
         current_stream = context.current_stream()
@@ -362,7 +380,7 @@ def post_backward(self) -> None:
         """Reduce gradients and return parameters to their sharded resting state."""
         self._reduce_gradient_groups()
         self._reshard_parameter_groups()
-        self._phase = FsdpModule.Phase.RESTING
+        self.phase = FsdpModule.Phase.RESTING
         torch.cuda.nvtx.range_pop()
 
     def _reduce_gradient_groups(self) -> None:

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

@@ -262,16 +262,16 @@ def test_fully_shard_activation_recompute_reshards_parameters(distributed_setup,
 
     # Backward completes each module before recomputing the previous one, so
     # every module-local phase must be cleared after its matching backward.
-    assert model._phase is FsdpModule.Phase.RESTING
-    assert model.fc1._phase is FsdpModule.Phase.RESTING
-    assert model.fc2._phase is FsdpModule.Phase.RESTING
+    assert model.phase is FsdpModule.Phase.RESTING
+    assert model.fc1.phase is FsdpModule.Phase.RESTING
+    assert model.fc2.phase is FsdpModule.Phase.RESTING
 
     # A second forward after backward runs in the forward phase again, so
     # forward-order prefetch resumes and the module phases return to resting.
     model(x).sum().backward()
-    assert model._phase is FsdpModule.Phase.RESTING
-    assert model.fc1._phase is FsdpModule.Phase.RESTING
-    assert model.fc2._phase is FsdpModule.Phase.RESTING
+    assert model.phase is FsdpModule.Phase.RESTING
+    assert model.fc1.phase is FsdpModule.Phase.RESTING
+    assert model.fc2.phase is FsdpModule.Phase.RESTING
 
 
 @pytest.mark.parametrize("set_to_none", [True, False])