diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index a240b148d62..8381f9a3a5c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -292,8 +292,9 @@ def redistribute( """Redistribute this buffer to ``new_placements``. This dispatcher supports the one-axis transitions: - Flat -> Replicate, Partial -> Replicate, Partial -> Flat, and - Replicate -> Flat. Other placement changes are intentionally unsupported. + Flat -> Replicate, Partial -> Replicate, Partial -> Flat, + Replicate -> Flat, and Replicate -> Partial. Other placement changes are + intentionally unsupported. """ new_placements = tuple(new_placements) if len(new_placements) != self.mesh.ndim: @@ -322,6 +323,23 @@ def redistribute( return self.reduce_scatter(axis, new_placement, out=out) if isinstance(old_placement, Replicate) and isinstance(new_placement, Flat): return self.scatter(axis, new_placement, out=out) + if isinstance(old_placement, Replicate) and isinstance(new_placement, Partial): + # Replicate and Partial share the same local layout, so relabel the + # buffer without communication. Value-preserving for AVG only -- the + # mean of identical per-rank locals is that value; SUM would need a + # 1/axis_size rescale, which no caller needs. + if new_placement.reduce_op != dist.ReduceOp.AVG: + raise NotImplementedError( + "Replicate -> Partial redistribute supports AVG only, got " + f"{new_placement.reduce_op!r}." + ) + if out is not None: + raise NotImplementedError( + "Replicate -> Partial redistribute does not support an out buffer." + ) + return DBuffer.from_local( + self.local_buffer, self.mesh, new_placements, self.layout.tensor_shapes + ) raise NotImplementedError( "Unsupported DBuffer placement transition on axis " f"{axis}: {old_placement!r} -> {new_placement!r}." @@ -452,7 +470,17 @@ def get_dtensor(self, index: int) -> DTensor: elif isinstance(placement, Flat): torch_placements.append(dist_tensor.Shard(0)) elif isinstance(placement, Partial): - raise ValueError("Partial DBuffer placements cannot be represented as DTensor.") + # main_grad backs .grad while it rests DP-outer-Partial between + # microbatches, so a Partial placement must round-trip to a DTensor. + if placement.reduce_op == dist.ReduceOp.AVG: + reduce_op = "avg" + elif placement.reduce_op == dist.ReduceOp.SUM: + reduce_op = "sum" + else: + raise ValueError( + f"Unsupported Partial reduce op for DTensor: {placement.reduce_op!r}." + ) + torch_placements.append(dist_tensor.Partial(reduce_op)) else: raise TypeError(f"Unsupported placement for DTensor conversion: {placement!r}.") 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 4a63ab8307a..125f8982744 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -341,7 +341,7 @@ def _reduce_gradient_groups(self) -> None: reduce_scatter_stream.wait_stream(current_stream) with torch.cuda.stream(reduce_scatter_stream): - group.reduce_partial_gradients(partial_grad) + group.reduce_partial_gradients(partial_grad, self.context.is_last_microbatch) @property def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]: 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 a710c4f07e1..6ea39031e60 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 @@ -150,13 +150,9 @@ def __init__( "main_grad is built from main_weight tensor shapes on the same mesh, " "and DBuffer layouts are deterministic from those shapes and mesh size." ) - if self.main_grad.placements != self.main_weight.placements: - raise ValueError( - "FSDP temporarily requires main_grad and main_weight to have the same " - "placements until HSDP/HFSDP support is implemented. " - f"Got main_grad placements {self.main_grad.placements} and " - f"main_weight placements {self.main_weight.placements}." - ) + # main_grad rests here (DP-outer-Partial for HSDP) between microbatches and + # is finalized to main_weight's placements after the last microbatch. + self._accumulation_placements = main_grad_placements sharded_parameters: list[nn.Parameter] = [] unsharded_parameters: list[nn.Parameter] = [] main_grad_dtype = self.main_grad.dtype if self.main_grad is not None else None @@ -251,6 +247,12 @@ def release_unsharded_storage(self) -> None: # so keep the shared storage-release path. self._unsharded_model_weight.release_storage() + def _install_sharded_grads(self) -> None: + """Point each sharded parameter's grad at main_grad's current DTensor view.""" + assert self.main_grad is not None + for index, sharded_parameter in enumerate(self.sharded_parameters): + sharded_parameter.grad = self.main_grad.get_dtensor(index) + def allocate_partial_grad_buffer(self) -> DBuffer: """Allocate the unreduced reduce-scatter input buffer.""" assert self.main_grad is not None @@ -279,8 +281,18 @@ def copy_gradients_to_partial_buffer(self, partial_grad: DBuffer) -> None: partial_grad.get_local_tensor(index).copy_(parameter.grad) parameter.grad = None - def reduce_partial_gradients(self, partial_grad: DBuffer) -> None: - """Reduce a packed partial gradient buffer into sharded parameter gradients.""" + def reduce_partial_gradients( + self, partial_grad: DBuffer, is_last_microbatch: bool = True + ) -> None: + """Reduce a packed partial gradient buffer into sharded parameter gradients. + + For HSDP main_grad rests DP-outer-Partial (Partial where main_weight is + Replicate) between microbatches, accumulating each backward through the + standard zero_grad contract; the last microbatch reduces the DP-outer axes, + finalizing main_grad to main_weight's placements so ``.grad`` is the fully + reduced gradient before ``optimizer.step()``. With every axis Flat (plain + DP) main_grad already rests finalized. + """ assert self.main_grad is not None def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool: @@ -295,10 +307,20 @@ def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool: raise RuntimeError("FSDP sharded gradients must be either all set or all None.") return has_any_grad - # zero_grad(set_to_none=True) clears sharded parameter grads, so the next + # zero_grad(set_to_none=True) clears sharded parameter grads, so this # backward can reduce directly into main_grad. zero_grad(set_to_none=False) # leaves sharded grads installed, so this backward accumulates into main_grad. has_sharded_grads = has_grad(self.sharded_parameters) + + # A non-accumulation main_grad means the previous step finalized it; this + # only happens on the first microbatch. Redistribute it back to the + # DP-outer-Partial accumulation placement -- a metadata relabel for HSDP, + # and a fresh reduce-scattered buffer for HFSDP in the future. + if self.main_grad.placements != self._accumulation_placements: + self.main_grad = self.main_grad.redistribute(self._accumulation_placements) + if has_sharded_grads: + self._install_sharded_grads() + can_reduce_into_main_grad = ( not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype ) @@ -321,10 +343,14 @@ def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool: self.main_grad.local_buffer.add_(reduced_grad.local_buffer) else: self.main_grad.local_buffer.copy_(reduced_grad.local_buffer) - if not has_sharded_grads: - for index, parameter in enumerate(self.sharded_parameters): - parameter.grad = self.main_grad.get_dtensor(index) + self._install_sharded_grads() + + if is_last_microbatch: + # Finalize the deferred DP-outer reduction (all-reduce for HSDP, + # reduce-scatter for HFSDP) and install the sharded parameter gradients. + self.main_grad = self.main_grad.redistribute(self.main_weight.placements) + self._install_sharded_grads() def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]: 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 6f4a1a36620..2804a462c82 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -6,6 +6,7 @@ import pytest import torch +import torch.distributed as dist from torch import nn from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor @@ -13,7 +14,9 @@ from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, + Partial, Placements, + Replicate, fully_shard, microbatch, ) @@ -100,6 +103,18 @@ def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) +def _hsdp_placements() -> Placements: + """HSDP: params/optimizer replicated across DP-outer (axis 0), sharded within + DP-inner (axis 1). main_grad rests [Partial, Flat] between microbatches and is + all-reduced to [Replicate, Flat] on the last microbatch.""" + return Placements( + dp_axes=[0, 1], + parameter=[Replicate(), Flat()], + gradient=[Partial(dist.ReduceOp.AVG), Flat()], + optimizer=[Replicate(), Flat()], + ) + + def _mb(num_bytes: int) -> str: return f"{num_bytes / 1024**2:.2f} MB" @@ -111,6 +126,17 @@ def _events_overlap(first, second) -> bool: ) +def _nccl_events(cuda_events, *name_fragments): + """CUDA NCCL events whose name contains any of ``name_fragments`` (case-insensitive).""" + return [ + event + for event in cuda_events + if "nccl" in event.name.lower() + and event.activity_type == "kernel" + and any(fragment in event.name.lower() for fragment in name_fragments) + ] + + @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.""" @@ -168,6 +194,147 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]: ) +@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): + """HSDP (DP-outer replicated, DP-inner sharded) training should match single-rank SGD. + + Gradients reduce-scatter within DP-inner every backward and accumulate into + main_grad; the DP-outer all-reduce runs only on the last microbatch, scoped + via ``microbatch(...)``. Every rank sees identical data, so the averaged + gradient equals the single-rank gradient and losses must match. Both + ``zero_grad`` modes are covered: ``set_to_none=True`` overwrites main_grad, + ``set_to_none=False`` accumulates into a zeroed main_grad. + """ + rank = distributed_setup.rank + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 4 or world_size % 2 != 0: + pytest.skip("This test requires an even number of at least 4 ranks for a 2-D DP mesh.") + + outer_size = 2 + inner_size = world_size // outer_size + mesh = init_device_mesh( + device.type, (outer_size, inner_size), mesh_dim_names=("dp_outer", "dp_inner") + ) + torch.manual_seed(1234) + dim = 8 + baseline = MultiChildModel(dim=dim, num_children=2).to(device) + model = MultiChildModel(dim=dim, num_children=2).to(device) + model.load_state_dict(baseline.state_dict()) + + # Shard the child layers, then the model, so the children share a root context + # and reduce through the overlap path instead of as independent roots. + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=_hsdp_placements()) + fully_shard(model, mesh=mesh, placements=_hsdp_placements()) + baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + + micro_batch_size = 2 + x = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + target = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + microbatches = tuple(zip(x.unbind(), target.unbind())) + + def train(model, optimizer, log_prefix) -> list[torch.Tensor]: + losses = [] + for step in range(5): + optimizer.zero_grad(set_to_none=set_to_none) + + for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): + is_last = microbatch_index == num_microbatches - 1 + with microbatch(model, is_last=is_last): + loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) + (loss / num_microbatches).backward() + losses.append(loss.detach()) + logger.debug( + "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s", + log_prefix, + rank, + step, + microbatch_index, + loss, + ) + + optimizer.step() + return losses + + baseline_losses = train(baseline, baseline_optimizer, "Baseline") + sharded_losses = train(model, optimizer, "HSDP") + + torch.testing.assert_close( + torch.stack(sharded_losses), + torch.stack(baseline_losses), + msg="HSDP losses did not match baseline losses.", + ) + + +def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup): + """HSDP reduce-scatters DP-inner every microbatch but all-reduces DP-outer once. + + ``fully_shard(model)`` makes the child units share a root context so their + reductions run through the overlap path rather than as independent roots. + Counting NCCL events over a multi-microbatch step, the DP-inner reduce-scatter + fires once per microbatch per group while the DP-outer all-reduce that + finalizes main_grad fires only on the last microbatch, so the reduce-scatter + count is exactly ``num_microbatches`` times the all-reduce count. This asserts + on event counts only, not numerics. + """ + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 4 or world_size % 2 != 0: + pytest.skip("This test requires an even number of at least 4 ranks for a 2-D DP mesh.") + + outer_size = 2 + inner_size = world_size // outer_size + mesh = init_device_mesh( + device.type, (outer_size, inner_size), mesh_dim_names=("dp_outer", "dp_inner") + ) + torch.manual_seed(1234) + dim = 8 + num_children = 2 + model = MultiChildModel(dim=dim, num_children=num_children).to(device) + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=_hsdp_placements()) + fully_shard(model, mesh=mesh, placements=_hsdp_placements()) + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + + num_microbatches = 3 + micro_batch_size = 2 + x = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + target = torch.randn(num_microbatches, micro_batch_size, dim, device=device) + microbatches = tuple(zip(x.unbind(), target.unbind())) + + def train_one_step() -> None: + optimizer.zero_grad(set_to_none=True) + for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): + is_last = microbatch_index == num_microbatches - 1 + with microbatch(model, is_last=is_last): + loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) + (loss / num_microbatches).backward() + optimizer.step() + + train_one_step() + torch.cuda.synchronize(device) + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + train_one_step() + torch.cuda.synchronize(device) + + cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"] + reduce_scatter_events = _nccl_events(cuda_events, "reducescatter", "reduce_scatter") + all_reduce_events = _nccl_events(cuda_events, "allreduce") + # One DP-outer all-reduce per parameter group -- each child layer plus the + # root unit's bias -- fired only on the last microbatch. Plain DP fires none. + assert len(all_reduce_events) == num_children + 1, [event.name for event in cuda_events] + # DP-inner reduce-scatter runs every microbatch; the DP-outer all-reduce runs + # only on the last, so the counts differ by exactly the microbatch factor. + assert len(reduce_scatter_events) == len(all_reduce_events) * num_microbatches, ( + f"Expected reduce-scatter ({len(reduce_scatter_events)}) to be {num_microbatches}x " + f"the DP-outer all-reduce count ({len(all_reduce_events)})." + ) + + def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): """An outer FsdpModule owns direct parameters but not nested child FsdpModule parameters.""" world_size = distributed_setup.world_size @@ -353,26 +520,9 @@ def train_one_iteration() -> None: # drop the CUDA events. torch.cuda.synchronize(device) - # Keep only real device kernels. NCCL also emits per-collective GPU user - # annotations (e.g. "nccl:_reduce_scatter_base") that the profiler reports as - # CUDA events. Filtering by activity type avoids counting those annotations as - # kernels without relying on their current naming convention. - cuda_events = [ - event - for event in prof.events() - if event.device_type.name == "CUDA" and event.activity_type == "kernel" - ] - all_gather_events = [ - event - for event in cuda_events - if "nccl" in event.name.lower() and "allgather" in event.name.lower() - ] - reduce_scatter_events = [ - event - for event in cuda_events - if "nccl" in event.name.lower() - and ("reducescatter" in event.name.lower() or "reduce_scatter" in event.name.lower()) - ] + cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"] + all_gather_events = _nccl_events(cuda_events, "allgather") + reduce_scatter_events = _nccl_events(cuda_events, "reducescatter", "reduce_scatter") # 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 = [