[AutoTP] Add support for vocab-parallel LM head and related configurations - #8309
[AutoTP] Add support for vocab-parallel LM head and related configurations#8309jinyouzhi wants to merge 6 commits into
Conversation
|
Hi @jinyouzhi when this PR is ready for review just ping me, thanks! |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
ab4139a to
8cf638a
Compare
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Require an explicit sequence-parallel group when the legacy wrapper gathers losses, and document the distinct backward conventions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8f03c27ea0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model)) | ||
| setattr(model, "ds_autotp_parsed", True) | ||
| finalize_autotp(attach_uc_metadata=True) |
There was a problem hiding this comment.
When neither a partition config nor a convertible HuggingFace TP plan exists, this branch only invokes replace_transformer_layer for parsed transformer-block classes; its training-mode set_lm_head path returns without touching the root output head. Consequently, finalize_autotp finds no VocabParallelLinear, so the documented autotp_size plus vocab_parallel_lm_head configuration silently leaves the full head replicated and never installs the distributed loss. Explicitly replace the model's output head before finalizing this fallback; an end-to-end training-path test would expose the no-op configuration.
AGENTS.md reference: AGENTS.md:L35-L36
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in a3b440523 . The heuristic AutoTP path now creates a dedicated AutoTP instance for the output head and explicitly replaces the single supported lm_head / embed_out before finalization. It also fails clearly when no supported head or multiple candidate heads are found, instead of silently leaving the full head replicated. I added a 2-rank end-to-end regression using a real Llama model through deepspeed.initialize , forward, distributed causal-LM loss, and backward.
| for _, module in self.module.named_modules(): | ||
| if isinstance(module, nn.Embedding) and getattr(module, "weight", None) is lm_head.weight: | ||
| raise ValueError("A no-gather vocab-parallel LM head requires untied embedding and output weights") |
There was a problem hiding this comment.
For the usual module order where a tied embedding precedes lm_head, an HF/custom plan can replace the embedding first; _slice_embedding creates a new Parameter, so this later identity scan no longer sees that the original head and embedding were tied. The configuration is then incorrectly accepted and silently breaks weight sharing instead of raising the documented error. Capture and validate candidate ties before traversal mutates either module; integration coverage using an actual tied model and embedding plan is also required.
AGENTS.md reference: AGENTS.md:L35-L36
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in a3b440523 . AutoTP now records the identities of vocab-head modules tied to embeddings when it is constructed, before traversal can replace an embedding parameter. _validate_untied_vocab_head() checks this original tie information as well as the current parameter identity. A regression test covers an embedding_rowwise plan followed by a colwise tied lm_head and verifies that the configuration is rejected before weight sharing can be broken.
Replace the output head in heuristic AutoTP and preserve original tied-weight evidence before embedding traversal mutates parameters. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Resolve and validate the heuristic output head before transformer partitioning, cache the candidate for replacement, and skip tied-weight scans when the feature is disabled. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Hi, this PR is ready for review now. Thank you for your patience! @delock During development, I identified and fixed some issues related to SP in a separate PR: #8457. |
2 participants
Motivation
Changes
This pull request adds support for vocabulary-parallel language modeling (LM) heads in DeepSpeed's tensor parallelism (TP) system. It introduces a new
VocabParallelLinearlayer that keeps logits sharded across TP ranks for untied LM heads and installs a compatible distributed loss automatically. The changes ensure correct partitioning, validation, and integration with the rest of the DeepSpeed engine and configuration system.Key changes include:
Vocabulary-parallel LM head support
vocab_parallel_lm_headconfiguration option inTPTrainingConfigand passed it through the relevant APIs to enable this feature. [1] [2] [3] [4] [5]auto_tp.pyto detect and correctly replace untied LM heads with the newVocabParallelLinearlayer, along with validation to ensure the vocab head is not tied to the embedding. [1] [2] [3] [4]New layer and loss integration
VocabParallelLinearlayer inlayers.py, which ensures logits remain sharded and provides partitioning metadata. [1] [2]Miscellaneous
deepspeed.sequencemodule's public API.These changes collectively enable efficient, distributed training of models with untied, sharded vocabulary projection heads and ensure compatibility with DeepSpeed's tensor parallelism ecosystem.
Limitations