external-cannbot-model-model-infer-parallel-impl

Model Parallel Partitioning Implementation

Prerequisites：

1. parallel_config has been determined by /model-infer-parallel-analysis and confirmed by the user
2. The target model already has complete single-card framework adaptation code (modeling + Runner + infer.py + infer.sh + YAML); this skill performs parallelization modifications on this basis

---

Scope of Application

• Parallel layer replacement of model code
• Communication group creation and management
• YAML configuration file generation
• Weight conversion script writing

Implementation Process

第一步：确认输入 + 选择参考模型
    ↓
第二步：通信组创建
    ↓
第三步：逐模块并行层替换
    ↓
第四步：Embedding / LMHead 并行
    ↓
第五步：YAML 配置生成
    ↓
第六步：权重处理
    ↓
第七步：验证

Prohibited: skipping the second step and directly replacing layers

---

Step 1: Confirm Inputs + Select Reference Model

1.1 Confirm parallel_config

Obtain the confirmed configuration from the orchestration layer (main agent prompt or user input):

parallel_config:
  attn_tp_size: {value}       # attn_dp_size = world_size // attn_tp_size
  dense_tp_size: {value}
  moe_tp_size: {value}        # moe_ep_size = world_size // moe_tp_size
  embed_tp_size: {value}
  lmhead_tp_size: {value}
  o_proj_tp_size: {value}     # MLA 模型需要，非 MLA 可省略
  cp_size: {value}            # 长序列 Prefill 需要，Decode 可省略
  kvp_size: {value}           # 超长序列需要，普通场景可省略

1.2 Select Reference Model

According to the target code implementation pattern, select the closest already-adapted model in the repository and read its modeling code:

Implementation Pattern	Reference Model	Focus Points
Standard GQA + pure TP	GPT-OSS	Basic usage of QKVParallelLinear
Standard GQA + MoE EP	Qwen3-MoE	MoE routing + npu_swiglu fusion
MLA + module-differentiated TP + EP	DeepSeek-R1 / V3.2	MLA projection partitioning, oproj independent group, EP dispatch/combine
MLA + differentiated TP + KVP/AFD	LongCat-Flash	oproj alignment constraints of KVP, data rearrangement between modules, AFD communication
MLA + CP + differentiated TP	Kimi-K2 / GLM-5	Creation of multiple comm groups (refer to infer.py), CP KV sharding

Must: read the reference model's modeling_*.py(parallel layer replacement) and infer.py(communication group creation) to understand the actual implementation method. For configuration values, refer to the model-infer-parallel-analysis skill's references/config-index.md。

Completion Criteria

• parallel_config confirmed
• Reference model selected and key code read

---

Configuration Parameters → Implementation Steps Quick Reference

According to the values of each parameter in parallel_config, determine which steps need to be executed:

Parameter Condition	Required Communication Groups (Step 2)	Required Code Modifications (Steps 3/4)
attn_tp_size > 1	attn_tp_group	Replace Attention QKV/O with ParallelLinear
attn_tp_size = 1 and world_size > 1	no Attention TP group (uses DP)	Attention does not need TP replacement
dense_tp_size > 1	dense_tp_group	Replace Dense FFN Gate/Up/Down
moe_tp_size > 1	moe_tp_group	Perform TP partitioning on MoE expert FFN
moe_tp_size = 1 and has MoE	moe_ep_group(requires HCCL group name)	MoE EP dispatch/combine implementation
o_proj_tp_size ≠ attn_tp_size	oproj_tp_group(independent group)	O_proj uses an independent communication group
embed_tp_size > 1	embed_tp_group	Replace Embedding with VocabParallelEmbedding
lmhead_tp_size > 1	lmhead_tp_group	Replace LMHead with ColumnParallelLinear
Adjacent modules have different TP degrees	—	Insert AllGather/ReduceScatter rearrangement between modules

tp_size = 1 modules do not require TP replacement and the corresponding communication group. First scan parallel_config and mark which modules need modification, then execute step by step.

---

Step 2: Communication group creation

Before replacing parallel layers, communication groups must first be created during model initialization.

2.1 Communication group retrieval

Current repository: all models pass **kwargs to pass in hccl_comm_dict(ordinary dict)

