Skip to content

Fix lost TP all-reduce results for non-contiguous tensors #8837

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jinyouzhi wants to merge 3 commits into
base: master
Choose a base branch
from
+45 −4
Open

Fix lost TP all-reduce results for non-contiguous tensors #8837

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0b7f6d4
Keep TP all-reduce results for non-contiguous tensors
jinyouzhi Oct 16, 2026
f8efe06
Merge branch 'master' into fix-tp-contiguous-allreduce
jinyouzhi Oct 17, 2026
7ebe0db
Merge branch 'master' into fix-tp-contiguous-allreduce
jinyouzhi Oct 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions deepspeed/module_inject/layers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,8 @@ def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, is_inferenc
if is_inference_mode:
dist.inference_all_reduce(input, group=group)
else:
dist.all_reduce(input.contiguous(), group=group)
input = input.contiguous()
dist.all_reduce(input, group=group)
return input

@staticmethod
Expand Down Expand Up @@ -276,8 +277,8 @@ def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, weight, bia
def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:

input, weight = ctx.saved_tensors
grad_input = grad_output.matmul(weight)
handle = dist.all_reduce(grad_input.contiguous(), group=ctx.group, async_op=True)
grad_input = grad_output.matmul(weight).contiguous()
handle = dist.all_reduce(grad_input, group=ctx.group, async_op=True)
grad_weight = grad_output.view(-1, grad_output.shape[-1]).t().matmul(input.view(-1, input.shape[-1]))
grad_bias = grad_output.sum(0) if ctx.use_bias else None
handle.wait()
Expand Down Expand Up @@ -310,7 +311,8 @@ def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
if not pg_utils.is_parallel(ctx.group):
return None, grad_output

dist.all_reduce(grad_output.contiguous(), group=ctx.group)
grad_output = grad_output.contiguous()
dist.all_reduce(grad_output, group=ctx.group)
return None, grad_output


Expand Down
39 changes: 39 additions & 0 deletions tests/unit/module_inject/test_tp_noncontiguous_reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch

import deepspeed.comm as dist
from deepspeed.accelerator import get_accelerator
from deepspeed.module_inject.layers import ColumnParallel, RowParallel

from unit.common import DistributedTest


class TestTensorParallelNonContiguousReduce(DistributedTest):
"""The TP all-reduce must reach the returned tensor even when it is not contiguous.

Reducing a temporary from ``x.contiguous()`` silently drops the result for non-contiguous ``x``.
"""
world_size = 2

def test_row_parallel_forward(self):
rank = dist.get_rank()
x = torch.full((3, 2), rank + 1.0, device=get_accelerator().device_name()).t()
assert not x.is_contiguous()

out = RowParallel.apply(dist.new_group([0, 1]), x, False)

assert torch.all(out == 3.0)

def test_column_parallel_backward(self):
group = dist.new_group([0, 1])
rank = dist.get_rank()
x = torch.ones(3, 2, device=get_accelerator().device_name(), requires_grad=True)

# The transpose makes the incoming gradient non-contiguous.
out = ColumnParallel.apply(group, x).t()

@chatgpt-codex-connector chatgpt-codex-connector Bot Oct 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Exercise the reduction through an AutoTP training loop

test_column_parallel_backward calls ColumnParallel.apply directly, but actual AutoTP training reaches this reduction through LinearLayer.forward's process-group, overlap, and deferred-collective routing, so this test can pass without showing that a transposed downstream loss receives the fix through the user-facing path. Add a minimal injected-model forward/backward test whose layer output is transposed before the loss; training-loop contract changes require integration coverage rather than only primitive unit coverage.

AGENTS.md reference: AGENTS.md:L76-L81

Useful? React with 👍 / 👎.

@jinyouzhi jinyouzhi Oct 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. I tried this, but an AutoTP-level test cannot reach the bug, so it would not guard the fix.

In-tree, these primitives only receive freshly computed tensors: LinearLayer / SubParamLinearLayer wrap their input in ColumnParallel, whose backward gradient comes from matmul's backward, and LinearAllreduce / VocabParallelEmbedding / SubParamLinearAllreduce feed RowParallel a matmul / embedding output. Those are contiguous whatever the caller does with the layer's input or output. I checked that transposed, expanded and permuted inputs and a transposed loss all still deliver a contiguous gradient to ColumnParallel.backward.

To confirm, I ran the suggested test (a LinearLayer -> LinearAllreduce stack over a 2-rank TP group, output transposed before the loss, input gradient compared with an unsharded reference) against the code before this fix:

  • integration test: passes (cannot detect the bug)
  • test_tp_noncontiguous_reduce.py: both tests fail (each rank keeps only its own value)

So the bug is latent for AutoTP layers today. It is reachable whenever these autograd functions get a non-contiguous tensor, e.g. a custom TP layer or a future caller, which is exactly the contract the primitive tests pin down. I kept the regression tests at that level, since a training-loop test would pass with or without the fix.

(out * (rank + 1)).sum().backward()

assert torch.all(x.grad == 3.0)
Loading