diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index 1e1a229a73a..03eeb3aafed 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/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: diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 863becf8b77..e0802c30068 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/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])