ZeRO 1/2: wait on all IPG-bucket producer streams in average_tensor (#8061)

[arunshar]

What

Fixes #8061. In ZeRO stage 1/2 with overlap_comm, average_tensor waits the reduction stream on only the current stream before reducing the contiguous IPG gradient bucket:

if self.overlap_comm:
    stream = self.reduction_stream
    if not get_accelerator().resolves_data_dependency():
        stream.wait_stream(get_accelerator().current_stream())   # only one stream

But the per-parameter gradient copies that fill the bucket (reduce_independent_p_g_buckets_and_remove_grads, the new_grad_tensor.copy_(...) into bucket.buffer[bucket.index]) can be issued on multiple streams. That is exactly the scenario #8061 reports under torch.compile, where gradient hooks run on different autograd streams and several device streams write slices into the same IPG buffer. Waiting on only the current stream lets the all-reduce read the bucket before the other producers finish, corrupting gradients (NaN loss from step 1).

Change

Implements the fix direction proposed in #8061 (record the streams used for the IPG bucket copies, then make the reduction stream wait on all of them):

• IPGBucket gains a copy_streams set, cleared in clear().
• After each copy into bucket.buffer[bucket.index], record current_stream() (overlap path only).
• average_tensor waits the reduction stream on every recorded producer stream, falling back to current_stream() when the set is empty (e.g. the extra-large-param path that reduces without a bucket copy).

The single-stream case is unchanged: when all copies are on one stream, copy_streams == {current_stream}, so the wait is identical to before, and there is no behavior change for the common path.

Tests / validation

• New CPU unit tests in tests/unit/v1/zero/test_overlap_comm_record_stream.py (fakes + monkeypatch, no GPU): the reduction stream waits on all producer streams, the empty-set fallback to current_stream, and IPGBucket.copy_streams reset. The two pre-existing tests in that file still pass (5 passed).
• pre-commit run --files is green on both changed files (yapf, flake8, check-torchdist, check-license, codespell).
• No-regression A/B on 2x A40 (seeded, identical data, ZeRO-2 + contiguous_gradients, overlap_comm on vs off, eager and torch.compile): before and after this change, suspect == baseline on every repeat with byte-identical final-param hashes, and the baseline-vs-baseline determinism gate passes. So the fix does not change results when grad-bucket copies stay on a single stream.

Honest scope on reproduction

I could not deterministically trigger the multi-stream NaN on the available hardware (A40, small MLP): neither the torch.compile A/B nor a synthetic two-stream microbenchmark surfaced it (PyTorch's caching allocator inserts implicit cross-stream syncs that mask the race in a microbenchmark, and a plain torch.compile(model) on this model kept the grad-bucket copies on one stream). This PR is therefore offered as the minimal correct synchronization for the clearly-missing producer-stream wait identified in average_tensor, validated for no-regression and with unit coverage; a reviewer with the original torch.compile multi-stream repro from #8061 can confirm the NaN is resolved.

Opened as a draft.