Split token_embs and lm_head weights #2252
Conversation
| 0, // only used for position encoding | ||
| token_num, | ||
| token_num, | ||
| 1, | ||
| hidden_units_, | ||
| stream_); | ||
| if (tensor_para_.world_size_ > 1) { | ||
| NcclGuard nccl_guard(tensor_para_, stream_); | ||
| ftNcclAllReduceSum(decoder_input, decoder_input, token_num * hidden_units_, tensor_para_, stream_); |
There was a problem hiding this comment.
Why use allreduce? Shouldn't we use allgather?
There was a problem hiding this comment.
If using allgather, the weight and buffer layout should change and after allgather it need another permute op.
And another reason which confuses me is in my test, the reduce sum is faster than allgather.
|
We need to benchmark the ar/ag case on different systems (NVLink/PCIe) first. |
|
|
@irexyc bus bandwidth of all-reduce and all-gather is computed differently. |
|
|
|
| // check converted file with tp | ||
| auto should_exist = prefix + "." + std::to_string(tensor_para_size - 1) + ".weight"; | ||
| auto should_not_exist = prefix + "." + std::to_string(tensor_para_size) + ".weight"; | ||
| if (!std::filesystem::exists(should_exist) || std::filesystem::exists(should_not_exist)) { |
There was a problem hiding this comment.
We can put TP check in turbomind.py
| embedding_table_size_ = vocab_size_padded_ * hidden_units_ / tensor_para_size_; | ||
| } | ||
| else { | ||
| embedding_table_size_ = vocab_size_padded_ * hidden_units_; |
There was a problem hiding this comment.
Is this case expected to work correctly?
There was a problem hiding this comment.
I can't find models that hidden size % 8 != 0
There was a problem hiding this comment.
I mean we just FT_CHECK the condition here and eliminate the else branch since it's not going to work when hidden_units_ % tensor_para_size_ != 0
Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily receiving feedbacks. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.
Motivation
Please describe the motivation of this PR and the goal you want to achieve through this PR.
Modification
Please briefly describe what modification is made in this PR.
BC-breaking (Optional)
Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.
Use cases (Optional)
If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.
Checklist