From e639bb548d3be1282657fbad8e258779158278ba Mon Sep 17 00:00:00 2001 From: Jianbin Chang Date: Thu, 30 Jul 2026 23:43:29 +0800 Subject: [PATCH] Fix FSDP activation recompute prefetch Signed-off-by: Jianbin Chang --- .../src/megatron_fsdp/experimental/module.py | 22 +++++++++--- .../distributed/mfsdp_v2/test_fully_shard.py | 36 +++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) 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 33e6985d8ed..6623af043c0 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/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 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 510690d09ec..ae1e68de732 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/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):