大语言模型微调优化技术研究

大语言模型微调优化技术研究

一、结论摘要

本文基于微软论文，针对大语言模型（LLM）微调过程中的资源预算、运行时间、模型大小和上下文长度等因素进行了详细研究。主要发现如下：

1. ZeRO-2 + LoRA：是可靠的默认优化选项，通常在内存使用和微调运行时之间提供最佳平衡。
2. 大型模型微调：通过正确的优化组合（如 ZeRO-3 + LoRA + GC），成功微调了高达 Falcon-180B 的模型。
3. 长上下文支持：Flash Attention 2 能显著降低较大上下文长度下的内存消耗和运行时间。
4. 资源受限环境：DeepSpeed ZeRO 对于微调数十亿参数模型至关重要，支持数据/模型并行及 CPU 卸载。

二、背景介绍

迁移学习为将大型语言模型专化到特定任务或领域提供了有效方法。然而，微调具有数十亿可训练参数的 LLMs 仍具挑战性，主要受限于巨大的内存需求。例如，在全浮点精度下微调一个 10 亿参数模型，每个 GPU 大约需要 24GB 的高带宽内存（HBM）。若无内存优化，微调过程中常会遇到内存不足（OOM）错误。

近年来提出了若干优化技术，但理解复杂且需实验验证。本文旨在简化用户微调过程，深入探讨内存与运行时间的权衡，并提供选择最佳优化配置的指导。

三、优化技术概述

3.1 梯度检查点 (Gradient Checkpointing)

梯度检查点通过在反向传播过程中重新计算激活来节省 GPU 内存，而非在前向传播中保留所有激活。最省内存的策略是在每 n 步保存检查点。这使得能够微调更大的 LLMs，仅增加约 20%-30% 的微调时间。

3.2 低秩适应 (LoRA)

LoRA 通过冻结预训练权重并注入可训练的秩分解矩阵，减少可训练参数数量。例如，使用 LoRA 微调 700 亿参数模型，可训练参数可减少至约 1.31 亿（占原始模型大小的 -0.19%），降低成本同时保持质量。

3.3 DeepSpeed ZeRO

ZeRO 冗余优化器通过分区模型状态（参数、梯度、优化器状态）跨数据并行进程，缓解并行局限性。提供三个优化阶段：

• ZeRO-1/2：分别实现 4 到 8 倍的内存减少。
• ZeRO-3：在大量 GPU 下可实现线性内存减少，但运行时开销较大。 结合 ZeRO-Offload 和 ZeRO-Infinity 可将优化器状态和模型参数卸载到 CPU，使少数 GPU 也能训练大型模型。

3.4 Flash Attention

Flash Attention 将注意力计算复杂度从二次降为线性。利用分块和重计算技术，最小化 SRAM 与 HBM 之间的读写操作。版本 2 进一步优化了非矩阵乘法操作，并沿序列长度维度并行化。

这些优化彼此正交，可组合使用以获得渐增的效益。

四、GPU 内存需求理论分析

训练期间 GPU 内存消耗主要受以下变量影响：

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

在混合精度训练中使用 Adam 优化器时，可训练参数、动量和方差以全精度保持。ZeRO-DP 阶段决定了哪些模型状态在 GPU 间分区。理论计算可作为微调前预测内存需求的粗略估计。启用 CPU 卸载优化器状态可将内存使用量减少至不启用时的四分之一。

五、实验与分析

5.1 实验设置

• 模型：LLaMA-2 (7B, 13B, 70B) 和 Falcon (180B)。
• 数据集：Sampled Samsum。
• 硬件：标准 ND40rs v2 (8xV100) 和 ND96amsr A100 v4 (8xA100)。
• 配置：混合精度，AdamW 优化器，序列长度 256，批次大小 8。
• ：Rank=64, Alpha=32。