近日复旦大学的研究团队发布了FULL PARAMETER FINE-TUNING FOR LARGE LANGUAGE MODELS WITH LIMITED RESOURCES。实现在有限资源下对大模型进行全参数微调。

论文：https://arxiv.org/abs/2306.09782

代码：https://github.com/OpenLMLab/LOMO

大型语言模型（LLMs）已经彻底改变了自然语言处理（NLP），但是训练LLMs需要大量的GPU资源。降低LLMs训练的门槛将鼓励更多研究人员参与，从而使学术界和社会受益。虽然现有的方法着重于参数高效微调，即微调或添加少量参数，但很少有人解决了有限资源下调整LLMs的全部参数的挑战。在本文中，我们提出了一种新的优化器LOw-Memory Optimization（LOMO），它将梯度计算和参数更新融合在一步中以减少内存使用。通过将LOMO与现有的内存节省技术集成，我们将内存使用降低到10.8％，与标准方法（DeepSpeed解决方案）相比。因此，我们的方法使单台机器上的65B模型的全参数微调成为可能，该机器配有8×RTX 3090，每个显存为24GB。

大型语言模型（LLMs）已经彻底改变了自然语言处理（NLP），展示了出色的能力，如出现和理解（Wei et al.，2022），推动了模型的大小变得越来越大。然而，使用数十亿个参数来训练这些模型，例如具有30B到175B个参数的模型，使NLP研究的门槛提高了。微调LLMs通常需要昂贵的GPU资源，例如8×80GB设备，这使得小型实验室和公司难以参与这个研究领域。

最近，一些参数高效微调方法（Ding等，2022），如LoRA（Hu等，2022）和Prefix-tuning（Li＆Liang，2021），为具有有限资源的LLMs提供了解决方案。然而，这些方法并没有提供适用于全参数微调的实际解决方案，而全参数微调被认为比参数高效微调更为强大（Ding等，2022; Sun等，2023）。在本研究中，我们旨在探讨在资源有限的情况下实现全参数微调的技术。

我们分析LLMs中内存使用的四个方面，即激活，优化器状态，梯度张量和参数，并在三个方面优化训练过程：1）我们从算法的角度重新思考了优化器的功能，并发现SGD是用于LLMs的全参数微调的很好的替代方法。这使我们能够删除整个优化器状态的部分，因为SGD不存储任何中间状态（Sec-3.1）。2）我们提出的优化器LOMO如图1所示，将梯度张量的内存使用降低到O（1），相当于最大梯度张量的内存使用（Sec-3.2）。3）为了稳定使用LOMO的混合精度训练，我们集成了梯度归一化，损失缩放和在训练期间转换某些计算以进行完全精度（Sec-3.3）。

我们的技术导致内存使用量等于参数使用量加上激活和最大梯度张量的使用量。我们将全参数微调的内存使用推向极限，使其仅相当于推理过程的使用量。这是因为前向+反向过程的内存使用量不应少于仅有前向过程的使用量。值得注意的是，在使用LOMO来节省内存时，我们确保微调过程保持不受影响，因为参数更新过程仍然等同于随机梯度下降（SGD）的过程。

我们通过实证评估了 LOMO 的内存和吞吐量表现，并且证明了 LOMO 的使用使得我们可以成功地使用8个 RTX 3090 GPU训练包含 65B 个模型参数的模型。为了验证我们提出技术的下游效果，我们将 LOMO 应用于 SuperGLUE 数据集收集的 LLMs 的全部参数调整中（Wang et al.，2019）。实验结果展示了 LOMO 为优化具有数十亿参数的 LLMs 的效率和有效性。总的来说，我们的贡献包括：

? 我们提供理论分析，证明 SGD 可以成功微调 LLM 的全部参数。之前可能阻碍 SGD 广泛使用的问题在 LLM 微调方面不再是严重问题。

? 我们提出了名为 LOw-Memory Optimization（LOMO）的低内存优化技术，可以显著减少 GPU 内存的使用且不影响精细微调流程。

? 通过内存使用和吞吐量表现的全面评估，我们实证验证了 LOMO 在有限资源下优化 LLMs 的有效性，并且得到了下游任务性能的支持。

本部分介绍在全参数微调过程中节省内存的相关技术，这些技术可以与 LOMO 有效地结合使用以进一步减少内存占用。

激活检查点技术（Activation Checkpointing）：

在标准反向传播中，所有正向过程中的激活变量都会被存储在内存中以计算梯度，这可能是巨大的内存开销，尤其是对于大型语言模型。另一种方法是，可以舍弃所有激活变量，并在计算梯度时重新计算这些变量以节省内存。但这可能会产生大量的额外计算代价。激活检查点技术同时考虑了内存使用和计算成本，提供了一种折中解决方案（Chen等人，2016）。计算图中策略性地选择的检查点节点的激活变量在正向阶段后保留在内存中，而其余节点的激活变量在梯度计算时最多重新计算一次。激活内存可以降低原始内存使用量的平方根，代价是需要多一次正向传递。

