文章目录

摘要
1、简介
2、相关工作
3、方法
4、实验
5、结论

摘要

论文链接：https://arxiv.org/pdf/2306.09782v1.pdf
大型语言模型(LLM)已经彻底改变了自然语言处理(NLP)，但需要大量的GPU资源进行训练。降低LLM的训练门槛将鼓励更多的研究人员参与，使学术界和社会都受益。虽然现有的方法主要集中在参数高效的微调上，即调整或增加少量参数，但很少有人能够在有限的资源下解决调优LLM的全部参数的挑战。在本文中，我们提出了一种新的优化器，低内存优化(LOMO)，它将梯度计算和参数更新融合在一起，以减少内存的使用。通过将LOMO与现有的内存节省技术集成，与标准方法(DeepSpeed解决方案)相比，我们将内存使用量减少到10.8%。因此，我们的方法能够在一台具有8×RTX 3090的机器上对65B模型进行全参数微调，每块显卡具有24GB内存。

1、简介

大型语言模型(LLM)已经彻底改变了自然语言处理(NLP)，展示了显着的能力，如涌现和grokking (Wei et al.， 2022)，推动模型规模变得越来越大。然而，用数十亿个参数训练这些模型，比如那些有30B到175B个参数的模型，提高了NLP研究的门槛。调优LLM通常需要昂贵的GPU资源，例如8×80GB设备，这使得小型实验室和公司很难参与这一领域的研究。

最近，参数高效的微调方法(Ding et al.， 2022)，如LoRA (Hu et al.， 2022)和Prefix-tuning (Li & Liang, 2021)，为资源有限的LLM的调优提供了解决方案。然而，这些方法并没有为全参数微调提供实用的解决方案，而全参数微调已被认为是比参数高效微调更强大的方法(Ding et al.， 2022;Sun et al.， 2023)。在这项工作中，我们的目标是探索在资源有限的情况下完成全参数微调的技术。

我们从激活、优化器状态、梯度张量和参数四个方面分析了LLM中的内存使用情况，并从三个方面对训练过程进行了优化:1)我们从算法的角度重新思考了优化器的功能，发现SGD在微调LLM的全参数方面是一个很好的替代品。这允许我们删除优化器状态的整个部分，因为SGD不存储任何中间状态(第3.1节)。2)我们提出的优化器LOMO如图1所示，将梯度张量的内存使用量减少到O(1)，相当于最大梯度张量的内存使用量(第3.2节)。3)为了稳定LOMO混合精度训练，我们在训练过程中集成了梯度归一化、损失缩放，并将某些计算转换为全精度(第3.3节)。
在这里插入图片描述

我们的技术导致内存使用等于参数加上激活和最大梯度张量的使用。我们将全参数微调的内存使用推向了一个极端，使其仅仅等同于推理的使用。这是因为前向传播+反向传播进程的内存使用量不应该少于单独前向传播进程的内存使用量。值得注意的是，当使用LOMO来节省内存时，我们确保微调过程不受影响，因为参数更新过程仍然相当于SGD。

我们通过经验评估了LOMO的内存和吞吐量性能，并表明LOMO的使用可以仅用8个RTX 3090 gpu成功训练65B模型。此外，为了验证我们提出的技术的下游性能，我们应用LOMO来调整SuperGLUE数据集集合上LLM的全部参数(Wang et al.， 2019)。实证结果表明，LOMO对具有数十亿参数的LLM进行优化的效率和有效性。总的来说，我们的贡献如下:

我们提供了一个理论分析，表明SGD可以成功地微调LLM的全部参数。以前阻碍SGD广泛使用的问题可能不再是微调LLM的严重问题。
我们提出低内存优化，命名为LOMO，以显着节省GPU内存使用，而不会损害微调过程。
通过对内存使用和吞吐量性能的全面评估，我们经验验证了LOMO在资源受限场景下优化LLM的有效性。下游任务的性能评估进一步支持了这一点。

2、相关工作

在本节中，我们将介绍在全参数调优期间有关内存节省技术的相关工作。这些技术可以有效地与LOMO结合使用，以进一步减少内存消耗。

