Hybrid LLM Parallelism

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$ 个 GPU 则有 $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)
  

References

• 揭秘 LLM 推理：全面解析 LLM 推理性能的关键因素
• EfficientML.ai Lecture, Fall 2023, MIT 6.5940
• Efficient Training on Multiple GPUs