昇腾MindSpeed分块专家计算Chunk MoE：让通信藏进计算、显存轻装上阵

随着MoE模型规模与序列长度持续增长，专家并行（EP）已成为训练扩容的关键路径，但All-to-All通信暴露和随token数线性增长的专家计算workspace，也在同时吞噬训练效率与显存空间。针对这一挑战，昇腾MindSpeed推出分块专家计算Chunk MoE：将MoE 输入沿token维切块，以相邻chunk的异步流水掩盖通信，并缩小单次专家计算及pair双缓冲的驻留规模。当前在Atlas 900 A3 SuperPoD上的实测中，当 num_chunks = 2时，相邻chunk的通信与计算重叠实现掩盖，减少训练耗时；num_chunks = 4时训练耗时与未掩盖方案基本持平，但专家激活显存峰值降低50%，为大规模MoE训练提供兼顾通信效率与显存优化的新选择。

核心原理：把“大块等待”拆成“小块流水”

标准 all2all dispatcher 一次处理 MoE 层的全部 token：先完成 permute 与 All-to-All 分发，再执行专家计算，最后通过 All-to-All 回收并 unpermute。这个流程不仅会产生与全量 token 成正比的临时 workspace，通信与计算的串行等待也会让宝贵的计算资源空转。

Chunk MoE 的解法可以概括为两件事：把大块拆小，把串行变流水。

一个通俗类比——让后厨与配送同时忙起来

把token看作餐厅订单，专家是后厨，All-to-All则是往返于餐桌与后厨之间的配送员。

传统方式像是一次性接收整批订单：配送员送完全部订单，后厨才开始集中处理；菜品做好后，再统一送回。订单越多，后厨需要预留的操作台越大，配送和烹饪之间的等待也越明显。

Chunk MoE则把订单分成若干小批：后厨烹饪当前批时，配送员已经在运送下一批；后厨开始处理下一批时，当前批成品又在同步回收。没有谁被省略，只是原本相互等待的环节开始并行运转，操作台也不必始终为整批订单保留。

一张专业时序图——相邻 Chunk 如何互相掩护

实际实现中，Chunk MoE首先将hidden_states、top_k_index和 top_k_weights沿token维切分为多个chunk，再由forward_pair将相邻两个chunk组成一对，依次完成permute、异步Dispatch All-to-All、专家计算、异步Combine All-to-All与unpermute。

Chunk MoE相邻chunk的通信与计算重叠时序

以当前块 c和下一块 n = c + 1为例，通信掩盖发生在三个阶段：

Expert(c)执行时，Dispatch(n)已在后台传输，当前计算掩盖下一块的分发通信。
Expert(n)执行时，Combine(c)同步回收，下一块计算掩盖当前块的回收通信。
Unpermute(c)还能进一步掩盖部分 Combine(n)，减少流水末尾的通信暴露。

通信并没有消失，而是被“藏”进了计算窗口；首个Dispatch(c)与末尾 Combine(n)仍可能暴露。反向过程由autograd镜像上述机制，梯度All-to-All 同样可以与相邻chunk的梯度计算重叠。

为了降低分块带来的额外开销，Chunk MoE还通过 batched_pre_dispatch_metadata对所有chunk只做一次 all_gather_into_tensor和一次主机拷贝，批量获得发送/接收splits、各rank token数及本地专家token数，避免逐块集合通信与主机同步。对于变长输入产生的空块，实现会保留集合通信和零token专家计算，维持autograd图与专家参数的梯度连接。

一组关键公式——显存收益为什么从3块开始显现

设单个 rank 的输入 token 数为 T，固定切分为 N块，则每块token数为：

T_chunk = ceil(T / N)

标准路径需要面向全量 T保留通信与专家中间量；Chunk MoE 的 pair 双缓冲主要驻留两个 chunk，因此与 chunk 相关的驻留规模可近似理解为：

W_pair = O(2 × T_chunk) = O(2 × ceil(T / N))

这也解释了不同分块数的收益差异：

N = 2：已经形成一个完整 pair，可以获得通信掩盖收益；但 2 × T_chunk约等于 T，显存峰值无明显下降。
N > 2：两个 chunk 的驻留规模开始小于整层规模，通信掩盖与显存优化可以兼得。
N 继续增大：单块规模进一步缩小，但集合通信、kernel 启动和空块处理开销也会随之增加，并非越多越好。

这里的 W_pair描述的是 chunk 相关通信与专家中间量的规模关系，并非端到端显存的严格上界；实际收益还受模型结构、EP 规模和其他驻留激活影响。压低专家激活峰值，主要通过开启 chunk_moe_recompute，重计算专家计算部分以换取更低的激活驻留来实现。

在当前实测中，N = 4实现了训练耗时基本持平、专家激活峰值显著降低的组合收益。具体表现会随模型与通信/计算占比变化，建议结合显存快照和 profiling 选择分块数。

技术收益和测试结果——

以Qwen3.5-35B模型为例，使用 coco数据packing到16k，训练迭代时间测试如下：

红色为moe_num_chunks =2时单步迭代时间-迭代步数曲线；绿色为moe_num_chunks =4时单步迭代时间-迭代步数曲线；蓝色为原始不开启chunkmoe时的单步迭代时间-迭代步数曲线；可以看到moe_num_chunks =2时，单步迭代时间相比原始配置明显下降；moe_num_chunks =4时，训练耗时与原始配置持平，但显存峰值下降50%，如下图所示，尖刺明显变小：

操作指南：三步开启Chunk MoE
第一步：准备环境

先安装MindSpeed MM配套环境，模型、权重、数据及其余训练配置沿用对应模型示例。

https://gitcode.com/Ascend/MindSpeed-MM/blob/master/docs/zh/guides/installation/install_guide.md

第二步：修改训练配置

在训练配置文件 examples/xxx/xxx_config.yaml的 parallel段中合并以下字段（xxx为相应模型名字），并将专家模块路径和 EP 规模调整为实际模型与集群配置：

parallel:  expert_parallel_size: 8  ep_plan:    apply_modules:      - model.language_model.layers.{*}.mlp.experts    dispatcher: chunkmoe    moe_num_chunks: 2    chunk_moe_recompute: false  # 可选：开启专家计算重计算

如果更关注专家激活峰值，可将chunk_moe_recompute设为4，chunk_moe_recompute设为 true；Attention、shared experts、gate 等其他部分如有需要，仍需分别配置重计算。

第三步：选择分块策略并启动训练

moe_num_chunks与 moe_chunk_size建议只设置一个正值；若两者同时大于 0，实际以 moe_num_chunks为准并输出告警。两者均为 0时不拆分为多个块。

完成配置后，沿用原模型训练脚本启动即可。建议从 num_chunks = 2起步，结合显存快照、训练耗时和 profiling 逐步调整。

结语：让每一块算力，都更少等待

Chunk MoE 不是简单地“把 token 切开”，而是以分块控制驻留规模、以成对流水压缩等待窗口、以批量元数据和自动微分保证工程可用性，让通信掩盖与显存优化在同一条 EP 路径上协同生效。

后续，MindSpeed MM将持续围绕 MoE 训练探索更精细的流水调度、更高效的通信掩盖与更极致的显存优化，为大规模模型训练释放更多算力空间。欢迎体验 Chunk MoE，并通过实际负载反馈最佳实践。

开源代码仓库：

https://gitcode.com/Ascend/MindSpeed-MM

路径：

mindspeed_mm/fsdp/distributed/expert_parallel/ep_chunkmoe.py

昇腾MindSpeed分块专家计算Chunk MoE：让通信藏进计算、显存轻装上阵

评论 (0)

相关文章