【LLM】对大语言模型微调优化的研究

一、结论写在前面

论文来自Microsoft。

微调大语言模型是用户尝试将其适应特定应用的流行选择。然而，微调这些模型是一项艰巨的任务，因为用户必须考虑多种因素，如资源预算、运行时间、模型大小和上下文长度等。一个具体挑战是微调过程对内存的需求很大，这限制了所需的硬件内存和能够处理的训练数据上下文长度。

论文分享了针对不同微调场景的各种微调优化的详细研究。特别是，论文评估了梯度检查点（Gradient Checkpointing）、低秩适应（Low Rank Adaptation）、DeepSpeed的ZeRO冗余优化器（ZeRO Redundancy Optimizer）和Flash Attention。

• 论文专注于内存和运行时间，研究了不同优化组合在微调阶段对GPU内存使用和执行运行时的影响。
• 论文提供了针对不同模型大小平衡内存和运行时的最佳默认优化建议。论文分享了有效策略，用于微调拥有数十亿或数百亿参数的非常大型模型，并在微调过程中实现大型上下文长度。
• 此外，论文提出了在GPU资源限制下进行微调的适当优化组合。

论文的结果表明：

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

二、论文的简单介绍

2.1 论文的背景

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

近年来，为优化GPU内存使用，提出了若干关键技术。然而，这些技术对许多用户来说理解复杂，且其应用往往需要通过实验来找到避免内存溢出错误的正确组合。这不仅需要人力资源（用于研究和实施这些优化），还需要计算资源（大量的GPU小时）来为特定任务选择合适的优化策略。

为了简化用户的微调过程，诸如Azure Machine Learning等平台提供了先进的GPU内存优化技术，使得即使使用少量GPU也能进行大型语言模型（LLM）的微调。

论文详细介绍了流行的LLM微调优化技术，深入探讨了内存与运行时间的权衡，提供了选择最佳优化配置的指导，并分享了论文的实验结果，这些结果可用于校准优化默认设置。

2.2 优化技术概述

论文简要概述了四种优化技术，这些技术可用于减少微调过程中的内存瓶颈。所有这些优化技术都可以被用户在其微调任务中，如Azure Machine学习和Hugging Face等平台上使用。

2.2.1 Gradient Checkpointing

梯度检查点（GC）（Chen et al., 2016）通过不在深度神经网络（DNN）前向传播（FP）过程中保留所有计算出的激活来审慎地利用GPU内存。相反，它在反向传播过程中重新计算许多激活，这有助于节省GPU内存。最节省内存的策略是在每n步保存检查点，其中n是DNN的层数（深度）。这种策略确保计算时间仍与DNN的深度成线性关系，同时将GPU内存需求从线性减少到DNN深度的平方根。因此，梯度检查点使得能够微调比仅增加20%-30%微调时间更大的LLMs。

2.2.2 LoRA

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

2.2.3 DeepSpeed

DeepSpeed的ZeRO冗余优化器（ZeRO）是一种内存优化技术，它提供了模型和数据并行性的好处，同时缓解了两者的局限性。ZeRO驱动的数据并行性（Zero-DP）将模型状态——参数、梯度和优化器状态——跨数据并行进程进行分区，并使用动态通信调度在进程之间共享必要的模型状态。

ZeRO-DP提供了三个优化阶段，这些阶段相继提供越来越大的内存减少，同时产生一些运行时开销。基于ZeRO-DP的微调可以分别实现第1阶段和第2阶段的4到8倍的内存减少，而对于第3阶段（假设有大量GPU）则可以达到线性内存减少。这确实是以运行时为代价的，尤其是对于ZeRO-DP第3阶段。

ZeRO优化通过包含ZeRO-Offload和ZeRO-Infinity（得到进一步增强，它们分别将优化器状态和模型参数卸载到CPU。虽然ZeRO-Offload在所有三个ZeRO-DP阶段都可用，但ZeRO-infinity仅适用于第3阶段。如果不使用DeepSpeed ZeRO这样的优化，许多大型模型（数十亿参数）实际上不可能仅使用少数几个V100或A100 GPU进行训练。    

2.2.4 快速注意力（Flash Attention）

快速注意力帮助实现注意力计算的复杂度从与序列/上下文长度相关的二次复杂度降为线性复杂度。它利用分块和重计算技术显著加速注意力计算。它明智地使用静态随机存取存储器（SRAM），这是GPU中最昂贵但容量最小的存储单元，最小化SRAM与高带宽存储器（HBM）之间的读写操作。在快速注意力版本2中，进一步优化了减少较慢的非矩阵乘法操作，并沿着序列长度维度（除了批次和头数维度外）并行化前向和后向传播。