class Model(nn.Module):
    def __init__(self, config, runner_settings, **kwargs):
        self.hccl_comm_dict = kwargs.get("hccl_comm_dict", None)
        self.attn_tp_group = self.hccl_comm_dict["attn_tp_group"]
        self.moe_ep_group = self.hccl_comm_dict["moe_ep_group"]
        self.embed_tp_size = runner_settings.get("embed_tp_size", 1)

Refactored version: through CommManager object to encapsulate communication groups

class Model(nn.Module):
    def __init__(self, config, infer_config, comm_manager):
        tp_rank = comm_manager.get_rank("attn_tp_group")
        ep_rank = comm_manager.get_rank("moe_ep_group")

Select the mode according to the target repository version. All models in the current repository use hccl_comm_dict.

2.2 Communication groups for module-level differentiated parallelism

When attn_tp ≠ dense_tp ≠ moe_tp, each module uses a different communication group:

# 示例：attn_tp=1, dense_tp=8, moe_tp=1
self.attn_tp_group = self.hccl_comm_dict.get("attn_tp_group", None)   # size=1
self.dense_tp_group = self.hccl_comm_dict.get("dense_tp_group", None) # size=8
self.moe_ep_group = self.hccl_comm_dict["moe_ep_group"]              # EP 通信组

Communication group creation: completed in the model's init_parallel_comm_group() method. Reference: in the modeling_deepseek.py of DeepSeek-R1 and Kimi-K2, this method is present.

2.3 Automatic DP size calculation

attn_dp_size = world_size // attn_tp_size
moe_dp_size = world_size // moe_tp_size
moe_ep_size = moe_dp_size  # EP size = DP size
embed_dp_size = world_size // embed_tp_size

Completion criteria

• All required communication groups have been created in infer.py
• The rank partitioning of each group is correct
• The DP size automatic calculation logic is correct

---

Step 3: Module-by-module parallel layer replacement

Parallel layer replacement also supports online weight sharding: ParallelLinear has built-in weight_loader(), according to tp_rank automatically loads shards. For weight handling, see Step 6.

3.1 Attention layer (when attn_tp_size > 1 )

QKV → QKVParallelLinear，O → RowParallelLinear, both use attn_tp_group. MLA model O_proj may use an independent oproj_tp_group。

3.2 Dense FFN layer (when dense_tp_size > 1 )

Gate/Up → ColumnParallelLinear，Down → RowParallelLinear, both use dense_tp_group。

3.3 MoE layer (must be handled when MoE is present)

• moe_tp_size > 1: expert FFN performs TP sharding + AllReduce
• moe_tp_size = 1: EP mode, Prefill uses re_routing + AllToAll, Decode uses dispatch/combine + AllToAll

MoE parallelism is tightly coupled with fused operators, see {file:./references/moe-parallel.md}(includes complete code and operator descriptions)

3.4 Inter-module data rearrangement (when adjacent modules have different TP degrees)

At boundaries, AllGather/ReduceScatter is needed for data rearrangement.

3.5 CP / KVP implementation

TODO: for CP see cann-recipes-infer/models/deepseek-v3.2-exp/, for KVP see cann-recipes-infer/models/longcat-flash/。

For code examples and data rearrangement examples of Attention/Dense FFN/Embed/LMHead, see {file:./references/code-examples.md}

Completion criteria

• The Attention layer QKV/O has been replaced with the parallel version
• The Dense FFN layer Gate/Up/Down has been replaced (if there is differentiated tp)
• The MoE layer has been implemented according to the selected mode (TP or EP Prefill/Decode)
• Inter-module data rearrangement has been correctly inserted (AllGather/ReduceScatter at boundaries with different TP degrees)
• Each module uses the correct communication group

---

Step 4: Embedding / LMHead parallelism

When embed_tp_size > 1 or lmhead_tp_size > 1 , parallelization is needed.

4.1 Embedding parallelism

from module.linear import VocabParallelEmbedding

# 按词表维度切分（参数为 tp_size + tp_rank，无 tp_group）
self.embed_tokens = VocabParallelEmbedding(
    config.vocab_size,
    config.hidden_size,
    self.padding_idx,
    torch.bfloat16,
    tp_size=self.embed_tp_size,
    tp_rank=dist.get_rank(self.hccl_comm_dict["embed_tp_group"]) if self.embed_tp_size > 1 else 0,
)

