轻松微调巨型AI模型：解密高效内存优化策略，让你的GPU事半功倍

Abstract

在用户尝试将大型语言模型适应特定应用时，微调这些模型是一种流行的选择。然而，微调这些模型是一项艰巨的任务，因为用户必须考虑多种因素，如资源预算、运行时、模型大小和上下文长度等。一个特定的挑战是微调对内存的消耗很大，这限制了所需的硬件内存和可以处理的训练数据的上下文长度。在这项工作中，作者分享了对不同微调场景下各种微调优化的详细研究。特别是，作者评估了梯度预训练权重、低秩适应、DeepSpeed的ZeRO冗余优化器和闪存注意力。在关注内存和运行时的同时，作者研究了不同的优化组合在微调阶段对GPU内存使用和执行运行时的影响。作者为平衡不同模型大小的内存和运行时提供了最佳的默认优化建议。作者分享了微调具有数十亿或数百亿参数的非常大的模型以及在微调期间启用大型上下文长度的有效策略。此外，作者还提出了在GPU资源限制下微调的适当优化组合。

1 Introduction

迁移学习为特定任务或领域定制大型语言模型（LLM）提供了一种非常有效的方法。在微调中，预训练的LLM会在特定任务的 数据集上进行进一步的训练，使它们能够将知识适应到特定任务。尽管这是一种强大的技术，但由于它们对内存的需求量很大，因此对具有数十亿可训练参数的LLM进行微调是具有挑战性的。例如，对一个拥有10亿参数的模型以全浮点精度（32位）进行微调大约需要每个GPU 24 GB的高带宽内存（HBM）（Fregly等人，2023年）。因此，在没有任何内存优化的情况下，在LLM微调过程中经常会遇到内存不足的错误。

在最近，已经提出了几种关键技术来优化GPU内存使用。然而，对于许多用户来说，它们可能难以理解，并且使用它们通常需要通过实验找到正确的组合以避免内存不足错误。这将需要人力资源（研究并运行这些优化），以及计算资源（大量的GPU小时），以为给定的任务选择正确的优化。为了便于用户进行微调，像Azure Machine Learning这样的平台提供了优化GPU内存使用的最先进技术，即使只有少数GPU也能进行LLM微调。在本文中，作者分享了关于微调LLM的流行优化技术的详细信息。作者深入探讨了内存和运行时之间的权衡，为选择最佳的优化配置提供指导，并分享作者的实验发现，这些发现可用于校准优化默认值。本文的目标如下：

• 通过实证结果展示哪些优化在实现内存有效微调的同时获得最佳训练时间方面效果良好。
• 在各种设置中为用户提供选择微调优化的指导。作者考虑的设置包括微调非常大的模型，具有长训练数据上下文的微调，以及在GPU限制（SKU或GPU数量）下的微调。

本文的其余部分组织如下。作者在第2节回顾现有的微调优化文献。第3节概述了常用的微调优化：梯度预训练权重，低秩适应，DeepSpeed ZeRO冗余优化器和Flash注意力。然后在第4节中描述了GPU内存使用的理论，并将理论值与作者的实证结果进行比较。在第5节中，作者讨论了作者的实验结果，并分析了一起使用微调优化时的影响。最后，在第6节中，作者分享关键结论和未来工作的方向。

2 Related Work

微调大型语言模型是一个计算密集和内存密集的过程。现有文献中提出的几种优化方法要么减少内存消耗，要么减少运行时间，或者在微调期间同时实现这两者。张量和流水线并行是模型并行技术的类型，它们使大型模型能够适应内存进行训练。在张量并行（Shoeybi等人，2020年）中，模型在GPU之间按张量方式进行分割。相反，流水线并行（Huang等人，2019年）将模型按层分布在GPU上。DeepSpeed ZeRO（Rajbhandari等人，2020年）通过在数据并行过程中划分模型状态，并在训练过程中需要时进行重构，集成了数据和模型并行。Fully Sharded Data Parallel（FSDP）（FairScale作者，2021年）与DeepSpeed ZeRO相似，并支持模型参数分片以训练大型模型。

