From 028f77a106f4a5cec9a3ab5ec86709d7bb5560e7 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Fri, 26 Jun 2026 20:23:42 +0000 Subject: [PATCH 1/8] Add FSDP all-gather stream overlap Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 2 + .../src/megatron_fsdp/experimental/module.py | 86 ++++++- .../experimental/parameter_group.py | 9 +- .../distributed/mfsdp_v2/test_context.py | 2 +- .../distributed/mfsdp_v2/test_fully_shard.py | 227 +++++++++++++++++- 5 files changed, 303 insertions(+), 23 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 0ab256a7ef4..0360284384d 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -24,6 +24,8 @@ from .module import FsdpContext, FsdpModule from .placement import Placements +__all__ = ["fully_shard", "microbatch"] + def fully_shard( module: nn.Module, 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 3f1d24b1517..3c97fda3242 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -14,6 +14,8 @@ """Module mixin for the minimal Megatron-FSDP path.""" +import dataclasses +from collections import deque from collections.abc import Callable from typing import Literal, cast @@ -26,23 +28,54 @@ from .placement import MeshAxis, Placements +@dataclasses.dataclass(frozen=True) +class DelayedRelease: + """A module whose unsharded storage can be released after its consumer event.""" + + consumer_event: torch.cuda.Event | None + module: "FsdpModule" + + class FsdpContext: - """Runtime state shared by one experimental FSDP subtree.""" + """Runtime state, stream, and release scheduler shared by one FSDP subtree.""" + allgather_stream: torch.cuda.Stream + delayed_releases: deque[DelayedRelease] # HFSDP/HSDP need explicit last-microbatch state. First-microbatch state is # unnecessary because it can be detected when ``model_weight``, after syncing # from ``main_weight``, has placements different from ``Placements.optimizer``. is_last_microbatch: bool root_module: "FsdpModule" - def __init__(self, root_module: "FsdpModule") -> None: + def __init__(self, device: torch.device, root_module: "FsdpModule") -> None: """Create rank-local runtime state for a root FSDP subtree. Args: + device: Device on which this context schedules communication. root_module: Outermost module that owns this context. """ self.root_module = root_module self.is_last_microbatch = True + self.delayed_releases = deque() + with torch.cuda.device(device): + self.allgather_stream = torch.cuda.Stream() + + def enqueue_release(self, module: "FsdpModule") -> None: + """Queue a module's unsharded storage for delayed release.""" + consumer_event = torch.cuda.current_stream(self.allgather_stream.device).record_event() + self.delayed_releases.append(DelayedRelease(consumer_event=consumer_event, module=module)) + + def drain_delayed_releases(self, target_length: int) -> None: + """Release queued module storages FIFO until the queue reaches ``target_length``.""" + if target_length < 0: + raise ValueError(f"target_length must be non-negative, got {target_length}.") + + while len(self.delayed_releases) > target_length: + delayed_release = self.delayed_releases.popleft() + with torch.cuda.stream(self.allgather_stream): + if delayed_release.consumer_event is not None: + self.allgather_stream.wait_event(delayed_release.consumer_event) + delayed_release.module.release_unsharded_storage() class FsdpModule: @@ -114,7 +147,7 @@ def _lazy_init_context(self) -> None: if self._context is not None: return - context = FsdpContext(root_module=self) + context = FsdpContext(device=self._parameter_groups[0].main_weight.device, root_module=self) for submodule_name, submodule in cast(nn.Module, self).named_modules(): if not isinstance(submodule, FsdpModule): continue @@ -172,33 +205,64 @@ def pre_forward(self) -> None: self._lazy_init_context() torch.cuda.nvtx.range_push(self._nvtx_label("forward")) self._ready_grad_parameters.clear() - for group in self._parameter_groups: - group.sync_model_weight_from_main_weight() - group.unshard_parameters() + if self.is_root(): + allgather_stream = self.context.allgather_stream + allgather_stream.wait_stream(torch.cuda.current_stream(allgather_stream.device)) + self._unshard_parameter_groups(sync_model_weight=True) + + def _unshard_parameter_groups(self, *, sync_model_weight: bool) -> None: + """Materialize full parameters for this FSDP unit.""" + self.context.drain_delayed_releases(target_length=1) + + allgather_stream = self.context.allgather_stream + current_stream = torch.cuda.current_stream(allgather_stream.device) + + with torch.cuda.stream(allgather_stream): + for group in self._parameter_groups: + if sync_model_weight: + # TODO: After NVIDIA/Megatron-LM#5411 lands, move this sync to the + # optimizer post-step hook instead of running it every microbatch. + group.sync_model_weight_from_main_weight() + group.unshard_parameters() + current_stream.wait_stream(allgather_stream) def post_forward(self) -> None: """Return parameters to their sharded resting state after forward compute.""" + self._reshard_parameter_groups() + self.context.enqueue_release(self) + if self.is_root(): + self.context.drain_delayed_releases(target_length=0) + torch.cuda.nvtx.range_pop() + + def _reshard_parameter_groups(self) -> None: for group in self._parameter_groups: group.reshard_parameters() - torch.cuda.nvtx.range_pop() def pre_backward(self) -> None: """Prepare full parameters for backward compute.""" torch.cuda.nvtx.range_push(self._nvtx_label("backward")) - for group in self._parameter_groups: - group.unshard_parameters() + self._unshard_parameter_groups(sync_model_weight=False) def post_backward(self) -> None: """Reduce gradients and return parameters to their sharded resting state.""" for group in self._parameter_groups: if group.requires_grad: group.reduce_gradients() - group.reshard_parameters() + self._reshard_parameter_groups() + self.context.enqueue_release(self) + if self.is_root(): + self.context.drain_delayed_releases(target_length=0) self._ready_grad_parameters.clear() torch.cuda.nvtx.range_pop() + def release_unsharded_storage(self) -> None: + """Release unsharded storage owned by this FSDP unit.""" + for group in self._parameter_groups: + group.release_unsharded_storage() + + @property def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]: - """Return parameter groups owned by this FSDP unit.""" + """Parameter groups owned by this FSDP unit.""" return self._parameter_groups def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 28145cf2e07..2c7aeaf5d06 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -232,12 +232,9 @@ def unshard_parameters(self) -> None: def reshard_parameters(self) -> None: """Install sharded DTensor parameters on the owning modules.""" self._switch_to_sharded_parameters() - # At post-backward time, replacing unsharded parameter .data with size-0 - # empty tensors would also be safe: autograd has consumed the saved - # forward views. That alternative is not much cleaner than releasing - # this storage, and splitting post-forward and post-backward reshard - # behavior would make the caller code less clean, so keep the shared - # storage-release path. + + def release_unsharded_storage(self) -> None: + """Release this group's full-parameter storage.""" self._unsharded_model_weight.release_storage() def reduce_gradients(self) -> None: diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_context.py b/tests/unit_tests/distributed/mfsdp_v2/test_context.py index 7235a094240..9ffd6bd8cdb 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_context.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_context.py @@ -83,7 +83,7 @@ def test_two_child_subtrees_then_parent_collapse_to_one_context(distributed_setu def test_sibling_roots_without_parent_keep_separate_contexts(distributed_setup): - """Independent FSDP roots should not share runtime state.""" + """Independent FSDP roots should not share runtime scheduling state.""" device = distributed_setup.device mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) 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 869903fd7fa..ff761bd43a0 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -9,6 +9,7 @@ from torch import nn from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor +from torch.profiler import ProfilerActivity, profile from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, @@ -48,6 +49,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.inner(x) + self.bias +class MultiChildModel(nn.Module): + """Model with direct parameters and multiple child FSDP units.""" + + def __init__(self, dim: int, num_children: int) -> None: + super().__init__() + self.bias = nn.Parameter(torch.ones(dim)) + self.layers = nn.ModuleList([nn.Linear(dim, dim, bias=False) for _ in range(num_children)]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run through every child layer with a root-owned bias.""" + x = x + self.bias + for layer in self.layers: + x = torch.relu(layer(x)) + return x + + class SaveNonLeafWeightView(torch.autograd.Function): """Autograd function that saves a non-leaf parameter view for backward.""" @@ -87,6 +104,13 @@ def _mb(num_bytes: int) -> str: return f"{num_bytes / 1024**2:.2f} MB" +def _events_overlap(first, second) -> bool: + return ( + first.time_range.start < second.time_range.end + and second.time_range.start < first.time_range.end + ) + + @pytest.mark.parametrize("num_microbatches", [1, 3]) def test_fully_shard_losses_match_baseline(distributed_setup, num_microbatches): """Minimal per-module FSDP training should match single-rank SGD.""" @@ -158,14 +182,207 @@ def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): fully_shard(model, mesh=mesh, placements=_flat_placements()) inner_names = [ - name for group in model.inner.parameter_groups() for name in group.parameter_names + name for group in model.inner.parameter_groups for name in group.parameter_names ] - outer_names = [name for group in model.parameter_groups() for name in group.parameter_names] + outer_names = [name for group in model.parameter_groups for name in group.parameter_names] assert inner_names == ["weight"] assert outer_names == ["bias"] +def test_forward_peak_memory_bounds_in_flight_child_all_gathers(distributed_setup): + """Forward peak memory should stay below three live child all-gathers.""" + rank = distributed_setup.rank + 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,)) + dim = 4096 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=4).to(dtype=dtype, device=device) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + x = torch.randn(2, dim, device=device, dtype=dtype) + with torch.no_grad(): + model(x) + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + + resting_allocated = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + with torch.no_grad(): + model(x) + torch.cuda.synchronize(device) + peak_delta = torch.cuda.max_memory_allocated(device) - resting_allocated + + child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() + bound_nbytes = 3 * child_weight_nbytes + + # A parent forward should keep one previous child unsharded until its compute + # stream consumer is safe, plus the current child being unsharded. The bound + # is looser than two child weights to avoid coupling this test to CUDA + # allocator granularity and small temporary buffers, while still catching + # delayed releases piling up across the four child layers. + assert peak_delta < bound_nbytes, ( + "FSDP forward peak memory exceeded the in-flight all-gather bound: " + f"rank={rank}, peak_delta={_mb(peak_delta)}, " + f"three_child_weights={_mb(bound_nbytes)}" + ) + + +def test_root_forward_returns_to_resting_memory(distributed_setup): + """Root forward should release child all-gather storage before returning.""" + rank = distributed_setup.rank + 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,)) + dim = 4096 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=2).to(dtype=dtype, device=device) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + x = torch.randn(2, dim, device=device, dtype=dtype) + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + resting_allocated = torch.cuda.memory_allocated(device) + + with torch.no_grad(): + output = model(x) + del output + torch.cuda.synchronize(device) + allocated_after_forward = torch.cuda.memory_allocated(device) + extra_allocated = allocated_after_forward - resting_allocated + child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() + + assert extra_allocated < child_weight_nbytes, ( + "Root forward did not return to resting memory after draining child releases: " + f"rank={rank}, extra_allocated={_mb(extra_allocated)}, " + f"one_child_weight={_mb(child_weight_nbytes)}" + ) + + +def test_root_backward_returns_to_resting_memory(distributed_setup): + """Root backward should release child all-gather storage before returning.""" + rank = distributed_setup.rank + 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,)) + dim = 4096 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=2).to(dtype=dtype, device=device) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + x = torch.randn(2, dim, device=device, dtype=dtype, requires_grad=True) + output = model(x) + loss = output.float().square().mean() + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + allocated_before_backward = torch.cuda.memory_allocated(device) + + loss.backward() + del loss, output + torch.cuda.synchronize(device) + allocated_after_backward = torch.cuda.memory_allocated(device) + extra_allocated = allocated_after_backward - allocated_before_backward + child_weight_nbytes = dim * dim * torch.empty((), dtype=dtype).element_size() + + assert extra_allocated < child_weight_nbytes, ( + "Root backward did not return to resting memory after draining child releases: " + f"rank={rank}, extra_allocated={_mb(extra_allocated)}, " + f"one_child_weight={_mb(child_weight_nbytes)}" + ) + + +def test_overlaps_all_gather_and_compute(distributed_setup): + """A shared root context should let child all-gathers overlap GEMM compute.""" + 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,)) + dim = 4096 + num_children = 4 + dtype = torch.bfloat16 + model = MultiChildModel(dim=dim, num_children=num_children).to(dtype=dtype) + placements = _flat_placements() + policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy) + fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + + x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True) + + def train_one_iteration() -> None: + model.zero_grad(set_to_none=True) + model(x).sum().backward() + + train_one_iteration() + torch.cuda.synchronize(device) + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + train_one_iteration() + prof.step() + torch.cuda.synchronize(device) + + cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"] + all_gather_events = [ + event + for event in cuda_events + if "nccl" in event.name.lower() and "allgather" in event.name.lower() + ] + compute_events = [ + event + for event in cuda_events + if any(token in event.name.lower() for token in ("gemm", "cutlass", "cublas")) + ] + assert all_gather_events, [event.name for event in cuda_events] + assert compute_events, [event.name for event in cuda_events] + + all_gather_streams = {event.device_resource_id for event in all_gather_events} + compute_streams = {event.device_resource_id for event in compute_events} + assert len(all_gather_streams) == 1 + assert all_gather_streams.isdisjoint(compute_streams) + + overlap_count = sum( + any(_events_overlap(all_gather_event, compute_event) for compute_event in compute_events) + for all_gather_event in all_gather_events + ) + # This profiles a full forward/backward iteration, so backward all-gathers are + # included in all_gather_events. The expected overlap count is from the forward + # child pipeline: each child after the first can all-gather while the previous + # child computes, giving num_children - 1 overlaps. Backward does not overlap + # in this all-gather-only path because gradient reduction is not delayed: + # each module synchronously reduces gradients in post_backward before autograd + # reaches the next module's pre_backward all-gather. The next PR addresses + # this by delaying gradient reduction. + expected_overlap_count = num_children - 1 + assert overlap_count >= expected_overlap_count, ( + f"Expected at least {expected_overlap_count} all-gather events to overlap compute, " + f"got {overlap_count}/{len(all_gather_events)}." + ) + + def test_frozen_parameter_group_does_not_allocate_main_grad(distributed_setup): """A non-trainable parameter group should not allocate persistent main gradients.""" world_size = distributed_setup.world_size @@ -179,7 +396,7 @@ def test_frozen_parameter_group_does_not_allocate_main_grad(distributed_setup): fully_shard(model, mesh=mesh, placements=_flat_placements()) - (group,) = model.parameter_groups() + (group,) = model.parameter_groups assert not group.requires_grad assert group.main_grad is None @@ -279,7 +496,7 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): fully_shard(model, mesh=mesh, placements=_flat_placements()) - (group,) = model.parameter_groups() + (group,) = model.parameter_groups full_weight = group.model_weight.allgather(0).get_local_tensor(0) assert full_weight.device.type == device.type torch.testing.assert_close(full_weight, expected_weight) @@ -296,7 +513,7 @@ def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup): model = NonLeafViewModel().to(device) fully_shard(model, mesh=mesh, placements=_flat_placements()) - group = model.parameter_groups()[0] + group = model.parameter_groups[0] x = torch.randn(8, device=device, requires_grad=True) loss = model(x).sum() From 8438bf1e6a764a5a5ed1e6b6f7762587ac707e50 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 4 Jul 2026 22:01:32 +0000 Subject: [PATCH 2/8] Remove fully_shard submodule export list Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/fully_shard.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 0360284384d..0ab256a7ef4 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py @@ -24,8 +24,6 @@ from .module import FsdpContext, FsdpModule from .placement import Placements -__all__ = ["fully_shard", "microbatch"] - def fully_shard( module: nn.Module, From b0dfb7d66585847a8b6da54bcbe8b34f31c622f1 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 4 Jul 2026 22:09:32 +0000 Subject: [PATCH 3/8] Format FSDP overlap files Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/__init__.py | 11 +---------- .../distributed/mfsdp_v2/test_fully_shard.py | 4 +--- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index bc9118598d1..5204a7645cf 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -18,13 +18,4 @@ from .fully_shard import fully_shard, microbatch from .placement import Flat, Partial, Placement, Placements, Replicate -__all__ = [ - "DBuffer", - "Flat", - "Partial", - "Placement", - "Placements", - "Replicate", - "fully_shard", - "microbatch", -] +__all__ = ["DBuffer", "Flat", "Partial", "Placement", "Placements", "Replicate", "fully_shard", "microbatch"] 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 ff761bd43a0..d00fe79d6b1 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -181,9 +181,7 @@ def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) fully_shard(model, mesh=mesh, placements=_flat_placements()) - inner_names = [ - name for group in model.inner.parameter_groups for name in group.parameter_names - ] + inner_names = [name for group in model.inner.parameter_groups for name in group.parameter_names] outer_names = [name for group in model.parameter_groups for name in group.parameter_names] assert inner_names == ["weight"] From f562558c4dbebb17dab87d5ba4c87ffe2ae750b8 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 6 Jul 2026 15:26:33 -0700 Subject: [PATCH 4/8] Synchronize inside profiler to capture FSDP overlap CUDA events The overlap test synchronized after the profiler context exited, so the profiler finalized its trace before in-flight device kernels completed and recorded no CUDA events. Move the synchronize inside the profiler window (matching test_symmetric_memory.py) and drop the no-op prof.step(). Signed-off-by: Jingyue Wu --- tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) 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 d00fe79d6b1..d5604304027 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -340,8 +340,11 @@ def train_one_iteration() -> None: with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: train_one_iteration() - prof.step() - torch.cuda.synchronize(device) + # Synchronize inside the profiler context so in-flight device kernels + # complete and get recorded before the profiler stops on __exit__. + # Synchronizing after the context would finalize the trace first and + # drop the CUDA events. + torch.cuda.synchronize(device) cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"] all_gather_events = [ From d59844b34d83212c0e3f2e102b51f4b5492c7e30 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 6 Jul 2026 22:29:19 +0000 Subject: [PATCH 5/8] Keep unsharded storage release comment Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/parameter_group.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 2c7aeaf5d06..e084bf5444e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -235,6 +235,12 @@ def reshard_parameters(self) -> None: def release_unsharded_storage(self) -> None: """Release this group's full-parameter storage.""" + # At post-backward time, replacing unsharded parameter .data with size-0 + # empty tensors would also be safe: autograd has consumed the saved + # forward views. That alternative is not much cleaner than releasing + # this storage, and splitting post-forward and post-backward reshard + # behavior would make the caller code less clean, so keep the shared + # storage-release path. self._unsharded_model_weight.release_storage() def reduce_gradients(self) -> None: From 039ceb8a53db96b32c309745cc7b1730b233a8c0 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 6 Jul 2026 22:38:02 +0000 Subject: [PATCH 6/8] Clarify unsharded storage release comment Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/parameter_group.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index e084bf5444e..eeec848416b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -235,12 +235,14 @@ def reshard_parameters(self) -> None: def release_unsharded_storage(self) -> None: """Release this group's full-parameter storage.""" - # At post-backward time, replacing unsharded parameter .data with size-0 - # empty tensors would also be safe: autograd has consumed the saved - # forward views. That alternative is not much cleaner than releasing - # this storage, and splitting post-forward and post-backward reshard - # behavior would make the caller code less clean, so keep the shared - # storage-release path. + # This method is shared by the post-forward and post-backward release + # paths. Post-forward must release storage because autograd may have + # saved forward views into the unsharded parameters. Post-backward could + # replace unsharded parameter .data with size-0 empty tensors, instead + # of releasing storage, because autograd has consumed those saved views. + # That alternative is not much cleaner, and splitting post-forward and + # post-backward reshard behavior would make the caller code less clean, + # so keep the shared storage-release path. self._unsharded_model_weight.release_storage() def reduce_gradients(self) -> None: From 5b8fa2962faaeaa241a8d9311bccf00915fd3820 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 6 Jul 2026 21:43:59 -0700 Subject: [PATCH 7/8] Match nvjet GEMM kernels in FSDP overlap test The name-based compute filter matched only gemm/cutlass/cublas, but on the CI H100 stack (torch 2.12 / CUDA 13) bf16 GEMMs are cuBLASLt nvjet_sm90_* kernels, so the filter found nothing and the test failed at the compute-event assertion. Add nvjet to the token set and rename compute_events -> gemm_events to reflect that it matches GEMM kernels by name. Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) 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 d5604304027..9703046b3ec 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -352,21 +352,25 @@ def train_one_iteration() -> None: for event in cuda_events if "nccl" in event.name.lower() and "allgather" in event.name.lower() ] - compute_events = [ + # GEMM device-kernel names vary across CUDA/cuBLAS versions and GPU archs + # (e.g. "*gemm*", "cutlass*", "cublas*", and cuBLASLt's Hopper "nvjet_sm90_*"). + gemm_events = [ event for event in cuda_events - if any(token in event.name.lower() for token in ("gemm", "cutlass", "cublas")) + if any( + token in event.name.lower() for token in ("gemm", "cutlass", "cublas", "nvjet") + ) ] assert all_gather_events, [event.name for event in cuda_events] - assert compute_events, [event.name for event in cuda_events] + assert gemm_events, [event.name for event in cuda_events] all_gather_streams = {event.device_resource_id for event in all_gather_events} - compute_streams = {event.device_resource_id for event in compute_events} + gemm_streams = {event.device_resource_id for event in gemm_events} assert len(all_gather_streams) == 1 - assert all_gather_streams.isdisjoint(compute_streams) + assert all_gather_streams.isdisjoint(gemm_streams) overlap_count = sum( - any(_events_overlap(all_gather_event, compute_event) for compute_event in compute_events) + any(_events_overlap(all_gather_event, gemm_event) for gemm_event in gemm_events) for all_gather_event in all_gather_events ) # This profiles a full forward/backward iteration, so backward all-gathers are From 9f5989cff6597b347d09a6926c05bf945697f210 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 8 Jul 2026 07:36:42 +0000 Subject: [PATCH 8/8] Fix FSDP overlap lint formatting Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/__init__.py | 11 ++++++++++- .../distributed/mfsdp_v2/test_fully_shard.py | 4 +--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index 5204a7645cf..bc9118598d1 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -18,4 +18,13 @@ from .fully_shard import fully_shard, microbatch from .placement import Flat, Partial, Placement, Placements, Replicate -__all__ = ["DBuffer", "Flat", "Partial", "Placement", "Placements", "Replicate", "fully_shard", "microbatch"] +__all__ = [ + "DBuffer", + "Flat", + "Partial", + "Placement", + "Placements", + "Replicate", + "fully_shard", + "microbatch", +] 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 9703046b3ec..229cd0bff4b 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -357,9 +357,7 @@ def train_one_iteration() -> None: gemm_events = [ event for event in cuda_events - if any( - token in event.name.lower() for token in ("gemm", "cutlass", "cublas", "nvjet") - ) + if any(token in event.name.lower() for token in ("gemm", "cutlass", "cublas", "nvjet")) ] assert all_gather_events, [event.name for event in cuda_events] assert gemm_events, [event.name for event in cuda_events]