Wait for all async all-gathers in coalesced ZeRO-3 parameter gather

[delock]

Problem

While investigating multi-rank CPU CI (#8381), TestUnmanagedGradientAccumulation::test_unmanaged_varying_backward_count[3] failed intermittently with parameters turned to NaN, and test_zero_coalesce_grad_reduction showed mismatches with denormal garbage values (4.8e+37, 1e-38) — the signature of uninitialized memory.

Root cause

_allgather_params_coalesced() launches one async all_gather per parameter, but only waits on the last handle before pointing param.data at the torch.empty flat buffers:

launch_handles[-1].wait()
param.data = gathered_tensor.narrow(...)
if not get_accelerator().resolves_data_dependency():
    get_accelerator().synchronize()

On CUDA this is safe by construction: the ops are enqueued from a single thread onto the same stream (FIFO), and torch.cuda.synchronize() waits for everything anyway. On gloo each async handle runs on an independent background thread with no ordering between handles — waiting for the last one says nothing about the earlier ones — and the CPU accelerator's synchronize() is a no-op. Reading param.data can therefore race the gather and expose uninitialized memory.

The corruption is timing-dependent (locally 9/20 runs fail; #8382's comparison-helper fix made the numeric assertions reachable), which is why it was never caught on GPU CI.

Fix

Wait on every handle. On CUDA the earlier handles are guaranteed complete by stream ordering, so the extra waits are immediate status checks with no real waiting; the change makes correctness independent of that structural coincidence.

Verification (CPU/gloo, world sizes 2–3)

Test	Before	After
test_unmanaged_varying_backward_count[3]	9/20 failed	0/20 failed
TestZero3ParamPartitioningBase / TestGradientAllreduceOp / TestUnmanagedGradientAccumulation (regression, 57 cases)	—	55 pass + 2 skip

GPU behavior is unchanged by construction (same buckets, same order, stream-FIFO + full synchronize); GPU CI will re-confirm.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: fe58c7ef6d

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This is a non-merge commit, but its commit message contains no Signed-off-by trailer, violating the repository's mandatory commit requirement. Recreate the commit with git commit --signoff before merging.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[pengdurice]

thanks for the fix!