diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index dc3014d72d2..38d84ffcce5 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -396,6 +396,11 @@ 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 + # Zero out grad_data since it was reused as the all-gather + # receive buffer. Without this, accumulation into main_grad + # (a view into grad_data) would start from the result of the + # latest parameter all-gather instead of zero. + bucket.grad_data.zero_() self.param_gather_handle = None else: # Standard distributed optimizer path: use _coalescing_manager. @@ -505,6 +510,11 @@ 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 + # Zero out grad_data since it was reused as the all-gather + # receive buffer. Without this, accumulation into main_grad + # (a view into grad_data) would start from the result of the + # latest parameter all-gather instead of zero. + bucket.grad_data.zero_() else: fp8_params = [] for bucket in self.buckets: diff --git a/tests/unit_tests/test_layer_wise_optimizer.py b/tests/unit_tests/test_layer_wise_optimizer.py index e43f2ab7d88..36e5fe11b67 100644 --- a/tests/unit_tests/test_layer_wise_optimizer.py +++ b/tests/unit_tests/test_layer_wise_optimizer.py @@ -821,6 +821,13 @@ def test_overlap_param_gather_multi_iteration(self): optimizer.step() model.start_param_sync(force_sync=True) + # Verify grad_data is zeroed after synchronous param_sync. + for bucket_group in model.bucket_groups: + for bucket in bucket_group.buckets: + assert torch.all( + bucket.grad_data == 0 + ), f"grad_data not zeroed after param sync at iteration {iteration}" + # Sync path: step (includes allgather) ref_optimizer.step() @@ -860,6 +867,13 @@ def test_overlap_param_gather_async_dispatch_and_finish(self): for bucket_group in model.bucket_groups: bucket_group.finish_param_sync(skip_next_bucket_dispatch=True) + # Verify grad_data is zeroed after asynchronous param_sync. + for bucket_group in model.bucket_groups: + for bucket in bucket_group.buckets: + assert torch.all( + bucket.grad_data == 0 + ), "grad_data not zeroed after finish_param_sync" + # Sync path: step (includes allgather) ref_optimizer.step()