PyTorch 详细学习笔记 第十六章:参数高效训练与大模型技巧
第十六章:参数高效训练与大模型技巧
目标:LoRA/Adapter、梯度检查点、FSDP/ZeRO 思想,显存优化与吞吐调参路径。
随着 GPT-3、Llama 等拥有数千亿参数的“巨兽”模型的出现,传统的全量微调(Fine-tuning)方法变得越来越遥不可及——它不仅需要庞大的 GPU 集群,还需要极高的成本。如何在消费级甚至单张专业级 GPU 上,“撬动”这些大模型,让它们为我们的特定任务服务?参数高效微调 (Parameter-Efficient Fine-Tuning, PEFT) 技术应运而生。本章将深入探讨 PEFT 的明星方法 LoRA,并介绍与之相辅相成的显存优化“三剑客”:梯度检查点 (Gradient Checkpointing)、FSDP 及其背后的 ZeRO 思想。我们将为你构建一个清晰的、在资源受限条件下训练大模型的“调参路径图”,让你能够有策略地组合这些技术,最大限度地压榨硬件潜能,实现“小马拉大车”的壮举。
知识卡片:核心概念 & API
| 核心概念 | API / 工具 | 简要说明 |
|---|---|---|
| 参数高效微调 (PEFT) | peft 库 (by Hugging Face) |
一系列旨在通过只微调模型一小部分参数(<1%)来达到或接近全量微调效果的方法。 |
| LoRA (低秩适应) | peft.LoraConfig, peft.get_peft_model |
PEFT 的主流方法。通过在原模型的权重矩阵旁边,插入两个小的、可训练的低秩矩阵(A 和 B)来模拟权重的更新,而原始权重保持冻结。 |
| Adapter (适配器) | peft.AdapterConfig |
另一种 PEFT 方法。在 Transformer 层的特定位置(如 FFN 之后)插入一个小的、瓶颈状的“适配器”模块,只训练该模块的参数。 |
| 梯度检查点 | model.gradient_checkpointing_enable() |
一种用计算换显存的技术。在前向传播时,不保存中间激活值,只在反向传播需要时重新计算它们。 |
| FSDP (完全分片数据并行) | torch.distributed.fsdp.FullyShardedDataParallel |
PyTorch 官方的、先进的分布式训练方案,是 ZeRO 思想的官方实现。 |
| ZeRO (零冗余优化器) | N/A (一种思想) | 由微软 DeepSpeed 提出的分布式训练优化思想。通过在多个 GPU 间分片 (Shard) 存储模型参数、梯度和优化器状态,极大地减少了单个 GPU 的显存冗余。 |
| 吞吐量 (Throughput) | samples/sec |
衡量训练效率的关键指标,表示每秒钟能够处理的样本数量。 |
16.1 “四两拨千斤”:参数高效微调 (PEFT)
全量微调一个 70 亿参数的 Llama 模型,即使使用 AdamW 优化器,也需要近 100GB 的 GPU 显存,这对于大多数研究者和开发者来说是不可承受的。PEFT 的核心动机就是:我们能否在冻结绝大部分预训练参数的前提下,通过只训练极少数(例如百万级别)的新增参数,来让模型适应新的任务?
16.1.1 LoRA (Low-Rank Adaptation) 的精妙之处
LoRA 的理论基础来自于一个观察:大型语言模型虽然参数众多,但它们在适应下游任务时,其权重的变化量(ΔW)是低秩 (Low-Rank) 的。这意味着,这个巨大的权重更新矩阵 ΔW,可以用两个更小的矩阵 B 和 A 的乘积 B*A 来很好地近似,其中 B 的形状是 (d, r),A 的形状是 (r, k),而 r (秩, rank) 是一个远小于 d 和 k 的超参数(例如 r=8 或 16)。
LoRA 的实现方式:
- 冻结原始权重: 原始的预训练权重矩阵
W(例如,Transformer 中query,key,value投影的全连接层权重) 被完全冻结,不参与训练。 - 注入可训练旁路: 在
W旁边,并联一个由A和B两个小矩阵组成的“旁路”。A通常用随机高斯分布初始化,B用零初始化。 - 计算流程: 对于输入
x,修改后的前向传播变为y = Wx + B(Ax)。 - 训练: 在训练过程中,只有矩阵
A和B的参数会被更新。由于r很小,A和B的总参数量非常小(通常不到原始模型参数的 0.1%)。 - 部署/推理: 训练完成后,我们可以将学习到的“变化量”
BA合并回原始权重W中,得到一个新的权重矩阵W' = W + BA。在推理时,我们只需要使用这个合并后的模型W',而无需保留A和B,这意味着 LoRA 在推理时不会引入任何额外的延迟!
使用 Hugging Face peft 库实现 LoRA:peft 库极大地简化了为任意 transformers 模型添加 LoRA 的过程。
最小可运行示例:为 Llama 模型添加 LoRA
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model, LoraConfig, TaskType
MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf" # 替换为你可访问的模型
# 需要 Hugging Face 登录和访问权限
# from huggingface_hub import notebook_login; notebook_login()
# 1. 加载原始模型和分词器
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# 2. 定义 LoRA 配置
# TaskType.CAUSAL_LM: 指定任务类型为因果语言模型
# r: LoRA 的秩
# lora_alpha: LoRA 缩放因子, 通常设为 r 的两倍
# target_modules: 指定要应用 LoRA 的模块名称,通常是 'q_proj', 'v_proj'
peft_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
inference_mode=False,
r=8,
lora_alpha=16,
lora_dropout=0.1,
target_modules=["q_proj", "v_proj"]
)
# 3. 将 LoRA 应用于模型
peft_model = get_peft_model(model, peft_config)
# 4. 打印可训练参数
peft_model.print_trainable_parameters()
# 输出会显示可训练参数仅占总参数的极小一部分
# trainable params: 4,194,304 || all params: 6,742,609,920 || trainable%: 0.0622
现在,你只需要像平常一样训练 peft_model 即可,peft 库会自动处理好梯度的计算和参数的更新。
16.2 “刀刃上的舞蹈”:极致显存优化
即使使用了 LoRA,对于非常大的模型(如 70B),或者在显存较小的 GPU 上(如 24GB),基础模型本身的加载仍然可能导致 OOM。这时,我们需要更激进的显存优化技术。
16.2.1 梯度检查点 (Gradient Checkpointing)
- 思想: 以时间换空间。
- 标准反向传播: 在前向传播时,
Autograd会保存大量的中间激活值(例如,每个 Transformer 层的输出),以便在反向传播时快速计算梯度。这些激活值是显存占用的主要来源之一。 - 使用梯度检查点: 当启用梯度检查点后,在前向传播时,大部分中间激活值不会被保存。当反向传播需要用到某个激活值来计算梯度时,它会重新从前向传播计算到那个点,以获取该激活值。
- 效果: 这种方法几乎可以将激活值占用的显存降低到接近零,代价是训练速度会变慢(大约 20-30%),因为需要额外的重计算。
- API: 在
transformers中,可以通过model.gradient_checkpointing_enable()一行代码启用。
16.2.2 ZeRO 思想与 FSDP 的实现
- 思想: 消除数据并行中的冗余。
- 标准数据并行 (DDP): 每块 GPU 都保存着完整的模型参数、梯度和优化器状态的副本。对于一个 7B 的
float32模型,仅参数就需要7 * 4 = 28GB显存,这在 4 块 GPU 上就是28 * 4 = 112GB的总显存占用,其中有112 - 28 = 84GB是冗余的。 - ZeRO (Zero Redundancy Optimizer):
- ZeRO-1: 分片优化器状态 (Optimizer States Sharding)。每块 GPU 只保存自己负责的那一部分参数的优化器状态。
- ZeRO-2: 在 ZeRO-1 的基础上,进一步分片梯度 (Gradients Sharding)。
- ZeRO-3: 在 ZeRO-2 的基础上,进一步分片模型参数本身 (Parameters Sharding)。这是最极致的优化,每块 GPU 在任何时刻都只持有模型的一小部分。当需要计算某个层时,它会通过
All-Gather操作从其他 GPU 获取该层所需的全部参数,计算完毕后立即释放。
- PyTorch FSDP (Fully Sharded Data Parallel): 这是 PyTorch 官方对 ZeRO-3 思想的高效实现。它与 DDP 的使用方式类似,但能够极大地降低单卡峰值显存。
FSDP 的优势:
- 它使得训练千亿级别参数的模型成为可能,因为单个 GPU 不再需要承载整个模型。
- 它与
gradient_checkpointing和AMP等技术可以完美结合。
使用 FSDP 需要更复杂的分布式环境设置,通常与 torch.distributed.run 结合使用。transformers 的 Trainer 和 accelerate 库对 FSDP 提供了很好的集成支持。
16.3 大模型训练的调参路径图
当你在资源受限的环境下,面对一个大模型,想要成功启动训练并获得最佳吞吐量,可以遵循以下系统性的调参路径:
目标: 在不 OOM 的前提下,尽可能地增大全局批次大小 (Global Batch Size),因为这通常能带来更好的收敛性和更高的训练吞吐量。
Step 0: 基础设置
- 模型加载: 使用
torch_dtype=torch.float16或torch.bfloat16加载模型。这能将模型参数的显存占用减半。bfloat16对训练更友好。 - AMP: 始终启用自动混合精度训练。
Step 1: 启用 PEFT (LoRA)
- 这是最优先的策略。它能将需要优化的参数量级从几十亿降低到几百万,从而极大地减少优化器状态和梯度占用的显存。
Step 2: 寻找最大单卡批次大小 (per_device_batch_size)
- 从
per_device_batch_size=1开始,逐步增大,直到发生 OOM。记录下不 OOM 的最大值。
Step 3: 启用梯度检查点
- 如果 Step 2 中
per_device_batch_size只能设为很小的值(如 1),导致 GPU 利用率很低,那么启用梯度检查点 (gradient_checkpointing_enable())。 - 启用后,重复 Step 2,你通常会发现最大单卡批次大小可以显著增加。
Step 4: 使用梯度累积 (gradient_accumulation_steps)
- 现在你已经有了最大的单卡批次大小,但可能全局批次大小仍然不够大(例如,你想达到 128,但 8 卡 GPU 每张最多只能跑
bs=4,全局只有 32)。 - 设置
gradient_accumulation_steps = 目标全局批次大小 / (GPU数量 * 单卡批次大小)。例如128 / (8 * 4) = 4。
Step 5: 终极方案 - FSDP/ZeRO
- 如果即使启用了梯度检查点,单卡批次大小仍然无法大于等于 1(即连一个样本都放不下),或者你希望进一步优化多卡训练的吞吐量,那么就必须使用 FSDP。
- FSDP 通过分片模型参数,可以让你在单卡上运行远超其显存容量的模型。
一个典型的 Hugging Face Trainer 配置:
from transformers import TrainingArguments
training_args = TrainingArguments(
output_dir="./results",
per_device_train_batch_size=4, # Step 2 & 3 的结果
gradient_accumulation_steps=4, # Step 4 的结果
gradient_checkpointing=True, # Step 3
optim="paged_adamw_8bit", # 使用 8-bit AdamW 进一步节省显存
learning_rate=2e-4,
fp16=True, # Step 0 (AMP)
# FSDP 配置 (Step 5)
# fsdp="full_shard auto_wrap",
# fsdp_transformer_layer_cls_to_wrap='LlamaDecoderLayer',
# ...
)
性能/踩坑清单
- LoRA
target_modules:target_modules的选择对最终性能有影响。对于 Llama-style 模型,通常选择["q_proj", "v_proj"]是一个不错的起点。有些研究表明,对所有nn.Linear层都应用 LoRA 效果可能更好。 - 合并 LoRA 权重: 在使用
peft训练后,如果你想进行部署或分享一个独立的、不需要peft库的模型,可以使用model = model.merge_and_unload()来将 LoRA 权重合并回原始模型中。 - FSDP 与
device_map冲突: FSDP 和device_map="auto"(一种简单的模型并行方式) 通常是互斥的。在使用 FSDP 时,每个 FSDP 进程应该只负责一块 GPU。 - 吞吐量 vs. 收敛速度: 所有的优化技巧都是在“时间”、“空间”和“模型性能”之间做权衡。例如,梯度检查点节省了显存但增加了计算时间;梯度累积模拟了大批次,但实际的参数更新频率变低了。最终的目标是找到一个能让模型在可接受的时间内收敛到理想性能的最佳组合。
- CPU Offload: 对于 ZeRO-3 和 FSDP,还可以选择将暂时不用的参数、梯度和优化器状态**卸载 (Offload)**到 CPU 内存中,这能让你在单 GPU 上运行惊人大小的模型,但代价是极大的训练速度下降,因为 CPU 和 GPU 之间的数据传输 (PCIe) 远慢于 GPU 内部通信 (NVLink)。
练习题
- LoRA 参数计算: 假设一个
nn.Linear层的权重矩阵W形状为(4096, 4096)。如果我们使用r=8的 LoRA,那么新增的可训练参数A(形状(8, 4096)) 和B(形状(4096, 8)) 的总参数量是多少?它与原始权重矩阵的参数量相比,比例是多少? - Adapter vs. LoRA: 阅读
peft库关于 Adapter 的文档。比较 Adapter 和 LoRA 在注入方式和计算流程上的异同。 - 梯度检查点效果测量: 使用一个中等大小的模型(如
gpt2-medium),编写一个脚本,测量在启用和禁用梯度检查点两种情况下,进行一次包含前向和反向传播的训练步骤所达到的峰值 GPU 显存占用。使用torch.cuda.max_memory_allocated()来获取峰值。 - FSDP 思想实验: 假设你正在使用 ZeRO-3 (FSDP) 在 4 块 GPU 上训练一个模型,该模型只有一个
nn.Linear层。请描述在前向传播和反向传播过程中,该层的权重矩阵、输入数据和梯度是如何在 4 块 GPU 之间被分片、收集 (All-Gather) 和减少 (Reduce-Scatter) 的。
延伸阅读
- PEFT & LoRA:
- Hugging Face
peftLibrary: https://github.com/huggingface/peft - 官方仓库,包含所有 PEFT 方法的实现和文档。 - LoRA Paper: “LoRA: Low-Rank Adaptation of Large Language Models” (by Hu et al., 2021) - LoRA 的原始论文。
- Hugging Face
- ZeRO & FSDP:
- PyTorch FSDP Documentation: https://pytorch.org/docs/stable/fsdp.html - 官方 FSDP 文档。
- DeepSpeed ZeRO Paper: “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models” (by Rajbhandari et al., 2019) - 提出 ZeRO 思想的开创性论文。
- Hugging Face FSDP Integration Guide: https://huggingface.co/docs/accelerate/fsdp - 在 HF 生态中使用 FSDP 的实用指南。
- 综合博客:
- “Illustrated Guide to Transformer Fine-tuning” (by various authors): 搜索相关主题,可以找到许多关于 PEFT 和大模型微调的优秀图文并茂的博客。
PyTorch 详细学习笔记 其他章节链接
以下是整个系列的20章目录,点击章节标题即可跳转阅读:。
- PyTorch 详细学习笔记 第一章:张量基础与运算模型
- PyTorch 详细学习笔记 第二章:Autograd 与计算图
- PyTorch 详细学习笔记 第三章:模型构建:
nn.Module的正确打开方式 - PyTorch 详细学习笔记 第四章:数据管道:
Dataset/DataLoader深入 - PyTorch 详细学习笔记 第五章:从零实现训练循环
- PyTorch 详细学习笔记 第六章:设备管理与性能基础
- PyTorch 详细学习笔记 第七章:调试与可解释性
- PyTorch 详细学习笔记 第八章:复现实验与工程可靠性
- PyTorch 详细学习笔记 第九章:可视化与实验管理
- PyTorch 详细学习笔记 第十章:计算机视觉入门实战
- PyTorch 详细学习笔记 第十一章:NLP 入门实战
- PyTorch 详细学习笔记 第十二章:时序与序列建模
- PyTorch 详细学习笔记 第十三章:生成式模型速览
- PyTorch 详细学习笔记 第十四章:推荐系统基础
- PyTorch 详细学习笔记 第十五章:并行与分布式训练
- PyTorch 详细学习笔记 第十六章:参数高效训练与大模型技巧
- PyTorch 详细学习笔记 第十七章:模型压缩与加速
- PyTorch 详细学习笔记 第十八章:导出、部署与推理性能
- PyTorch 详细学习笔记 第十九章:工程化最佳实践与项目骨架
- PyTorch 详细学习笔记 第二十章:融会贯通:从学习者到实践者的蜕变
更多推荐
所有评论(0)