与并行策略相反，参数高效微调（PEFT）方法（Ding等人，2023年；Han等人，2024年）在训练期间修改一小部分模型参数，而将其他参数冻结。这种可训练参数总数的减少随后降低了内存占用，同时仍能实现与完全微调相当的性能。低秩适应（Hu等人，2021年）是一种广泛使用的PEFT技术，其中在训练过程中更新模型参数的低秩张量公式。前缀调整（Li和Liang，2021年）和提示调整（Lester等人，2021年）是PEFT方法的例子，其中在模型或输入 Level 注入可训练的嵌入，而整个模型被冻结。PEFT方法可以轻松与DeepSpeed ZeRO等并行技术集成，从而在微调期间提高内存和运行时减少的收益。

某些优化旨在解决微调期间特定内存瓶颈。在反向传播过程中，梯度计算可能会消耗大量GPU内存，因为需要维护内存中的大量激活。梯度预训练权重（Sohoni等人，2019年）仅保存前向传播中的少数激活，并在反向传播过程中重新计算大多数激活，以减少内存使用，代价是增加运行时间。当上下文长度非常大时，注意力计算可能会形成内存瓶颈。Flash Attention（Dao等人，2022年；Dao，2023年）方法专注于优化注意力计算操作期间消耗的内存。梯度预训练权重和Flash Attention都可以与PEFT和DeepSpeed ZeRO结合使用。

鉴于有各种各样的优化方法和不同的用例，如微调非常大的（数十亿参数）模型，在较低资源GPU如ND40（Azure，2022年）上微调，或在大型上下文长度上微调，选择合适的优化组合进行微调变得很重要。并非所有的优化，即使它们彼此兼容，也可能是必要的，因为它们常常在内存、运行时间和准确性之间进行权衡。为企业微调LLM的实用指南（J等人，2024年）是在没有利用DeepSpeed等数据和模型并行框架的情况下生成的，并没有提供理解GPU内存消耗的理论基础。尽管存在差异，但那些指南检查了QLoRA（Dettmers等人，2023年）可以被视为对作者研究的补充。据作者所知，很少有研究提供针对特定但常见的用例量身定制的微调指导，并探索最先进的优化之间的相互作用。

3 Overview of Optimization Techniques

作者提供了四种优化技术的简要概述，这些技术可用于在微调过程中减少内存瓶颈。所有这些优化都可以被用户应用于在Azure机器学习和Hugging Face等平台上进行微调任务。

Gradient Checkpointing

梯度预训练权重（GC）（Chen等人，2016年）通过在深度神经网络（DNN）的前向传播（FP）过程中不保存所有计算的激活值，而是有选择地使用GPU内存。它反而在后向传播过程中重新计算许多激活值，这有助于节省GPU内存。最节省内存的策略是每隔步保存一次预训练权重，其中n是DNN的层数（深度）（Sohoni等人，2019年；Bulatov，2018年）。这种策略确保了计算时间仍然与DNN的深度线性扩展，同时将GPU内存需求从线性减少到深度的平方根。因此，梯度预训练权重允许微调比额外增加20%-30%的微调时间所能实现的更大的LLM（Bulatov，2018年）。

Low Rank Adaptation

低秩适应（LoRA）减少了可训练参数的数量，从而在微调过程中降低了GPU内存的需求。它是通过冻结预训练模型的权重，并在Transformer架构的每个选定的层中注入可训练的秩分解矩阵来实现的（Hu等人，2021年）。可训练参数的数量减少了几个数量级，降低了微调的成本，同时保持了结果的质量。例如，使用LoRA微调并将秩设置为64时，一个700亿参数模型的训练参数数量减少到大约1.31亿个参数（约为原始模型大小的0.19%）。

DeepSpeed

