Imported from HC-Skadi/ascendc (agent-skills/model-agent-skills/verification/model-infer-parallel-impl/SKILL.md). Install upstream with npx skills add HC-Skadi/ascendc --skill model-infer-parallel-impl. Copyright stays with the author.

模型并行切分实施

前置条件：

1. parallel_config 已由 /model-infer-parallel-analysis 确定并经用户确认
2. 目标模型已有完整的单卡框架适配代码（modeling + Runner + infer.py + infer.sh + YAML），本 skill 在此基础上做并行化改造

适用范围

• 模型代码的并行层替换
• 通信组创建和管理
• YAML 配置文件生成
• 权重转换脚本编写

实施流程

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

禁止：跳过第二步直接替换层

第一步：确认输入 + 选择参考模型

1.1 确认 parallel_config

从编排层（主 agent prompt 或用户输入）获取已确认的配置：

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 选择参考模型

根据目标代码实现模式，选择仓库中最接近的已适配模型阅读其 modeling 代码：

必须：读取参考模型的 modeling_*.py（并行层替换）和 infer.py（通信组创建），了解实际实现方式。 配置数值参考 model-infer-parallel-analysis skill 的 references/config-index.md。

完成标志

• parallel_config 已确认
• 参考模型已选定并阅读关键代码

配置参数 → 实施步骤速查

根据 parallel_config 中各参数值，确定需要执行哪些步骤：

tp_size = 1 的模块不需要 TP 替换和对应通信组。先扫一遍 parallel_config，标注哪些模块需要改，再逐步执行。

第二步：通信组创建

并行层替换前，必须先在模型初始化中创建通信组。

2.1 通信组获取

当前仓库：所有模型通过 **kwargs 传入 hccl_comm_dict（普通 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)

重构版：通过 CommManager 对象封装通信组

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")

根据目标仓库版本选择模式。当前仓库所有模型使用 hccl_comm_dict。

2.2 模块级差异化并行的通信组

当 attn_tp ≠ dense_tp ≠ moe_tp 时，各模块使用不同的通信组：

# 示例：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 通信组

通信组创建：在模型的 init_parallel_comm_group() 方法中完成。 参考：DeepSeek-R1、Kimi-K2 的 modeling_deepseek.py 中均有此方法。

2.3 DP 大小自动计算

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

完成标志

• 所有需要的通信组已在 infer.py 中创建
• 各组的 rank 划分正确
• DP size 自动计算逻辑正确

第三步：逐模块并行层替换

并行层替换同时支持了权重在线切分：ParallelLinear 内置 weight_loader()，按 tp_rank 自动加载切片。权重处理详见第六步。

3.1 Attention 层（当 attn_tp_size > 1 时）

QKV → QKVParallelLinear，O → RowParallelLinear，均使用 attn_tp_group。 MLA 模型 O_proj 可能使用独立的 oproj_tp_group。

3.2 Dense FFN 层（当 dense_tp_size > 1 时）

Gate/Up → ColumnParallelLinear，Down → RowParallelLinear，均使用 dense_tp_group。

3.3 MoE 层（有 MoE 时必须处理）

• moe_tp_size > 1：专家 FFN 做 TP 切分 + AllReduce
• moe_tp_size = 1：EP 模式，Prefill 用 re_routing + AllToAll，Decode 用 dispatch/combine + AllToAll

MoE 并行与融合算子紧密耦合，详见 {file:./references/moe-parallel.md}（含完整代码和算子说明）

3.4 模块间数据重排（当相邻模块 TP 度不同时）

边界处需要 AllGather/ReduceScatter 做数据重排。

3.5 CP / KVP 实施

TODO: CP 见 cann-recipes-infer/models/deepseek-v3.2-exp/，KVP 见 cann-recipes-infer/models/longcat-flash/。

Attention/Dense FFN/Embed/LMHead 的代码示例和数据重排示例见 {file:./references/code-examples.md}

完成标志

