大模型分布式训练并行技术分享

目前业内解决大模型问题，基本以多节点、分布式方案为主。分布式方案具体的实施时，又分为数据并行、参数并行、流水线并行等，针对具体的业务场景采取合适的并行方案方可带来更高的效率。
后续结合业内主流的分布式框架，具体介绍各种并行的思路以及可能带来的收益。

数据并行

一些基础知识的补充：
Pytorch DDP分布式细节分享

ZeRO（零冗余优化）

零冗余优化的核心思想：用通信换显存，数据算完即废，等需要的时候，再同步过来。
从效果来说，零冗余优化属于数据并行+张量并行，从根本来说属于数据并行。
模型在训练时需要的显存大小，假设模型的参数W大小是phi，以byte为单位，存储如下：

结论：优化器、模型参数、梯度是占据显存的主要数据。
将优化器、模型参数、梯度等数据进行切分可达到不同程度的显存优化，可分为zero1、zero2、zero3

zero1（优化器切分）

由于每块GPU上只保管部分optimizer states，因此只能将相应的W（蓝色部分）进行更新；需要对W做一次All-Gather，从别的GPU上把更新好的部分W取回来，额外产生单卡通讯量phi。

zero2（优化器+梯度切分）

• 对梯度做一次Reduce-Scatter，保证每个GPU上所维持的那块梯度是聚合梯度。单卡通讯量phi。
• 每块GPU用自己对应的O和G去更新相应的W。更新完毕后，每块GPU维持了一块更新完毕的W。同理，对W做一次All-Gather，将别的GPU算好的W同步到自己这来。单卡通讯量phi。

zero3（优化器+梯度+参数切分）

• 做forward时，对W做一次All-Gather，取回分布在别的GPU上的W，得到一份完整的W，单卡通讯量phi 。forward做完，立刻把不是自己维护的W抛弃。
• 做backward时，对W做一次All-Gather，取回完整的W，单卡通讯量phi。backward做完，立刻把不是自己维护的W抛弃。
• 做完backward，算得一份完整的梯度G，对G做一次Reduce-Scatter，从别的GPU上聚合自己维护的那部分梯度，单卡通讯量phi。聚合操作结束后，立刻把不是自己维护的G抛弃。

优化效果：

用1.5倍的通讯开销，换回近60倍的显存
基于zero的实现的工具有：

• 微软Deepspeed
• Pytorch fsdp（1.11+）

参考论文：
zero-deepspeed.pdf

模型并行

在数据并行训练中，一个明显的特点是每个 GPU 持有整个模型权重的副本，这就带来了冗余问题。如果将模型参数、优化器等分割在一个设备整列，将有效缓解显存的压力和副本冗余。
模型并行，主流上分为张量并行和流水线并行。
张量并行为层内并行，对模型 Transformer 层内进行分割、流水线为层间并行，对模型不同的 Transformer 层间进行分割。

张量并行（TP）

张量并行可视为层内并行，可分为按行进行切分和按列进行切分，分别对应行并行（Row Parallelism）与列并行（Column Parallelism）。

受 GSPMD、Oneflow 和 TF DTensor 的启发，PyTorch 从 2.0.0 开始引入 DTensor，通过DTensor抽象，我们可以无缝构建张量并行。
参考论文：
Megatron-LM 1D 2020-03-13

流水线并行（PP）

经典的流水线并行范式有Google推出的Gpipe，和微软推出的PipeDream。两者的推出时间都在2019年左右，大体设计框架一致。主要差别为：在梯度更新上，Gpipe是同步的，PipeDream是异步的。
多维混合并行
在进行上百亿/千亿级以上参数规模的超大模型预训练时，通常会组合多种并行技术一起使用。
常见的组合方式：
DP+PP
3D 并行（DP + PP + TP）
ZeRO-DP + PP + TP