support context parallel #3951
support context parallel #3951
Conversation
* [Fix] device args in chat cli when using pytorch engine * [Fix] change device into device_type in chat cli
|
may resolve the build error on windows platform |
| const int tp_rank_; | ||
| const DataType data_type_; | ||
| const bool debug_; | ||
| const bool is_driver_; |
There was a problem hiding this comment.
Consider renaming is_driver_ to be more specific. The current name is vague - what exactly is it driving or controlling?
| tp_rank_(model->tp_rank_), | ||
| data_type_(data_type), | ||
| debug_(isDebug()), | ||
| is_driver_(param.attn_tp_rank == 0 && param.attn_cp_rank == 0), |
There was a problem hiding this comment.
At current setting, this is the same as tp_rank_ == 0, is_driver_ is not needed.
lmdeploy/turbomind/turbomind.py
Outdated
|
|
||
| self._postprocess_config(tm_model.tm_config, engine_config) | ||
|
|
||
| print(yaml.safe_dump(self.config_dict)) |
| // for context parallel, we use symm_alloc_ and both prefill and decode stage have reduce process | ||
| // w/o context parallel, we use common alloc and only decode stage has reduce process | ||
| // perhaps it would be more appropriate to put this buffer in the unified_attention_layer. | ||
| Allocator alloc = param_.attn_cp_size > 1 ? symm_alloc_ : core::Context::alloc(kDEVICE); |
There was a problem hiding this comment.
This will create a new allocator which is not needed in the case. Use core::Context::device_alloc() to get the device allocator in current context.
| if (qi_begin + qi < qi_end && ri == 0 && check_h(hi)) { | ||
| params.partial_M[index] = M; | ||
| params.partial_L[index] = L; | ||
| params.partial_ML[index * 2] = M; |
There was a problem hiding this comment.
make the partial_ML a pointer to float2 so that load / store can be vectorized.
| @@ -0,0 +1,20 @@ | |||
| // Copyright (c) OpenMMLab. All rights reserved. | |||
|
|
|||
| #include "src/turbomind/models/llama/cp_utils.h" | |||
There was a problem hiding this comment.
move cp_utils.* to kernels/attention
| } | ||
| } | ||
|
|
||
| int cp_quo, cp_rem; |
There was a problem hiding this comment.
use expressive names, e.g. local_ti and local_ti_rank
| } | ||
|
|
||
| const int tile_count = cdiv(std::min(params.max_k_len, params.window_size), Kernel::CTA_S); | ||
| const int max_cp_k_len = (params.max_k_len + params.cp_size - 1) / params.cp_size; |
| }(); | ||
|
|
||
| const int tile_count = cdiv(std::min(params.max_k_len, params.window_size), Kernel::CTA_S); | ||
| const int max_cp_k_len = (params.max_k_len + params.cp_size - 1) / params.cp_size; |
| const int qi = offset.y / CTA_H; | ||
| const int ti = history_len; | ||
|
|
||
| int cp_quo, cp_rem; |
| }); | ||
| } | ||
|
|
||
| const bool separate_reduce = need_separate_reduce(cta_map.split_count()); |
There was a problem hiding this comment.
This code path can be removed.
|
|
||
| Impl::Merge(frag_O, frag_M, frag_L, params.inv_sqrt_dh, storage); | ||
|
|
||
| if (params.sinks && iter_end == tile_count) { |
There was a problem hiding this comment.
attention sink should be applied to cp rank 0 ONLY
4 participants
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