DeepSpeed的ZeRO冗余优化器（ZeRO）（Rajbhandari等人，2020年）是一种内存优化技术，它兼具模型并行和数据并行的优点，同时减轻了两者的局限性。ZeRO驱动的数据并行（Zero-DP）将模型状态——参数、梯度以及优化器状态——在数据并行进程间进行分割，并使用动态通信计划在进程间共享必要的模型状态。Zero-DP提供了三个优化阶段，这些阶段逐步提供越来越大的内存减少，同时带来一些运行时开销。基于ZeRO-DP的微调可以实现4倍（第一阶段）到8倍（第二阶段）的内存减少，而对于第三阶段，在大量GPU的情况下，可以实现线性内存减少。这当然是以运行时为代价的，特别是对于ZeRO-DP的第三阶段。通过包括ZeRO-Offload和ZeRO-Infinity（Ren等人，2021年；Rajbhandari等人，2021年）进一步增强了ZeRO优化，它们分别允许将优化器状态和模型参数卸载到CPU。尽管ZeRO-Offload在所有三个ZeRO-DP阶段都可用，但ZeRO-Infinity仅适用于第三阶段。如果不使用类似DeepSpeed ZeRO的优化，许多大型模型（数十亿个参数）实际上仅使用少量V100或A100 GPU是不可能训练的。

Flash Attention

闪存注意力帮助实现了与序列/上下文长度成线性而非二次复杂度的注意力计算（Dao等人，2022年）。它利用分块和重新计算技术显著加快注意力计算速度。它明智地使用了静态随机访问存储器（SRAM），这是GPU内最昂贵但体积最小的存储单元，最小化SRAM与HBM之间的读写操作。在闪存注意力版本2中，有进一步的优化措施，以减少较慢的非矩阵乘法操作，并沿着序列长度维度（除了批处理和头数维度）并行化前向和后向传递（Dao，2023年）。

所有这些优化的一个关键特点是它们相互正交。因此，它们可以被组合起来，为用户带来逐步增加的内存和运行时效益。

4 Theoretical Analysis of GPU Memory Requirements

在本节中，作者的目标是帮助读者基本了解在训练过程中影响GPU内存消耗的计算。以下三个变量主要在训练期间影响GPU内存：

1. 模型状态——包括模型参数、梯度以及优化器状态
2. 激活——代表中间计算结果
3. 临时缓冲区和碎片

模型状态内存 ：作者采用（Rajbhandari等人，2020年）中的模型状态内存计算来估计（A）。在这里，作者简要讨论图1中显示的模型状态内存计算。在混合精度训练中，使用Adam优化器时，可训练参数、动量和方差以全精度（32位）维护。假设参数，内存总计为字节。因此，K=12是在内存计算中使用的特定于Adam优化器的常数乘法项。此外，前向和反向传递期间，参数和梯度以半精度维护，这会额外增加字节的内存需求。ZeRO-DP阶段确定哪些模型状态（参数、梯度、优化器状态）跨GPU分割，从而节省内存。

激活内存 ：在没有模型并行性的情况下，每层 Transformer 的激活内存可以用论文（Korthikanti等人，2022年）第4节的方程1来估计。作者依赖相同的计算来确定（B）的激活内存需求的近似上限。每层的激活内存可以计算为：

这里，是序列长度，是微批次大小，是隐藏维度，是注意力头数。总的激活内存然后计算为：Transformer 层数每层激活内存。

表1展示了作者在ZeRO-DP阶段对LLaMA-2-7B模型进行微调时，理论上预期和实际GPU内存分配的比较结果。作者使用ZeRO-1、ZeRO-2和ZeRO-3的表示分别指代ZeRO-DP阶段1、2和3。理论内存是通过将模型状态和激活的内存相加来计算的。作者使用以下模型默认参数：Transformer 层数=32，隐藏维度大小=4096，注意力头数=32来理论计算总激活内存为1.48 GB。实际GPU内存是指在微调期间分配的峰值GPU内存。作者展示了有无CPU卸载优化器状态的实际内存。实验设置细节在5.1小节中描述。

