大模型训练中优化策略（数据并行、模型并行、ZeRO等）

大耳朵爱学习

于 2024-06-07 09:59:27 发布

阅读量566

点赞数 10

文章标签： python 算法人工智能学习 agi 深度学习 ai

本文链接：https://blog.csdn.net/2401_85379281/article/details/139519248

版权

GPU 显存分析

GPU显存分布.png

在微调时，模型显存占用主要包括模型参数、参数梯度、优化器和中间结果四个部分。

对于一个 6B 参数量的模型，它的模型参数占用为：

6×109×4(FP32)10243≈22GB\frac{6 \times 10^9 \times 4(FP32)}{1024^3} \approx 22GB102436×109×4(FP32)≈22GB

将模型参数视为基准，模型梯度占用量与模型参数相同。

优化器主采用 Adam Optimizer ，它核心计算公式如下：

mt=β1mt−1+(1−β1)gtvt=β2vt−1+(1−β2)gt2m_t = \beta_1 m_{t-1} + (1- \beta_1) g_t \\ v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2mt=β1mt−1+(1−β1)gtvt=β2vt−1+(1−β2)gt2

由于需要保存 m 和 v，而 m 和 v 规模与参数梯度相同，因此优化器需要两倍显存容量。

同时，在计算中得到的中间结果需要保存在显存中，以便反向传播时计算梯度。对于每一个中间结果，其数据形状为 [Batch, SeqLen, Dim]。

Collective Operations

为了节省显存，可以将模型或者数据分配到不同的显卡上，显卡之间有如下几种 [Collective Operations]

Broadcast

广播.png

The Broadcast operation copies an N-element buffer on the root rank to all ranks.

广播操作将一张显卡上数据广播到所有显卡。

AllReduce、Reduce、ReduceScatter

The AllReduce operation is performing reductions on data (for example, sum, min, max) across devices and writing the result in the receive buffers of every rank.

The Reduce operation is performing the same operation as AllReduce, but writes the result only in the receive buffers of a specified root rank.

The ReduceScatter operation performs the same operation as the Reduce operation, except the result is scattered in equal blocks between ranks, each rank getting a chunk of data based on its rank index.

AllReduce 操作将所有显卡上数据进行聚合如求和、取最大值或取最小值，并将结果写入所有显卡。

Reduce 只会将结果写入一张显卡。

ReduceScatter 则将结果分散在所有显卡中。

AllGather

The AllGather operation gathers N values from k ranks into an output of size k*N, and distributes that result to all ranks.

AllGather 操作会收集所有显卡数据，并写入所有显卡中。

数据并行

数据并行是将数据分成若干份，装载到不同节点上进行计算。

数据并行.png

数据并行计算流程如下：

有个参数服务器保存模型参数。
参数被复制到不同的设备中，构成若干 replicas。每个 replica 处理一部分数据，进行前向传播和反向传播。
每个设备得到梯度进行 Reduce 操作，得到最终梯度，并按照这个梯度更新参数服务器中的模型参数。
在后向传播时，每计算完一层的梯度，就可以进行 Reduce 操作，提高并行性。

分布式数据并行

分布式数据并行.png

分布式数据并行中不存在参数服务器，其计算流程如下：

每个 replica 都保存模型参数，但是分别计算部分数据，进行前向传播和反向传播。
每个设备都得到梯度后进行 AllReduce 操作，将梯度写入所有设备，每个设备根据自己的优化器和梯度更新参数。

分布式数据并行中，每个设备显存占用情况如图：

分布式数据并行显存占用.png

其中每个设备仍需要保存模型参数、梯度和优化器参数。

模型并行

由于模型越来越大，单个设备保存模型参数、梯度和优化器越来越难。因为深度学习主要是矩阵计算，而矩阵计算可以分块计算，因此可以将模型参数拆成若干份，每份单独计算，以减少显存占用。

yA=WA×BXB=[WAn×B(1);WAn×B(2);WAn×B(3)]XB=[WAn×B(1)XB;WAn×B(2)XB;WAn×B(3)XB]y_A=W_{A \times B} X_B \\ =[W^{(1)}_{\frac{A}{n} \times B};W^{(2)}_{\frac{A}{n} \times B};W^{(3)}_{\frac{A}{n} \times B}] X_B \\ =[W^{(1)}_{\frac{A}{n} \times B} X_B;W^{(2)}_{\frac{A}{n} \times B} X_B;W^{(3)}_{\frac{A}{n} \times B} X_B]yA=WA×BXB=[WnA×B(1);WnA×B(2);WnA×B(3)]XB=[WnA×B(1)XB;WnA×B(2)XB;WnA×B(3)XB]