激活检查点。在普通的反向传播期间，来自前向传播的所有激活都存储在内存中以计算梯度。这可能是一个很大的内存开销，特别是对于大型语言模型。或者，可以丢弃所有激活，并根据需要重新计算它们以进行梯度计算，以节省内存。然而，这可能会导致大量的额外计算成本。激活检查点(或梯度检查点)考虑了内存使用和计算成本，提供了一个折衷的解决方案(Chen et al.，2016)。计算图中策略选择的检查点节点的激活在向前传递后保留在内存中，而剩余节点的激活最多重新计算一次。激活内存可以减少到原始数量的平方根，代价是一个额外的向前通过。

混合精度训练。由于混合精度训练能够加快训练速度并减少内存占用，它已经成为训练大型语言模型的一种流行方法(Narayanan等人，2021;Rajbhandari等人，2020)。通过对参数、激活值和梯度采用半精度存储，混合精度训练可以在前向和后向传播期间实现高通量计算。为了保持稳定性和模型精度，micicikevicius等人(2018)提出了三种技术，涉及使用全精度权重副本、损失缩放和全精度执行特定算术操作。

扫描二维码关注公众号，回复： 15424729 查看本文章

异构训练系统。多项研究(Rhu等人，2016;王等，2018;Ren等人，2021a)试图通过利用异构内存(如CPU和NVMe内存)来减少GPU内存消耗。L2L (Pudipeddi et al.， 2020)采用层到层策略，其中只将特定层计算所需的张量传输到GPU内存，而其余张量保留在CPU内存中。ZeRO-Offload (Ren et al.， 2021b)是ZeRO-2的扩展(Rajbhandari et al.， 2020)，它在CPU内存中保留梯度和优化器状态，并通过CPU计算更新参数。张量和计算操作根据数据流图分配给gpu或cpu。ZeRO-Infinity (Rajbhandari等人，2021)，是ZeRO3上ZeRO-Offload的后续改进(Rajbhandari等人，2020)，使模型大小能够进一步扩展。分割后的模型状态和其他张量不仅可以卸载到CPU内存，还可以卸载到NVMe，以充分利用异构架构。

3、方法

3.1、重新思考optimizer的功能

优化器状态占用了用于训练LLM的大部分内存。Adam (Kingma & Ba, 2015)等现代优化器存储的中间状态是参数大小的两倍。随着参数大小的增加，优化器状态成为内存使用的主导项。

3.1.1、使用SGD

尽管Adam在训练深度模型方面取得了巨大的成功，但我们提出了一个问题“我们可以使用更便宜的优化器来微调LLM吗?”我们的答案是SGD，最基本的优化器。幸运的是，当我们限制范围时，发现这是一个可接受的LLM微调解决方案。

先前的工作经常讨论SGD的三个挑战:1)大曲率损失面，2)局部最优，3)鞍点(Ruder, 2016;Sun等人，2020)。现代优化器在处理1)问题方面已经显示出有效性，并且在某些情况下可以缓解2)和3)。然而，当我们将范围限制在微调LLM时，这三个挑战可能会有所不同。

平滑损失平面。一个重要的假设是LLMs的参数空间非常平滑，对参数的微小扰动不会过多改变损失。有实证结果和理论分析支持这一假设(Hao等人，2019)。如果我们相信较大的模型具有更平滑的损失曲面，我们可以得出结论:1)问题不是问题，因为LLM的损失曲面不应该具有较大的曲率。请注意，这只适用于我们教授LLM基于自然语言的任务(或基于代码的，如果用代码预训练)。与预训练任务无关的合成损失函数确实会面临大曲率问题。

在这里插入图片描述

局部最优就足够了。微调的目标是使LLM适应新的任务和域，而不显著改变模型本身。因此，局部最优往往是一个足够好的解决方案，而有限的训练数据(与预训练语料库相比)使得模型很难达到远处的全局最优。

远鞍点。类似地，对于一个常见的NLP任务，LLM的初始点应该在山谷中。如果模型是用指令(任务)进行预训练的，那么这种现象可能会明显得多，因为我们有更多的机会找到与新任务相似的预训练任务。鞍点通常出现在脊上，与谷有一段距离，因此如果我们将参数改变得离预训练值不太远，我们可能不会遇到鞍点问题。

