CANNBot MoE并行实施指南

MoE 并行实施指南【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsEP 改造同时包含并行基础设施通信组、权重加载和 MoE 算子链替换routing → GMM → finalize。本文档覆盖 EP 实施的完整内容。MoE TP 模式moe_tp_size 1适用于小规模部署如 16 卡纯 TP。每个专家的 FFN 做 TP 切分。expert_output torch_npu.npu_grouped_matmul( [hidden_states], [self.expert_weights], group_listexpert_tokens, group_type0 ) dist.all_reduce(expert_output, opdist.ReduceOp.SUM, groupself.moe_tp_group)通信组moe_tp_group参考cann-recipes-infer/models/qwen3_moe/models/modeling_*.pyMoE EP 模式moe_tp_size 1适用于大规模部署。专家分布到各卡通过 AllToAll 交换 token。EP Prefilldouble-routing 模式# 来源models/deepseek-v3.2-exp/models/modeling_deepseek.py # Step 1: 门控 路由初始化 topk_weight, topk_idx, _ torch_npu.npu_moe_gating_top_k( logits, kself.top_k, biasself.gate.e_score_correction_bias.float(), ...) topk_idx topk_idx.to(torch.int32) expanded_x, expanded_row_idx, tokens_per_expert, pertoken_scale \ torch_npu.npu_moe_init_routing_v2( hidden_states.view(-1, h), expert_idxtopk_idx, active_numtopk_idx.shape[0] * topk_idx.shape[1], expert_numnum_experts, expert_tokens_num_type1, # 1count 模式 expert_tokens_num_flagTrue, active_expert_range[0, num_experts], quant_mode-1 # BF16: -1, W8A8: 1 ) # Step 2: AllToAll dispatch expert_input dist.all_to_all_single(..., groupself.moe_ep_group) # Step 3: EP 重路由按本地专家重排 hidden_states, gathered_scale, gathered_ids_unsort, tokens_per_local_expert \ torch_npu.npu_moe_re_routing( gathered_tokens, tokens_per_expert_group.view(self.moe_ep_size, -1), per_token_scalesgathered_pertoken_scale ) # Step 4: 专家计算GMM expert_output experts(hidden_states, tokens_per_local_expert, ...) # Step 5: 恢复顺序 AllToAll combine new_x torch.index_select(expert_output, 0, gathered_ids_unsort.float().argsort().int()) output dist.all_to_all_single(..., groupself.moe_ep_group) # Step 6: 最终聚合 hidden_states torch_npu.npu_moe_finalize_routing( output, skip1shared_expert_output, scalestopk_weight.to(output.dtype), expanded_src_to_dst_rowexpanded_row_idx, drop_pad_mode2 )EP Decodedispatch/combine 模式硬件约束npu_moe_distribute_dispatch_v2每卡最多支持 24 个 expert。当experts_per_rank 24时如 256/EP832需要改用 Prefill 的 double-routing 路径。# 来源models/deepseek-v3.2-exp/models/modeling_deepseek.py # 适用条件experts_per_rank 24 # Step 1: MC2 分发融合 AllToAll token 分组 output torch_npu.npu_moe_distribute_dispatch_v2( xhidden_states.view(-1, h), expert_idstopk_ids, **dispatch_kwargs # 包含 group、moePara 等配置 ) expand_x, dynamic_scale, expand_idx, expert_token_num, ep_recv, tp_recv output[:6] # Step 2: 专家计算GMM expert_output experts(expand_x, expert_token_num, ...) # Step 3: MC2 聚合融合 AllToAll 加权聚合 hidden_states torch_npu.npu_moe_distribute_combine_v2( expand_xexpert_output, shared_expert_xshared_expert_output, expert_idstopk_ids, assist_info_for_combineexpand_idx, expert_scalestopk_weight.float(), ep_send_countsep_recv, tp_send_countstp_recv, **combine_kwargs )EP Decode 回退experts_per_rank 24当experts_per_rank 24无法使用 dispatch_v2 时Decode 也使用 double-routing 路径同 Prefill用all_to_all_single替代 MC2 融合通信。性能略低但无 expert 数量限制。通信组moe_ep_group需要 HCCL group name因为 NPU dispatch 算子要求 参考cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py关键算子速查算子功能阶段约束npu_moe_init_routing_v2路由TopK 选专家 权重计算Prefill Decodenpu_moe_re_routingEP 重路由按 ep_size 重新分配Prefill EPnpu_moe_distribute_dispatch_v2MC2 Dispatch融合 AllToAll 分组Decode EP每卡 ≤24 expertsnpu_moe_distribute_combine_v2MC2 Combine融合 AllToAll 聚合Decode EP每卡 ≤24 expertsnpu_moe_finalize_routing最终聚合含 shared expert skipPrefillnpu_grouped_matmul批量专家计算GMM全部npu_moe_gating_top_k门控sigmoid/noaux 打分DeepSeek 系列npu_moe_gating_top_k_softmax门控softmax 打分Qwen3-MoE 等Shared Expert 处理部分 MoE 模型有 Shared Expert如 DeepSeek其计算独立于 EP routing但需要与 MoE 输出做残差加法TP 模式通过finalize_routing的skip1参数融合hidden_states torch_npu.npu_moe_finalize_routing( expert_output, skip1shared_expert_output, ...)EPTP 模式MC2通过combine_v2的shared_expert_x参数融合hidden_states torch_npu.npu_moe_distribute_combine_v2( expand_xexpert_output, shared_expert_xshared_expert_output, ...)参考cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.pyMoE block 实现EP 负载均衡EPLBmodel_config: perfect_eplb: True开启后框架会重新分配 expert 到各 rank确保负载均衡。需配合对应 routing 算子参数。通信组注意事项通信组获取方式因仓库版本不同详见 SKILL.md 第二步。以下注意事项按版本标注。通用moe_ep_group创建时需返回 HCCL group namereturn_nameTrue因为 NPU dispatch 算子要求EP 模式下 Prefill 和 Decode必须使用不同的 routing 路径double-routing vs dispatch/combine需要is_prefill分支AllToAll 的 input/output splits 取决于各卡的 token 分布是动态的Identity expert无 FFN 权重的 routing weight 可能是 FP32来自 router与 hidden_states 运算前需 cast 到 BF16重构版 CommManager 特有CommManager 对moe_ep_group使用延迟初始化_deferred_ep模型__init__时获取的 group/group_name 是 None 占位符权重加载后才调用initialize_deferred_groups()创建真实 HCCL 组。模型代码中需要运行时实时获取 group不能缓存__init__时的值当前仓库 hccl_comm_dict 模式moe_ep_group通过hccl_comm_dict[moe_ep_group]获取在模型实例化时已创建完毕无延迟初始化问题group name 通过hccl_comm_dict.get(moe_ep_group_name)获取量化模式下的 EP MoE不同量化模式影响 routing 和 GMM 的参数选择BF16init_routing(quant_mode-1) → grouped_matmul(BF16→BF16) → npu_swiglu → grouped_matmul(BF16→BF16)W8A8init_routing(quant_mode1) → grouped_matmul(INT8→INT32) → npu_dequant_swiglu_quant → grouped_matmul(INT8→BF16)W8A8C8init_routing(quant_mode1) → grouped_matmul(INT8→BF16, with scale) → npu_swiglu_clip_quant → grouped_matmul(INT8→BF16, with scale)【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考