SuperAPI AI 工具集

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

编程与开发

查看项目

更新于 10-15 17:19

技能介绍

用途:在昇腾 NPU 上对已具备单卡框架适配代码(modeling、Runner、infer.py、infer.sh、YAML)的 PyTorch 模型实施并行切分,前置条件是 parallel_config 已由 /model-infer-parallel-analysis 确定并经用户确认。

主要能力与实施流程分七步:确认输入并选择最接近的参考模型(如 GPT-OSS、Qwen3-MoE、DeepSeek-R1、LongCat-Flash、Kimi-K2);创建通信组,当前仓库模型通过 **kwargs 传入 hccl_comm_dict,并支持 attn_tp、dense_tp、moe_tp 等模块级差异化通信组;逐模块替换并行层,Attention 的 QKV 用 QKVParallelLinear、O 用 RowParallelLinear,Dense FFN 的 Gate/Up 用 ColumnParallelLinear、Down 用 RowParallelLinear,MoE 按 moe_tp_size 选择 TP 或 EP 模式(Prefill 用 re_routing + AllToAll,Decode 用 dispatch/combine);Embedding 使用 VocabParallelEmbedding,LMHead 使用 ColumnParallelLinear;按场景生成独立 YAML 配置;权重处理可选在线切分(enable_online_split_weight: True)或离线转换(utils/weight_convert.sh);最后做配置校验与功能验证。

适用场景:需在昇腾 NPU 上对模型推理做 TP、EP、CP、KVP 等并行改造的工程实现,尤其是 GQA、MoE、MLA 等结构。

限制与注意:禁止跳过通信组创建直接替换层;相邻模块 TP 度不同需插入 AllGather/ReduceScatter;embed_tp_size 需满足框架约束且不小于 attn_tp_size;offline 预切权重与 parallel_config 绑定,改配置必须重新转换。CP 与 KVP 部分资料标注为 TODO,具体实现暂未提供。