这些优化中的一个关键特点是它们彼此正交。因此，它们可以组合起来为用户带来渐增的内存和运行时效益。

2.3 理论分析GPU内存需求

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

A. 模型状态 - 包括模型参数、梯度和优化器状态

B. 激活 - 表示中间计算结果

C. 临时缓冲区和碎片

模型状态内存：论文使用Rajbhandari等人（2020）中的模型状态内存计算来估计(A)。这里论文简要讨论图1中所示的模型状态内存计算。在混合精度训练中使用Adam优化器时，可训练参数、动量和方差以全精度（32位）保持。假设有0个参数，内存总计为(0 + 0 + 0) x 4 = 120字节。因此，K-12是Adam优化器特有的常数优化器状态：可训练参数 + 动量 + 方差梯度模型参数    

图1：使用Adam优化器在混合精度设置下微调0参数模型时的模型状态内存。模型状态包括优化器状态、梯度和模型参数。在未启用任何优化的情况下，总模型状态内存总计达到160字节

在内存计算中使用的乘数项。此外，在前向和后向传播过程中，参数和梯度以半精度维护，这导致额外的内存需求为20 + 20 = 40字节。ZeRO-DP阶段决定了哪些模型状态（参数、梯度、优化器状态）在GPU之间分区，从而节省内存。

激活内存：在没有模型并行性的情况下，每个transformer层的激活内存可以使用论文第4节中的公式1（Korthikanti等人，2022）进行估计。论文依赖相同的计算来确定（B）激活内存需求的近似上限。每层的激活内存可以表示为：

其中，s是序列长度，b是微批次大小，h是隐藏维度，a是注意力头的数量。总激活内存计算为transformer层数 × 每层激活内存。

表1：ZeRO-DP三个阶段理论与实测GPU内存分配（以GB为单位）的比较。微调模型为LLaMA-2-7B，使用8个A100 GPU（每个具有80 GB HBM）。理论估计可用于在运行微调作业之前近似实测内存需求    

表 

 展示了论文在比较理论上预期的和经验上分配给LLaMA-2-7B微调的GPU内存方面的结果，涉及ZeRO-DP的不

同阶段。论文使用ZeRO-1、ZeRO-2和ZeRO-3分别表示ZeRO-DP的第1、2和3阶段。理论内存是通过汇总模型状态和激活的内存来计算的。论文使用以下模型默认参数：transformer层数=32，隐藏维度大小h=4096，注意力头数a-32，理论上计算出总激活内存为1.48 GB。经验GPU内存指的是微调过程中分配的峰值GPU内存。论文展示了有无CPU卸载优化器状态的经验内存。

论文的结果表明，理论计算可以作为微调过程中GPU内存消耗（无CPU卸载）的粗略估计。在运行微调任务之前计算这一估计有助于预测内存需求并有效规划资源分配。论文注意到，启用CPU卸载优化器状态可以将内存使用量减少至不启用卸载时的四分之一。因此，CPU卸载可以允许对具有数十亿参数的模型进行微调时，对GPU和CPU资源进行高效利用，这是一个内存密集型过程。

需要注意的是，在没有ZeRO-DP优化的情况下，使用传统数据并行性对70亿参数模型进行完全微调将需要超过112（16 x 7）GB的GPU内存。这样的内存需求远远超过了A100s提供的80 GB GPU HBM。

2.4 实验

论文首先分析微调优化对内存和运行时的影响。论文的分析结果用于推荐一组优化，这些优化可以作为平衡默认值，以优化微调期间的内存和运行时。随后，论文深入研究大型模型（具有数十亿参数）的微调，并探索使这种微调成为可能的优化。接着，论文检查Flash Attention 2在LLMs微调长上下文数据中的作用。最后，论文研究在资源受限的GPU（特别是V100s）上进行微调的具体情况。论文的目标是识别能够在各种模型大小和上下文长度下，甚至在有限资源下实现高效微调的优化。

2.4.1 设置

论文在Causal Language Modeling任务上对LLaMA-2（7B、13B、70B）和Falcon（180B）系列的模型进行微调。以下是论文的实验设置的详细信息：

• 数据：所有实验中使用了Sampled Samsum数据集（Gli et al., 2019）进行微调。
• 计算：使用了标准ND4Ors v2（8xV100）（Azure, 2022）和标准ND96amsr A100 v4（8xA100）（Azure, 2024）作为GPU计算资源。
• 优化器：所有实验使用混合精度设置和AdamW优化器（β1=0.9，β2=0.99），配合线性调度器和学习率4e-4。
• 序列长度、批次大小和周期：序列长度通常设置为256。在探索序列长度影响的研究中，通过填充扩展输入序列的长度。所有实验中使用的有效批次大小为8。注意，可以通过梯度累积来增加有效批次大小，同时实现对可用GPU内存的高效利用。所有模型都进行单周期微调。
• LoRA：在使用LoRA的实验中，LoRA的rank设置为64，alpha设置为32。
• CPU卸载：除非明确指出，论文使用ZeRO-Offload在微调过程中将优化器状态和计算卸载到CPU。这有助于使用可用的GPU和CPU的计算及内存资源训练数十亿参数的模型。

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

