模型训练-3D并行

佛系调参

已于 2023-07-03 21:28:51 修改

阅读量1k

点赞数

文章标签：深度学习自然语言处理人工智能语言模型

于 2023-06-27 22:07:15 首次发布

本文链接：https://blog.csdn.net/yzy__zju/article/details/131425945

版权

文章介绍了数据并行、模型并行两种主要的并行计算方法以及ZeRO框架的三个阶段，详细阐述了它们如何通过分布式计算来优化显存使用，减少内存占用。ZeRO从不同层面进行显存管理优化，到第三阶段能将四种显存占用降至显卡数分之一。流水线并行则通过将模型层分配到不同显卡实现显存减载，但可能导致计算效率下降。

摘要由CSDN通过智能技术生成

1. 数据并行（Data Parallel）

1. 数据并行（Data Parallel）

1. 数据并行（Data Parallel）

1.1常规数据并行

有一张显卡（例如rank 0显卡）专门用于存储模型参数、梯度信息和更新模型参数。将训练数据分成多份（份数等于显卡数量），每张卡上的模型参数相同，进行前向和反向传播后，每张卡上都计算得到对应部分数据的梯度，然后对多张卡上的梯度进行reduce操作，将平均后的梯度结果存放在专门的显卡上，然后在专门的显卡上利用优化器进行参数更新。最后将更新后的参数再broadcast到所有显卡上，重复上述过程

1.2 distributed data parallel（分布式数据并行）

区别：不需要专门的参数服务器

初始时每张显卡上都有相同的模型参数，同样将训练数据均分成多份，每张卡上利用单独一小份的数据进行前向和反向得到梯度，然后将多张卡上的梯度参数all reduce到所有的显卡上，这样每张显卡上的梯度信息也是完全一致的，同时优化器的历史信息数据也是完全一致的，这样便可以在每一张显卡上单独进行参数更新，并且能够保证每张卡上更新后的模型参数也是完全一致的。

总结：每张卡上单独进行参数更新

1.3 数据并行带来的显存优化效果

transformer中，显卡上存储的模型中间结果（即是每一层的输入，也可以理解成上一层的输出）的维度是[batch, Len, Dim]，多卡数据并行后，每张显卡上存储的模型中间结果的维度变成[batch/卡数, Len, Dim]

2. 模型并行

2.1 原理

将参数矩阵（全连接层）按照行进行切分，使得在每张显卡上分别进行计算（前提：需要保证多张显卡上的输入XB是一样的，即多张显卡采用同一batch的输入数据），然后再将计算结果进行拼接

2.2 模型并行带来的显存优化结果

模型参数、梯度和优化器参数，三大部分的显存都降低为原来的1/显卡数量

中间结果显存没有变化，因此不采用数据并行的话可能显存仍然很大导致在一张显卡上放不下。

3. ZeRO框架

3.1 ZeRO1

是基于数据并行的做法，但与原始的数据并行做法不同，ZeRO第1阶段和数据并行的区别在于

（1）每张卡上前向和反向得到梯度之后，采用Reduce Scatter得到每张显卡上对应部分（1/3）的平均梯度

（2）优化器参数也是采用对应部分的（1/3），然后去更新模型的部分参数（1/3）

（3）采用All gather操作将多张显卡上的模型参数拼接在一起

核心：每张显卡只获取一部分的梯度，只进行一部分的参数更新

3.2 ZeRO2

与第一阶段不同，不需要全部反向传播完成之后再更新模型参数（反向传播完成之后再删除gradient），而是修改成：当计算出网络最后一层的梯度后，就采用reduce scatter得到最后一层参数的gradient*，然后进行参数更新，并且可以删去最后一层各张显卡上的gradient，然后对倒数第2层进行同样操作，依此类推。