Repository navigation
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
dist.all_reduce(x.contiguous()) reduces a temporary copy when x is not contiguous, so the result is lost. Keep the contiguous tensor and return it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0b7f6d4d96
ℹ️ 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".
| 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() |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
This branch has not been deployed
1 participant
Problem
dist.all_reducereduces its argument in place. Three tensor-parallel autograd functions indeepspeed/module_inject/layers.pycalldist.all_reduce(x.contiguous(), ...)and then returnx. Whenxis not contiguous,.contiguous()returns a copy. The reduction lands in that copy and is discarded, so the TP sum is silently skipped:RowParallel.forward(training path): the forward output is not reduced.ColumnParallel.backward: the input gradient is not reduced, e.g. when the output is transposed before the loss.AsyncColumnParallel.backward:grad_inputcomes frommatmuland is normally contiguous, but it uses the same pattern, so it is fixed the same way.In-tree AutoTP layers only hand these functions
matmul/ embedding outputs and their gradients, which are contiguous, so the bug is latent there today. It is hit by any caller that passes a non-contiguous tensor, such as a custom TP layer. A training-loop test throughLinearLayerpasses with or without this fix, so the regression tests target the autograd functions directly.Fix
Assign the contiguous tensor first, reduce it, and return it.
Tests
tests/unit/module_inject/test_tp_noncontiguous_reduce.py(2 ranks) feeds a transposed input toRowParalleland a transposed gradient toColumnParallel, and checks that the result is the sum over ranks (1 + 2 = 3).test_tp_shard.py(17 passed).Run on 2× Intel Arc Pro B70 (XPU, XCCL), torch 2.14.0.dev20260804+xpu.
Contribution checklist
Related work. A search of open and closed PRs found no existing fix for this.
Sibling paths checked.
git grep 'all_reduce(<x>.contiguous()'overdeepspeed/finds only the three runtime sites fixed here, all inmodule_inject/layers.py. The fourth match,ColumnParallel.symbolic, only runs during ONNX tracing and is left unchanged. Theinference_all_reducebranch ofRowParallelpasses the tensor itself, so it is not affected.Commands and results. Hardware and software: 2× Intel Arc Pro B70 (XPU, XCCL), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu.
test_tp_shard.py: 17 passed).Numerics, performance, memory. Nothing changes for contiguous inputs, which is the case for every in-tree AutoTP layer:
.contiguous()returns the same tensor. For non-contiguous inputs, the result is now the reduced tensor, which is the intended behavior.Compatibility matrix. Validated: TP=2, fp32, at the autograd-function level. ZeRO stage and offload do not interact with this code. Not validated: CUDA, HPU.
AI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.