混合精度训练（Mixed-Precision Training）

由于其加速训练速度和降低内存占用的能力，混合精度训练已成为训练大型语言模型的流行方法（Narayanan等，2021; Rajbhandari等，2020）。通过为参数，激活和梯度使用半精度存储，混合精度训练在正向传播和反向传播期间实现高吞吐计算。为了保持稳定性和模型准确性，Micikevicius等（2018）提出了三种技术，包括使用完整精度权重副本，损失缩放和在完整精度下执行特定算术运算。

异构训练系统（Heterogeneous Training System）

多项研究（Rhu等，2016; Wang等，2018; Ren等，2021a）尝试通过利用异构内存（如CPU和NVMe内存）来减少GPU内存消耗。L2L（Pudipeddi等，2020）采用逐层策略，仅将计算特定层所需的张量传输到GPU内存，而将其余张量保留在CPU内存中。ZeRO-Offload（Ren等，2021b）是ZeRO-2（Rajbhandari等，2020）的扩展，将梯度和优化器状态预留在CPU内存中，并通过CPU计算更新参数。张量和计算操作根据数据流图分配给GPU或CPU。ZeRO-Infinity（Rajbhandari等，2021）是ZeRO-Offload在ZeRO 3（Rajbhandari等，2020）上的后续进展，可进一步扩展模型大小。分区模型状态和其他张量不仅可以卸载到CPU内存中，还可以卸载到NVMe以充分利用异构架构。

优化器状态占用了用于训练LLM的大部分内存。像Adam（Kingma＆Ba，2015）这样的现代优化器存储中间状态，其大小是参数的两倍。随着参数的增加，优化器状态成为内存使用的主导术语。

虽然Adam在训练深度模型方面取得了巨大成功，但我们问一个问题“我们可以使用更便宜的优化器来微调LLM吗？”我们的答案是最基本的优化器SGD。幸运的是，我们发现当我们限定范围时，它是微调LLM的可接受解决方案。以前的研究经常讨论SGD的三个挑战：1）大曲率损失表面，2）局部最优解，以及3）鞍点（Ruder，2016; Sun等，2020）。现代优化器已经显示出处理1）问题的有效性，并且在某些情况下可以减轻2）和3）问题。然而，当我们将范围限定为微调LLM时，这三个挑战可能会有所不同。

更平滑的损失面（Smoother loss surface）

一个重要的假设是LLM的参数空间非常平滑，对参数进行小的扰动不会显着改变损失。有实证结果和理论分析支持这一假设（Hao等，2019）。如果我们相信更大的模型具有更平滑的损失曲面，那么我们可以得出结论：由于LLM的损失表面不应具有很大的曲率，因此1）的问题不是问题。请注意，仅当我们教LLM基于自然语言的任务（或者像以前一样用代码预训练）时，才有效。与预训练任务无关的合成损失函数确实会面临大曲率问题。

局部最优解已经足够（Local optimum is good enough）

微调的目标是将LLM调整到新任务和领域中，并且不会显着改变模型本身。因此，局部最优解通常是足够好的解决方案，并且有限的训练数据（与预训练语料库相比）使其难以将模型推向遥远的全局最优解。

算法1： LOMO中的融合更新

远离的鞍点（Distant saddle points）

同样，对于常见的NLP任务，LLM的初始点应该在一个山谷中。如果使用指令（任务）预先训练模型，则现象可能更加明显，因为我们有更多机会找到与新任务类似的预先训练任务。鞍点通常出现在山脊上，并且与山谷相距甚远，因此，如果我们不将参数更改得太远离预训练值，则可能不会遇到鞍点问题。然而，无法保证SGD相比现代优化器是一个强大的优化器。我们的直觉是创建一个简单且实用的微调LLM解决方案，并确定其缺点以持续改进。

正如我们上面提到的，假设LLMs的损失面是平滑的，而较大的批量大小表示更强的训练稳定性，因此我们认为使用SGD优化器对LLMs进行微调是稳定的。这也解释了SGD在小模型上失败但在大模型上工作的原因。

梯度张量表示参数张量的梯度，并具有与参数相同的大小，从而产生大量的内存开销。像PyTorch（Paszke等，2017）等现代深度学习训练框架为所有参数存储梯度张量。存储梯度张量有两个原因：计算优化器状态和规范化梯度。由于我们采用SGD作为优化器，在没有依赖于梯度的优化器状态的情况下，我们有一些替代梯度规范化的方法。因此，我们提出了LOw-Memory Optimization（LOMO），如算法1所示，将梯度计算和参数更新融合在一步中，以避免存储任何梯度张量。