然而，与现代优化器相比，不能保证SGD是一个强大的优化器。我们的直觉是为微调LLM创建一个简单实用的解决方案，并识别其缺陷以不断改进它。

3.1.2、隐式BatchSize

除了上述定性讨论之外，本文想用SGD对微调LLM的稳定性进行更深入的分析。假设我们有一个带参数 $\boldsymbol{\theta}$ 的预训练模型 $f(\cdot)$ ，一个训练集 $\mathcal{D}=\left\{d_{1}, d_{2}, \cdots, d_{n}\right\}$ ，和一个损失函数l，有两个数据点的批次的SGD的一步更新可以是:
$\boldsymbol{\theta}^{\prime}=\boldsymbol{\theta}-\alpha\left[\nabla \mathcal{L}\left(d_{i}, f\left(d_{i}, \boldsymbol{\theta}\right)\right)+\nabla \mathcal{L}\left(d_{j}, f\left(d_{j}, \boldsymbol{\theta}\right)\right)\right]$

其中 $\alpha$ 是学习率， $d_{i}$ ， $d_{j}$ 是两个不同的训练样本。
接下来，在这两个训练样本 $d_{i}$ ， $d_{j}$ 上更新SGD的两个步骤依次为:
$\begin{array}{l} \boldsymbol{\theta}_{\mathbf{1}}=\boldsymbol{\theta}-\alpha \nabla \mathcal{L}\left(d_{i}, f\left(d_{i}, \boldsymbol{\theta}\right)\right) \\ \boldsymbol{\theta}_{\mathbf{2}}=\boldsymbol{\theta}_{\mathbf{1}}-\alpha \nabla \mathcal{L}\left(d_{j}, f\left(d_{j}, \boldsymbol{\theta}_{\mathbf{1}}\right)\right) \end{array}$
根据微分均值定理，我们有
$\begin{array}{c} \mathcal{L}\left(d_{j}, f\left(d_{j}, \boldsymbol{\theta}_{\mathbf{1}}\right)\right)=\mathcal{L}\left(d_{j}, f\left(d_{j}, \boldsymbol{\theta}\right)\right)+\mathcal{L}\left(d_{j}, \xi\right)\left(f\left(d_{j}, \boldsymbol{\theta}_{\mathbf{1}}\right)-f\left(d_{j}, \boldsymbol{\theta}\right)\right), \\ \boldsymbol{\theta}_{\mathbf{2}}=\boldsymbol{\theta}-\alpha \nabla \mathcal{L}\left(d_{i}, f\left(d_{i}, \boldsymbol{\theta}\right)\right)-\alpha \nabla \mathcal{L}\left(d_{j}, f\left(d_{j}, \boldsymbol{\theta}\right)\right)-\alpha \nabla\left[\mathcal{L}\left(d_{j}, \xi\right)\left(f\left(d_{j}, \boldsymbol{\theta}_{\mathbf{1}}\right)-f\left(d_{j}, \boldsymbol{\theta}\right)\right)\right], \\ \boldsymbol{\theta}_{\mathbf{2}}=\boldsymbol{\theta}-\alpha\left[\nabla \mathcal{L}\left(d_{i}, f\left(d_{i}, \boldsymbol{\theta}\right)\right)+\nabla \mathcal{L}\left(d_{j}, f\left(d_{j}, \boldsymbol{\theta}\right)\right)\right]-\alpha \nabla\left[\mathcal{L}\left(d_{j}, \xi\right)\left(f\left(d_{j}, \boldsymbol{\theta}_{\mathbf{1}}\right)-f\left(d_{j}, \boldsymbol{\theta}\right)\right)\right] \end{array}$

其中 $\xi$ 是 $f (d j ， θ)$ 和 $f (d j ， θ 1)$ 之间的一个点，我们可以看到方程6减去方程1等于 $α \nabla [L (d j ， ξ) (f (d j ， θ 1) - f (d j ， θ))]$ 。假设损失曲面足够平滑，这一项可以忽略不计。它表明，在平滑的损失表面上使用SGD优化器可能意味着更大的批量大小。

正如我们上面提到的，我们有理由假设LLMs的损失曲面是平滑的，并且更大的批量大小意味着更强的训练稳定性，因此我们认为，使用SGD优化器的LLMs的微调过程是稳定的。这也解释了为什么SGD在小型模型上失败，但在大型模型上有效。