作者的结果显示，理论上计算可以作为一个粗略估计微调过程中（没有CPU卸载）的GPU内存消耗。在运行微调工作之前计算这个估计可以帮助预测内存需求并有效地规划资源分配。作者注意到，CPU卸载优化器状态可以将内存使用降到没有卸载时的四分之一。因此，CPU卸载可以允许高效地利用GPU和CPU资源进行微调数十亿参数的模型，这是一个内存密集型过程。请注意，如果没有任何ZeRO-DP优化，使用传统数据并行主义对70亿参数模型进行完全微调将需要超过112（16 7）GB的GPU内存。这种内存需求远远超过了A100s上可用的80 GB GPU HBM。

5 Experiments

作者从这个部分开始，分析微调优化对内存和运行时间的影响。作者的分析结果用于推荐一组可以在微调过程中平衡内存和运行时间优化的默认设置。随后，作者深入探讨微调大型模型（拥有数十亿个参数）并探索使这种微调成为可能的优化方法。接着，作者检查Flash Attention 2在微调LLM（大型语言模型）处理长上下文数据时的作用。最后，作者研究在资源受限的GPU上进行微调的特殊情况，特别是V100型号。作者的目标是确定可以使各种模型大小和上下文长度在有限资源下高效微调的优化方法。

Setup

作者微调了来自LLaMA-2（7B、13B、70B）和Falcon（180B）系列的模型，以应对因果语言建模任务。以下是作者实验设置的细节：

• 数据 ：在所有实验中，作者都使用了抽样的Samsum数据集（Gliwa等人，2019）进行微调。
• 计算资源 ：作者使用了Standard_ND40rs_v2（8xV100）（Azure，2022）和Standard_ND96amsr_A100_v4（8xA100）（Azure，2024）作为GPU计算资源。
• 优化器 ：所有实验都使用混合精度设置和AdamW优化器（=0.9和=0.99），并采用线性调度器和学习率为4e-4。
• 序列长度、批处理大小和时期数 ：序列长度通常设置为256。在作者探索序列长度影响的实验中，通过填充来延长输入序列的长度。所有实验中使用的有效批处理大小为8。请注意，可以使用梯度累积来增加有效批处理大小，同时有效地利用可用的GPU内存。所有模型都经过单个时期的微调。
• LoRA ：在采用LoRA的实验中，LoRA等级设置为64，alpha设置为32。
• CPU卸载 ：除非特别说明，作者在微调过程中使用ZeRO-Offload将优化器状态和计算卸载到CPU上。这有助于使用可用GPU和CPU的计算和内存资源训练数十亿参数的模型。

在作者的实验中，GPU内存使用量是通过微调过程中所有使用的GPU上的峰值GPU内存分配来衡量的。

Best Default Optimizations to Balance Memory and Runtime

作者已经证实，对于一个相对较小的70亿参数模型，在没有优化措施的情况下，即使使用A100显卡，也无法进行微调。因此，作者使用了DeepSpeed ZeRO，它提供了模型和数据并行。作者研究了一组微调优化措施，这些措施应在微调框架中默认启用。作者选择这些优化的标准是基于它们与ZeRO-DP结合时在内存使用和运行时间之间取得最佳平衡的能力。

作者在一个节点（8xA100显卡）上对LLaMA-2-7B进行了五种不同配置的微调：(a)不启用任何优化（b）启用梯度预训练权重（c）启用LoRA（d）启用闪存注意力2（e）不启用CPU卸载。每种配置都启用了所有三个ZeRO-DP阶段。作者选择A100显卡进行这些实验，因为它们具有80GB的大内存容量。这种充足的内存使作者能够灵活地通过打开和关闭来实验各种优化，同时避免了内存不足的失败。

