热搜词:

训练MoE足足提速70%!华为只用了3招

Scaling Law 之下,MoE(混合专家)如今已经成为各大模型厂商扩展模型能力的制胜法宝。

不过,在高效实现模型参数规模化的同时,MoE 的训练难题也日益凸显:

训练效率不足,甚至一半以上训练时间都浪费在 " 等待 " 上。

现在,为了突破 MoE 的训练瓶颈,华为出手了:

构建了一套名为Adaptive Pipe & EDPB的优化方案,开启 " 上帝视角 ",让 MoE 面临 " 交通拥堵 " 的训练集群,实现无等待流畅运行。

MoE 大规模训练难题:一半以上的训练时间在等待?

实践已经表明,MoE 模型训练集群的效率面临两方面挑战:

首先,是专家并行引入了计算和通信等待。

当模型规模较大时,需要切分专家到不同设备形成并行(EP),这就引入额外 All-to-All 通信。

与此同时,MoE 层绝大部分 EP 通信与计算存在时序依赖关系,一般的串行执行模式会导致大量计算单元空闲,等待通信。

其次,负载不均会引入计算和计算等待。

MoE 算法核心是 " 有能者居之 ",在训练过程中会出现部分热专家被频繁调用,而冷专家使用率较低的情况。

同时,真实训练数据的长度不一,不同的模型层(如稀疏层、嵌入层等)的计算量也存在明显差异,造成不同卡之间计算也在互相等待。

用一个形象点的说法就是,MoE 训练系统就像一个存在局部交通阻塞的城区,面临两大核心问题:

人车混行阻塞:所有车辆(计算)与行人(通信)在红绿灯交替通行,互相等待。

车道分配僵化:固定划分的直行、左转车道就像静态的专家分配,导致热门车道(热专家)大排长龙,而冷门车道(冷专家)闲置。

针对以上问题,华为团队打造了 " 智慧化交通 " 设施:

首先,建造 " 行人地下通道 "(通信掩盖技术),彻底分离人车动线,使计算不再等待通信。

其次,部署 " 智能可变车道 "(动态专家路由),根据实时车流(数据分布)动态调整车道功能,让闲置的左转车道也能分担直行压力,实现负载均衡。

这套组合方案既解决了资源分配不均的问题,又消除了通信阻塞的瓶颈,就像为城市交通装上了 " 智慧大脑 ",让每个方向的通行效率都得到最大化提升。

DeployMind 仿真平台,小时级自动并行寻优

具体来说,华为首先构建了名为 DeployMind 的仿真平台,它是一个基于昇腾硬件训练系统的 " 数字孪生 " 平台,通过计算 / 通信 / 内存三维度的多层级建模、昇腾硬件系统的高精度映射、全局化算法加速运行等技术,能在1 小时内模拟百万次训练场景,实现 MoE 模型多样化训练负载的快速分析和自动找到与集群硬件规格匹配的最优策略选择。

在训练实践验证中,该建模框架可达到 90% 精度指标,实现低成本且高效的最优并行选择。

针对 Pangu Ultra MoE 718B 模型,在单卡内存使用约束下,华为通过 DeployMind 以训练性能为目标找到了 TP8/PP16/VPP2/EP32(其中 TP 只作用于 Attention),这一最适合昇腾集群硬件规格的并行方案,综合实现计算、通信、内存的最佳平衡。

通信掩盖 >98%,让计算不再等待通信

华为还提出了一套名为 Adaptive Pipe 的通信掩盖框架。在 DeployMind 仿真平台自动求解最优并行的基础上,采用层次化 All-to-All 降低机间通信和自适应细粒度前反向掩盖,实现通信几乎 " 零暴露 "。

层次化专家并行通信

针对不同服务器之间通信带宽低,但机内通信带宽高的特点,华为创新地将通信过程拆成了两步走:

第一步,让各个机器上 " 位置相同 " 的计算单元联手,快速地从所有机器上收集完整的数据块(Token);

第二步,每台机器内部先对数据块进行整理,然后利用机器内部的高速通道,快速完成互相交换。