具体来说，我们可以将基本的梯度下降表示为grad=?L/?p，p=p - lr * grad，这是一个两步处理过程，首先计算梯度，然后将其更新到参数。融合版本是p=p - lr * ?L/?p。关键思想是在计算参数梯度时立即更新参数，从而不必在内存中存储梯度张量。这可以通过将钩子函数注入到反向传播中来实现。PyTorch提供了相关的API用于注入钩子函数，但我们不能使用当前的API实现精确的即时更新。相反，我们最多只存储一个参数的梯度，并随着反向传播逐个更新每个参数。我们的方法将梯度的内存使用量从存储所有参数梯度降至仅存储一个参数的梯度。

大部分 LOMO 训练所使用的内存和 parameter-efficient fine-tuning（PEFT）方法的相同，这表明将 LOMO 与这些方法结合只会导致梯度占用的内存增加很少，从而可以针对 PEFT 方法调整更多参数。

梯度归一化和剪裁是处理梯度爆炸和梯度消失问题的基本工具（Chen等人，2018），但它们的计算过程需要使用所有参数的梯度张量。我们在这里提出两种替代方案：

• 根据值剪裁梯度张量，而不是根据范数。
• 在额外的遍历中计算梯度范数。

根据值剪裁梯度张量是一种简单但有效的解决方案，可以在梯度范数接近之前防止梯度爆炸。剪裁值的主要问题是截断某些梯度元素可能会改变梯度张量的方向。例如，一个二维向量[1.3, 0.8]及其剪裁版本[1.0, 0.8]（剪裁为1.0）表示不同的方向。我们的经验是，在学习率较高时，根据值剪裁的效果较差，因为在这种情况下剪裁发生得更频繁。然而，对于中等和小学习率，根据值剪裁效果很好。注意，学习率的大小很大程度上取决于任务和数据，但一般来说，我们建议在学习率小于1e-3时使用根据值剪裁的方法。

我们的方法不能直接计算梯度范数，因为我们在反向传播过程中更新参数，所以在更新某个特定参数时，我们不知道其他参数的范数。然而，我们可以引入额外的遍历来计算和累积每个参数的梯度范数，从而需要两个反向传播过程，一个用于计算梯度范数，一个用于更新参数。内存使用量保持不变，但速度会受到牺牲。

一个有争议的解决方案。我们当前的训练框架根据所有参数计算梯度范数，并需要两个反向传播过程。节省额外的反向传播过程的一种解决方案是使用一组参数近似梯度张量的范数，例如，相邻的层。然而，这种方法实际上是有偏差的，因为它导致不同参数的更新步长不同。在更新过程中，根据梯度范数，参数会乘以一个比例因子。由于不同参数组之间的梯度范数不同，这样的近似会导致比例因子的差异。尽管存在这个限制，这种分组梯度剪裁方法可以被视为根据梯度范数向不同参数组应用动态学习率。Sun等人（2020）认为在随机梯度下降（SGD）中，并不总是适合为所有参数使用相同的学习率，因此我们相信我们的方法也具有进一步受益SGD的潜力。我们将探索作为一个引人注目的未来方向。

混合精度训练常用于加速训练过程。为了减轻精度下降，我们使用动态损失缩放并将某些计算转换为全精度。损失缩放的方法对于在FP16训练中防止下溢是至关重要的，在反向传播之前通过特定因子放大损失，并通过相同因子减小梯度。

在这种情况下，我们将动态损失缩放与LOMO结合在一起，动态调整整个训练过程中的缩放因子。如果在指定数量的反向传播过程中没有发生溢出，缩放因子将加倍。否则，将跳过此步骤并将缩放因子减半。这个过程类似于梯度归一化时遇到的情况。在反向传播完成之前，我们无法确定是否会发生溢出。因此，我们执行两个反向传播过程：第一个遍历用于检测溢出，第二个遍历用于在未检测到溢出时更新参数。这两个动态损失缩放的反向传播过程可以与梯度归一化同时执行。为了有效地更新参数和处理梯度，例如归一化和缩放，这些计算中的梯度及其关联的参数将转换为全精度。

我们首先在不同设置下分析模型状态和激活的内存使用情况。如表1所示，在训练LLaMA-7B模型时，与AdamW优化器（Loshchilov和Hutter，2019）相比，LOMO优化器的内存占用从102.20GB减少到14.58GB，与SGD相比从51.99GB减少到14.58GB。这种显著的内存使用减少主要归因于梯度和优化器状态的内存需求降低。因此，在训练过程中，内存主要被参数占用，与推断过程中的内存使用相称。

