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