昇腾+鲲鹏联手上大招!华为爆改MoE训练,吞吐再飙升20%,内存省70%

Wait 5 sec.

最近,华为在MoE训练系统方面,给出了MoE训练算子和内存优化新方案:三大核心算子全面提速,系统吞吐再提20%,Selective R/S实现内存节省70%。在通往更强大的AI路上,MoE已成为科技巨头另一个首选路径。只要Scaling Law没有失效,大模型的参数规模依旧不断扩大,由此AI智能水平才能不断攀升。凭借独特的架构设计,MoE正以前所未有的参数规模,成为突破大规模模型训练的算力瓶颈的关键路径之一。然而,如何将MoE潜力真正转化为高效的训练实践,一直是业界探索的难题。此前,华为曾通过Adaptive Pipe&EDPB框架,实现了集群级高效分布式计算,让通信和计算能完美并行,提高训练集群效率。本次,华为通过昇腾与鲲鹏算力的深度协同,进一步实现了训练算子计算效率和内存利用率大幅提升。他们从单节点视角出发,深入到NPU和CPU内部,对算子计算、下发、训练内存使用等进行细粒度拆解。令人惊喜的是,结果显示,MOE训练在之前的基础上,吞吐又提升了20%,内存占用降低了70%。这不仅是一次技术突破,更是引领MoE训练的风向标。MoE训练单节点效率提升挑战现实中,MoE模型的训练并非易事。因其框架的复杂性,除了集群分布式训练效率优化外,在单节点训练效率提升方面有两大核心挑战。算子计算效率低,存在等待首先,硬件核心计算单元,如Cube利用率不足,存在冗余操作和可优化的数据流水搬运,进而拖累了整体计算吞吐。同时,专家路由机制导致算子下发频繁,且容易中断。因为,复杂的专家路由机制增加了算子下发的调度压力,如同高速闸道入口既窄还有频繁红灯,形成了Host-Bound瓶颈。「昂贵的」NPU内存永远不够用为了扩展模型参数量以提高模型智能水平,MoE模型参数量动辄千亿甚至万亿。加之训练过程中前向传播累积的海量激活值,让内存资源显得捉襟见肘。如果太过挤压内存,还容易引发NPU内存溢出(OOM),造成训练中断。因此,NPU内存优化是大规模MoE训练永恒的主题。针对这些难题,华为给出了业界最完整的解决方案。昇腾算子计算加速,训练吞吐飙升15%毋庸置疑,只有更高算力的利用率,才能将AI系统的效用最大化。除了并行策略、通算掩盖等框架层优化方法,硬件亲和算子优化,也是进一步获得潜在性能优化的新路径。在MoE模型中,最「吃时间」的几个核心算子有:融合算子FlashAttention、基础算子MatMul,以及负责数据重排与反重排的Vector(矢量)算子。这些算子,占据了总计算耗时75%以上。从数学等价和昇腾硬件亲和角度出发,华为提出了「瘦身术」、「均衡术」、「搬运术」三大优化策略。得益于这些方案,MoE模型计算冗余消除,数据流效率提升,同时计算单元间数据移动减少,充分发挥出昇腾的硬件能力。在Pangu Ultra MoE 718B模型训练实践中,三大算子加速实现整体训练吞吐量提升15%。FlashAttention「瘦身术」(RECT-FA)针对FA算子,华为团队优化了计算顺序,进而消除了冗余计算,进一步让FA内部处理流水排布更紧密。它能支持原生非对齐计算,直接处理Key/Query维度不匹配场景,省去填充与切片开销。利用昇腾片上缓存原位累加技术,可基于数学等价计算消除旋转位置编码中复杂的拼接操作(如图1所示)。 图1:基于数学等价计算消除拼接开销原理通过核间高效同步与缓存智能预搬运技术,实现FA内部计算步骤的高效流水线衔接,消除等待阻塞。通过这三点优化,实现FA前/反向性能分别提升50%/30%;MatMul(矩阵乘法)「均衡术」(AscendP-MM)针对矩阵乘法算子,华为优化了双级数据流水搬运问题,如下图2所示。在数据从通用内存传输至L1缓存时,基于更小的L0缓存容量调整搬运量,从而更早启动从L1至L0的搬运,以及后续计算。通过矩阵子块斜向分配计算核,降低并发冲突,提升数据在芯片内部的流水效率。最终,实现了Cube(核心计算单元)的利用率提升10%。图2:MatMul算子数据流水优化前后示意图Vector算子「搬运术」(VectorSort)针对Vector算子,华为充分利用了昇腾芯片Vector指令特性,融合了多个细粒度小算子,以降低内存搬运耗时。通过等价变换(如图3所示),减少了重排与反重排操作中数据的反复搬运。上述优化操作的效果立竿见影,VectorSort类算子性能直接飙升3倍以上。图3:VectorSort重排操作优化前后流水示意图昇腾鲲鹏协同再加速,吞吐提升至20%,内存节省70%通过昇腾和鲲鹏的高效协同,华为的研究团队实现了算子下发几乎「零」等待,内存节省70%,其中关键在于Host-Device协同的算子下发优化和Selective R/S精准内存手术两项创新。鲲鹏Host-昇腾Device协同优化算子下发 等效计算的同步消除算子下发中断通常是因为Host需要等待Device返回结果(即同步),就如同高速有车需要逆行回到闸道入口,会阻塞所有后续算子的下发。在Host或Device上就地等效计算,避免了数据逆行,实现算子下发无同步。图4展示Token分发预处理过程的同步消除结果,同步消除后,算子下发不再中断,提高Device上算力利用率。图4:Token分发结果预处理同步消除前后对比重排下发序规避空闲等待通过对模块的细粒度分离并重排序(图5所示),如同在绿灯后闸道入口优先放行起步快的小车。这让Device上尽快接收到计算任务,避免过长时间的空闲等待,实现单次Host-Bound从2.1ms降低至0.6ms,超过70%的降幅。图5:重排下发序规避空闲等待昇腾+鲲鹏亲和提升下发速度为了发挥算子下发的极致性能,可以通过任务绑核来控制CPU端算子任务的处理器亲和性,将主要任务绑定在最亲和的核上,并隔离其余任务对算子下发线程的影响。这如同加宽闸道,有助于提升算子的下发速度。采用自定义粗粒度绑核方式(每NPU绑24核),完全消除系统型持续Host-Bound。通过上述Host-Device协同优化,有效缓解了同步型与持续型Host-Bound瓶颈。在Pangu Ultra MoE 718B模型训练实践中,团队实现了算子下发几乎零等待(free时间占比