[pengdurice]

Part of: #8672 (Q4 2026 roadmap), "Enable training of large scale models".

Summary

AutoTP training today runs the norms, residual adds and other token-wise ops of every decoder layer on the whole
sequence on every tensor-parallel rank, keeps their activations at full size, and makes one all-reduce per
column-parallel linear in backward plus one per row-parallel linear in forward (7 per Qwen3 layer), all on the compute
stream. This PR adds tensor_parallel.sequence_sharding (default false): the hidden states between the attention
and MLP blocks are split along the sequence over the TP ranks, as in Megatron-LM's sequence parallelism.

"tensor_parallel": {
    "autotp_size": 4,
    "sequence_sharding": true
}

• Each block all-gathers its input; its row-parallel linear reduce-scatters its output instead of all-reducing it.
• In backward, one reduce-scatter per block sums the input gradients of all its column-parallel linears.
• Norms and residual adds run on 1/tp of the tokens, and their activations are 1/tp the size.
• sequence_sharding_overlap (default true): the column-parallel linears that read a block input run in one
  autograd node; backward reduce-scatters their summed input gradient while the weight-gradient GEMMs run; only this
  rank's part of the input is saved for backward and gathered again there.

The name avoids "sequence parallel", which in DeepSpeed means Ulysses (sequence_parallel_size), a different
scheme; the two are not combined.

Motivation

Today AutoTP training makes 7 all-reduces per Qwen3 layer on the compute stream (5 column-parallel input gradients in
backward, 2 row-parallel outputs in forward). On Qwen3-32B with activation recompute on 2 nodes, AutoTP got 62% slower
from TP 2 to TP 8 (Megatron-LM: 19%), and its best TP layout was 1.42x slower than Megatron-LM's best.

