当前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等三种流水线调度算法,用户可根据使用场景进行自由选择。
结语
本期为大家介绍了MindSpeed细粒度DualPipe通信掩盖特性,在训练性能加速上,MindSpeed将持续探索演进,欢迎访问MindSpeed开源社区
https://gitee.com/ascend/MindSpeed/
参考链接:
[1] :https://zhuanlan.zhihu.com/p/26915547331?utm_medium=social&utm_psn=1878740300547221291&utm_source=wechat_session&utm_id=0
「免责声明」:以上页面展示信息由第三方发布,目的在于传播更多信息,与本网站立场无关。我们不保证该信息(包括但不限于文字、数据及图表)全部或者部分内容的准确性、真实性、完整性、有效性、及时性、原创性等。相关信息并未经过本网站证实,不对您构成任何投资建议,据此操作,风险自担,以上网页呈现的图片均为自发上传,如发生图片侵权行为与我们无关,如有请直接微信联系g1002718958。
