使用 trl 和 DeepSpeed 进行分布式 SFT 第 2 部分：本地扩展
简介
在本系列的第一部分中，我们介绍了使用 trl 设置本地 SFT 实验的基础知识。我们学习了如何为 trl 的 SFTTrainer 格式化数据集，并对其进行预处理以适应所需的结构。
现在，是时候迈出下一步了。在这篇文章中，我们将重点关注如何扩展 SFT 设置以处理更大的任务。具体来说，我们将探讨如何在单节点、多 GPU 环境中微调 LLM。在此过程中，我们将讨论优化技术，以减少内存使用、加快训练速度并支持更大模型的微调。让我们开始吧！
先决条件
要遵循本教程，您需要一台配备多个 NVIDIA GPU 的机器。确保 GPU 通过高速互连连接，以最大限度地减少通信开销。作为参考，我使用 8 块 NVIDIA V100 SXM2 GPU 运行了此实验。
重要考量
GPU 架构：虽然我使用 V100 GPU 运行了此实验，但强烈推荐使用 Ampere 或 Hopper 等较新的架构。这些 GPU 提供了高级功能，例如支持更高效的精度类型和改进的通信速度。此外，像 flash-attention 这样的技术仅与 Ampere 或更新的 GPU 兼容。
互连质量：使用
nvidia-smi topo -m验证 GPU 通信带宽。糟糕的互连在训练期间可能成为瓶颈。
此外，您还需要安装以下依赖项：
datasets
torch
transformers
trl
调整超参数
BAAI/Infinity-Instruct 提供了几个官方微调模型，包括 Llama3.1-70B、mistral-7B、Qwen2-7B 和 Yi-1.5-9B。他们还慷慨地分享了这些模型的训练细节。
对于本教程，我们将使用 Qwen2-7B 的超参数作为参考。这些超参数在 trl 中的训练参数对应关系如下：
- epoch:
--num_train_epochs - lr:
--learning_rate - lr_warmup_steps:
--warmup_steps - lr_decay_style:
--lr_scheduler_type（与min_lr一起设置为cosine_with_min_lr。可用的调度器选项可以在此处找到。） - min_lr:
--lr_scheduler_kwargs（设置为"{\"min_lr\": 0}"。这个参数没有明确的文档；我是通过这个 PR 和这个测试用例发现的。） - weight_decay:
--weight_decay - adam_beta1:
--adam_beta1 - adam_beta2:
--adam_beta2 - clip_grad:
--max_grad_norm
另一个值得一提的参数是 global_batch_size，它没有直接在训练脚本中设置。全局批大小由以下等式确定：
global_batch_size = per_device_train_batch_size * gradient_accumulation_steps * numGPUs.
例如，如果我们的目标全局批大小为 528 并且我们使用 8 个 GPU，则本地批大小（每个 GPU）将为
528 / 8 = 66.
如果我们每个 GPU 的每个批次可以容纳 2 个样本，那么我们可以将 per_device_train_batch_size 设置为 2，将 gradient_accumulation_steps 设置为 33。
另一个重要考虑因素是训练精度。现代 GPU（Ampere 系列或更新版本）支持 bf16 和 tf32，而旧版 GPU 仅支持 fp16 和 fp32。微调时，请确保精度与基础模型匹配。具体来说，如果基础模型是使用 bf16 训练的，请避免使用 fp16。有关更多详细信息，请参阅此 PR。
您可以通过 config.json 文件中的 torch_dtype 字段找到基础模型的数据类型。因此，如果您正在微调 bf16 模型但无法访问 Ampere 或更新的 GPU（就像我一样），最好暂时坚持使用 fp32。
现在我们已经介绍了基本的超参数和注意事项，接下来让我们看看一些优化技术，它们将有助于提高训练效率和资源利用率。
梯度累积
您可能已经注意到我使用 per_device_train_batch_size 和 gradient_accumulation_steps 来计算本地批大小。梯度累积允许您在更新模型之前在多个小批量上累积梯度。当所需的批大小超出硬件的内存容量时，此技术特别有用。
作为一般准则：
- 使用适合您的显存的最大
per_device_train_batch_size。 - 如有必要，调整
gradient_accumulation_steps以达到您的目标批大小。
通过这种方式，您可以有效地模拟更大的批大小而不会遇到内存限制
梯度检查点
梯度检查点是一种内存优化技术，通过权衡计算来减少内存使用。在训练过程中，大部分内存用于存储中间激活以进行反向传播。梯度检查点通过选择性地保存一部分激活并在反向传播期间重新计算其余部分来减少这种内存使用。
注意：根据 https://pytorch.ac.cn/docs/stable/checkpoint.html
目前有两种可用的检查点实现，由
use_reentrant参数决定。建议您使用use_reentrant=False。
您可以阅读该部分，以更深入地了解两种实现之间的差异。
在撰写本文时，transformers 库 (v4.48.1) 默认使用可重入实现。要使用非可重入版本，您必须显式传递以下参数：
--gradient_checkpointing_kwargs "{\"use_reentrant\": false}"
ZeRO
并行训练任务有几种方法，包括**数据并行 (DP)**、**张量并行 (TP)**、**流水线并行 (PP)**、**零冗余优化器 (ZeRO)**、**序列并行**和**专家并行**。有关这些方法的详细概述，我建议您查看这篇优秀资源。
在本教程中，我们将重点关注 **ZeRO**，它比传统的 DP 提供更高的效率，而无需修改训练代码。
ZeRO (零冗余优化器) 是一种通过减少内存使用来扩展训练的强大技术。如果您是 ZeRO 的新手，请查看原始论文或这篇详细文章。
ZeRO 有三个阶段，每个阶段都针对不同的内存节省方面：
- 阶段 1：减少优化器状态内存。
- 阶段 2：进一步减少梯度内存。
- 阶段 3：完全划分模型状态，以增加通信开销为代价实现最高的内存节省。
阶段 3 提供最大的内存效率，但如果 GPU 间通信不够快，可能会显著减慢训练速度。作为一般准则：
- 从阶段 2 开始。
- 仅当阶段 2 仍导致 CUDA OOM 时才尝试阶段 3。
话虽如此，始终值得在您的设置上测试阶段 2 和阶段 3，以确定哪一个在您的硬件上表现更好。
在本教程中，我们将使用 ZeRO 的官方实现——DeepSpeed。要使用 DeepSpeed，您需要先安装它。DeepSpeed 提供可以预安装或 JIT 编译的 C++/CUDA 操作。如果您选择预安装选项，请参阅此文档。我们将通过运行以下命令使用 JIT 方法安装 DeepSpeed：
pip install deepspeed
Hugging Face transformers 库内置了对 DeepSpeed 的支持。您可以通过在训练脚本中使用 --deepspeed 标志指定 DeepSpeed 配置文件来启用它。有关更多信息，请参阅transformers 中的 DeepSpeed 文档。
Liger 内核
Liger Kernel 是 LinkedIn 开发的一系列 Triton 内核，旨在减少内存使用并提高训练吞吐量。最棒的是它不需要复杂的配置，使其成为您设置的轻松补充。要安装它，请运行：
pip install liger-kernel
安装后，在训练脚本中添加 --use_liger 标志，您将自动节省显存，无需任何额外设置或麻烦。这是一种在不牺牲性能的情况下优化训练的直接方法。
样本打包
大型模型在 GPU 上训练以利用其并行性。然而，在语言模型的上下文中，我们在文本序列上进行训练，每个样本的长度各不相同。
处理变长序列的传统方法是将每个样本填充到与批处理中最长的样本匹配。虽然这确保了统一的输入维度，但由于填充，也导致了大量的内存浪费。
样本打包通过将较短的样本组合成一个序列来解决这个问题。这项技术允许更有效地利用 GPU 内存，减少浪费并可能加快训练速度。
虽然概念很简单，但正确实现它可能是我这次实验中最具挑战性的任务之一。
乍一看，trl 通过简单地传递一个参数来支持打包数据集。然而，经过进一步调查，我发现该实现可能不适合我的需求。正如此问题中指出的那样，注意力掩码没有得到正确处理，这可能导致序列之间注意力存在潜在的交叉污染。下图清楚地说明了这个问题。左侧是使用 --packing 的结果，右侧是打包样本的正确方法：
经过进一步深入研究，我发现，至少目前，“正确的打包方式”仅支持 Flash Attention。如果您无法访问 Ampere 或更新的 GPU，您可能需要坚持使用传统的填充方法。
然而，如果您有幸拥有这些 GPU，您可以按照这篇博客文章在训练期间启用样本打包。请注意，我尚未亲自验证这种方法。此外，截至撰写本文时，有一些与此功能相关的 PR 尚未发布（例如此 PR）。要访问此功能，您可能需要从源代码安装 trl 和 transformers：
pip install git+https://github.com/huggingface/trl
pip install git+https://github.com/huggingface/tranformers
分布式训练
通过所有优化到位，我们现在可以跨多个 GPU 扩展我们的 SFT 实验。为此，我们可以使用像 torchrun、deepspeed 或 accelerate 这样的工具。我个人更喜欢 torchrun，因为它简单易用。
通过运行以下命令，我们可以将训练作业分发到多个 GPU：
哦，别忘了设置 wandb 进行日志记录——我们现在正在进行适当的微调！😉
sft2.sh
torchrun \
--nproc_per_node 8 \
sft.py \
--model_name_or_path Qwen/Qwen2.5-3B \
--dataset_name BAAI/Infinity-Instruct \
--dataset_config 0625 \
--do_train \
--learning_rate 1e-5 \
--lr_scheduler_type cosine_with_min_lr \
--lr_scheduler_kwargs "{\"min_lr\": 0}" \
--warmup_steps 40 \
--weight_decay 0.0 \
--max_grad_norm 1.0 \
--adam_beta1 0.9 \
--adam_beta2 0.95 \
--per_device_train_batch_size 11 \
--gradient_accumulation_steps 6 \
--gradient_checkpointing \
--gradient_checkpointing_kwargs "{\"use_reentrant\": false}" \
--num_train_epochs 3 \
--use_liger \
--deepspeed ./ds-config.json \
--output_dir /tmp/Qwen2.5-3B-Infinity-Instruct-0625 \
--report_to wandb \
--run_name my-second-sft-exp
ds-config.json
{
"fp16": {
"enabled": false
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},
"zero_optimization": {
"stage": 2,
"overlap_comm": false,
"allgather_bucket_size": 5e8,
"reduce_bucket_size": "auto",
"allgather_partitions": true,
"reduce_scatter": true,
"contiguous_gradients": true,
"round_robin_gradients": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
多亏了所有的优化，我能够微调一个 3B 模型，而不是第一部分中使用的 0.5B 模型。
在 V100 上完成训练确实花费了大量时间（大约 133 小时），所以我强烈建议使用现代 GPU 并启用 Flash Attention 和样本打包以获得更好的性能。
评估
现在训练已经完成，重要的是评估所有操作是否都正确完成。
快速检查模型性能的方法是与它进行交互。您可以参考官方 SFT 模型的快速入门部分进行尝试。以下是我刚刚微调的模型进行交互的示例：
Me: Give me a short introduction to large language model.
AI: A large language model (LLM) is a type of artificial intelligence (AI) that is designed to understand and generate human language. These models are trained on vast amounts of text data, allowing them to learn patterns, structures, and nuances
（请注意：AI 的响应由于我设置的 max_new_tokens 被截断，但您可以看到模型正在适当响应。）
虽然直接交互对于快速检查很有用，但正式评估对于更严格的验证至关重要。评估 LLM 是一个相当广泛的话题，我只在这里分享一些技巧。
有几个可用于评估 LLM 的框架，选择最佳框架具有挑战性，并且比较不同框架的结果有时会导致不公平的结论。
一个著名的评估平台是 Open LLM Leaderboard，它根据 LLM 的评估结果对其进行排名。Open LLM Leaderboard 使用 lm-evaluation-harness 作为其后端。通过使用相同的工具，您可以确保与排行榜中的模型进行公平比较。因此，对于本教程，我将使用 lm-evaluation-harness 来运行与 Open LLM Leaderboard 中相同的评估，以评估我刚刚微调的模型。
lm-evaluation-harness 集成了 Open LLM Leaderboard 中使用的所有任务。要评估您的模型在这些任务上的表现，您可以运行以下命令：
lm_eval \
--model hf \
--model_args pretrained=$MODEL_YOU_WANT_TO_EVAL \
--tasks leaderboard
MATH-hard 任务不可用
然而，截至撰写本文时，competition_math 数据集由于法律问题目前不可用。因此，我们需要跳过依赖此数据集的 MATH-hard 任务。您可以修改脚本以包含除 leaderboard_math_hard 之外的所有其他任务：
lm_eval \
--model hf \
--model_args pretrained=$MODEL_YOU_WANT_TO_EVAL \
--tasks leaderboard_bbh,leaderboard_gpqa,leaderboard_ifeval,leaderboard_mmlu_pro,leaderboard_musr
评估代码生成
除了排行榜评估之外，如果您有兴趣评估您的模型在代码生成任务（例如 humaneval）上的表现，请记住通常需要执行生成的代码才能评估其正确性。由于执行 LLM 生成的代码可能存在风险，大多数框架将默认中止此类任务。要允许在评估期间执行代码，您需要将 HF_ALLOW_CODE_EVAL 设置为 1 并在评估命令中包含 --confirm_run_unsafe_code 参数：
lm_eval \
--model hf \
--model_args pretrained=$MODEL_YOU_WANT_TO_EVAL \
--tasks leaderboard_bbh,leaderboard_gpqa,leaderboard_ifeval,leaderboard_mmlu_pro,leaderboard_musr \
--confirm_run_unsafe_code # Add this line
结论
在这篇文章中，我们涵盖了从基本设置到在单节点、多 GPU 环境中扩展大型语言模型的高级技术。通过利用 DeepSpeed 和 trl，即使在原本无法支持此类模型的硬件上，我们也可以高效地微调 Qwen2-3B 及更大型的模型。我还将微调后的模型上传到了 Hugging Face 模型中心，您可以亲自尝试一下：https://huggingface.co/jlzhou/Qwen2.5-3B-Infinity-Instruct-0625。
在本系列的下一部分中，我们将探索跨多个节点的分布式训练，解决跨不同机器的多个 GPU 的更复杂设置。敬请期待！