4.2 LMHead parallelism

self.lm_head = ColumnParallelLinear(
    config.hidden_size,
    config.vocab_size,
    tp_size=self.lmhead_tp_size,
    tp_rank=dist.get_rank(self.hccl_comm_dict["lmhead_tp_group"]) if self.lmhead_tp_size > 1 else 0,
    # gather_output 由框架处理
)

Completion criteria

• Embedding has been parallelized according to embed_tp_size
• LMHead has been parallelized according to lmhead_tp_size
• Constraint checks passed

---

Step 5: YAML configuration generation

Generate an independent YAML configuration file for each deployment scenario.

5.1 Configuration template

model_name: "{model_name}"
model_path: "/path/to/weights"
exe_mode: "eager"                    # 初始用 eager，后续可切 ge_graph
world_size: {W}

model_config:
  enable_pa: True
  pa_block_size: 128
  enable_weight_nz: True
  with_ckpt: True
  enable_online_split_weight: True   # 推荐：运行时自动按 tp_rank 切分权重，改配置无需重新转换
  enable_multi_streams: False
  enable_profiler: False
  # MoE 专用
  moe_chunk_max_len: {1024 for Decode, 65536 for Prefill}
  perfect_eplb: False

data_config:
  dataset: "default"
  input_max_len: {根据场景}
  max_new_tokens: {根据场景}
  batch_size: {根据显存估算}

parallel_config:
  attn_tp_size: {value}
  dense_tp_size: {value}
  moe_tp_size: {value}
  embed_tp_size: {value}
  lmhead_tp_size: {value}
  o_proj_tp_size: {value}     # MLA 模型需要
  cp_size: {value}            # 长序列 Prefill 需要
  kvp_size: {value}           # 超长序列需要

5.2 Naming convention

config/
├── {model_name}_rank_{W}_{W}ep_decode.yaml          # Decode 纯 EP
├── {model_name}_rank_{W}_{tp}tp_prefill.yaml        # Prefill 纯 TP
├── {model_name}_rank_{W}_densetp{n}_ep{m}.yaml      # 混合模式
└── ci/
    └── {model_name}_ci.yaml                          # CI 测试用

5.3 CI configuration (optional)

The repository has a config/ci/ directory storing simplified configurations for automated testing. If CI coverage is needed, refer to existing configurations to generate them.

Completion Criteria

• Each deployment scenario has an independent YAML file
• Configuration file naming complies with the specification

---

Step 6: Weight Handling

After parallel layer replacement, you need to ensure weights can be correctly loaded onto each card.

6.1 Online Weight Splitting (Recommended)

Enable enable_online_split_weight: True, and at runtime the framework automatically loads the corresponding slice by rank through each module's weight_loader() . All ranks read the same complete checkpoint, and each retains only the part needed by its own card.

TP Weight Loading：ColumnParallelLinear / RowParallelLinear 's weight_loader takes the corresponding column/row slice by tp_rank .

EP Weight Loading(MoE models):FusedMoEGMM 's weight_loader filters experts by ep_rank —retaining only expert weights within the [ep_rank * experts_per_rank, (ep_rank+1) * experts_per_rank) range and discarding the rest.load_weights() needs to generate a global expert mapping through make_expert_params_mapping(num_experts=...) and pass them in one by one to weight_loader(expert_id=...)。

