1.背景近年来，随着“大模型”概念的引入，深度学习模型变得越来越大。如何训练这些大型模型已成为亟待解决的工程问题。最初的视觉模型只有几百兆的参数，但现在语言模型拥有数百亿、数千亿甚至数万亿的参数并不罕见。如此庞大的参数量会消耗巨大的存储空间。下表显示了当前大模型的参数（以Float32计算）以及对应的存储空间。
目前最好的Nvidia GPU显卡只有40G显存容量。显然，将大型模型塞进一张显卡是不现实的。本质上，所有大型模型都是以分布式方式训练的。目前的分布式训练中，常用的有数据并行、模型并行和管道并行。就计算效率而言，数据并行远远优于模型并行和管道并行。然而，数据并行消耗的视频内存量最大，因为它需要整个网络在一个GPU 上运行。在模型训练过程中，除了参数之外，还有很多地方需要存储空间，这进一步增加了训练大型模型时的内存消耗。那么你是否好奇，如何训练这么大的深度学习模型呢？
2.深度学习中的视频内存使用。在讨论如何进行大规模训练之前，我们先详细了解一下网络中的显存使用情况。通常在深度学习训练过程中，涉及的内存使用包括：网络参数、梯度、激活值、激活值的梯度、优化器状态信息，如果使用混合精度[6]训练，还有备份参数。 (master_weight)等。这里需要指出的是，旧的caffe框架中并没有对激活值的梯度进行优化，它与激活值占用相同的空间。不过在tensorflow、pytorch等框架中，已经做了很好的优化，所以激活值的梯度实际上并没有占用大量的显存空间。另外，很多朋友可能对之前的显存使用类型（网络参数、梯度、激活值、激活值的梯度）比较清楚，但是对于优化器的状态信息以及混合精度的备份参数（master_weight）[ 6]训练）不是很清楚，这里稍微解释一下。
2.1 SGD优化器在简单的SGD优化器中，使用以下公式来更新参数：
2.2 Moment SGD优化器但通常我们不会直接使用SGD来更新参数，而是对梯度进行滑动平均，然后进行更新。即使用Moment SGD优化器，计算公式如下：
编辑切换到主页
由于需要进行计算，因此需要保存在显存中。它是优化器的状态信息。它的大小与梯度一致，因此与参数大小一致。例如，如果参数大小是100亿，那么优化器缓存信息也是100亿。
2.3 ADAM优化器很多情况下，我们也会使用ADAM优化器来更新参数，ADAM会使用梯度的一阶矩估计和二阶矩估计。公式如下：
编辑切换到中心
同样的，既然要使用计算，要使用计算，那么和两个变量也需要保存在显存中。它们的大小也与梯度一致，因此也与参数一致。例如，如果参数大小为100亿，则总大小为200亿。
2.4 混合精度[6] 有时为了提高效率，我们会使用混合精度[6]进行训练。为了抵消参数更新时float16的舍入误差，混合精度[6]训练会额外保存一份FP32参数。用于参数更新，称为master-weights，因此会占用较多显存空间。
编辑切换到中心
2.5 激活值优化和参数优化上面分析了深度学习网络的显存占用，在不同的架构中，各部分占用显存的比例是不同的。在CNN中，激活值通常占据大部分显存空间。 MLP/Transformer等结构是模型的参数和参数的梯度，优化器的状态信息占据较大比例。因此，对于不同的网络结构，优化策略是不同的。对于CNN 网络，通常优化的重点是激活值。 MLP/Transformer更关注与网络参数相关的内存优化。因此，对于两种不同类型的网络，需要分别采用激活值优化策略和参数优化策略。
3.激活值优化策略3.1 以时间换空间正如前面分析的，CNN网络的显存消耗主要是激活值和激活值的梯度。随着输入分辨率的提高和批量大小的扩大，激活值和激活值梯度的内存使用量将按平方因子增加。因此，大型CNN模型的训练主要集中在激活值的优化上。这里我们介绍一种通过增加计算时间来减少内存空间占用的方法：“次线性内存优化[5]”。首先我们看一下传统的深度学习流程，如下图所示。前向计算后，保存所有激活值，如图中的a1、a2、a3、a4。反向计算时，根据之前的激活值计算各层的梯度，包括激活值梯度和参数梯度。
编辑切换到中心
“次线性记忆优化[5]”的深度学习过程与传统的深度学习过程有些不同，如下图所示。在前向计算中，为了减少内存消耗，一些激活值，例如a2，会被选择性地丢弃。逆向计算时，按常规方法进行计算。当激活值缺失时，例如计算layer3的参数梯度时，需要使用激活值a2，但a2已被丢弃。这时，反向传播就会暂停。再次进行最短路径的正向计算，根据a1计算a2（如果a1也被丢弃，则继续向前查找）。然后继续原来的反向传播。
编辑切换到中心
可以看出“亚线性内存优化[5]”完全采用了时间换空间的策略。那么你好奇它为什么有这样的名字吗？其实很简单理解。假设每一层的激活值相同，那么整个网络的激活值与层数线性相关。层的深度增加数倍，激活值的占用扩大数倍。使用“亚线性内存优化[5]”策略后，显存占用与层数增加的关系不再是线性的，而是次线性的，因此被称为“亚线性内存优化[5]”。事实上，如何丢弃激活值的选择会极大地影响最终的网络性能。丢弃的激活值需要满足前向计算简单（减少重新计算耗时）和激活值占用空间大的特点。比如BN层的计算非常简单，重新计算基本不耗时（相对于卷积），所以更适合丢弃。
3.2 低精度训练此外，还可以采用低精度训练。这样，激活值和激活值的梯度就以更小的数据格式存储。可以大大减少激活值的存储空间。例如，当使用混合精度进行训练时[6]，所有激活值都使用float16的格式。与原来的float32相比，内存占用直接减少了一半。
4.参数优化策略上面介绍的激活值优化策略适用于CNN等结构。对于目前的一些大型模型来说，其显存占用主要集中在参数以及参数相关的显存占用上。例如参数的梯度、优化器的状态信息等。因此，如何优化与参数相关的内存使用对于大型模型的训练来说更为重要。
4.1 ZeRO[4]数据并行原理微软开源的DeepSpeed训练框架采用了一种名为ZeRO[4]的显存优化技术，称为零冗余优化技术。本质上，它是一种数据并行分布式训练策略，重点优化数据并行中的内存使用问题。在ZeRO[4]数据并行中，虽然每个GPU都有完整的网络，但每个GPU只保存一部分参数、梯度和优化器状态信息，从而可以将参数、梯度和优化器状态信息平均分配给多个GPU。对于参数视图较大的网络，视频内存的减少将是巨大的。但由于分布式存储参数，也会导致通信量的增加。
4.2 传统数据并行流程首先简单回顾一下传统数据并行流程，包括正向计算、反向计算和参数更新。假设总共有两个GPU参与训练。前向流程如下图所示。在初始阶段，每个GPU都被初始化为相同的参数，并分为互斥的训练子集。每个站点的GPU独立完成所有层的前向计算。
编辑切换到中心
其逆过程如下图所示。初始阶段，通过loss对最后一个激活值进行求导，得到了激活值的梯度。然后每个GPU独立地执行所有层的逆向计算。
编辑切换到中心
参数更新流程如下图所示。首先对所有GPU的梯度进行归约操作（平均值），然后针对每个GPU独立更新参数。 （由于初始参数值相同，梯度下降后也相同，所以最终每个GPU上更新的参数也相同）
编辑切换到中心
4.3 ZeRO[4]数据并行流程ZeRO[4]数据并行有多个级别，分别是os级别（只优化优化器状态）、os+g级别（优化优化器状态+梯度）、os+g+p级别（优化优化器状态+梯度+参数）。我们直接分析最优化的os+g+p流程。首先看前向计算过程。初始阶段，每个GPU只保存W/GPU_NUM的参数。这里假设网络有两层，总共有两个GPU参与并行。因此，GPU1只保存了layer1的参数w1，GPU2只保存了layer2的参数w2。在进行layer1的前向计算之前，由于GPU2没有layer1的参数，因此需要对w1进行参数分配。然后进行layer1的前向计算。同样，在进行layer2的前向计算之前，需要对layer2的参数w2进行一次分布，然后完成layer2的前向计算。整体示意图如下图所示。
编辑切换到中心
我们来分析一下ZeRO[4]的数据并行反向传播过程。同样，在反向计算该层之前，需要对参数进行分布。然后进行反向传播计算。完成反向传播后，会有一个梯度收集过程。例如GPU2需要保存w2对应的梯度g2，因此所有其他GPU都将g2梯度发送给GPU2。获得GPU2上各个GPU的g2梯度后，进行约简操作并保存，得到g2~。其他GPU会删除w2、g2。然后重复这个过程，直到所有层都完成反向传播计算。示意图如下所示。
编辑切换到中心
最后我们来分析一下ZeRO[4]数据并行性的参数更新过程。由于反向传播时已经进行了梯度降低操作，ZeRO[4]数据并行性可以直接更新优化器的状态信息，进而更新参数。原理图如下图所示。
4.4 ZeRO[4]数据并行通信开销从前面的ZeRO[4]数据并行流程可以看出，ZeRO[4]数据并行在os+g+p层面进行了优化（优化优化器状态+梯度+参数）时，会有两种参数分布（一种正向计算，一种反向计算）和一种梯度集合。传统的数据并行只需要执行一次梯度下降。那么ZeRO[4]数据并行的通信消耗会是传统数据并行的3倍？事实上，在传统的数据并行中，虽然只需要梯度约简操作，但由于每个GPU都需要获取约简后的梯度，因此使用了all-reduce通信原语。在ZeRO[4]数据并行中，虽然有3个数据传输，但只需要一对多的分布参数或多对一的梯度收集，使用广播和聚集的通信原语。广播和聚集的通信消耗基本相同，大约是all-reduce的一半。因此，最终在os+g+p级别通信开始时的ZeRO[4]数据并行度是原始数据并行度的1.5倍，而不是3倍。当使用os+g级优化或者os级优化时，通信消耗相当于原来的数据并行度。关于分布式数据并行中的通信原语和通信消耗，我会另外写一篇文章来分析。
4.5 ZeRO[4]论文原图分析最后我们来分析ZeRO[4]论文中最nuble的图片。如何将传统数据并行需要120G显存的模型变成只需要1.9G。没看到你不知道，但是当你看到的时候真的很震惊！在分析这张图之前，需要知道一个前提，那就是都是基于混合精度[6]训练，采用adam的优化策略。
编辑切换到中心
首先，这意味着模型参数大小为7.5B，即75亿个参数。由于所有训练均基于混合精度训练，因此参数和梯度均使用float16 存储。一个参数占用两个字节，参数和渐变的显存占用是其两倍。 K表示的梯度状态信息和混合精度master_weight的内存占用是多少倍。由于梯度状态信息和混合精度master_weight必须使用float32存储，即一个参数占用4字节存储空间，而adam中有两个状态信息，分别是一阶矩估计和二阶矩梯度的估计，所以K=(2 + 1)*4=12 次。 Nd=64表示使用64个GPU进行零数据并行训练。首先看第一行的Baseline。采用传统的数据并行方式，每个GPU 的内存消耗为：然后看第二行，使用os级别的优化，那么参数和梯度大小都没有改变，优化器状态+master_weight均匀分配到所有GPU上，所以每个GPU的内存消耗是31.4G；同样分析第三个OK，使用os+g级别优化。由于梯度数据均匀分布到所有GPU上，因此每个GPU的内存消耗为16.6G。归根结底，采用了os+g+p级别的优化，参数也均匀的分配到所有GPU上，所以最终每个GPU的显存消耗为1.9G。从上面的分析可以看出，采用os+g+p级别优化时，每个GPU的显存消耗是传统数据并行的1/Nd。用白话来说，就是使用更多的GPU可以减少显存消耗。多少次了。
5.总结以上是大模型训练中显存使用的一些优化措施。包括针对激活值优化的策略和针对参数优化的策略。正是有了这些强有力的工程手段，才使得大型模型的训练成为可能。但由于作者水平有限，时间仓促，难免有错误，希望读者不吝指正。感激的。
6、参考[1]DeepSpeed的ZeRO系列：将显存优化进行到底
https://zhuanlan.zhihu.com/p/513571706
[2] 2022年，大模型能走多远？
https://www.51cto.com/article/697186.html
[3]ZeRO+DeepSpeed:微软发布的高效大规模训练套件（含详细的分布式训练流程）
https://zhuanlan.zhihu.com/p/108571246
[4] Rajbhandari S、Rasley J、Ruwase O 等人。 ZeRO: 针对训练万亿参数模型的内存优化[C]//SC20: 高性能计算、网络、存储和分析国际会议。 2020.
[5] 陈天琪，徐兵，张驰远，卡洛斯·格斯特林。使用次线性内存成本训练深度网络。 arXiv 预印本arXiv:1604.06174，2016。
[6] Micikevicius, P.Narang, S.Alben, J.Diamos, G.Elsen, E.Garcia, D. 等人。 （2017）。混合精准训练。
