[draft] support async param gather in layer-wise optimizer

megatron/core/distributed/param_and_grad_buffer.py

@@ -11,6 +11,7 @@
 
 import torch
 from torch.distributed import _coalescing_manager
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 import megatron.core.nccl_allocator as nccl_allocator
 from megatron.core import parallel_state
@@ -108,6 +109,12 @@ def __init__(
             self.param_to_index[param] = (offset, offset + param.numel())
             offset += param.numel()
 
+        self.lw_params_list = None
+        self.lw_param_flat_sizes = None
+
+    def set_lw_params_list(self, lw_params_list: List[List[torch.nn.Parameter]]):
+        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 _ParamAndGradBucketGroup:
     """
@@ -250,20 +257,41 @@ def start_param_sync(self, force_sync: bool = False):
         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_layer_wise_optimizer:
+                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,
                     )
-                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,
-                )
+            else:
+                assert async_op, "Layer-wise optimizer requires overlap_param_gather=True"
+                for bucket in self.buckets:
+                    src = (
+                        _flatten_dense_tensors(bucket.lw_params_list[self.intra_distributed_optimizer_instance_rank])
+                        if len(bucket.lw_params_list[self.intra_distributed_optimizer_instance_rank]) > 0
+                        else torch.empty(0, device=bucket.param_data.device, dtype=bucket.param_data.dtype)
+                    )
+                    # TODO(deyuf): reuse param_data as the gather tensor buffer to avoid unnecessary memory allocation
+                    bucket.lw_gather_tensor_list = [
+                        torch.empty(size, device=bucket.param_data.device, dtype=bucket.param_data.dtype)
+                        for size in bucket.lw_param_flat_sizes
+                    ]
+                    torch.distributed.all_gather(
+                        bucket.lw_gather_tensor_list,
+                        src,
+                        group=self.intra_distributed_optimizer_instance_group,
+                        async_op=True,
+                    )
+
         if async_op:
             self.param_gather_handle = cm
         else:
@@ -328,6 +356,17 @@ 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 self.ddp_config.use_layer_wise_optimizer:
+                for bucket in self.buckets:
+                    # unflatten and copy gathered params for each rank i
+                    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()
             else:
                 fp8_params = []
                 for bucket in self.buckets:

megatron/core/optimizer/layer_wise_optimizer.py

@@ -45,6 +45,8 @@ def __init__(
         config: OptimizerConfig,
         pg_collection: Optional[ProcessGroupCollection] = None,
         init_state_fn_list: Optional[List[Callable]] = None,
+        model_chunks: Optional[List] = None,
+        async_allgather: Optional[bool] = False,
     ) -> None:
         """
         Initialize LayerWiseDistributedOptimizer.
@@ -58,6 +60,13 @@ def __init__(
 
         self.pg_collection = pg_collection
         self.shard_params(optimizers)
+
+        # set bucket lw params list for async allgather
+        self.async_allgather = async_allgather
+        if self.async_allgather:
+            assert model_chunks is not None, "model_chunks must be provided if async_allgather is True"
+            self.set_bucket_lw_params_list(model_chunks)
+
         if init_state_fn_list:
             assert len(init_state_fn_list) == len(
                 optimizers
@@ -143,6 +152,28 @@ def shard_params(self, optimizers):
         if expt_dp_size == 1 or len(self.expt_dp_params_list[0]) == 0:
             self.expt_dp_params_list = None
 
+    def set_bucket_lw_params_list(self, model_chunks: List[MegatronModule]):
+        for model_chunk in model_chunks:
+            for group in model_chunk.bucket_groups:
+                for bucket in group:
+                    # find all params that belong to this bucket and keep their sharding structure
+                    bucket_params_list = [[] for _ in range(get_pg_size(self.pg_collection.dp_cp))]
+                    for bucket_list, full_params_list in zip(bucket_params_list, self.dp_cp_params_list):
+                        for param in full_params_list:
+                            if param in bucket.params:
+                                bucket_list.append(param)
+                    bucket.set_lw_params_list(bucket_params_list)
+            # do the same for expert parallel
+            for group in model_chunk.expert_parallel_bucket_groups:
+                for bucket in group:
+                    # find all params that belong to this bucket and keep their sharding structure
+                    bucket_params_list = [[] for _ in range(get_pg_size(self.pg_collection.expt_dp))]
+                    for bucket_list, full_params_list in zip(bucket_params_list, self.expt_dp_params_list):
+                        for param in full_params_list:
+                            if param in bucket.params:
+                                bucket_list.append(param)
+                    bucket.set_lw_params_list(bucket_params_list)
+
     @torch.no_grad()
     def allgather_params(self) -> None:
         """All-gather updated params from all ranks."""
@@ -223,8 +254,9 @@ def step(self):  # type: ignore[no-untyped-def]
         """step function for layer-wise optimizer."""
         update_successful, grad_norm, num_zeros_in_grad = super().step()
 
-        # All gather updated params.
-        self.allgather_params()
+        # All gather updated params. If async_allgather is True, the allgather is done in the forward pre-hook.
+        if not self.async_allgather:
+            self.allgather_params()
 
         return update_successful, grad_norm, num_zeros_in_grad