[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 9 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>
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.
Thank you for your suggestions! You‘re absolutely right, this PR is far from the target. I've refined the PR description, and I'll submit a separate PR for those optimizations. |
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Suggest profile GPU foot print with the separate PR. Online softmax and in-place computate could save memory for OPSD teacher TP sharding. |
Great! list the keypoints into an issue to track #8173 |
3 participants
Changes
Enable HF
colwise_gather_outputintp_planthrough AutoTP, allowing vocab parallelism for an untiedlm_head.Known Limitations
lm_head. Tiedlm_headfalls back to the legacy implementation.