这种分层设计的巧妙之处在于,它把每个数据块最多的复制分发操作都限制在单台机器内部的高速网络上完成,而在跨机器传输时,每个数据块只需要发送一份拷贝,相比传统 All-to-All 通信加速 1 倍。

也就是说,有效通过减少跨机通信,提升了集群的通信速度。

自适应细粒度前反向掩盖

在 DualPipe 掩盖框架的基础上,华为基于虚拟流水线并行技术,实现了更精密的调度,即 Adaptive Pipe。

相比 DualPipe,Adaptive Pipe仅利用一份权重,不仅将流水线并行所需的内存占用减半,有效降低了计算 " 空泡 ",释放了流水线的峰值性能潜力;同时,该策略能够额外实现与分层通信的完美协同,无缝覆盖机间与机内两层通信的掩盖。

在这种层次化通信和细粒度计算通信切分调度优化下,Adaptive Pipe 可实现 98% 以上的 EP 通信掩盖,让计算引擎不受通信等待的束缚。

克服负载不均,训练再加速 25%

由于 MoE 模型训练过程中天然存在的负载不均问题,集群训练效率时高时低,华为团队还提出了EDPB 全局负载均衡,实现专家均衡调度。

在最优并行和通信掩盖基础上,EDPB 再取得了 25.5% 的吞吐提升收益。

△集群 P2P 通信分析对比

所谓 EDPB,E 是专家预测动态迁移。

MoE 模型训练中,设备间的专家负载不均衡如同 " 跷跷板 " ——部分设备满载运行,另一些却处于 " 半休眠 " 状态。团队提出了基于多目标优化的专家动态迁移技术,让专家在分布式设备间 " 智能流动 "。

该技术主要有三个特点:

预测先行,让专家负载 " 看得见未来 ":预测负载趋势,实现 " 计算零存储开销,预测毫秒级响应 ";

双层优化,计算与通信的黄金分割点:提出节点 - 设备双层贪心优化架构,在让计算资源 " 齐步走 " 的同时,给通信链路 " 减负 ";

智能触发,给专家迁移装上 " 红绿灯 ":设计分层迁移阈值机制,通过预评估迁移收益动态决策,实现专家迁移的智能触发。

△基于专家动态迁移的 EP 间负载均衡整体框架图

D 是数据重排 Attention 计算均衡。

在模型预训练中普遍采用数据拼接固定长度的策略,但跨数据的稀疏 Attention 计算量差异显著,会引入负载不均衡问题,导致 DP 间出现 " 快等慢 " 的资源浪费。

为解决这一问题,华为团队提出了一种精度无损的动态数据重排方案,其核心在于:通过线性模型量化单样本计算耗时,在严格保持训练精度无损下,批次内采用贪心算法构建最小化耗时的数据重排,实现负载均衡。

P 是虚拟流水线层间负载均衡。

MoE 模型通常采用混合结构,Dense 层、MTP 层、输出层所在的 Stage 与纯 MoE 层所在的 Stage 负载不均,会造成的 Stage 间等待。

华为团队提出虚拟流水线层间负载均衡技术,将 MTP 层与输出层分离,同时将 MTP Layer 的 Embedding 计算前移至首个 Stage,有效规避 Stage 间等待问题,实现负载均衡。

△基于异构模块设计的 VPP 并行负载均衡系统端到端 72.6% 训练吞吐提升

在 Pangu Ultra MoE 718B 模型的训练实践中,华为团队在 8K 序列上测试了 Adaptive Pipe & EDPB 吞吐收益情况。

实验结果显示,在最优并行策略的初始性能基础上,华为这套 " 通信掩盖 + 动态专家迁移 " 的优化方案,能实现系统端到端 72.6% 的训练吞吐提升。

总而言之,华为的这套打法可以说是为大模型训练优化提供了关键路径。感兴趣的小伙伴可以再通过完整技术报告深入了解——

技术报告地址:

https://gitcode.com/ascend-tribe/ascend-training-system/tree/main/DistributedOptimization

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法!

—  完  —

点亮星标

科技前沿进展每日见