模型并行.png

其计算流程如下：

将参数矩阵分成若干子矩阵，分发到不同设备中。
每个设备计算不同矩阵，然后将结果收集起来。

模型并行后，显存占用如下：

模型并行显存占用.png

由于每个设备处理所有数据，因此中间结果都会保存在所有设备中。

ZeRO

在分布式数据并行中，最后梯度更新在不同设备进行的操作相同，多个设备中参数相同，梯度相同，优化器状态相同，存在大量冗余。

ZeRO-1 对优化器状态进行分片。

ZeRO-1 计算流程如下：

每个 replica 处理一部分数据输入。
独立进行前向传播。
独立进行反向传播。
得到完整梯度后进行 ReduceScatter，每个 replica 得到对应梯度。
每个 replica 更新梯度对应的部分参数。
使用 AllGather 同步更新所有参数。

ZeRO-2 计算流程与1基本相同，ZeRO-2在后向传播时，每计算一层梯度，就可以使用 ReduceScatter 进行同步，提高并行度。同时由于不需要完整计算梯度之后进行 ReduceScatter，每个 replica 只需要保存部分梯度即可。

ZeRO-3 在 2 的基础上，将模型参数进行分片。

ZeRO-3 计算流程如下：

每个 replica 处理一部分输入。
前向传播时，当需要别的层参数，使用 AllGather 获取。
反向传播时，当需要别的层参数时，使用 AllGather 获取，同时计算出每一层梯度时，使用 ReduceScatter 分发到对应 replica。
每个 replica 用于部分优化器参数和梯度，进行对应参数更新。

不同 ZeRO 对应的显存占用情况：

ZeRO显存占用.png

流水线并行

将模型一层一层分开，不同层放入不同 GPU 进行计算。个人理解与模型并行不同的是，模型并行保留从头到尾每一层的部分参数，输入可以计算出结果。流水线并行需要等前一层计算完毕才能进行计算。

流水线并行.png

流水线并行显存分析：

流水线并行显存分析.png

混合精度

FP16 相较于 FP32 计算更快，同时占用更少的显存。但同时 FP16 表示的范围小，可能产生溢出错误。

特别的，在权重更新时 gradient * lr 导致下溢出。

混合精度训练的思路在优化器中保留一份 FP32 格式的参数副本，而模型权重、梯度等数据在训练中都是用 FP16 来存储。

混合精度.png

优化器中参数更新在 FP32 格式下保证精度，之后转换为 FP16 格式。

Checkpointing

由于模型反向传播需要中间结果计算梯度，大量中间结果占用大量显存。

Checkpointing 思路是保存部分隐藏层的结果（作为检查点），其余的中间结果直接释放。当反向传播需要计算梯度时，从检查点开始重新前向传播计算中间结果，得到梯度后再次释放。

那么，我们该如何学习大模型？

作为一名热心肠的互联网老兵，我决定把宝贵的AI知识分享给大家。至于能学习到多少就看你的学习毅力和能力了。我已将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

一、大模型全套的学习路线

学习大型人工智能模型，如GPT-3、BERT或任何其他先进的神经网络模型，需要系统的方法和持续的努力。既然要系统的学习大模型，那么学习路线是必不可少的，下面的这份路线能帮助你快速梳理知识，形成自己的体系。

L1级别:AI大模型时代的华丽登场

L2级别：AI大模型API应用开发工程

L3级别：大模型应用架构进阶实践

L4级别：大模型微调与私有化部署

一般掌握到第四个级别，市场上大多数岗位都是可以胜任，但要还不是天花板，天花板级别要求更加严格，对于算法和实战是非常苛刻的。建议普通人掌握到L4级别即可。

以上的AI大模型学习路线，不知道为什么发出来就有点糊，高清版可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。那以下这些PDF籍就是非常不错的学习资源。

四、AI大模型商业化落地方案

作为普通人，入局大模型时代需要持续学习和实践，不断提高自己的技能和认知水平，同时也需要有责任感和伦理意识，为人工智能的健康发展贡献力量。