Skip to content

Fix wgrad race condition when using double buffers. #5222

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
cspades merged 3 commits into from
Jun 10, 2026
Merged

Fix wgrad race condition when using double buffers. #5222

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3527f99
Fix wgrad race condition when using double buffers.
cspades Jun 9, 2026
e2e5663
Lint lint lint
cspades Jun 9, 2026
b5a6cf8
Merge branch 'main' into cye/double-buf-wgrad-race
cspades Jun 9, 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
18 changes: 5 additions & 13 deletions megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,15 +340,8 @@ def __init__(
self._init_fsdp_param_and_grad_buffer()
self._register_fsdp_hooks(self.module)
self.microbatch_count = 0

# Add a reference from the distributed parameters to self for API
# accessibility, e.g. when attaching MegatronFSDP scheduled ops
# to the distributed optimizer.step() and optimizer.zero_grad().
self.is_param_fsdp_distributed = False
self._replace_param_with_distributed_if_needed()
for param in self.module.parameters():
# Attach MegatronFSDP reference to the parameter.
setattr(param, "_megatron_fsdp_model", self)

def _check_module_parameter_types(self):
"""
Expand Down Expand Up @@ -604,12 +597,7 @@ def _grad_acc(param):
# If TransformerEngine gradient accumulation is fused, then param.get_main_grad()
# already holds the wgrad and param.grad_added_to_main_grad=True.
if not param.grad_added_to_main_grad:
# Get `main_grad` will allocate bucket, check that the currently
# used main_grad buffer does not exceed the scope of two FSDP Unit
# Modules, i.e., the buffer limit imposed by double-buffer allocator.
if self.ddp_config.fsdp_double_buffer:
self.grad_reduce_pipeline._enforce_double_buffer_limit([group_id])

# Allocate a unsharded gradient buffer.
param.main_grad = param.get_main_grad()
if param.grad is not None:
if self.report_nan_in_param_grad:
Expand Down Expand Up @@ -1359,6 +1347,8 @@ def _replace_param_with_distributed_if_needed(self):
# DTensor parameter is managed by Megatron FSDP.
if not hasattr(dist_param, "__fsdp_param__"):
dist_param.__fsdp_param__ = True
if not hasattr(dist_param, "_megatron_fsdp_model"):
dist_param._megatron_fsdp_model = self
_replace_module_parameter(self.module, name, dist_param)

# Handle shared weights
Expand All @@ -1371,6 +1361,8 @@ def _replace_param_with_raw_if_needed(self):

for name, _ in self.module.named_parameters():
assert name in self.raw_param, f"Raw parameter {name} not found in module."
if not hasattr(self.raw_param[name], "_megatron_fsdp_model"):
self.raw_param[name]._megatron_fsdp_model = self
_replace_module_parameter(self.module, name, self.raw_param[name])

# Handle shared weights
Expand Down
20 changes: 13 additions & 7 deletions megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -906,7 +906,7 @@ def __init__(
# Build the data parallel buffer index, which contains information
# on where each parameter / gradient tensor will be stored in this
# distributed buffer.
(self.item_index_map, self.bucket_index, self.shard_bucket_index) = (
self.item_index_map, self.bucket_index, self.shard_bucket_index = (
build_data_parallel_buffer_index(
[to_local_if_dtensor(p).shape for p in self.params],
self.dp_rank,
Expand Down Expand Up @@ -1767,7 +1767,7 @@ def __init__(
)

# Get the parameter groups.
(self.parameter_groups, self.param_to_param_group, self.bucket_to_bucket_group) = (
self.parameter_groups, self.param_to_param_group, self.bucket_to_bucket_group = (
_get_parameter_groups(module, bucketing_policy, meta_device_init_fp8_params)
)
self._init_each_parameter_group_buffers(meta_device_init_fp8_params)
Expand Down Expand Up @@ -2685,18 +2685,23 @@ def _alloc(dtype, size):
p._item_id = item_id

def main_grad_getter(p):
# Get gradient buffer and item ID.
gbuf = p._gbuf
item_id = p._item_id
# Need to free a bucket for the incoming bucket if double-buffering.
p._megatron_fsdp_model.grad_reduce_pipeline._enforce_double_buffer_limit(
[gbuf.bucket_id]
)
# Make sure main_grad memory is allocated when initially accessed.
# When gradients are sharded, we can pre-allocate a communication
# bucket to avoid casting to a communication data-type. Otherwise,
# return the item backed by the main gradient buffer required to
# support un-sharded gradient accumulation at high precision.
bucket = p._gbuf.fetch_bucket(
bucket = gbuf.fetch_bucket(
dtype=(
self.mp_policy.grad_comm_dtype if p._gbuf.is_data_distributed else None
)
)
gbuf = p._gbuf
item_id = p._item_id
# View it as p.shape so you can insert the param.grad into
# the bucket seamlessly.
return gbuf.get_item_from_bucket(bucket, item_id).view(
Expand Down Expand Up @@ -2909,6 +2914,7 @@ def set_param_attribute():
"is_embedding_or_output_parameter",
"is_embedding_parameter",
"_tensor_parallel_mode",
"_megatron_fsdp_model",
]:
if hasattr(orig_param, attr_name):
setattr(param, attr_name, getattr(orig_param, attr_name))
Expand Down Expand Up @@ -3558,7 +3564,7 @@ def _enforce_double_buffer_limit(self, add_buckets):
for _, _, bucket_id in reversed(self.grad_reduce_queue):
fsdp_unit_id = param_groups[bucket_id].fsdp_unit_id
double_buf_units.add(fsdp_unit_id)
if len(double_buf_units) > 1:
if len(double_buf_units) > 2:
keep_n -= 1

with torch.cuda.stream(self.rs_stream):
Expand Down Expand Up @@ -3940,7 +3946,7 @@ def reset(self, preserve_non_fsdp_units: bool = True):
UserWarning,
)
while len(self.param_gather_event_map) > 0:
(bucket_id, bwd) = next(iter(self.param_gather_event_map))
bucket_id, bwd = next(iter(self.param_gather_event_map))
self.wait_bucket_ready(bucket_id, bwd)

for bucket_id in range(self.num_buckets):
Expand Down
Loading