From 3527f994e5fb123f4434a6436340bea429b09870 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Mon, 8 Jun 2026 17:57:05 -0700 Subject: [PATCH 1/2] Fix wgrad race condition when using double buffers. Signed-off-by: Cory Ye --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 18 +++++------------- .../src/megatron_fsdp/param_and_grad_buffer.py | 14 ++++++++++---- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 37777e7581b..e90e3879dcb 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -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): """ @@ -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: @@ -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 @@ -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 diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index 690ec263890..fcd5707b742 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py @@ -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( @@ -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)) @@ -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): From e2e566365b6eebfc3de56977ad310059782b6a55 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Mon, 8 Jun 2026 19:38:51 -0700 Subject: [PATCH 2/2] Lint lint lint Signed-off-by: Cory Ye --- .../fsdp/src/megatron_fsdp/param_and_grad_buffer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index fcd5707b742..ee33df209ea 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py @@ -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, @@ -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) @@ -3946,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):