在大规模语言模型(LLM)的训练中,高效利用计算资源、降低通信开销以及维持负载均衡是亟待解决的关键问题。尤其在面对超大规模模型和海量数据时,传统训练方法往往难以应对。DeepSeek 团队在这一领域取得了突破性进展,其中 DualPipe 和 EPLB 作为两项核心技术,为优化大规模模型训练提供了创新解决方案。
DualPipe 是一种创新的双向流水线并行算法。 它通过在流水线的两端同时注入微批次,实现了前向和反向传播的完全重叠,从而大幅减少了流水线空闲时间(Pipeline Bubble),显著提高了计算资源的利用率。
EPLB(Expert Parallelism Load Balancer)则是一种专家并行负载均衡算法。 通过冗余专家策略和分组限制专家路由,优化了专家并行(EP)中的负载分配,确保不同 GPU 之间的负载均衡,提高训练效率。
官方开源代码链接:
流水线并行根据执行的策略,可以分为 F-then-B 和 1F1B 两种模式。
F-then-B 模式,先进行前向计算,再进行反向计算。该模式由于缓存了多个 micro-batch 的中间变量和梯度,显存的实际利用率并不高。
1F1B(One Forward pass followed by One Backward pass)模式是一种前向计算和反向计算交叉进行的方式。此模式下,前向计算和反向计算交叉进行,可以及时释放不必要的中间变量(如下图所示,F42 在计算前,F41 的反向 B41 已经计算结束,即可释放 F41 的中间变量,从而 F42 可以复用 F41 中间变量的显存。)
下文讲述的朴素流水线并行以及 GPipe 都是 F-then-B 模型,而 PipeDream 和 ZBPP 则是 1F1B 模式。
朴素流水线并行是实现流水线并行训练的最直接的方法。我们将模型按照层间切分成多个部分(Stage),并将每个部分(Stage)分配给一个 GPU。
假设有 K K K 块GPU,而单块GPU上做一次forward和backward的时间为: t f + t b t_f + t_b tf+tb。则:
因此,bubble的时间复杂度为: O ( K − 1 K ) O(\frac{K-1}{K}) O(KK−1),当K越大(即GPU的数量越多时),空置的比例接近1,即GPU的资源基本都被浪费。
微批次(MicroBatch)流水线并行与朴素流水线几乎相同,但它通过将传入的小批次(minibatch)分块为微批次(microbatch),并人为创建流水线来解决 GPU 空闲问题,从而允许不同的 GPU 同时参与计算过程。
假设有 K K K 块GPU,每个mini-batch划分为 M M M个,单块GPU上对一个microbatch做一次forward和backward的时间为: t f i + t b i t_{fi} + t_{bi} tfi+tbi。则:
因此,bubble的时间复杂度为: O ( K − 1 K + M − 1 ) O(\frac{K-1}{K+M-1}) O(K+M−1K−1),当 M ≫ K M\gg K M≫K的时候,这个时间可以忽略不计。
Gpipe 的流水线的问题:
为此,微软 DeepSpeed 提出的 PipeDream(1F1B 策略),针对这些问题进行了改进。其解决思路就是努力减少每个 activation 的保存时间,即这就需要每个微批次数据尽可能早的完成后向计算,从而让每个 activation 尽可能早释放。
相比 GPipe,表面上看 PipeDream 在Bubble率上并没有优化,Bubble 时间仍然为: O ( K − 1 K + M − 1 ) O(\frac{K-1}{K+M-1}) O(K+M−1K−1)。但在节省了显存之后,在设备显存一定的情况下,就可以通过增大 M 的值来降低Bubble率了。
深入阅读:https://zhuanlan.zhihu.com/p/653860567
(ZBPP)Zero Bubble Pipeline Parallelism 提出一种新的调度方法,实现了近似零流水线空闲。这一改进背后的关键思想是将反向计算分为两个部分,一个计算输入的梯度,另一个计算参数的梯度。
图中 F 表示前向传递,B 和 W 分别表示计算相对于输入 x 和层级参数 W 的梯度。
ZBPP核心原理是通过在预热阶段引入更多的前向计算(F)来填补第一个反向传播(B)之前的气泡,并在尾部重新排序权重更新(W),将流水线的布局从梯形变为平行四边形,从而消除了流水线中的所有气泡。
DualPipe 的架构基于 Transformer 框架,并针对流水线并行进行了深度优化。其在ZB-PP划分粒度的基础上,DualPipe做了针对前向传播与后向传播都需要计算的情况,做了更细致的拆分。
DualPipe的核心思想是在一对单独的前向和反向块中重叠计算和通信。具体来说,每个计算块被划分为四个部分:Attention、All-to-All Dispatch、MLP 和 All-to-All Combine。对于后向传播块,Attention 和 MLP 进一步细分为输入梯度计算(Backward for Input)和权重梯度计算(Backward for Weights)。
如上图所示,对于一对前向和反向块,DualPipe通过重新排列这些组件,并手动调整专用于通信与计算的GPU SM(Streaming Multiprocessors)的比例,可以确保在执行过程中,全对全和PP通信都可以完全隐藏。
完整的DualPipe调度如图所示。它采用双向流水线调度,同时从流水线的两端输入微批次,并且可以完全重叠大量的通信。这种重叠还确保了随着模型进一步扩展,只要保持恒定的计算与通信比,仍然可以在节点间使用细粒度的专家,并实现接近零的全对全通信开销。
DualPipe的优势:
DualPipeV 是由 DualPipe 通过“一分为二”的方法推导出的一种简洁的 V 形调度。如下图所示:
流水线气泡和内存使用对比如下:(基于相同数量的 PP 阶段)
方法 | 气泡大小 | 每设备参数量 | 每设备激活量 | 设备数量 |
---|---|---|---|---|
1F1B | (PP-1)(+) | 1× | PP | PP |
ZB1P | (PP-1)(+-2) | 1× | PP | PP |
DualPipe | (PP/2-1)(&+-3) | 2× | PP+1 | PP |
DualPipeV | (PP/2-1)(&+-3) | 2× | PP+1 | PP/2 |
PP 表示流水线并行阶段的数量(偶数)。 表示正向块的执行时间, 表示完整反向块的执行时间, 表示“反向权重”块的执行时间,& 表示两个相互重叠的正向和反向块的执行时间。
专家并行(Expert Parallelism,EP) 是一种在混合专家模型(Mixture of Experts,MoE)中采用的训练方法。它将模型中的不同“专家”分配到不同的计算设备上,以提高计算效率。(有关DeepSeekMoE架构解析在我的DeepEP学习笔记中已经分析过,不了解的读者可以去看看)
在专家并行中,负载均衡是一个关键挑战,因为不同专家分配在不同GPU上,负载往往不平均——有的专家经常被大量样本命中,成为热点,而有的却很少工作。这导致某些GPU忙不过来、成为瓶颈,而其他GPU又闲置浪费。这种木桶短板效应,限制了整体的吞吐。
在DeepSeek-V3论文3.4.1节中提到:
To achieve load balancing among different experts in the MoE part, we need to ensure that each GPU processes approximately the same number of tokens. To this end, we introduce a deployment strategy of redundant experts, which duplicates high-load experts and deploys them redundantly. The high-load experts are detected based on statistics collected during the online deployment and are adjusted periodically (e.g., every 10 minutes). After determining the set of redundant experts, we carefully rearrange experts among GPUs within a node based on the observed loads, striving to balance the load across GPUs as much as possible without increasing the cross-node all-to-all communication overhead. For the deployment of DeepSeek-V3, we set 32 redundant experts for the prefflling stage. For each GPU, besides the original 8 experts it hosts, it will also host one additional redundant expert.
简单来说,冗余专家策略就是复制负载较重的专家。然后将这些复制的专家打包到 GPU 上,以确保不同 GPU 之间的负载平衡,从而提高模型训练的效率和稳定性。
EPLB冗余专家策略的关键点:
EPLB负载均衡算法包含两种策略,用于不同的场景。
以下代码展示了一个两层 MoE 模型的示例,每层包含 12 个专家。我们在每层引入 4 个冗余专家,总共 16 个副本被放置在 2 个节点上,每个节点包含 4 个 GPU。
import torch
import eplb
weight = torch.tensor([[ 90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
[ 20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27]])
num_replicas = 16
num_groups = 4
num_nodes = 2
num_gpus = 8
phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
print(phy2log)
# 输出:
# tensor([[ 5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1],
# [ 7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1]])
由分层负载均衡策略生成的输出,表示为下图所示的专家复制和放置计划。
类型 | 优点 | 缺点 | 适用场景 |
---|---|---|---|
静态分配 | 实现简单 | 资源浪费严重 | 小规模模型 |
动态路由 | 适应性较强 | 通信开销大 | 中等规模模型 |
EPLB策略 | 资源利用率最大化 | 需要冗余计算资源 | 超大规模MoE模型 |
总的来说,DualPipe是一种创新的双向流水线并行算法,旨在通过完全重叠前向和后向计算-通信阶段来提高模型训练效率,减少流水线气泡。它通过对称调度微批次,实现了高效的计算资源利用。EPLB则是为解决专家并行中负载不均衡问题而设计的负载均衡器,通过动态调整专家分布和冗余专家策略,最大化资源利用率,适用于超大规模MoE模型。这两种技术都显著提升了分布式训练的性能和效率,对于处理大规模机器学习任务具有重要意义。