MindSpeed 细粒度 DualPipe 通信掩盖，实现 AllToAll 和 EP 通信双掩盖，内存优化 40%

2025-04-01
浙江
本文字数：1607 字
阅读完需：约 5 分钟

当前 DeepSeek 等 MoE(Mixture of Experts)混合专家模型已逐步成为 LLM(大语言模型)领域的主流架构,但 MoE 模型训练面临几个关键问题,如内存占用大、通信时延高、负载不均衡、算力利用率低等,其中 All2All 通信耗时占比高和显存开销大是显著瓶颈。MindSpeed 基于昇腾硬件特点实现了 DualPipe 通信掩盖,创新提出重计算细粒度通信掩盖和 Zero-Memory 内存优化技术,缓解显存瓶颈的同时保持计算效率,并支持兼容多种流水线调度算法,供用户灵活选择。

MoE 模型的 All2All 通信耗时和显存占用面临的挑战

在 MoE 模型每一层的前向和反向计算中,各有两次 All2All 通信,共计四次 All2All 通信,称为 Dispatch(F), Combine(F), Combine(B), Dispatch(B)。当模型专家数量较多时需要进行 EP(专家并行)切分,EP 切分涉及跨节点服务器传输数据,通信耗时受到节点间带宽限制,大规模集群中 AllToAll 通信耗时占单次训练迭代时间高达 50%。

MoE 模型虽比稠密模型模型激活参数更小,但显存占用更高,首先所有专家参数都需要被保存,总参数量增大,单卡显存占用是稠密模型的 4~8 倍;其次激活值膨胀,显存需求增加 30%以上,加剧显存压力。

昇腾 MindSpeed 细粒度 DualPipe 通信掩盖

MindSpeed 基于昇腾硬件特点实现 DualPipe 方案,达成细粒度的 1F1B 阶段 All2All 通信全掩盖、Warmup/Cooldown 阶段 All2All 自掩盖的效果,创新性地提出重计算细粒度通信掩盖和 ZeroMemory 内存优化技术,掩盖了重计算阶段通信,缓解显存瓶颈的同时保持计算效率。同时,MindSpeed 的 DualPipe 通信掩盖特性兼容 VPP、Dualpipe、Cut-in-half[1]等三种流水线调度算法,用户根据使用场景灵活选择。

01 基于昇腾硬件特点实现 DualPipe 通信掩盖,达成 1F1B 阶段 All2All 通信 100%掩盖的同时,Warmup 和 Cooldown 阶段额外掩盖 50% EP 通信

在 1F1B 阶段,通过将计算和通信任务分离为两条独立的处理流,实现细粒度的调度与编排,通过拆分模型计算图中的数据流与梯度流,完成通信效率与计算资源的双重优化‌(dw 分离技术)提升通信掩盖比例,实现 All2All 通信的 100%掩盖。在 1F1B 基础上,MindSpeed 在 Warmup/Cooldown 阶段,利用共享专家的计算掩盖了部分的 Dispatch(F)/Combine(B) 通信,采用无依赖关系的专家反向计算的 dw 进一步掩盖了部分 Dispatch(B) 通信,在 1F1B 掩盖方案的基础上额外掩盖 50%的 EP 通信。

掩盖计算和流水掩盖示意图

02 支持重计算细粒度通信掩盖和 Zero-Memory/Gelu 重计算技术,实现内存优化 40%

针对 MoE 模型训练中高显存占用的问题,MindSpeed 创造性地提出了重计算细粒度通信编排和 ZeroMemory 技术,通过细粒度联合编排重计算以及前向和反向过程中所有的计算通信隐藏重计算的通信开销,进而丢弃前向计算中每层的 GroupedMatmul 融合算子的输入,从而实现内存优化。该特性在 DeepSeek V3 模型上每层能节省接近 200M 的激活值开销,叠加 Gelu 重计算等技术,整体内存优化 40%,且性能几乎无损。

重计算细粒度通信掩盖示意图

03 MindSpeed 的 DualPipe 通信掩盖特性兼容多种流水线调度算法,用户可灵活选择

DualPipe 的掩盖收益与流水线调度策略并不是强绑定关系,这种掩盖的使能前提是在流水线策略中当前可做两个无依赖关系的 MicroBatch(微批次)计算。同时,针对社区提出的 DualPipe 的改进算法 Cut-In-Half[1]调度(只取 DualPipe 镜像对称结构的一半,呈 V 字排布,相比 DualPipe 在同 PP 并行度下算法启动规模和静态显存减半,消除冗余参数且 bubble 保持一致)与 DualPipe 类似。针对 VPP(Virtual Pipeline Parallelism,虚拟流水线并行)的 1F1B 阶段以跨 MicroBatch 的前反向计算为主,稍改调度也能实现 EP 通信掩盖。因此,MindSpeed 的通信掩盖技术可兼容 Dualpipe、Cut-in-half、VPP 等三种流水线调度算法,用户可根据使用场景进行自由选择。