[AutoTP] Complete uneven sharding and universal checkpoint support - #8185

Open

jinyouzhi wants to merge 28 commits into

deepspeedai:masterfrom

jinyouzhi:uneven

Commits on Aug 2, 2026

Enhance AutoTP Uneven Column Parallel Gathered Output
jinyouzhi
committed
2936b72
Complete AutoTP uneven sharding for row parallelism and checkpoints

jinyouzhi
and
Copilot
committed
847c4e4
Fix tied parameter shape collection in universal checkpoints
jinyouzhi
committed
8acd22a
make GatherFromTensorParallelRegion partition sizes recomputed locally
jinyouzhi
committed
be4c6d7
Remove the diagnostics or change as logger.info
jinyouzhi
committed
790dab9
Freeze AutoTP partition sizes at layer construction
jinyouzhi
committed
000c92f

Commits on Aug 3, 2026

Merge remote-tracking branch 'origin/master' into uneven
jinyouzhi
committed
f3bf28d

Commits on Aug 4, 2026

Add shape consistency guard for merging tied parameters in pipeline replicas
jinyouzhi
committed
73f1faa
Keep the sub-grain remainder when sharding a dimension
jinyouzhi
committed
ad5c1af
Cleanup the remaining AutoTP lm_head replacement traces
jinyouzhi
committed
ad58bc3
fix yapf formatting
jinyouzhi
committed
1e7cde3
Add support for uneven sub-parameter shard widths in AutoTP
jinyouzhi
committed
5bb9a43
Enhance AutoTP checkpoint handling for uneven sub-parameter sizes
jinyouzhi
committed
e497ab2
Allow excess TP ranks for attention
jinyouzhi
committed
0d41b1c
Support uneven AutoTP sharding
jinyouzhi
committed
8445a00
Merge remote-tracking branch 'origin/master' into HEAD
jinyouzhi
committed
000a46a
Size AutoTP shards by the tensor parallel group rank

jinyouzhi
and
Copilot
committed
b0ef6c6
Partition and gather fused layers by their frozen widths

jinyouzhi
and
Copilot
committed
692d54e
Shard attention head helpers by the tensor parallel group

jinyouzhi
and
Copilot
committed
e41e43e
Convert and restore AutoTP shards at their recorded widths

jinyouzhi
and
Copilot
committed
c371163

Commits on Aug 6, 2026

Merge branch 'master' into uneven
jinyouzhi
authored
f6a03ac
add tp3 round trip test for uneven train & ucp scenario

committed
48f438e
uc_info only needs to be computed when tp_degree>1
jinyouzhi
committed
622faac

Commits on Aug 8, 2026

Merge branch 'master' into uneven
delock
authored
f5db391