3.2、LOMO:低内存优化

梯度张量表示参数张量的梯度，与参数具有相同的大小，导致很大的内存开销。像PyTorch这样的现代深度学习训练框架(Paszke等人，2017)存储了所有参数的梯度张量。存储梯度张量有两个原因:计算优化器状态和归一化梯度。

由于我们使用SGD作为优化器，因此没有依赖梯度的优化器状态，我们有一些梯度归一化的替代方案。因此，本文提出了低内存优化(LOMO)，如算法1所示，将梯度计算和参数更新融合在一个步骤中，以避免存储任何梯度张量。

详细地说，我们可以将vanilla梯度下降表示为 $\operatorname{grad}=\frac{\partial \mathcal{L}}{\partial p}$ ，这是一个两步过程，首先计算梯度，然后将其更新为参数。融合版本是 $\operatorname{grad}$ 。

其关键思想是在计算梯度时立即更新参数，这样我们就不会将梯度张量存储在内存中。这可以通过向反向传播注入挂钩函数来实现PyTorch为注入钩子函数提供了相关的api，但我们无法使用当前api实现准确的立即更新。相反，我们在内存中最多存储一个参数的梯度，并随着反向传播逐个更新每个参数。该方法将梯度的内存使用从存储所有参数的梯度减少到只存储一个参数的梯度。

大多数LOMO内存使用与参数高效微调(PEFT)方法一致，这表明将LOMO与这些方法相结合只会引入梯度占用的内存的少量增加。这使得可以调优PEFT方法的更多参数。

3.3、使用LOMO稳定训练

3.3.1、梯度归一化和裁剪的替代方法

梯度归一化和裁剪是处理梯度爆炸和消失问题的必要工具Chen et al.(2018)，但它们的计算过程需要使用所有参数的梯度张量。本文在这里提出两种选择:

根据梯度张量的值而不是范数裁剪梯度张量。
在一个额外的通道中计算梯度规范。

在梯度范数逼近之前，根据梯度张量的值裁剪梯度张量是解决梯度爆炸的一种简单而有效的方法。按值裁剪的主要问题是截断一些梯度元素可能会改变梯度张量的方向。例如，一个二元向量[1.3,0.8]和它的裁剪版本1.0,0.8表示不同的方向。我们的经验是，当学习率较高时，按值裁剪的表现较差，因为在这种情况下截断发生得更频繁。然而，按值裁剪在中、小学习率下表现良好。请注意，学习率的规模在很大程度上取决于任务和数据，但一般来说，我们建议对学习率小于1e−3的值进行裁剪。

我们的方法不能直接计算梯度范数，因为我们随着反向传播更新参数，所以我们在更新某个参数时不知道其余参数的范数。但是，我们可以引入一个额外的通道来计算和累积每个参数的梯度范数，从而产生两个反向传递，一个用于计算梯度范数，另一个用于更新参数。内存使用保持不变，但牺牲了速度。

一个有争议的解决方案。我们当前的训练框架基于所有参数计算梯度范数，并需要两次反向传递。保存额外反向传递的一种解决方案是用一组参数近似梯度张量的范数，例如相邻层。然而，这种方法确实是有偏差的，因为它会导致不同参数的更新步长不同。在更新时，根据梯度范数乘以一个比例因子。由于参数组之间的梯度范数不同，这种近似会导致尺度因子的差异。尽管有这个限制，这种分组梯度裁剪方法可以被认为是根据其梯度范数对不同的参数组应用动态学习率。Sun等人(2020)认为，在SGD中对所有参数使用相同的学习率并不总是合适的，因此我们相信我们的方法也有可能进一步使SGD受益。我们将探索作为一个令人信服的未来方向。

3.3.2、缓解精度下降

混合精度训练通常用于加快训练过程。为了缓解精度的下降，利用动态损失缩放，并将某些计算过渡到全精度。损失缩放的方法对于FP16训练期间防止下溢至关重要，在反向传递之前用特定的因子放大损失，并通过相同的因子减少梯度。