• Attention 层 QKV/O 已替换为并行版本
• Dense FFN 层 Gate/Up/Down 已替换（如有差异化 tp）
• MoE 层已按选定模式实现（TP 或 EP Prefill/Decode）
• 模块间数据重排已正确插入（不同 TP 度边界的 AllGather/ReduceScatter）
• 各模块使用了正确的通信组

第四步：Embedding / LMHead 并行

当 embed_tp_size > 1 或 lmhead_tp_size > 1 时需要并行化。

4.1 Embedding 并行

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 并行

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 由框架处理
)

完成标志

• Embedding 已按 embed_tp_size 并行化
• LMHead 已按 lmhead_tp_size 并行化
• 约束检查通过

第五步：YAML 配置生成

为每种部署场景生成独立的 YAML 配置文件。

5.1 配置模板

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 命名规范

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 配置（可选）

仓库有 config/ci/ 目录存放精简配置用于自动化测试。如需 CI 覆盖，参考已有配置生成。

完成标志

• 每种部署场景有独立的 YAML 文件
• 配置文件命名符合规范

第六步：权重处理

并行层替换后，需要确保权重能正确加载到各卡。

6.1 在线权重切分（推荐）

启用 enable_online_split_weight: True，框架在运行时通过各模块的 weight_loader() 自动按 rank 加载对应切片。所有 rank 读同一份完整 checkpoint，各自只保留本卡需要的部分。

TP 权重加载：ColumnParallelLinear / RowParallelLinear 的 weight_loader 按 tp_rank 取对应列/行切片。

EP 权重加载（MoE 模型）：FusedMoEGMM 的 weight_loader 按 ep_rank 过滤专家——只保留 [ep_rank * experts_per_rank, (ep_rank+1) * experts_per_rank) 范围内的专家权重，丢弃其他。load_weights() 中需通过 make_expert_params_mapping(num_experts=...) 生成全局专家映射，逐个传入 weight_loader(expert_id=...)。

适配要点（详见 docs/common/online_split_weight_guide.md）：

• 模型类须实现 load_weights() 方法，遍历权重文件并匹配到各模块的 weight_loader()
• MoE 模型必须调用 make_expert_params_mapping 生成专家权重映射
• MergedColumnParallelLinear（如 gate+up 合并）需要特殊的 weight_loader 处理 slice 顺序
• process_weights_after_loading() 处理权重转置和 NZ 格式转换

6.2 离线权重转换（备选）

若模型未实现 online split，或需要预切权重用于离线部署：

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

输出为 rank_0/ ~ rank_N/ 目录结构，每个 rank 只包含该卡需要的权重切片。

注意：offline 预切的权重与 parallel_config 绑定。改了配置必须重新转换。

参考实现：cann-recipes-infer/models/deepseek_r1/utils/convert_model.py

完成标志

• 权重加载方式已确定（online split / offline convert）
• 若 online split：load_weights() 和 weight_loader() 已实现
• 若 offline convert：转换脚本已编写或复用，输出目录结构正确

第七步：验证

7.1 配置校验（优先复用现有校验函数）

# 当前仓库没有统一的 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 功能验证

1. 确认推理实际加载的是修改后的代码（检查模型注册表、import 路径、日志中的模块路径等，确保运行时代码路径与修改路径一致）
2. 修改 infer.sh 中 YAML_FILE_NAME 指向目标配置
3. 执行 bash infer.sh
4. 检查 Prefill + Decode 推理成功（无 crash）
5. 检查各 Rank 输出形状一致
6. 如加载权重：检查输出文本合理性

每种配置独立验证，通过后再验下一个。

7.3 权重加载验证

加载权重验证时，在 YAML 中设置 with_ckpt: True + model_path，确认各 rank 权重加载无报错。

完成标志

• 配置校验脚本通过
• 至少一种配置的 Prefill + Decode 验证通过
• 使用 enable_profiler: True 运行一次，生成 profiler 数据供后续策略校准
• 验证结果已输出

常见错误

仓库参考实现索引