在Zero-DP + LoRA配置中，ZeRO-1和ZeRO-2阶段的内存使用和运行时间大致相似。当使用LoRA（秩为64）时，7B模型的可训练参数减少到大约3360万。可训练参数的优化器状态和梯度占GPU内存的一小部分。ZeRO-1和ZeRO-2中的内存消耗主要由总模型参数（70亿）组成，这些参数在半精度下占14GB的内存。2. 配置（a）和（d）在ZeRO-DP阶段的性能如预期那样变化，即内存消耗减少，运行时间从阶段1增加到2再到3。对于配置（b）和（e），ZeRO-1的运行时间略高于ZeRO-2，并且内存使用随ZeRO-DP阶段的增加而减少，符合预期趋势。3. 在配置（e）中禁用CPU卸载可以提供最快的运行时间，但代价是要求多2-4倍的GPU内存。

作者得出结论，ZeRO-2和LoRA的组合是一个很好的默认选择，因为它能够在内存使用和运行时间之间保持平衡。此外，ZeRO-2可以处理比ZeRO-1更大的模型，这有助于支持更广泛的模型尺寸范围。对于某些用例，例如复杂的问答任务或构建聊天机器人，用户可能希望对模型进行完全微调。因此，使用LoRA的选择权在用户手中。因此，作者建议在微调过程中默认激活ZeRO-2。

Fine-tuning Large Models

对大型模型如LLaMA-2 70B或Falcon 180B进行微调通常会在没有启用正确的优化措施时遇到内存不足（OOM）错误。在这种情况下，启用ZeRO-3是一个可行的选择，尽管这会以增加运行时间为代价。让作者来分析一下在ZeRO-3下，跨个节点、每个节点个GPU进行微调的模型状态内存项：

这里，表示总的可训练参数（以十亿计）。假设一个标准的配置，每个节点的=8个GPU。这使得模型状态内存为。这意味着通过扩展节点数量可以有效抵消模型尺寸的增加。例如，对=100B参数模型在=5个节点上进行完全微调，模型状态内存占用40 GB，可以适应具有80 GB HBM的A100s。实际上，通常使用LoRA和ZeRO-3的CPU卸载来微调具有数千亿个参数的大型模型，这进一步将内存需求降低到可管理的规模。

作者能够使用ZeRO-3 + LoRA的组合在一个带有8xV100 GPU（32 GB HBM）的单个节点上微调LLaMA-2 70B。表2显示在作者的实验中，70B模型的实际GPU内存消耗约为15.54 GB。此外，可以看出对于较小的模型，如LLaMA-2 7B和13B，在使用ZeRO-3 + LoRA时，GPU内存利用率不足。因此，作者的研究显示，对于微调具有数十亿或数百亿个参数的大型模型，启用ZeRO-3是不可或缺的。此外，与LoRA结合使用时，实际中的总体内存需求显著减少。对于较小的模型，ZeRO-3 + LoRA显然是过度的，因为它可能导致GPU利用率不足以及运行时间增加。作者建议读者参考5.5小节，作者在其中探讨了针对不同模型尺寸的适当优化选择。

Long Context Fine-tuning

上下文长度在微调LLM时是一个关键因素。特别是在使用包含非常长文本序列的数据集进行微调时更是如此。如果没有Flash-Attention 2（FA2），注意力计算与序列长度的关系是二次方扩展的。由于V100 GPU不支持FA2，使用更长的上下文长度对GPU内存消耗的影响可能是显著的。为了研究在微调过程中长上下文长度的影响，作者绘制了在使用和不使用FA2的A100s以及不使用FA2的V100s上，不同上下文长度下的内存和运行时间。具体来说，作者在8xV100和8xA100 GPU上分别使用1024、2048和4096的上下文长度微调LLaMA-2 70B模型。在微调期间，除了ZeRO-3之外，还启用了LoRA和梯度预训练权重技术。

