Contents
Distributed LLM Training
Comparing Different Parallelizations
- Data Parallelism. high utilization (DP 训练模型 GPU 利用率较高), high memory cost (即使是 ZeRO-3 也只是将优化器参数、梯度和模型参数切分到不同设备上,仍然无法降低激活值显存,而如果像 ZeRO-R 一样利用 gradient checkpointing 将激活值切分到不同设备上,又会有额外的计算开销,还会导致每层在反向传播前都需要先进行一次同步操作,进行 all-gather,会带来额外的通信开销,如果是大规模训练只用 ZeRO-R,那么每次反向传播还需要), low communication (仅需在反向传播时 all-reduce 汇聚权重梯度,并且梯度反传和梯度聚合还可以 overlap 进一步降低通信开销)
- Pipeline Parallelism. low utilization (PP 在训练时存在 bubble,GPU 利用率不高), low memory cost, medium communication (PP 只需要在相邻 stage 间前向/反向传播时进行通信,但通信量随 bsz 增大而增大)
- Tensor Parallelism. high utilization (当通信带宽不是瓶颈时,TP 会有比较高的 GPU 利用率), low memory cost, high communication (TP 的每个 Transformer block 中前向/反向各需要两次 all-reduce,并且和 PP 一样,通信量也是随 bsz 增大而增大,因此通常只会在单机内使用 TP)
3D Parallelism (DP + PP + TP)
- 上图中,同一颜色的 GPU 代表是在同一台机器上。首先在每台机器内部做 TP (如果一个层在单卡上放不下就要先用 TP 切分),然后跨机做 PP (如果一个模型在 TP 后仍然单卡放不下就要再用 PP 继续切分),把模型划分到 16 张卡上,然后再用总共 32 卡做 DP
- 需要注意的是,DP 配合 PP 使用时通常只会使用 ZeRO-1 (optimizer sharding). 考虑 ZeRO-1 和 PP 同时使用的情况,同一 stage 上所有 micro-batch 都反向传播完毕后就可以得到整个 batch 的梯度,之后可以统一地采用 reduce-scatter 汇聚梯度用于优化器参数更新,而如果使用的是 ZeRO-2,模型梯度也被切分到了不同 GPUs 上,这意味着每个 micro-batch 结束后都需要进行 reduce-scatter 汇聚梯度,极大地增加了通信量;此外,PP 本来就会降低单卡上模型参数梯度的显存占用,这样 ZeRO-2 本身带来的收益也不高;同理,ZeRO-3 也一般不会和 PP 结合使用,因为使用 ZeRO-3 将会带来更多的跨节点通信量
How to Auto Parallelize
- 结合 DP + PP + TP 有非常多种 3D 并行方案,其实可以自动搜寻最佳并行方案
- Zheng, Lianmin, et al. “Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning.” 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). 2022.
Distributed LLM Inference
- 对于 LLM 推理来说,常用的分布式推理方法包括数据并行 (Data Parallelism, DP)、张量并行 (Tensor Parallelism, TP) 和流水线并行 (Pipeline Parallelism, PP);在推理阶段,因为不会更新模型参数,DP 不再有意义,多实例和并发推理在某种程度上也可以理解为数据并行;因此下面主要讨论 TP 和 PP 的优劣以及组合应用
Comparing Different Parallelizations
- TP (Intra-op). TP 可以有效降低时延,缓解显存的限制,从而能够使用更大的 batch size,提高吞吐量;而 TP 的局限性在于其通信量比较大,可能成为另一种形式的访存瓶颈,一次推理时一个 Transformer block 内就需要做两次 all-reduce,在 decoding 阶段,通常需要使用 continuous batching 来提升 GPU 利用率,随着 batch size 增加,MHA 和 FFN 的 kernel 计算时延不会明显增加,而 all-reduce 的通信量却线性增加,相应的通信时延基本也线性增加,以至于 all-reduce 通信可能进一步成为瓶颈。此外,continuous batching 的方式会希望组合尽可能大的 batch size,也就意味着 GPU 上同一时间可能只有一个 CUDA 计算流,当 all-reduce 通信的时候并没有其他流的计算可以 overlap,导致通信开销难以掩盖。而 GPU 间的通信时延与 GPU 之间的互联方式有关,比如节点内采用 PCIE 或 NVLink 互联,跨节点采用 IB 网卡互联等会导致时延有数倍的差距,因此 TP 一般只会在单机内部使用
- PP (Inter-op). 相比 TP,PP 中模型按层切分,一次 forward,
N
N
N 个 GPU 则有
N
N
N 次 GPU 通信,而且是两个 GPU 间,不是 all-reduce,PP 的通信量会小很多。然而,PP 并不能降低单个请求的推理时延 (甚至增加),对于时延敏感的服务不一定合适;此外,不同 requests 具有不同的序列长度,最后一个 PP 还要承担 Word Embedding 相关计算等原因也可能导致负载不均衡,随着 GPU 数量的增加,bubble 时间也会增加,调度策略也会更加复杂
How to Auto Parallelize
- TP + PP. TP 和 PP 有各自的优势,因此通常也会将其结合起来进行混合并行推理,以便充分发挥它们各自的优势。AlpaServe 提出自动混合并行的调度策略,如下图可以看出,在没有 AlpaServe 的情况下,4 PP + 4 TP 相比 16 PP、8 PP + 2 TP 和 2 PP + 8 TP 获得了更优的结果,而 AlpaServe 比 4 TP + 4 PP 获得了更优的结果 (其中 SLO 为 Service Level Objective,CV 为 Coefficient of Variance)