2.4.2 最佳默认优化以平衡内存和运行时间

论文已经确定，对于相对较小的7B参数模型，即使使用A100，不进行任何优化进行微调也是不可行的。因此，论文使用DeepSpeed ZeRO，它提供了模型和数据并行性。论文检查了在微调框架中应默认启用的微调优化集合。论文选择这些优化的标准是基于它们在ZeRO-DP结合下，能够在内存使用和运行时间之间达到最佳平衡的能力。

论文在一个节点（8x A100 GPU）上对LLaMA-2-7B进行了五种不同配置的微调：(a) 未启用优化 (b) 启用梯度检查点 (c) 启用LoRA (d) 启用Flash Attention 2 (e) 未启用CPU卸载。每种配置都启用了所有三个ZeRO-DP阶段。论文选择A100 GPU进行这些实验，因为它们具有80 GB的大内存容量。这种充足的内存使论文能够通过开启和关闭各种优化来灵活地进行实验，同时避免内存不足的失败。    

图2：LLaMA-2 7B在不同优化配置下，使用ZeRO-1、ZeRO-2和ZeRO-3时的GPU内存使用情况和微调运行时间。使用ZeRO-2结合LoRA提供了最佳的内存使用与运行时间平衡

图2展示了每种配置和ZeRO-DP阶段下的GPU内存使用和运行时间。根据结果，论文得出以下结论：

1.从运行时间和内存使用的角度来看，最佳的优化组合是配置(c)，即ZeRO-DP + LoRA。由于LoRA减少了总的可训练参数，它节省了内存并缩短了运行时间。在Zero-DP + LoRA配置中，ZeRO-1和ZeRO-2阶段的内存使用及运行时间大致相似。当使用LoRA（rank=64）时，可训练参数减少到约3360万个，对于7B模型而言。优化器状态和可训练参数的梯度占据了GPU内存的极小部分。在ZeRO-1和ZeRO-2中，内存消耗主要由总模型参数（70亿）主导，这些参数在半精度下占用14GB内存。

1.配置(a)和(d)在ZeRO-DP阶段的表现符合预期，即从阶段1到2到3，内存消耗下降，运行时间增加。对于配置(b)和(e)，ZeRO-1的运行时间略高于ZeRO-2，内存使用遵循随着Zero-DP阶段增加而减少的预期趋势。

2.在(e)中禁用CPU卸载提供了最快的运行时间，但需要2-4倍更多的GPU内存。

论文得出结论，ZeRO-2与LoRA的结合是维持内存使用和运行时间平衡的优秀默认选择。此外，ZeRO-2能处理比ZeRO-1更大的模型，这有助于支持更广泛的模型尺寸范围。对于某些用例，如复杂的问答任务或构建聊天机器人，用户可能更喜欢完全微调模型。因此，是否使用LoRA由用户自行决定。因此，论文建议在微调过程中默认激活ZeRO-2。    

2.4.3 大型模型的微调

微调大型模型如LLaMA-2 70B或Falcon 180B时，如果没有启用正确的优化设置，常常会遇到内存不足（OOM）错误。在这种情况下，启用ZeRO-3是一个可行的替代方案，尽管它会增加运行时间。论文来检查在ZeRO-3中，经过个节点和每个节点个GPU微调后的模型状态内存项：

这里，表示总可训练参数（以十亿计）。假设一个标准的配置，每个节点有

个GPU，这将随着模型尺寸的增加而增加。例如，在个节点上完全微调一个

亿参数的模型，将占用40GB的模型状态内存，这可以在具有80GB HBM的A100上容纳。在实践中，具有数百亿参数的大型模型通常与LoRA一起微调，并使用ZeRO-3与CPU卸载。这进一步降低了内存需求，使其达到可管理的规模。             

论文能够在单个节点上使用8xV100 GPU（32 GB HBM）结合ZeRO-3 + LoRA对LLaMA-2 70B进行微调。表显示，在论文的实验中，70B模型的实际GPU内存消耗约为15.54 GB。此外，可以看出，对于LLaMA-2 7B和13B这样的较小模型，使用ZeRO-3 + LoRA会导致GPU内存未充分利用。因此，论文的研究表明，对于微调具有数十亿或数百亿参数的大型模型，启用ZeRO-3是必不可少的。此外，当与LoRA结合使用时，实际的总体内存需求显著减少。

