From bec1750499504bee799704becf89fd3514e1932d Mon Sep 17 00:00:00 2001 From: Deyu Fu Date: Mon, 9 Mar 2026 17:16:33 +0800 Subject: [PATCH 1/2] reuse grad buffer for param allgather --- .../core/distributed/param_and_grad_buffer.py | 52 ++++++++----------- 1 file changed, 21 insertions(+), 31 deletions(-) diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index d24bd639bc1..bb5cb086a1b 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -117,7 +117,6 @@ def __init__( self.layerwise_params_list = None self.layerwise_param_flat_sizes = None self.layerwise_gather_list = None - self._layerwise_src_buffer = None def set_layerwise_params_list(self, layerwise_params_list: List[List[torch.nn.Parameter]]): """Set per-rank parameter lists for layer-wise async all-gather. @@ -334,44 +333,38 @@ def start_param_sync(self, force_sync: bool = False): param_dtype = bucket.params_list[0].dtype if max(bucket.layerwise_param_flat_sizes) == 0: - # All ranks have empty params for this bucket — skip. bucket.layerwise_gather_list = None continue - # Flatten local params. Detach from the autograd graph because - # start_param_sync can be called during the forward pass (where - # autograd is active) and all_gather will write into gather_list - # entries in-place. local_size = bucket.layerwise_param_flat_sizes[local_rank] + total_gather_size = sum(bucket.layerwise_param_flat_sizes) + + # Reuse grad_data as the all_gather receive buffer; it is idle + # during forward and grad_dtype.element_size >= param_dtype. + reuse_buf = bucket.grad_data.view(param_dtype) + assert reuse_buf.numel() >= total_gather_size + + # Partition reuse_buf into contiguous per-rank receive slices. + gather_list = [] + offset = 0 + for i in range(dp_size): + size = bucket.layerwise_param_flat_sizes[i] + gather_list.append(reuse_buf[offset : offset + size]) + offset += size + local_slot_view = gather_list[local_rank] + + # Flatten local params and copy into the local rank's slot. + # Detach from autograd since start_param_sync may be called + # during the forward pass where autograd is active. if local_size > 0: flat_local_params = _flatten_dense_tensors( bucket.layerwise_params_list[local_rank] ).detach() - else: - flat_local_params = torch.empty( - 0, device=bucket.grad_data.device, dtype=param_dtype - ) - # Keep flat_local_params alive until the async operation completes. - bucket._layerwise_src_buffer = flat_local_params - - # Allocate per-rank receive buffers with actual sizes (no padding). - # Reuse flat_local_params for local_rank's slot to avoid an extra allocation. - gather_list = [] - for i in range(dp_size): - if i == local_rank: - gather_list.append(flat_local_params) - else: - gather_list.append( - torch.empty( - bucket.layerwise_param_flat_sizes[i], - device=flat_local_params.device, - dtype=flat_local_params.dtype, - ) - ) + local_slot_view.copy_(flat_local_params) bucket.layerwise_gather_list = gather_list work = torch.distributed.all_gather( - gather_list, flat_local_params, group=group, async_op=async_op + gather_list, local_slot_view, group=group, async_op=async_op ) if async_op and work is not None: layerwise_work_handles.append(work) @@ -392,7 +385,6 @@ def start_param_sync(self, force_sync: bool = False): for updated_p, model_p in zip(updated_params, params): model_p.data.copy_(updated_p) bucket.layerwise_gather_list = None - bucket._layerwise_src_buffer = None self.param_gather_handle = None else: # Standard distributed optimizer path: use _coalescing_manager. @@ -502,7 +494,6 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False): for updated_p, model_p in zip(updated_params, params): model_p.data.copy_(updated_p) bucket.layerwise_gather_list = None - bucket._layerwise_src_buffer = None else: fp8_params = [] for bucket in self.buckets: @@ -700,7 +691,6 @@ def free_overlap_buffers(self): self.param_gather_handle = None for bucket in self.buckets: bucket.layerwise_gather_list = None - bucket._layerwise_src_buffer = None def register_grad_ready( self, param: torch.nn.Parameter, force_all_reduce: Optional[bool] = False From 89a38b3918781d96f71d1d77a8e3e3681f7ca5c7 Mon Sep 17 00:00:00 2001 From: Deyu Fu Date: Mon, 9 Mar 2026 20:00:14 +0800 Subject: [PATCH 2/2] fix unit test --- tests/unit_tests/distributed/test_param_and_grad_buffer.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/unit_tests/distributed/test_param_and_grad_buffer.py b/tests/unit_tests/distributed/test_param_and_grad_buffer.py index 48f815531c5..26572820377 100644 --- a/tests/unit_tests/distributed/test_param_and_grad_buffer.py +++ b/tests/unit_tests/distributed/test_param_and_grad_buffer.py @@ -407,7 +407,6 @@ def test_bucket_group_clears_buffers(self): # Simulate buffers that would be allocated by start_param_sync. for bucket in bg.buckets: bucket.layerwise_gather_list = [torch.empty(8), torch.empty(8)] - bucket._layerwise_src_buffer = torch.empty(16) bg.free_overlap_buffers() @@ -415,9 +414,6 @@ def test_bucket_group_clears_buffers(self): assert ( bucket.layerwise_gather_list is None ), "layerwise_gather_list should be None after free_overlap_buffers" - assert ( - bucket._layerwise_src_buffer is None - ), "_layerwise_src_buffer should be None after free_overlap_buffers" Utils.destroy_model_parallel() @@ -446,7 +442,6 @@ def test_bucket_group_noop_when_no_buffers(self): assert bg.param_gather_handle is None for bucket in bg.buckets: assert bucket.layerwise_gather_list is None - assert bucket._layerwise_src_buffer is None # Should not raise. bg.free_overlap_buffers()