ZeRO-1/2: synchronize IPG buffer reuse per buffer

[almightychang]

Summary

Fixes #8364.

ZeRO-1/2 currently protects IPG ping-pong buffer reuse by making the current/autograd stream wait on the entire reduction stream in average_tensor(). That reverse dependency also waits for unrelated work already queued on the reduction stream, serializing the next backward bucket behind the preceding collective.

This change scopes that dependency to the physical IPG buffer being reused:

• preserve the producer-stream to reduction-stream waits added for multi-stream gradient copies;
• record reduction completion per physical IPG buffer after its final reduction-stream consumer;
• wait on that event immediately before the same buffer is overwritten;
• initialize and release event state together with the physical buffers.

Oversized gradients and non-contiguous fallback reductions do not participate in physical IPG buffer reuse. Accelerators that resolve data dependencies implicitly retain their existing path.

Stream ordering: before and after

Time flows left to right. P0/P1 are producer copies, R0/R1 are reductions, and E0 is the completion event for physical buffer 0.

Before: wait on the whole reduction-stream tail

current/autograd  [ P0 ][ backward + P1 ][ wait for reduction_stream ........ ][ next backward ]
                              |                         ^
                              | producer dependency     | broad reverse wait
                              v                         |
reduction_stream         [ wait P0 ][ R0 collective ...][ wait P1 ][ R1 ... ]

Result: preparing or enqueueing the next bucket can stall behind the preceding collective.

After: wait only before reusing the same physical buffer

current/autograd  [ P0 ][ backward + P1 ][ more backward ][ wait E0 ][ overwrite buffer 0 ]
                     |          |                              ^
                     v          v                              | buffer-0 reuse only
reduction_stream  [ wait P0 ][ R0 ... ][ record E0 ][ wait P1 ][ R1 ... ]

buffer index            0                    1                       0 (reused)

Result: producer-to-reduction safety remains, while unrelated backward work overlaps R0.

Correctness invariants

• A reduction never reads a buffer before all producer copies complete.
• A producer never overwrites a physical buffer before that buffer's previous reduction-stream consumers complete.
• Backward compute for a different ping-pong buffer does not wait on the reduction stream's unrelated tail.
• Completion state is isolated by communication dtype and physical buffer index, and survives logical bucket clears until reuse.

Tests / validation

• pre-commit run --files deepspeed/runtime/zero/stage_1_and_2.py tests/unit/v1/zero/test_overlap_comm_record_stream.py
• Local focused suite: 14 passed; CUDA distributed cases skipped on the CPU-only host.
• Modal 2x H200 stress: 10 repeated overlap/non-overlap comparisons across three bucket sizes and delayed reduction-stream work. Losses, full gradients, and final parameters matched.
• Modal 2x H200 related regression selection: overlap on/off, contiguous/non-contiguous gradients, ZeRO stages 1/2/3, multi-bucket overflow, gradient accumulation, unused parameters, and shared backward graphs. The initial run passed 44 cases and exposed one over-constrained new test assertion; after narrowing it to the actual oversized-gradient invariant, the corrected case passed.
• Bucket boundary coverage exercises parameter size at bucket size - 1, equal to bucket size, and bucket size + 1. Bucketed gradients record and wait for reuse; oversized gradients bypass reuse waits.

The correctness suite deliberately delays reduction-stream work to amplify reuse races. I have not repeated the issue's original 8x H200 profiler experiment in this PR; the performance evidence and expected recoverable interval are documented in #8364.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: db34ce6a29

ℹ️ 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".

[almightychang]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Keep them coming!

Reviewed commit: 9dc5ec5977

ℹ️ 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".