[AutoTP] Add support for vocab-parallel LM head and related configurations - #8309

Open

jinyouzhi wants to merge 6 commits into

deepspeedai:masterfrom

jinyouzhi:vocab_parallel_ce

Commits on Sep 8, 2026

Add vocab-parallel LM loss

jinyouzhi
and
Copilot
committed
8cf638a
Fix vocab-parallel integration

jinyouzhi
and
Copilot
committed
bf8b7e2
Preserve sequence loss gradients

jinyouzhi
and
Copilot
committed
8f03c27
Document sequence loss gradients

jinyouzhi
and
Copilot
committed
97078ce
Fix vocab-parallel heuristic integration

jinyouzhi
and
Copilot
committed
a3b4405
Validate vocab head before AutoTP

jinyouzhi
and
Copilot
committed
a47d7ab