优化器状态 图2显示，使用AdamW优化器进行LLaMA-7B训练时，内存的相当大比例（73.7%）被分配给了优化器状态。这是混合精度训练方法的结果，其中在权重更新过程中保留了权重、动量和方差的全精度副本。用SGD优化器替换AdamW优化器可以有效减少内存中优化器状态的比例，从而减轻GPU内存的使用（从102.20GB减少到51.99GB）。这种减少是因为SGD优化器不需要存储全精度的动量和方差。对于LOMO来说，参数更新和反向传播被融合为一步，进一步消除了对优化器状态内存的需求。

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

激活 在训练一个包含512×8个标记的7B模型时，对于激活而言需要大量的内存。LOMO与激活内存减少技术（如激活检查点）兼容。通过将激活检查点与LOMO集成，由于激活的内存占用可以从45.61GB减少到1.79GB。

我们评估LOMO相对于AdamW和SGD的吞吐性能。实验在一台配备8个RTX 3090 GPU的服务器上进行，通过PCIe主板互连。序列长度和批大小分别设置为1024和1。吞吐量以每个GPU每秒处理的标记数量（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上遇到了内存不足的问题，而LOMO在只使用4个GPU时表现良好。

最后，我们成功使用8个RTX 3090 GPU训练了65B模型，实现了4.93 TGS的吞吐量。利用这样的服务器配置和LOMO，处理包含512个标记的1000个样本的训练过程大约需要3.6小时。

为了评估LOMO在微调大型语言模型中的有效性，我们进行了一系列广泛的实验。我们将LOMO与Zero-shot和LoRA两种方法进行了比较，其中Zero-shot不需要微调，LoRA是当前最流行的参数高效微调技术之一（详见Hu等人，2022）。LoRA对密集层进行了重新参数化，并且只在更新低秩矩阵时引入了延迟。

我们使用SuperGLUE数据集集合来评估模型性能，特别关注RTE（Dagan等人，2005），BoolQ（Clark等人，2019），WSC（Levesque等人，2012），WIC（Pilehvar和Camacho-Collados，2019），MultiRC（Khashabi等人，2018）和COPA（Roemmele等人，2011）。由于运行大型语言模型的计算成本较高，我们按照MeZO（Malladi等人，2023）的方法从训练集中随机抽取1000个训练数据和从验证集中随机抽取1000个测试数据，并报告使用相同随机种子获得的最佳结果。我们的实验使用了与MeZO相同的提示，并在附录A中详细说明了超参数。在推断过程中，我们将不同的标签或候选项插入到提示中，并计算每个标签的平均对数似然。选择得分最高的标签作为模型的答案。我们使用准确度作为评估指标。

LOMO与Zero-shot和LoRA在下游任务上的性能对比结果如表3所示。根据结果，我们得出以下观察结论。

LOMO在Zero-shot方面表现优势显著。在所有六个数据集和模型大小上，LOMO始终比Zero-shot取得更好的结果，使用LLaMA-13B平均增益超过20个点。尽管先前的研究展示了大型语言模型在Zero-shot设置中的惊人能力，但微调仍然为特定的下游任务提供了显著的性能增强。实验结果证实了LOMO在优化不同规模的大型语言模型方面的有效性。

LOMO通常在大多数实验中优于LoRA。我们展示了相对于LoRA，LOMO提供了强大的性能，例如，使用LLaMA-13B平均增益2.8个点。这表明相对于参数高效微调，模型性能更多地受益于全参数微调，因为前者调整了更多参数。LOMO在性能和效率之间取得了良好的平衡，使其成为微调的竞争选择。

在某些情况下，LOMO的性能不如LoRA。一个可能的原因是我们使用的训练集相对较小，对于大型模型的全参数微调可能不足够。此外，LoRA和LOMO采用了不同的模型架构。具体而言，LoRA为模型调优提供了一种捷径，这在某些情况下具有优势。实际上，这两种方法并不矛盾或相互排斥。在下一个子部分中，我们验证了将LoRA与LOMO结合使用不会损害模型性能，并且在大多数情况下可以提高性能。

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

LOMO和LoRA在基本上是相互独立的。为了验证这一说法，我们使用LLaMA-13B在BoolQ和MultiRC数据集上进行实验。结果如图3所示。我们发现，LOMO无论在LoRA取得的更高结果上如何，都一致提高了LoRA的性能。这表明LOMO和LoRA采用的不同微调方法是相辅相成的。具体而言，LOMO专注于微调预训练模型的权重，而LoRA则微调额外的模块。因此，LOMO不会损害LoRA的性能，而是为下游任务提供更好的模型调优。

本文介绍了LOw-Memory Optimization（LOMO），这是一种针对资源有限的大型语言模型的全参数微调的新优化器。我们已经展示了在配备RTX 3090等消费级GPU的服务器上微调65B模型的可行性。通过分析LOMO的内存使用情况，进行吞吐量测试，并在SuperGLUE上进行实验，我们展示了其有效性和潜在影响。

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