Fix sequence overlap backward gradient permutation (#8342)
## Description
Sequence-parallel overlap backward must wait for the asynchronous Q/K
all-to-all and return the post-permuted gradients before the Q/K
producers run backward. Returning those gradients from hooks on the
producers' parent nodes instead can change gradient sizes for
transpose/view inputs or bypass the producers' backward transformations.
Apply @tohtana's patch from
https://github.com/taking-lying-flat/DeepSpeed/pull/1, preserving his
authorship. Register hooks directly on the Q/K tensors and return the
completed, post-permuted gradients. Also use `ctx.stream is not None` so
the stream check does not invoke CUDA stream equality operators.
The included regression test compares synchronous and overlap gradients
for transpose-produced Q/K/V using a mocked two-rank all-to-all with the
real reshape, permutation, and autograd paths.
## Validation
- Included regression test: `.venv-cherry/bin/python -m pytest --forked
-q tests/unit/v1/sequence_parallelism/test_overlap_grad_permutation.py`
— 1 passed.
- Additional local regression matrix: 10 passed across CPU and CUDA,
covering scaling, transpose, view, shared packed-QKV unbind, and leaf
tensors. The asynchronous collective mock defers writing its output
until `wait()`; gradients match synchronous execution exactly. CUDA
cases use real CUDA streams; multi-GPU NCCL was not tested locally.
- Pre-commit checks passed for both changed files.
- Restored the missing sign-off on the earlier test-removal commit; all
PR commits now include sign-offs.
Merged current master to resolve the stream-check conflict and follow
the test-directory migration. Removed the test mock for
`get_num_kv_heads`, which upstream deleted. All validation above passed
again after integration; CUDA hardware: NVIDIA RTX A1000 Laptop GPU (4
GB), PyTorch 2.13.0+cu130.
---------
Signed-off-by: taking-lying-flat <1615405@qq.com>
Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
Co-authored-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
