From c486e79687340fd50a630e3163a7f0ea2fb6a4c7 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Mon, 4 May 2026 14:30:39 -0700 Subject: [PATCH] Zero grad_data after layerwise param all-gather to prevent gradient corruption The layerwise param sync path reuses grad_data as the all-gather receive buffer. After copying gathered params to model params, grad_data was not zeroed, so it still contained the all-gather results. Since main_grad is a view into grad_data, accumulation into main_grad would start from the result of the latest parameter all-gather instead of zero, corrupting gradients when using --overlap-param-gather with the layerwise distributed optimizer (e.g., Muon). Fix: zero bucket.grad_data after copying gathered params in both the synchronous (start_param_sync) and async (finish_param_sync) paths, matching the existing pattern used for the mxfp8 path. Add grad_data zeroing assertions to test_overlap_param_gather_multi_iteration and test_overlap_param_gather_async_dispatch_and_finish. Co-Authored-By: Claude Opus 4.6 --- megatron/core/distributed/param_and_grad_buffer.py | 10 ++++++++++ tests/unit_tests/test_layer_wise_optimizer.py | 14 ++++++++++++++ 2 files changed, 24 insertions(+) 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()