add --overlap-param-gather support for layer-wise optimizer. lots of unit tests.

megatron/core/distributed/distributed_data_parallel.py

@@ -236,7 +236,10 @@ def _allocate_buffers_for_parameters(
 
             # Set `next_param_gather_bucket_group` for different bucket groups by iterating through
             # buckets in reverse order (since all-gathers happen in reverse order of buckets).
-            if self.ddp_config.use_distributed_optimizer and self.ddp_config.overlap_param_gather:
+            # Note: overlap_param_gather covers both the distributed optimizer and the
+            # layer-wise optimizer cases; the latter sets overlap_param_gather=True
+            # without use_distributed_optimizer.
+            if self.ddp_config.overlap_param_gather:
                 num_bucket_groups = len(bucket_groups)
                 for i in range(1, num_bucket_groups):
                     bucket_groups[num_bucket_groups - i].next_param_gather_bucket_group = (
@@ -344,9 +347,10 @@ def unmap_weight_tensor(m):
                     grad_acc.register_hook(self._make_backward_post_hook(param))
                     self.grad_accs.append(grad_acc)
 
-        self.use_forward_hook = (
-            self.ddp_config.use_distributed_optimizer and self.ddp_config.overlap_param_gather
-        )
+        # Note: overlap_param_gather covers both the distributed optimizer and the
+        # layer-wise optimizer cases; the latter sets overlap_param_gather=True
+        # without use_distributed_optimizer.
+        self.use_forward_hook = self.ddp_config.overlap_param_gather
         self.remove_forward_pre_hook_handles = {}
         if self.use_forward_hook:
             self.enable_forward_pre_hook()
@@ -534,6 +538,11 @@ def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
         for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
             bucket_group.finish_grad_sync(force_all_reduce=force_all_reduce)
 
+    def free_overlap_buffers(self):
+        """Free overlap param-gather GPU buffers across all bucket groups."""
+        for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+            bucket_group.free_overlap_buffers()
+
     def scale_gradients(self, scaling_factor: float):
         """Scale all gradients inside the buffers by `scaling_factor`."""
         for buffer in self.buffers + self.expert_parallel_buffers:

megatron/core/distributed/param_and_grad_buffer.py

@@ -10,6 +10,7 @@
 from typing import Dict, List, Optional
 
 import torch
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 from torch.distributed import _coalescing_manager
 
 import megatron.core.nccl_allocator as nccl_allocator
@@ -112,6 +113,42 @@ def __init__(
             global_start, global_end, _ = param_index_map[param]
             self.param_to_index[param] = (global_start - offset, global_end - offset)
 
+        # Layer-wise optimizer attributes for async param gather.
+        self.lw_params_list = None
+        self.lw_param_flat_sizes = None
+        self.lw_gather_tensor_list = None
+        self._lw_src_buffer = None
+
+    def set_lw_params_list(self, lw_params_list: List[List[torch.nn.Parameter]]):
+        """Set per-rank parameter lists for layer-wise async all-gather.
+
+        Args:
+            lw_params_list: List of param lists, one per rank in the DP group.
+        """
+        self.lw_params_list = lw_params_list
+        self.lw_param_flat_sizes = [
+            sum([p.numel() for p in param_list]) for param_list in lw_params_list
+        ]
+
+
+class _LayerWiseAllGatherHandle:
+    """Handle for multiple async broadcast operations used by the layer-wise optimizer.
+
+    Wraps the list of work handles returned by per-rank broadcast calls so that
+    a single wait() call waits on all outstanding operations.
+    """
+
+    def __init__(self, handles):
+        self.handles = handles
+
+    def wait(self):
+        # All broadcasts are on the same NCCL communicator stream, so NCCL
+        # guarantees in-order completion. Waiting on only the last handle is
+        # sufficient and avoids intermediate CUDA stream synchronizations that
+        # can cause timing-dependent deadlocks across ranks.
+        if self.handles:
+            self.handles[-1].wait()
+
 
 class _ParamAndGradBucketGroup:
     """
@@ -138,11 +175,13 @@ def __init__(
         self.buckets = buckets
         self.ddp_config = ddp_config
 
-        if self.ddp_config.use_distributed_optimizer:
+        # overlap_param_gather covers the layer-wise optimizer case, which sets
+        # overlap_param_gather=True without use_distributed_optimizer.
+        if self.ddp_config.use_distributed_optimizer or self.ddp_config.overlap_param_gather:
             self.intra_distributed_optimizer_instance_group = collective_group
             self.intra_distributed_optimizer_instance_size = collective_group_size
             self.intra_distributed_optimizer_instance_rank = collective_group.rank()
-        else:
+        if not self.ddp_config.use_distributed_optimizer:
             self.data_parallel_group = collective_group
 
         # State for bookkeeping: params is the set of parameters this bucket group is
@@ -262,7 +301,9 @@ def start_param_sync(self, force_sync: bool = False):
             force_sync (bool, optional): force synchronous collective regardless of
                 other settings if true.
         """
-        assert self.ddp_config.use_distributed_optimizer
+        # overlap_param_gather covers the layer-wise optimizer case, which sets
+        # overlap_param_gather=True without use_distributed_optimizer.
+        assert self.ddp_config.use_distributed_optimizer or self.ddp_config.overlap_param_gather
 
         if force_sync:
             if self.param_gather_handle is not None:
@@ -273,33 +314,123 @@ def start_param_sync(self, force_sync: bool = False):
             assert self.param_gather_handle is None
 
         async_op = self.ddp_config.overlap_param_gather and not force_sync
-        # Coalesce communication kernels across buckets in the bucket group.
-        with _coalescing_manager(
-            self.intra_distributed_optimizer_instance_group, async_ops=async_op
-        ) as cm:
-            for idx, bucket in enumerate(self.buckets):
-                if self.cached_param_buffer_shard_list[idx] is None:
-                    self.cached_param_buffer_shard_list[idx] = shard_buffer(
-                        bucket.param_data, self.intra_distributed_optimizer_instance_size
+
+        if not self.ddp_config.use_distributed_optimizer:
+            # Layer-wise optimizer path: use per-rank broadcasts for
+            # variable-size param gather.
+            #
+            # Each rank may own a different number of params per bucket, so
+            # lw_param_flat_sizes can vary across ranks. NCCL's ncclAllGather
+            # requires uniform send sizes, so we cannot use all_gather directly.
+            #
+            # Instead, we issue dp_size broadcasts per bucket: for each rank i,
+            # all ranks call broadcast with rank i as the source. On the source
+            # rank the buffer is the flattened local params; on other ranks the
+            # buffer is a correctly-sized receive allocation. This avoids padding
+            # and uses only collectives (no P2P send/recv, which can deadlock
+            # with subsequent collectives on the same NCCL communicator).
+            #
+            # Memory cost: sum(lw_param_flat_sizes) elements per bucket (the
+            # receive buffers for each remote rank, plus the local flattened src).
+            dp_size = self.intra_distributed_optimizer_instance_size
+            local_rank = self.intra_distributed_optimizer_instance_rank
+            group = self.intra_distributed_optimizer_instance_group
+            lw_work_handles = []
+            for bucket in self.buckets:
+                # Use param dtype (e.g., bf16), NOT grad dtype (which may be fp32
+                # when grad_reduce_in_fp32 is enabled). All ranks must use the same
+                # dtype for broadcast buffers, and ranks with local params will have
+                # bf16 src tensors from _flatten_dense_tensors.
+                param_dtype = bucket.params_list[0].dtype
+                src = (
+                    _flatten_dense_tensors(bucket.lw_params_list[local_rank])
+                    if len(bucket.lw_params_list[local_rank]) > 0
+                    else torch.empty(
+                        0, device=bucket.grad_data.device, dtype=param_dtype
                     )
-                local_data_view = self.cached_param_buffer_shard_list[idx][
-                    self.intra_distributed_optimizer_instance_rank
-                ]
-                dist_all_gather_func(
-                    bucket.param_data,
-                    local_data_view,
-                    group=self.intra_distributed_optimizer_instance_group,
-                    async_op=async_op,
                 )
-        if async_op:
-            self.param_gather_handle = cm
+                # Keep src alive until the async operation completes.
+                bucket._lw_src_buffer = src
+
+                if max(bucket.lw_param_flat_sizes) == 0:
+                    # All ranks have empty params for this bucket — skip.
+                    bucket.lw_gather_tensor_list = [
+                        torch.empty(0, device=src.device, dtype=src.dtype) for _ in range(dp_size)
+                    ]
+                    continue
+
+                # Allocate per-rank receive buffers (actual sizes, NO padding).
+                gather_list = []
+                for i in range(dp_size):
+                    if i == local_rank:
+                        gather_list.append(torch.empty(0, device=src.device, dtype=src.dtype))
+                    else:
+                        gather_list.append(
+                            torch.empty(
+                                bucket.lw_param_flat_sizes[i], device=src.device, dtype=src.dtype
+                            )
+                        )
+                bucket.lw_gather_tensor_list = gather_list
+
+                # Broadcast each rank's params to all other ranks.
+                for i in range(dp_size):
+                    if bucket.lw_param_flat_sizes[i] == 0:
+                        continue
+                    src_global = torch.distributed.get_global_rank(group, i)
+                    if i == local_rank:
+                        buf = src
+                    else:
+                        buf = gather_list[i]
+                    work = torch.distributed.broadcast(
+                        buf, src_global, group=group, async_op=async_op
+                    )
+                    if async_op and work is not None:
+                        lw_work_handles.append(work)
+
+            if async_op:
+                self.param_gather_handle = _LayerWiseAllGatherHandle(lw_work_handles)
+            else:
+                # Synchronous: unflatten and copy gathered params immediately.
+                for bucket in self.buckets:
+                    for idx, (flat_params, params) in enumerate(
+                        zip(bucket.lw_gather_tensor_list, bucket.lw_params_list)
+                    ):
+                        if len(params) == 0 or idx == local_rank:
+                            continue
+                        updated_params = _unflatten_dense_tensors(flat_params, params)
+                        for updated_p, model_p in zip(updated_params, params):
+                            model_p.data.copy_(updated_p)
+                    bucket.lw_gather_tensor_list.clear()
+                    bucket._lw_src_buffer = None
+                self.param_gather_handle = None
         else:
-            # When using `_coalescing_manager`, even if a synchronous op (async_op=False) is used,
-            # `cm` is not None, which is different from when `_coalescing_manager` is not used in
-            # which case the torch.distributed._all_gather_base() will return None. In order to
-            # maintain consistency with prior code, we need to manually set communication handle to
-            # None.
-            self.param_gather_handle = None
+            # Standard distributed optimizer path: use _coalescing_manager.
+            # all_gather_into_tensor writes directly into a contiguous output buffer and
+            # does not need a copy-back step, so coalescing works correctly.
+            with _coalescing_manager(
+                self.intra_distributed_optimizer_instance_group, async_ops=async_op
+            ) as cm:
+                for idx, bucket in enumerate(self.buckets):
+                    if self.cached_param_buffer_shard_list[idx] is None:
+                        self.cached_param_buffer_shard_list[idx] = shard_buffer(
+                            bucket.param_data, self.intra_distributed_optimizer_instance_size
+                        )
+                    local_data_view = self.cached_param_buffer_shard_list[idx][
+                        self.intra_distributed_optimizer_instance_rank
+                    ]
+                    dist_all_gather_func(
+                        bucket.param_data,
+                        local_data_view,
+                        group=self.intra_distributed_optimizer_instance_group,
+                        async_op=async_op,
+                    )
+            if async_op:
+                self.param_gather_handle = cm
+            else:
+                # When using `_coalescing_manager`, even if a synchronous op
+                # (async_op=False) is used, `cm` is not None. Manually set to None for
+                # consistency with prior code.
+                self.param_gather_handle = None
         self.param_gather_dispatched = True
 
     def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
@@ -317,7 +448,9 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
             skip_next_bucket_dispatch (bool, optional): if true, dispatch next
                 bucket's communication if available.
         """
-        assert self.ddp_config.use_distributed_optimizer
+        # overlap_param_gather covers the layer-wise optimizer case, which sets
+        # overlap_param_gather=True without use_distributed_optimizer.
+        assert self.ddp_config.use_distributed_optimizer or self.ddp_config.overlap_param_gather
         assert self.ddp_config.overlap_param_gather
 
         # If current bucket's param AG has not been dispatched, dispatch it now (e.g., first
@@ -355,6 +488,23 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
                     # correspond to multiple param buffers. If we zero out the entire grad buffer,
                     # it would clear the data of those param buffers that have not yet completed AG.
                     bucket.param_data.zero_()
+            elif not self.ddp_config.use_distributed_optimizer:
+                for bucket in self.buckets:
+                    # Unflatten and copy gathered params for each rank.
+                    for idx, (flat_params, params) in enumerate(
+                        zip(bucket.lw_gather_tensor_list, bucket.lw_params_list)
+                    ):
+                        # Skip local params and empty tensors.
+                        if (
+                            len(params) == 0
+                            or idx == self.intra_distributed_optimizer_instance_rank
+                        ):
+                            continue
+                        updated_params = _unflatten_dense_tensors(flat_params, params)
+                        for updated_p, model_p in zip(updated_params, params):
+                            model_p.data.copy_(updated_p)
+                    bucket.lw_gather_tensor_list.clear()
+                    bucket._lw_src_buffer = None
             else:
                 fp8_params = []
                 for bucket in self.buckets:
@@ -539,6 +689,22 @@ def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
         self.grad_reduce_handle.wait()
         self.grad_reduce_handle = None
 
+    def free_overlap_buffers(self):
+        """Free GPU buffers used by overlap param gather.
+
+        Waits on any pending param all-gather handle, then releases the
+        per-bucket temporary buffers so that the CUDA memory allocator can
+        reclaim them.  Called before async checkpoint saves to avoid OOM in
+        the persistent checkpoint worker process.
+        """
+        if self.param_gather_handle is not None:
+            self.param_gather_handle.wait()
+            self.param_gather_handle = None
+        for bucket in self.buckets:
+            if bucket.lw_gather_tensor_list is not None:
+                bucket.lw_gather_tensor_list.clear()
+            bucket._lw_src_buffer = None
+
     def register_grad_ready(
         self, param: torch.nn.Parameter, force_all_reduce: Optional[bool] = False
     ):