[Feature] Prefill Context Parallel (PCP) basic support (#28718)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Signed-off-by: FENP <yuanyongjie.yyj@antgroup.com> Signed-off-by: LookAround <lixushi@huawei.com> Signed-off-by: Jingchun Gao <gaojingchun1@huawei.com> Signed-off-by: zhenwenqi2024 <zhenwenqi_2022@qq.com> Co-authored-by: FENP <yuanyongjie.yyj@antgroup.com> Co-authored-by: LookAround <lixushi@huawei.com> Co-authored-by: Jingchun Gao <gaojingchun1@huawei.com> Co-authored-by: zhenwenqi2024 <zhenwenqi_2022@qq.com> Co-authored-by: Jingchun Gao <63247409+gjc0824@users.noreply.github.com>
This commit is contained in:
@@ -13,6 +13,7 @@ from vllm.config import CacheConfig, VllmConfig
|
||||
from vllm.distributed import (
|
||||
get_dp_group,
|
||||
get_ep_group,
|
||||
get_pcp_group,
|
||||
get_pp_group,
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
@@ -322,10 +323,12 @@ class GptOssModel(nn.Module):
|
||||
|
||||
# In MoE, we need to flatten the tensor parallel size across the data
|
||||
# parallel size when EP is disabled.
|
||||
tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
|
||||
tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
|
||||
tp_size=get_tensor_model_parallel_world_size(),
|
||||
dp_size=get_dp_group().world_size,
|
||||
dp_rank=get_dp_group().rank_in_group,
|
||||
pcp_size=get_pcp_group().world_size,
|
||||
pcp_rank=get_pcp_group().rank_in_group,
|
||||
)
|
||||
|
||||
intermediate_size = self.config.intermediate_size
|
||||
@@ -507,10 +510,12 @@ class GptOssModel(nn.Module):
|
||||
|
||||
# In MoE, we need to flatten the tensor parallel size across the data
|
||||
# parallel size when EP is disabled.
|
||||
tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
|
||||
tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
|
||||
tp_size=get_tensor_model_parallel_world_size(),
|
||||
dp_size=get_dp_group().world_size,
|
||||
dp_rank=get_dp_group().rank_in_group,
|
||||
pcp_size=get_pcp_group().world_size,
|
||||
pcp_rank=get_pcp_group().rank_in_group,
|
||||
)
|
||||
|
||||
intermediate_size = self.config.intermediate_size
|
||||
|
||||
Reference in New Issue
Block a user