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

[cjxjxjx]

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():

DeepSpeed/deepspeed/runtime/zero/stage_1_and_2.py

Lines 1119 to 1122 in b214589

	new_grad_tensor = bucket.buffer[bucket.index].narrow(0, bucket.elements, param.numel())
	new_grad_tensor.copy_(
	grad_reduc.view(-1) if not self.zenflow else grad_reduc.permute(
	*reversed(range(grad_reduc.ndim))).contiguous().view(-1))

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:

DeepSpeed/deepspeed/runtime/zero/stage_1_and_2.py

Line 1575 in b214589

self.average_tensor(bucket.buffer[bucket.index].narrow(0, 0, bucket.elements), comm_dtype)

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:

DeepSpeed/deepspeed/runtime/zero/stage_1_and_2.py

Lines 1230 to 1238 in b214589

	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()

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

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:

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:

• Fix overlap communication of ZeRO stage 1 and 2 #5606 fixed the bidirectional wait between current/default stream and reduction_stream in average_tensor().
• fix #7188  #7371 / [BUG]when use 'overlap_comm:True' w/ 'contiguous_gradients:True', grad_norm is NaN #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:

# 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.

[hwchen2017]

Can you provide an example to reproduce? Thanks!

[cx2009]

Is this happening in production or just local testing?

[cjxjxjx]

Can you provide an example to reproduce? Thanks!

Thanks for looking into this. @hwchen2017

This happens in our real training workload, not just a toy local test. The original model/data are not publicly shareable, and the issue is hard to reduce to a small standalone script because it depends on the interaction between:

• ZeRO stage 1/2
• overlap_comm: true
• contiguous_gradients: true
• reduce_scatter: true
• bf16
• torch.compile
• autograd / compiled-backward hooks issuing contiguous-gradient bucket copies on multiple device streams

A naive small Python reproducer usually does not trigger the same timing, because all gradient hooks may run on the same current stream.

What I can provide now is the evidence matrix from our run:

• torch.compile = false: no NaN
• overlap_comm = false: no NaN
• adding a device-wide synchronize before the ZeRO reduction reads the IPG bucket: no NaN
• adding stream waits from the actual IPG bucket copy streams to reduction_stream: no NaN
• without the additional stream dependency: loss becomes NaN very early, from step 1

The profiler trace shows that the copy_ operations that fill the same contiguous gradient/IPG bucket are issued on multiple streams, while average_tensor() only waits for the stream current when average_tensor() is entered. Therefore reduction_stream may read the bucket before all producer streams for that bucket have completed.

[cjxjxjx]

Is this happening in production or just local testing?

It is from a real training workload / production-like training setup. The run is local to our cluster, but the failure is not from a synthetic unit test. @cx2009