图3展示了Flash-Attention 2对GPU内存和运行时间的影响。作者的结果显示，随着上下文长度增加到4096，出现了预期的趋势。使用Flash-Attention 2相比于不使用，GPU内存使用量和运行时间显著减少。如图3所示，在两个A100 GPU上比较使用和不使用Flash-Attention 2的运行，或者在A100和V100 GPU上的运行，这一趋势是一致的。作者的结论如下：由于具有较大HBMs的高性能GPU（如A100s），即使没有Flash-Attention 2也能支持相对较长的上下文微调。然而，在兼容的GPU架构上激活Flash-Attention 2是最佳选择，因为它可以提供额外的内存节省和减少的运行时间。这在训练包含长文本数据的模型，且允许使用较大上下文长度进行微调时尤其重要。

Fine-tuning under GPU Resource Constraints

在本节中，作者讨论了在有限资源下微调大型语言模型（参数量在7B到180B之间）。作者定义资源限制如下：

• 低资源GPU配备了有限的HBM。
• 可用GPU数量有限。作者在五种不同的上下文长度上微调LLaMA-2 7B、13B、70B和Falcon-180B模型：[256、512、1024、2048、4096]。遵循资源限制规范，作者的实验是在配备32GB HBM的标准_ND40rs_v2（8xV100 GPU）上进行的。唯一偏离此设置的是Falcon-180B模型。由于其体积庞大，在微调期间作者使用两个节点（16xV100 GPU）。图4说明了为了满足GPU内存需求同时优化运行时间，应该启用的一组优选优化。请注意，虽然可能存在其他允许成功微调的组合，但作者选择那些可以最小化微调时间的组合。例如，用户可以选择启用所有优化，但通常它们并不都是必需的，可能会不必要地增加微调时间。

作者结果的关键洞察如下：

1. 即使是拟合参数量小至7B的模型，结合模型和数据并行性也是必要的，除非作者利用量化。
2. 随着上下文长度或模型大小的增加，转向更高的ZeRO-DP阶段以支持微调变得重要。这一趋势在表中清晰可见，从左到右或从上到下移动对应于更高的ZeRO-DP阶段。
3. 梯度预训练权重（GC）是一种有效的节省内存的优化技术，特别是对于大型模型。如图4所示，为LLaMA-2 13B和70B启用GC支持了从512到4096的上下文长度，否则这是不可能的。同样，对于在所有考察的上下文长度上微调Falcon-180B，GC是必不可少的。

尽管V100 GPU不支持Flash-Attention 2，但作者的实验结果表明，在支持的架构上应始终启用它。值得注意的是，图4中概述的配置可以应用于比V100具有更高HBM的GPU上的微调。虽然这些配置可能不是所有GPU的最优选择，但作者的发现可以扩展到其他GPU，需要进一步研究。

6 Conclusion and Future Work

在本文中，作者分析了最先进的微调优化方法及其在微调过程中对内存和运行时间的影响。作者的结果显示，ZeRO-2 + LoRA是一种可靠的默认优化方法，通常在内存使用和微调运行时间之间提供了最佳的平衡。通过使用正确的优化组合，例如ZeRO-3 + LoRA + GC，作者成功地对大小达到Falcon-180B的模型进行了微调。鉴于在GPU受限环境下微调LLM的需求，作者引入了一个跨越模型大小上下文长度的优化矩阵。该优化矩阵可以指导用户选择正确的优化组合，以实现对特定用例的微调。作者得出结论，像DeepSpeed ZeRO这样的优化对于微调数十亿参数的模型至关重要，因为它们支持数据并行和模型并行，以及CPU卸载。结合DeepSpeed ZeRO，选择额外的优化措施以避免内存不足失败同时平衡微调运行时间变得必要。

以下方向可作为未来工作的考虑：1)融入4位和8位量化分析（Dettmers等人，2023年）2)探索微调小型语言模型（SLMs）如Phi-2（Abdin等人，2023年）和Phi-3-mini（Abdin等人，2024年）的最佳配置，以及3)研究具有高达128K上下文长度的微调策略。

参考

[1].A Study of Optimizations for Fine-tuning Large Language Models.