在这种情况下，将动态损失缩放与LOMO集成，在整个训练过程中动态调整缩放因子。如果在指定的向后传递次数内没有发生溢出，则比例因子将加倍。否则，这一步被删除，比例因子减半。这个过程与梯度归一化过程中遇到的情况类似。在反向操作完成之前，不知道是否会发生溢出。因此，我们执行了两次反向传递:第一次传递确定是否有溢出，第二次传递在没有检测到溢出时更新参数。动态损失缩放的这两个反向过程可以与梯度归一化同时执行。为了在归一化和缩放等操作中有效地更新参数和处理梯度，这些计算中将梯度及其相关参数转换为完全精度。

4、实验

在本节中，我们从内存配置、吞吐量和下游性能三个方面对所提方法进行了评估。如果没有进一步解释，所有实验都是用LLaMA模型进行的(Touvron等人，2023)，范围从7B到65B。
在这里插入图片描述

4.1、内存配置

首先分析了不同设置下训练过程中模型状态和激活的内存使用情况。如表1所示，在训练lama - 7b模型时，与AdamW优化器相比(Loshchilov & Hutter, 2019)， LOMO优化器的使用导致内存占用从102.20GB大幅减少到14.58GB，与SGD相比，内存占用从51.99GB减少到14.58GB。内存使用的显著减少主要是由于梯度和优化器状态的内存需求减少。因此，训练过程中的内存主要被参数占用，与推理过程中的内存使用量相称。
在这里插入图片描述

优化器状态。图2说明了将AdamW优化器用于lama - 7b训练(一种广泛采用的配置)，产生了分配给优化器状态的大量内存(73.7%)。这个结果是混合精度训练方法的结果，在这种方法中，权重、动量和方差的全精度副本被维护在优化器状态中，以进行权重更新。将AdamW优化器替换为SGD优化器可以有效降低优化器状态在内存中的百分比，从而减少GPU内存使用量(从102.20GB降低到51.99GB)。这种减少是由于SGD优化器不需要存储全精度动量和方差。对于LOMO，参数更新和反向更新融合在一个步骤中，进一步消除了对优化器状态记忆的需要。
在这里插入图片描述

梯度。在使用LOMO的训练过程中，参数在接收到梯度后立即更新，然后从内存中丢弃梯度。因此，梯度内存消耗的上界由与最大大小的参数矩阵相关联的梯度决定。这种方法大大减少了内存使用，几乎减少了参数的大小。

激活。在一个批次中训练带有512×8标记的7B模型需要大量的内存来激活。LOMO与激活检查点等减少激活内存的技术兼容。通过将激活检查点与LOMO集成，可以将由激活引起的内存占用从45.61GB减少到1.79GB。

4.2、吞吐量

我们评估了LOMO与AdamW和SGD的吞吐量性能。实验在配备8个RTX 3090 GPU的服务器上进行，通过PCIe主板互连。序列长度和批量大小分别设置为1024和1。吞吐量是根据每GPU每秒处理的token数量(TGS)来衡量的，并且参数划分是使用ZeRO-3实现的(Rajbhandari等人，2020)。

对于7B模型，LOMO表现出显著的吞吐量，超过AdamW和SGD约11倍。这种显著的改进可以归因于LOMO在单个GPU上训练7B模型的能力，从而减少了GPU间的通信开销。SGD的吞吐量略高于AdamW，这可以归因于SGD排除了动量和方差的计算。

至于13B模型，由于内存限制，无法在现有的8块RTX 3090 GPU上使用AdamW进行训练。在LOMO需要模型并行的情况下，LOMO在吞吐量方面仍然优于SGD。这一优势归功于LOMO的内存高效特性，以及在相同设置下只需要两个gpu就可以训练模型，从而减少了通信成本和更大的吞吐量。此外，在训练30B模型时，SGD在8个RTX 3090 GPU上遇到内存不足(OOM)问题，而LOMO在仅4个gpu上表现良好。

最后，使用8块RTX 3090 GPU成功训练了65B模型，取得了4.93 TGS的吞吐量。利用这样的服务器配置和LOMO，在1000个样本上的训练过程需要大约3.6小时，每个样本包含512个令牌。

4.3、下游性能

