ZeRO-1/2 average_tensor serializes backward compute behind the preceding reduction

Description

ZeRO stage 1/2 with overlap_comm=true still serializes the autograd/current stream behind the reduction stream at the start of every average_tensor() call:

if self.overlap_comm:
    stream = self.reduction_stream
    if not get_accelerator().resolves_data_dependency():
        for producer_stream in producer_streams:
            stream.wait_stream(producer_stream)
        get_accelerator().current_stream().wait_stream(stream)

The producer-to-reduction dependency is required. However, the reverse dependency makes the current stream wait for all previously queued work on reduction_stream, including the preceding bucket's collective. When bucket N+1 becomes ready, backward compute is therefore serialized behind bucket N's collective.

This confines overlap to the interval between launching bucket N and preparing bucket N+1. It is especially costly when collective residence is dominated by rank arrival skew rather than data transfer.

Environment

• Reproduced on DeepSpeed 0.18.7
• Confirmed present by source inspection in 0.19.6 and current master
• PyTorch 2.7.0
• 8x H200 on one NVLink node
• ZeRO stage 2
• bf16 communication
• overlap_comm=true
• contiguous_gradients=true
• reduce_scatter=true
• Hugging Face Trainer
• global batch 512, gradient accumulation 2
• 32 samples per rank per micro-batch
• 6,913,191,280 total parameters
• 2,376,274,560 trainable parameters
• 4,536,916,720 frozen parameters
• no CPU or parameter offload

Profile evidence

The following values are means per rank per micro-batch from eight Kineto traces:

Metric	Time
Compute kernel union	489.0 ms
NCCL kernel union	364.0 ms
Compute/NCCL overlap	75.4 ms
NCCL pure transfer estimate	29.9 ms
Wait inside collectives	334.1 ms
NCCL outside compute	288.6 ms

The traces show the following consistent timeline for non-final gradient buckets:

1. Bucket N launches on reduction_stream.
2. Backward compute resumes and prepares bucket N+1.
3. At or immediately before bucket N+1 launches, the compute stream becomes idle while bucket N's NCCL kernel is still resident.
4. Compute resumes only after that preceding collective completes.

Interval attribution assigns NCCL-outside-compute time as follows:

• 212.2 ms (73.5%): the compute stream is stalled on the reduction stream's tail
• 61.6 ms (21.3%): the final bucket is flushed after backward
• 14.8 ms: other collectives

An event replay that keeps NCCL on its own stream, preserves collective order and measured duration, but removes this premature reverse dependency bounds recoverable time at approximately 204 ms per rank per micro-batch. The result is consistent across ranks.

The effect is not explained by transfer bandwidth or SM contention in this profile. NCCL kernels launch with grid [16, 1, 1] on 132 SMs. Only about 29.9 ms of 364.0 ms NCCL residence is attributable to transfer. Rank arrival skew dominates: the slowest rank's compute is 652 ms versus a 520 ms rank mean. This makes the reverse dependency particularly visible: one late rank extends the collective, and every rank's compute stream is then held behind that tail at the next bucket enqueue.

Why this is distinct from existing fixes

• #5606 intentionally added the reverse wait to prevent reuse of an IPG ping-pong buffer before its preceding reduction completed. It also noted a performance gap for short-compute workloads.
• #7805 removed redundant waits after bucket flush and repaired ping-pong index reset, but left the reverse wait inside average_tensor().
• #8061 / #8080 correctly extended the producer-to-reduction dependency to all streams that copied gradients into the IPG bucket. That correctness dependency must remain.

This issue is not proposing removal of producer dependencies. It concerns the location and scope of the reverse dependency that protects buffer reuse.

Expected behavior

Launching reduction for bucket N should not block unrelated backward compute that produces bucket N+1.

The current/autograd stream should wait for bucket N's reduction only when the same physical IPG ping-pong buffer is about to be reused or overwritten, not whenever average_tensor() is called for the next bucket.

Proposed design direction

Track reduction completion per IPG buffer index:

1. Keep recording all producer streams for each bucket, as implemented by #8080.
2. Make reduction_stream wait on those producers before reading the bucket.
3. Record a reduction-complete event for the physical ping-pong buffer after all reduction-stream reads/copies using that buffer have been queued.
4. Before writing into a buffer index again, make the producer/current stream wait on that buffer's completion event.
5. Do not make the current stream wait on the entire reduction stream when merely enqueueing the next bucket.

This should preserve the buffer-reuse safety intended by #5606 and the multi-producer correctness of #8080 while allowing backward compute to overlap the preceding collective.

Special paths that likely need separate handling include oversized parameters, contiguous_gradients=false, CPU offload, MoE process groups, and the final bucket flush.

Validation requirements

A safe change should include:

• regression coverage for #5606's ping-pong buffer reuse race;
• regression coverage for #8061's multi-producer-stream case;
• rapid reuse of both IPG buffer indices;
• gradient, loss, and final-parameter parity against overlap_comm=false;
• a CUDA profiler test or reproducible benchmark showing that compute does not stop at the next bucket enqueue while the preceding collective is resident.

I can prepare a draft PR implementing per-buffer completion events if this direction is acceptable to the maintainers.

Start in the average_tensor() implementation and trace how reduction_stream, producer streams, and the IPG ping-pong buffer indices are managed; review #5606, #7805, #8061, and #8080 for the relevant correctness constraints. The issue proposes per-buffer reduction-completion events, but calls out special paths including oversized parameters, CPU offload, MoE process groups, and final bucket flush. Validate buffer-reuse and multi-producer races, compare gradients, loss, and parameters with overlap_comm=false, and profile that next-bucket compute overlaps the preceding collective.