From fe58c7ef6dbe45439852c89d27ad74bdfd71795c Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Tue, 1 Sep 2026 15:01:15 +0800 Subject: [PATCH] Wait for all async all-gathers in coalesced ZeRO-3 parameter gather _allgather_params_coalesced() launched one async all_gather per parameter but only waited on the last handle before pointing param.data at the torch.empty flat buffers. On CUDA the trailing torch.cuda.synchronize() and same-stream ordering make that safe. On gloo each async handle is independent and the CPU accelerator's synchronize() is a no-op, so the earlier gathers could still be in flight when the gathered parameters were read, exposing uninitialized memory (NaN / denormal garbage) through param.data. The corruption is timing-dependent: local reproduction of tests/unit/v1/zero/test_zero_user_backward.py::TestUnmanagedGradientAccumulation::test_unmanaged_varying_backward_count[3] failed 9/20 runs without this change and 0/20 with it (after #8382's comparison-helper fix made the numeric assertion reachable). Signed-off-by: Guokai Ma Signed-off-by: Ma, Guokai --- deepspeed/runtime/zero/partition_parameters.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index b7dc0ef78db0..f7d0d2e3ef7a 100644 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -2088,8 +2088,10 @@ def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False): launch_quantize_handles.append(quant_handle) launch_handles.append(h) - # Wait ensures the operation is enqueued, but not necessarily complete. - launch_handles[-1].wait() + # gloo handles are independent and the CPU synchronize() below is a + # no-op, so every handle must be waited on, not just the last one. + for handle in launch_handles: + handle.wait() if quantize: for quant_handle in launch_quantize_handles: quant_handle.wait()