Fix SequenceTiledCompute backward for empty trailing shards - #8434
Fix SequenceTiledCompute backward for empty trailing shards#8434taking-lying-flat wants to merge 1 commit into
Conversation
Signed-off-by: taking-lying-flat <1615405@qq.com>
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. |
![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: 781f2e2cc6
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| with torch.no_grad(): | ||
| shard_step = math.ceil(seqlen / shards) | ||
| ctx.shard_step = shard_step |
There was a problem hiding this comment.
The reviewed commit is a non-merge commit, but its message contains no Signed-off-by trailer, violating the repository's mandatory commit policy and risking rejection by DCO/CI checks; recreate the commit using --signoff.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
1 participant
Description
sequence_tiled_computecan finish forward and then raiseIndexErrorin backward whentorch.chunkreturns fewer pieces than requested. Forseqlen=5, shards=4, forward calls the function on slices of lengths[2, 2, 1, 0], while backward creates only three chunks and still indexes four. This also occurs without ZeRO.Save the forward slice step and reuse it to slice every sharded input during backward. Clamp trailing empty-slice offsets to the sequence length so the gradient
narrow()calls remain valid. This preserves the existing forward calls, requested shard count, reduction behavior, and final-shard gradient signaling.The regression test runs two SGD steps against an untiled PyTorch MLP and compares outputs, input gradients, parameter gradients, and updated weights. It covers batch sizes 1/2,
(seqlen, shards)pairs(5, 4),(2, 4),(7, 4), and(8, 4), multiple sharded inputs, and tensor/sum/mean outputs. The callback supports empty inputs, as required by the existing forward behavior.Validation
Hardware: NVIDIA RTX A1000 Laptop GPU, 4 GiB, driver 595.84; Python 3.12.13, PyTorch 2.13.0+cu130, CUDA 13.0.
IndexError; the 12 control variants passed.DS_ACCELERATOR=cuda python -m pytest -q tests/unit/ulysses_alst/test_tiled_compute.py -k sequence_tiled_compute_shard_boundaries --tb=short: 24 passed in FP32, including both optimizer steps.TestTiledComputeZeRO-2 cases: 2 passed on the GPU.reduce_bucket_size=65536override: 2 passed. The override is not included in this PR.pre-commit run --files deepspeed/runtime/sequence_parallel/ulysses_sp.py tests/unit/ulysses_alst/test_tiled_compute.py: passed.python ci/tests_fetcher.py --base origin/master: no affected tests in its configuredtests/unit/v1scope; the ALST tests above were run explicitly.