[BUG] ZeRO stage 1/2 overlap_comm only waits current stream, but contiguous gradient bucket copies may come from multiple streams

## Describe the bug On current `master` (`b2145896daeddefcb70d9301c04c4d0a03b71bb6` as of 2026-06-12), the ZeRO stage 1/2 overlap communication path appears to assume that the stream current at `average_tensor()` time is the same stream that produced all contiguous gradient bucket writes. This assumption can be false when the gradient hooks are triggered from different autograd/compiled-backward streams. In a profiler trace, the `copy_` operations that write parameter gradients into the IPG buffer are distributed across multiple device streams, while `average_tensor()` only makes the reduction stream wait for `get_accelerator().current_stream()` at the time `average_tensor()` is called. As a result, `reduction_stream` may start reading the IPG buffer (`div_`, `narrow`, `cat/concat`, all-reduce/scatter path) before all streams that wrote slices of that same bucket have completed. ## Relevant code on master The bucket copy is done in `reduce_independent_p_g_buckets_and_remove_grads()`: https://github.com/deepspeedai/DeepSpeed/blob/b2145896daeddefcb70d9301c04c4d0a03b71bb6/deepspeed/runtime/zero/stage_1_and_2.py#L1119-L1122 ```python new_grad_tensor = bucket.buffer[bucket.index].narrow(0, bucket.elements, param.numel()) new_grad_tensor.copy_( grad_reduc.view(-1), non_blocking=self.device == get_accelerator().device_name()) ``` The bucket is later reduced here: https://github.com/deepspeedai/DeepSpeed/blob/b2145896daeddefcb70d9301c04c4d0a03b71bb6/deepspeed/runtime/zero/stage_1_and_2.py#L1575 ```python self.average_tensor(bucket.buffer[bucket.index].narrow(0, 0, bucket.elements), comm_dtype) ``` But `average_tensor()` only waits for the current stream at the moment it is called: https://github.com/deepspeedai/DeepSpeed/blob/b2145896daeddefcb70d9301c04c4d0a03b71bb6/deepspeed/runtime/zero/stage_1_and_2.py#L1230-L1238 ```python def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dtype): if self.overlap_comm: stream = self.reduction_stream if not get_accelerator().resolves_data_dependency(): stream.wait_stream(get_accelerator().current_stream()) get_accelerator().current_stream().wait_stream(stream) else: stream = get_accelerator().current_stream() ``` This is not sufficient if the same `bucket.buffer[bucket.index]` was filled by `copy_` calls issued on multiple earlier streams. ## Observed behavior Configuration characteristics: - ZeRO stage 1/2 code path - `overlap_comm: true` - `contiguous_gradients: true` - `reduce_scatter: true` - bf16 training - `torch.compile` enabled - observed on a PrivateUse1/MLU backend Symptoms: - loss becomes `NaN` very early, starting from step 1 in our run - disabling `torch.compile` avoids the issue - disabling `overlap_comm` avoids the issue - adding a device-wide synchronize in the DeepSpeed reduction path avoids the issue, but is too heavy - adding stream synchronization from the actual gradient-copy streams to `reduction_stream` also avoids the issue Profiler trace observation: - the gradient bucket `copy_` operations corresponding to the DeepSpeed contiguous-gradient path are emitted on more than one device stream - later, `average_tensor()` / concat / reduction reads the same IPG bucket on the reduction stream - the existing wait only covers the stream returned by `get_accelerator().current_stream()` when `average_tensor()` is entered, not all streams that previously wrote slices into the bucket Conceptually, this can happen as: ```text stream A: copy grad slice A -> ipg bucket stream B: copy grad slice B -> ipg bucket stream C: calls average_tensor(); current_stream == stream C DeepSpeed today: reduction_stream waits stream C Missing dependency: reduction_stream should also wait stream A and stream B before reading the bucket ``` ## Why this is different from existing overlap_comm fixes This is related to, but not fully covered by, previous fixes: - #5606 fixed the bidirectional wait between current/default stream and `reduction_stream` in `average_tensor()`. - #7371 / #7188 addressed a synchronization issue around double IPG-buffer swapping. The remaining issue here is narrower: even on current master, `average_tensor()` still assumes that one current stream represents all producers of the current contiguous gradient bucket. With compiled autograd / torch.compile, gradient hooks can run under different current streams, so the bucket can have multiple producer streams. ## Possible fix direction One conservative fix is to record the unique streams used for the IPG bucket copy operations, then make `reduction_stream` wait for those streams before `average_tensor()` reads the bucket. Sketch: ```python # after new_grad_tensor.copy_(...) self._record_ipg_copy_stream(bucket) # inside average_tensor(), after switching to reduction_stream self._wait_for_ipg_copy_streams(bucket) ``` Where `_record_ipg_copy_stream()` stores unique `get_accelerator().current_stream()` values for the current IPG bucket, and `_wait_for_ipg_copy_streams()` calls `current_stream.wait_stream(copy_stream)` for each recorded stream not already covered by the existing `stream.wait_stream(current_stream)`. This avoids a device-wide synchronize and only adds stream dependencies for actual producer streams of the bucket. ## Expected behavior `overlap_comm: true` with `contiguous_gradients: true` should not allow the reduction stream to read the IPG bucket before all streams that wrote gradient slices into that bucket have completed.