表2：在8xV100 GPU上使用ZeRO-3 + LoRA对LLaMA-2模型进行微调的GPU内存使用量和时间

对于较小模型，ZeRO-3 + LoRA显然过于冗余，因为它可能导致GPU未充分利用以及运行时间增加。    

2.4.4 长上下文微调

              图3：LLaMA-2 70B在不同上下文长度下，使用和不使用Flash-Attention 2对GPU内存使用和微调时间的影响。在A100上启用Flash-Attention 2显著降低了较大上下文长度（如4096)的内存消耗和运行时间

上下文长度是微调LLMs时的关键因素。特别是在使用包含非常长文本序列的数据集进行微调时尤为如此。没有Flash-Attention 2（FA2），注意力计算随着序列长度的平方增长而增加。由于V100 GPU不支持FA2，因此在使用较长上下文长度时，对GPU内存消耗的影响可能很大。为了研究在微调过程中长上下文长度的影响，论文绘制了在A100（启用和未启用FA2）和V100（未启用FA2）上不同上下文长度的内存和运行时间。

具体来说，论文在8xV100和8xA100 GPU上分别使用1024、2048和4096的上下文长度对LLaMA-2 70B模型进行微调。在微调过程中启用了LoRA和梯度检查点技术，并基于ZeRO-3进行优化。

图3展示了Flash-Attention 2对GPU内存和运行时间的影响。论文的结果显示，随着上下文长度增加，至4096时，使用Flash-Attention 2相比于不使用，GPU内存使用和运行时间有显著减少。如图3所示，在两块A100 GPU上或分别在A100和V100 GPU上比较使用和不使用Flash-Attention 2的运行时，这一趋势是一致的。

论文的结论如下：高性能GPU如A100，由于其更大的HBM，可以在没有Flash-Attention 2的情况下支持相对较长的上下文微调。然而，在兼容的GPU架构上激活Flash-Attention 2是最优的，因为它允许额外的内存节省和减少运行时间。当使用允许较大上下文长度微调的模型训练长文本数据时，这一点变得尤为重要。    

2.4.5 在GPU资源限制下的微调

图4：使用V100 GPU时，不同大小LLMs微调的最佳配置。所有LLaMA-2实验均使用8x V100进行，而Falcon-180B需要16xV100。由于V100不支持Flash-Attention 2，因此未包含在内

本节讨论在资源有限的情况下对LLMs（参数范围从7B到180B）进行微调。论文将资源限制定义如下：

•可用的低资源GPU具有有限的HBM。

•可用的GPU数量较少。

论文对LLaMA-2 7B、13B、70B和Falcon-180B模型在五种不同的上下文长度上进行了微调。遵循资源约束规范，论文的实验在一台标准ND4Ors v2（8xV100 GPU）上进行，配备32 GB HBM。唯一例外的是Falcon-180B模型，由于其庞大的规模，论文在微调过程中使用了两个节点（16x V100 GPU）。图4展示了为满足GPU内存需求同时优化运行时间应启用的首选优化集合。

需要注意的是，虽然可能存在其他允许成功微调的组合，但论文选择了那些最小化微调时间的组合。例如，用户可以选择启用所有优化，但通常并非所有优化都是必要的，且可能会不必要地增加微调时间。

论文的结果中的关键见解包括：

• 即使是为了在有限数量的GPU上适配像7B参数这样的小型模型，结合模型和数据并行是先决条件，除非论文利用量化技术。
• 随着上下文长度或模型大小的增加，转向更高阶段的ZeRO-DP（零冗余优化器-数据并行）以促进微调变得重要。这一趋势在表格中清晰可见，从左到右或从上到下的移动对应于更高阶段的ZeRO-DP。
• 梯度检查点（GC）是一种有效的内存节省优化技术，尤其适用于大型模型。如图4所示，为LLaMA-2 13B和70B启用GC使得支持的上下文长度从512扩展到4096，否则这是不可能的。类似地，GC对于在论文检查的所有上下文长度上启用Falcon-180B的微调至关重要。
• 尽管Flash-Attention 2在V100 GPU上不支持，但论文的实验结果表明，在支持的架构上应始终启用它。值得一提的是，图中概述的配置可用于在具有比V100更高HBM的GPU上进行微调。虽然这些配置可能不是所有GPU的最佳选择，但论文的发现可以通过进一步的研究扩展到其他GPU。

论文标题：A Study of Optimizations for Fine-tuning Large Language Models

论文链接：​​​https://arxiv.org/pdf/2406.02290​​    

本文转载自 ​​AI帝国​​，作者： 无影寺