[AutoTP] Enalbe HF colwise_gather_output to support lm_head replace#8146
[AutoTP] Enalbe HF colwise_gather_output to support lm_head replace#8146jinyouzhi wants to merge 5 commits into
colwise_gather_output to support lm_head replace#8146Conversation
386d2bf to
67285d6
Compare
![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: 386d2bf3bf
ℹ️ 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".
2ab4f47 to
c63d2d4
Compare
…iguration(`colwise_gather_output`/`colwise_rep`) - Introduced `gather_output` option in `TPLayerSpec` for column-parallel layers. - Added validation to prevent using gathered output with tied embeddings. - Updated `LinearLayer` to handle gathered output during forward pass. - Enhanced documentation to reflect new gathered column parallelism capabilities. - Added tests for gathered column parallelism scenarios and configurations. Signed-off-by: iLeGend <824040212@qq.com>
…e tp_plan extraction logic Signed-off-by: iLeGend <824040212@qq.com>
… for fallback behavior Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
| """Create column-parallel layer (AllReduce in backward).""" | ||
| if spec.gather_output and self.mp_size is not None and self.mp_size > 1: | ||
| output_dim = module.weight.shape[0] | ||
| if output_dim % self.mp_size != 0: |
There was a problem hiding this comment.
Does this condition mean TP on lm_head is only possible when vocab size divisible by number of TP ranks?
PKUWZP
left a comment
PKUWZP
left a comment
There was a problem hiding this comment.
Thanks for submitting the PR, a couple of suggestions:
-
I don't see this PR leverages online softmax and in-place computation, this is the common optimization that saves GPU memory significantly.
-
We need to run some profiling to measure the GPU footprint.
Happy to work closely together to have a separate PR for this.
|
|
||
| # Column-parallel layer with replicated output (e.g., an untied LM head) | ||
| TPLayerSpec( | ||
| patterns=[".*\\.lm_head$"], |
There was a problem hiding this comment.
pattern should be changed to *lm_head\.weight$ in order to be matched.
| """Create column-parallel layer (AllReduce in backward).""" | ||
| if spec.gather_output and self.mp_size is not None and self.mp_size > 1: | ||
| output_dim = module.weight.shape[0] | ||
| if output_dim % self.mp_size != 0: |
There was a problem hiding this comment.
can you also add a test to test checkpoint are correct in uneven case? i.e. even number vocab size, with 3 ranks.
3 participants
Change
Enable HF
colwise_gather_outputin tp_plan via AutoTP, which support vocab parallel forlm_head.For now, only available for untied lm head, tied fallback legacy.
Plan
Here is a plan to implement the issue "Extend AutoTP capabilities by integrating Liger Kernel (sharding LM heads and adopting Online Softmax)" in #8104
lm_head(untied part) via AutoTP with HFcolwise_gather_output.gather_output=Falseand parallel CE loss feature (integrating Liger kernel).