为评估LOMO在微调大型语言模型方面的有效性，进行了一组广泛的实验。将LOMO与另外两种方法进行了比较，一种是不需要微调的Zero-shot，另一种是目前最流行的参数高效微调技术之一LoRA。如(Hu等人，2022)所述，LoRA重新参数化密集层，只更新低秩矩阵，而在推理过程中没有引入延迟。

我们使用SuperGLUE数据集来评估模型性能，特别是RTE (Dagan等人，2005年)、BoolQ (Clark等人，2019年)、WSC (Levesque等人，2012年)、WIC (Pilehvar和camachoc - collados, 2019年)、MultiRC (Khashabi等人，2018年)和COPA (Roemmele等人，2011年)。考虑到运行大型语言模型相关的高计算成本，本文跟随MeZO (Malladi等人，2023)从训练集随机采样1000个训练数据和从验证集随机采样1000个测试数据，并报告使用相同的随机种子获得的最佳结果。实验中使用的提示与MeZO相同，超参数的详细信息在附录a中。在推理过程中，我们将不同的标签或候选标签插入提示符，并计算每个标签的平均对数似然。分数最高的标签被选择为模型的答案。为了评估性能，我们使用准确性作为评估指标。

4.3.1、主要结果

与零样本和LoRA相比，LOMO的下游性能如表3所示。基于结果，我们得到以下观察结果。
在这里插入图片描述

LOMO的性能明显优于Zero-shot。在所有6个数据集和模型大小上，LOMO始终取得了优于零样本的结果，使用LLaMA-13B平均提高了20多个点。虽然之前的研究已经展示了大型语言模型在零样本设置中的令人印象深刻的能力，但微调仍然可以为特定的下游任务带来显著的性能提升。实验结果证实了LOMO在优化不同规模的大型语言模型方面的有效性。

在大多数实验中，LOMO通常优于LoRA。与LoRA相比，LOMO提供了强大的性能，例如，使用LLaMA13B导致平均提高2.8个点。这表明，与参数高效的微调相比，全参数微调对模型性能的好处更多，因为前者调整了更多的参数。LOMO在性能和效率之间取得了很好的平衡，使其成为微调的一个有竞争力的选择。

在某些情况下，LOMO的表现比LoRA更差。一个可能的原因是我们使用的训练集相对较小，这可能不足以对大型模型进行全参数微调。此外，LoRA和LOMO采用不同的模型架构。具体来说，LoRA为模型调优提供了一种快捷方式，这在某些情况下可能是有利的。实际上，这两种方法并不冲突或互斥。在下一小节中，我们将验证LoRA与LOMO结合不会损害模型性能，并且在大多数情况下会带来性能提升。

LOMO可以有效扩展到650亿个参数模型。尽管所有实验都是在配备8 × RTX 3090的单机上进行的，但即使在65个参数的规模上，LOMO也始终表现出强大的性能。这进一步支持了LOMO在资源受限场景下优化LLM的有效性。

4.3.2、LoRA与LOMO

LOMO和LoRA基本上是相互独立的。为了验证这一结论，使用LLaMA-13B在BoolQ和MultiRC数据集上进行了实验。结果如图3所示。发现LOMO始终提高了LoRA的性能，而不管LoRA取得了更高的结果。这表明LOMO和LoRA采用的不同微调方法是互补的。具体来说，LOMO专注于微调预训练模型的权重，而LoRA则调整其他模块。因此，LOMO不会影响LoRA的性能;相反，它有助于更好地对下游任务进行模型调优。
在这里插入图片描述

5、结论

本文提出低内存优化(LOMO)，一种新的优化器，旨在促进资源有限的大型语言模型的全参数微调。我们已经演示了在配备了消费级gpu(如RTX 3090)的服务器上微调65B模型的可行性。通过分析LOMO的内存使用情况、进行吞吐量测试以及在SuperGLUE上进行实验，展示了其有效性和潜在影响。

展望未来，我们未来的工作旨在进一步降低训练大型语言模型所需的资源门槛，从而使这些模型得到更广泛的获取和采用。在使用LOMO进行训练时，目前大部分内存被参数占用。因此，一个有希望的方向是探索参数量化技术，它可以显著减少内存使用。研究LOMO的更多适用场景，并深入研究优化大型语言模型的理论分析，对推动该领域的发展具有重要价值。

loss surface：损失平面

资源有限的大型语言模型的全参数微调