Fix gradient corruption with layerwise param all-gather overlap

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:

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()