Skip to content

reuse grad buffer for layer-wise param allgather #3751

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
FDecaYed merged 4 commits into from
Apr 14, 2026
Merged

reuse grad buffer for layer-wise param allgather #3751

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bec1750
reuse grad buffer for param allgather
FDecaYed Mar 9, 2026
89a38b3
fix unit test
FDecaYed Mar 9, 2026
f6730fc
Merge branch 'main' into deyuf/reuse_grad_buf
FDecaYed Mar 9, 2026
dd22546
Merge branch 'main' into deyuf/reuse_grad_buf
FDecaYed Apr 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 21 additions & 31 deletions megatron/core/distributed/param_and_grad_buffer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,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.
Expand Down Expand Up @@ -343,44 +342,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)
Expand All @@ -401,7 +394,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.
Expand Down Expand Up @@ -511,7 +503,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:
Expand Down Expand Up @@ -720,7 +711,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 _copy_back_extra_main_grads(self):
"""
Expand Down
5 changes: 0 additions & 5 deletions tests/unit_tests/distributed/test_param_and_grad_buffer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -448,17 +448,13 @@ 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()

for bucket in bg.buckets:
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()

Expand Down Expand Up @@ -487,7 +483,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()
Expand Down
Loading