DeepSpeed

5f9f48f0 - [AutoTP] Implement VocabParallelEmbedding and enhance tied model support (#8582)

[AutoTP] Implement VocabParallelEmbedding and enhance tied model support (#8582) Implement PR-C in #8173 This pull request introduces support for vocabulary-parallel (vocab-parallel) embedding layers that are tied to vocab-parallel output heads in DeepSpeed's AutoTP module. It ensures correct sharding and replacement of tied embedding/output head pairs, adds a new `VocabParallelEmbedding` layer, and improves the handling and detection of tied embeddings in both heuristic and HuggingFace tp_plan-driven partitioning paths. Key changes include: **Vocab-parallel embedding support and replacement logic:** * Added detection and handling of tied vocab-parallel embedding layers, ensuring that tied embeddings are replaced with `VocabParallelEmbedding` and share the same sharded weight as their tied output head. This avoids double-partitioning and maintains correct weight sharing. (`auto_tp.py`, `layers.py`) [[1]](diffhunk://#diff-3e2a27cf6940a22f2abf4c9e19d007d3824b6a97afb8d8cc028c7efa9ddd7e44L231-R246) [[2]](diffhunk://#diff-3e2a27cf6940a22f2abf4c9e19d007d3824b6a97afb8d8cc028c7efa9ddd7e44L554-R595) [[3]](diffhunk://#diff-3e2a27cf6940a22f2abf4c9e19d007d3824b6a97afb8d8cc028c7efa9ddd7e44R735-R740) [[4]](diffhunk://#diff-0df5862062e8a6030b3b4c689914f06701a7023014f94d353cbd756e8a898776R1057-R1150) **New layer implementation:** * Introduced the `VocabParallelEmbedding` class, which shards embedding weights along the vocabulary dimension and supports tying to a vocab-parallel output head, mirroring Megatron's approach. (`layers.py`) [[1]](diffhunk://#diff-0df5862062e8a6030b3b4c689914f06701a7023014f94d353cbd756e8a898776R1057-R1150) [[2]](diffhunk://#diff-0df5862062e8a6030b3b4c689914f06701a7023014f94d353cbd756e8a898776L30-R30) **Partitioning and plan conversion improvements:** * Updated the tp_plan converter to clarify that "embedding_rowwise" now allows tied vocab-parallel embedding/output head pairs to be sharded together, instead of skipping them entirely. (`tp_plan_converter.py`) [[1]](diffhunk://#diff-58109fbb6814482eae2a122647ad8ca8b2d862f06885419f1848479ffc91f605L33-R35) [[2]](diffhunk://#diff-58109fbb6814482eae2a122647ad8ca8b2d862f06885419f1848479ffc91f605L63-R67) **Engine and configuration integration:** * Modified the engine's AutoTP partitioning logic to detect when a HuggingFace `tp_plan` requests vocab-parallel embedding for a tied output head, enabling vocab-parallel sharding and distributed loss automatically. Added stricter checks to ensure all partitioning paths agree on vocab-parallel head requirements. (`engine.py`) [[1]](diffhunk://#diff-e6635a81d2c2bf0938b5f83b1c4945e0f344e0106191a6960df8ee4aa64cc55fL1050-R1050) [[2]](diffhunk://#diff-e6635a81d2c2bf0938b5f83b1c4945e0f344e0106191a6960df8ee4aa64cc55fL1087-R1089) [[3]](diffhunk://#diff-e6635a81d2c2bf0938b5f83b1c4945e0f344e0106191a6960df8ee4aa64cc55fR1103-R1121) [[4]](diffhunk://#diff-e6635a81d2c2bf0938b5f83b1c4945e0f344e0106191a6960df8ee4aa64cc55fL1120-R1151) [[5]](diffhunk://#diff-e6635a81d2c2bf0938b5f83b1c4945e0f344e0106191a6960df8ee4aa64cc55fL1162-R1185) **Validation and warnings:** * Improved warning and error messages to clarify when embedding or output head partitioning specs are overridden due to vocab-parallel tying, and removed obsolete validation logic that previously forbade tied output/embedding weights for vocab-parallel heads. (`auto_tp.py`) [[1]](diffhunk://#diff-3e2a27cf6940a22f2abf4c9e19d007d3824b6a97afb8d8cc028c7efa9ddd7e44L575-R625) [[2]](diffhunk://#diff-3e2a27cf6940a22f2abf4c9e19d007d3824b6a97afb8d8cc028c7efa9ddd7e44L599) --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>