Adaptation Points (see docs/common/online_split_weight_guide.md）：

• The model class must implement the load_weights() method to iterate over weight files and match them to each module's weight_loader()
• MoE models must call make_expert_params_mapping to generate the expert weight mapping
• MergedColumnParallelLinear(such as gate+up merging) requires special weight_loader handling of the slice order
• process_weights_after_loading() handles weight transposition and NZ format conversion

6.2 Offline Weight Conversion (Alternative)

If the model has not implemented online split, or pre-split weights are needed for offline deployment:

bash utils/weight_convert.sh \
    --input_path /path/to/origin \
    --output_path /path/to/output \
    --world_size {W} \
    --quant_mode {w8a8/w8a8c8/...}

The output is a rank_0/ ~ rank_N/ directory structure, where each rank contains only the weight slices needed by that card.

Note: offline pre-split weights are bound to parallel_config. If the configuration is changed, conversion must be redone.

Reference implementation:cann-recipes-infer/models/deepseek_r1/utils/convert_model.py

Completion Criteria

• The weight loading method has been determined (online split / offline convert)
• If online split:load_weights() and weight_loader() have been implemented
• If offline convert: the conversion script has been written or reused, and the output directory structure is correct

---

Step 7: Verification

7.1 Configuration Validation (Prefer Reusing Existing Validation Functions)

# 当前仓库没有统一的 scripts/validate_config.py。
# 优先复用 executor/utils/common_utils.py 中的 check_common_parallel_settings，
# 再调用模型自己的 model_setting.py / infer.py 中的 check_parallel_settings / check_vars。
assert world_size % attn_tp_size == 0
assert world_size % moe_tp_size == 0
assert world_size % embed_tp_size == 0
assert world_size % lmhead_tp_size == 0
assert num_attention_heads % attn_tp_size == 0
assert num_key_value_heads % attn_tp_size == 0  # GQA
assert num_experts % ep_size == 0               # MoE
assert embed_tp_size >= attn_tp_size
assert embed_tp_size % attn_tp_size == 0

7.2 Functional Verification

1. Confirm that inference actually loads the modified code (check the model registry, import paths, module paths in logs, etc., to ensure the runtime code path matches the modified path)
2. Modify infer.sh in YAML_FILE_NAME to point to the target configuration
3. Run bash infer.sh
4. to check that Prefill + Decode inference succeeds (no crash)
5. Check that the output shapes of each Rank are consistent
6. If loading weights: check the reasonableness of the output text

Verify each configuration independently, and verify the next one only after passing.

7.3 Weight Loading Verification

When verifying weight loading, set with_ckpt: True + model_pathin the YAML to confirm that weight loading for each rank reports no errors.

Completion Criteria

• The configuration validation script passes
• At least one configuration passes Prefill + Decode verification
• Run once using enable_profiler: True to generate profiler data for subsequent policy calibration
• The verification results have been output

---

Common Errors

Error Pattern	Root cause	Prevention
Skip communication group creation and directly replace layers	Group cannot be found at runtime	Step 2 must be completed before step 3
All modules use the same tp_group	attn_tp / dense_tp / moe_tp are not distinguished	Use the configuration-to-step quick reference table to check the source of each group item by item
In EP mode, Prefill/Decode use the same code	The routing operators differ between the two phases	Refer to the Prefill/Decode branches in moe-parallel.md
TP degrees differ between modules but data rearrangement is missing	Tensor shapes do not match between adjacent modules	The "adjacent modules have different TP degrees" row in the quick reference table indicates this step
embed_tp_size < attn_tp_size	Framework constraint: embed output must be consumable by attn	Checked by the step 7 validation script
Weight loading shape mismatch	parallel_config was changed but the weights were not reprocessed	Step 6 confirms the weight processing method

---

Repository reference implementation index

Implementation pattern	Reference file	Search keyword
TP linear layer replacement	cann-recipes-infer/models/gpt_oss/models/modeling_*.py	QKVParallelLinear
Communication group creation	cann-recipes-infer/models/kimi-k2-thinking/models/modeling_*.py	init_parallel_comm_group
Communication group validation	cann-recipes-infer/models/kimi-k2-thinking/infer.py	check_parallel_settings
MoE EP Prefill	cann-recipes-infer/models/deepseek_r1/models/modeling_*.py	npu_moe_re_routing
MoE EP Decode	cann-recipes-infer/models/deepseek_r1/models/modeling_*.py	npu_moe_distribute_dispatch
Embed/LMHead parallelism	cann-recipes-infer/models/kimi-k2-thinking/models/modeling_*.py	VocabParallelEmbedding
Independent oproj_tp configuration	cann-recipes-infer/models/longcat-flash/models/modeling_*.py	oproj_tp
Weight conversion script	cann-recipes-infer/models/deepseek_r1/utils/	weight_convert
YAML multi-scenario configuration	cann-recipes-infer/models/deepseek_r1/config/	decode/prefill separation