This PR is one of four. The others add copy-engine collectives for sequence sharding, vocab_parallel_embedding, and
ZeRO-1/2 overlap_param_gather. The table shows where DeepSpeed ends up with them: the Qwen2.5-72B and Qwen3-32B runs
use all four options, the Qwen3-8B run uses only overlap_param_gather. Each framework ran at its best measured
setting, all three in one allocation per model, 8 x H200 per node, random weights, 8,192-token sequences, bf16
compute, fp32 master weights and Adam states, bf16 gradient reduction, no recompute; seconds per step, mean of steps
4-13; MFU at 989 TFLOPS per GPU. DeepSpeed runs the Hugging Face model with Liger kernels for RMSNorm, rotary
embedding and SwiGLU (applied by our benchmark script, not part of these PRs; without them DeepSpeed was 4.02 against
Megatron-LM's 3.32 s per step on Qwen3-8B without TP). Megatron-LM uses Transformer Engine's fused kernels, TorchTitan
its compile regions.

Model, GPUs	DeepSpeed (AutoTP + these options)	Megatron-LM (main, Transformer Engine 2.19)	TorchTitan (main)
Qwen2.5-72B, 32	15.18 s, 50.3% (TP 8, micro-batch 2)	15.02 s, 50.9%	15.39 s, 49.6%
Qwen3-32B, 16	13.71 s, 52.6% (TP 4, micro-batch 2)	14.11 s, 51.1%	13.66 s, 52.8%
Qwen3-8B, 8	3.21 s, 54.3% (ZeRO-1 without TP)	3.22 s, 54.1%	3.18 s, 54.9%

AutoTP before these options: Qwen2.5-72B 25.22 s (micro-batch 1; micro-batch 2 did not fit), Qwen3-32B 20.54 s
(TP 8, micro-batch 1).

Results

Measured with this branch, 8 x H200 per node, bf16 compute (torch autocast), fp32 master weights and Adam states,
8,192-token sequences, no activation recompute, vocab_parallel_lm_head with the Liger vocabulary-parallel loss, and
Liger kernels for RMSNorm, rotary embedding and SwiGLU (applied by our benchmark script, not part of this PR). Master
and this branch ran in one allocation for each model. Step time is the mean of steps 4-13 of 13.

Setup	Master	This PR	Peak memory (rank 0), master -> this PR
Qwen3-8B, 1 node, AutoTP 2 x ZeRO-1 over 4, micro-batch 2, 32 sequences per step	3.701 s	3.524 s (4.8% less)	87.0 -> 78.0 GiB
Qwen3-8B, 1 node, AutoTP 8, micro-batch 4, 32 sequences per step	5.814 s	5.030 s (13.5% less)	80.8 -> 49.5 GiB
Qwen2.5-72B, 4 nodes (EFA), AutoTP 8 x ZeRO-1 over 4, 64 sequences per step	27.00 s at micro-batch 1; micro-batch 2 runs out of GPU memory	17.43 s at micro-batch 2 (35.4% less, 1.55x)	126.1 -> 114.1 GiB

The loss the benchmark script prints after the last step (global_loss) is within 2.5e-4 of master's in all three
setups.

Design

• Where it applies: apply_autotp_sequence_sharding runs in finalize_autotp after AutoTP has replaced the
  linears. A block is the parent module of a row-parallel linear (self_attn, mlp). Forward pre-hooks on the blocks
  gather their inputs; a pre-hook on the first decoder layer splits the hidden states and a forward hook on the last
  gathers them again, so the final norm, output layer and loss see the whole sequence (combine with
  vocab_parallel_lm_head to keep those sharded too).
• Checks at deepspeed.initialize(): every TP linear in the decoder layers belongs to a block; the blocks are in
  one numbered layer list and every layer has one; no column-parallel layer gathers its output; no DeepCompile
  autotp pass; no AutoTP layer class with its own forward (LinearAllreduceWithReplicatedInput,
  LmHeadLinearAllreduce, SubParamLinearAllreduce, SubParamLinearLayer), which would keep its all-reduce. The
  sequence length must be divisible by the TP size (checked in the forward pass).
• Token layout: each rank holds one contiguous chunk of the flattened batch x sequence tokens, so the collectives
  work on contiguous buffers and need no reordering copy at batch > 1. The ops between blocks act on each token
  alone, so any split is valid.
• Replicated parameters inside the decoder layers (norm weights, row-parallel biases) see only this rank's tokens;
  a gradient hook sums them over the TP group (skipped where AutoTP already registered one).
• ZeRO-3: partitioned weights use the plain path (separate collectives), not the one-autograd-node overlap.
• Core-file footprint: layers.py gets one branch in LinearLayer.forward and one in LinearAllreduce.forward,
  engine.py four lines in finalize_autotp; the rest is a new module.

Duplicate check

Searched issues and pull requests on 2026-10-18 for "AutoTP sequence parallel", "tensor parallel sequence parallel
reduce-scatter" and sequence_sharding: none adds sequence sharding to AutoTP. DeepSpeed's sequence-parallel work
(Ulysses, ALST, AutoSP) splits attention over ranks of its own and is a different scheme. Searched again on
2026-10-19: still none.

Related open pull request: #8851 (AutoTP mode per model) changes add_bias(output, self.bias) in
LinearAllreduce.forward to pass self.is_training_mode(). This PR changes the lines just above it, so whichever of
the two merges second needs a one-line rebase. If #8851 merges first, the new add_bias call in this PR's
LinearLayer.forward branch should pass self.is_training_mode() too.

The change

File	Change
deepspeed/module_inject/autotp_sequence_sharding.py (new, 381 lines)	the autograd functions (ScatterToSequenceShards, GatherFromSequenceShards, GatherToTensorParallelBlock, ReduceScatterFromTensorParallelBlock, ColumnGroupFromSequenceShards), the block hooks and apply_autotp_sequence_sharding
deepspeed/module_inject/layers.py (+19, -2)	the branches in LinearLayer.forward and LinearAllreduce.forward; two class attributes on TensorParallel_Layer
deepspeed/runtime/engine.py (+5)	the call in finalize_autotp
deepspeed/runtime/tensor_parallel/config.py (+15)	sequence_sharding, sequence_sharding_overlap
docs/_pages/config-json.md, docs/_tutorials/autotp-training.md	the options and a "Sequence Sharding" section
tests/unit/v1/autotp/test_autotp_sequence_sharding.py (new)	the tests below

Tests

TestAutoTPSequenceSharding (2 ranks, AutoTP 2, a 2-layer Qwen3 built from a config, untied; ZeRO 1 and 3; modes
plain and overlap): 3 SGD steps match the unsharded model's losses (within 1e-5), the decoder norms see half of
the sequence on each rank, and under ZeRO-1 a decoder norm weight matches the unsharded model after the steps (its
summed gradient). TestAutoTPSequenceShardingRejects: a sequence length not divisible by TP raises.

Hardware: 1 node, 4 x H200 (driver 580.126.09), torch 2.8.0+cu126, transformers 5.15.1, pytest 9.1.1, pytest-forked
1.7.5.

cd tests && pytest --forked unit/v1/autotp/test_autotp_sequence_sharding.py unit/v1/autotp/test_autotp_training.py \
    unit/v1/autotp/test_tp_plan_real_models.py unit/module_inject/test_tp_partition_config_path.py
# 103 passed, 2 skipped (the existing BF16-optimizer overflow skips), 59 deselected (sequential) in 13:24
pytest --forked -m sequential unit/v1/autotp/test_autotp_training.py
# 59 passed in 12:49

In the study, with this code under another name: losses equal TP 1 within 1.9e-6 in fp32 over 5 AdamW steps at batch 1
and 2 (within 1.1e-3 in bf16, 9.0e-4 with autocast and the vocabulary-parallel loss); every parameter gradient equals
the unsharded model's within a relative 4.3e-7 to 1.1e-6 (fp32, one backward, ZeRO-0).

pre-commit run --files passes on every changed file.

Compatibility

Cell	Validated
AutoTP 2, ZeRO 1 and 3, fp32 (unit tests)	yes
AutoTP 2-8, ZeRO-1, bf16 + torch autocast, Qwen3-8B / 32B, Qwen2.5-72B, 1-4 nodes	yes: this branch on Qwen3-8B (AutoTP 2 and 8, 1 node) and Qwen2.5-72B (AutoTP 8, 4 nodes); Qwen3-32B in the study tree only
ZeRO 2	not run in the unit tests; ZeRO-2 runs in the study's 4-GPU equivalence script
fp16	not run
Activation checkpointing	not run (the study runs used none)
Checkpoint save / load	not run; the option changes no parameter or optimizer state
Ulysses / ALST sequence parallelism together with it	not supported, not checked at initialization
Pipeline parallel, CPU offload, torch.compile, the DeepCompile autotp pass	not run; the DeepCompile pass is rejected
Models other than Qwen3 / Llama-style decoders	not run; the checks above reject unsupported layouts
Accelerators other than CUDA	not run; uses deepspeed.comm collectives only

AI assistance

This change, its tests and this description were prepared with AI assistance (Claude).

🤖 Generated with Claude Code