在单个 GPU 上高效训练的方法和工具¶
本指南展示了实用的技术，用于通过优化内存利用、加速训练或两者来提高模型训练的效率。如果您想了解在训练过程中如何利用 GPU，请首先参考 模型训练解剖学 概念指南。本指南侧重于实用技术。
如果您有一台多 GPU 的机器，这些方法仍然有效，此外您还可以利用 多 GPU 部分 中列出的其他方法。
在训练大型模型时，需要同时考虑两个方面：
- 数据吞吐量/训练时间
- 模型性能
最大化吞吐量（每秒样本数）可以降低训练成本。这通常通过尽可能利用 GPU 并填满 GPU 内存来实现。如果所需的批量大小超过了 GPU 内存的限制，可以使用梯度累加等内存优化技术来帮助解决。
然而，如果首选的批量大小适合内存，就没有必要应用内存优化技术，因为它们可能会减慢训练速度。仅仅因为可以使用大批量，并不意味着应该这样做。在超参数调优过程中，您应该确定哪个批量大小效果最好，然后相应地优化资源。
本指南中涵盖的方法和工具可以根据它们对训练过程的影响进行分类：
| 方法/工具 | 提高训练速度 | 优化内存利用 |
|---|---|---|
| 批量大小选择 | 是 | 是 |
| 梯度累加 | 否 | 是 |
| 梯度检查点 | 否 | 是 |
| 混合精度训练 | 是 | 可能* |
| torch_empty_cache_steps | 否 | 是 |
| 优化器选择 | 是 | 是 |
| 数据预加载 | 是 | 否 |
| DeepSpeed ZeRO | 否 | 是 |
| torch.compile | 是 | 否 |
| 参数高效微调 (PEFT) | 否 | 是 |
* 注：当使用混合精度与小模型和大批量时，会有一定的内存节省；但对于大模型和小批量，内存使用会更大。
您可以结合上述方法以获得累积效果。无论您是使用 Trainer 进行模型训练，还是编写纯 PyTorch 循环，您都可以 使用 🤗 Accelerate 配置这些优化。
如果这些方法没有带来足够的收益，您可以探索以下选项：
最后，如果所有方法仍然不够，即使切换到服务器级 GPU（如 A100），也可以考虑转向多 GPU 设置。所有这些方法在多 GPU 设置中仍然有效，并且您可以利用 多 GPU 部分 中概述的额外并行技术。
批量大小选择¶
要实现最佳性能，首先要确定合适的批量大小。建议使用 2^N 大小的批量和输入/输出神经元数量。通常它是 8 的倍数，但具体取决于所使用的硬件和模型的数据类型。
参考 NVIDIA 推荐的 输入/输出神经元数量 和 批量大小，这些涉及完全连接层（参与 GEMMs（通用矩阵乘法））。
Tensor Core 要求 定义了基于数据类型和硬件的乘数。例如，对于 fp16 数据类型，建议使用 8 的倍数，除非是 A100 GPU，这种情况下使用 64 的倍数。
对于较小的参数，还应考虑 维度量化效应。这里会发生切片，合适的乘数可以显著加速。
梯度累加(Gradient Accumulation)¶
梯度累加 方法旨在分小批计算梯度，而不是一次性计算整个批量的梯度。这种方法通过多次前向和后向传递模型来迭代计算梯度，并在过程中积累梯度。一旦积累了足够数量的梯度，就执行模型的优化步骤。通过使用梯度累加，可以将 有效批量大小 增加到超过 GPU 内存容量的限制。然而，需要注意的是，梯度累加引入的额外前向和后向传递可能会减慢训练过程。
您可以通过在 TrainingArguments 中添加 gradient_accumulation_steps 参数来启用梯度累加：
training_args = TrainingArguments(per_device_train_batch_size=1, gradient_accumulation_steps=4, **default_args)
在上面的例子中，您的有效批量大小变为 4。
或者，使用 🤗 Accelerate 来获得对训练循环的完全控制。有关 🤗 Accelerate 示例，请参阅 本指南中的进一步说明。
虽然建议尽可能充分利用 GPU，但过多的梯度累加步骤会导致更明显的训练减速。考虑以下示例。假设 per_device_train_batch_size=4 不使用梯度累加达到了 GPU 的极限。如果您希望以 64 的批量进行训练，不要将 per_device_train_batch_size 设为 1 并将 gradient_accumulation_steps 设为 64。相反，保持 per_device_train_batch_size=4 并将 gradient_accumulation_steps 设为 16。这样可以获得相同的有效批量大小，同时更好地利用可用的 GPU 资源。
有关更多信息，请参阅 RTX-3090 和 A100 的批量大小和梯度累加基准测试。
梯度检查点(Gradient Checkpointing)¶
即使将批量大小设为 1 并使用梯度累加，某些大型模型仍可能面临内存问题。这是因为在前向传递期间保存所有激活以便在后向传递期间计算梯度会占用大量内存。另一种方法是在需要时丢弃激活并重新计算，但这会引入大量的计算开销并减慢训练过程。
梯度检查点 提供了一种折衷方案，通过在计算图中战略性地保存部分激活，使得只需重新计算一部分激活即可完成梯度计算。有关梯度检查点的深入解释，请参阅 这篇文章。
要在 Trainer 中启用梯度检查点，可以将相应的标志传递给 TrainingArguments：
training_args = TrainingArguments(
per_device_train_batch_size=1, gradient_accumulation_steps=4, gradient_checkpointing=True, **default_args
)
或者，使用 🤗 Accelerate - 请参阅 本指南中的进一步说明。
尽管梯度检查点可以改善内存效率，但它会使训练速度减慢约 20%。
混合精度训练(Mixed precision training)¶
混合精度训练 是一种通过使用低精度数值格式来优化模型训练计算效率的技术。传统上，大多数模型使用 32 位浮点精度（fp32 或 float32）来表示和处理变量。然而，并非所有变量都需要如此高的精度来获得准确的结果。通过将某些变量的精度降低到更低的数值格式（如 16 位浮点数（fp16 或 float16）），我们可以加快计算速度。在这种方法中，一些计算以半精度进行，而另一些则仍以全精度进行，因此称为混合精度训练。
最常见的混合精度训练是使用 fp16（float16）数据类型，但某些 GPU 架构（如 Ampere 架构）提供 bf16 和 tf32（CUDA 内部数据类型）。有关这些数据类型差异的详细信息，请参阅 NVIDIA 博客。
fp16¶
混合精度训练的主要优势在于以半精度（fp16）保存激活。尽管梯度也以半精度计算，但在优化步骤中会转换回全精度，因此在这里不会节省内存。虽然混合精度训练可以加快计算速度，但也会导致 GPU 内存使用增加，尤其是对于小批量。这是因为模型现在在 GPU 上同时存在 16 位和 32 位精度（1.5 倍于原始模型）。
要启用混合精度训练，可以设置 fp16 标志为 True：
training_args = TrainingArguments(per_device_train_batch_size=4, fp16=True, **default_args)
如果您更喜欢使用 🤗 Accelerate，请参阅 本指南中的进一步说明。
BF16¶
如果您有访问 Ampere 或更新硬件的权限，可以使用 bf16 进行混合精度训练和评估。虽然 bf16 的精度不如 fp16，但其动态范围更大。在 fp16 中，最大的数字是 65504，任何大于该值的数字都会导致溢出。bf16 数字可以达到 3.39e+38，与 fp32 相当，因为两者都使用 8 位来表示数值范围。
您可以在 🤗 Trainer 中启用 BF16：
training_args = TrainingArguments(bf16=True, **default_args)
TF32¶
Ampere 硬件使用一种名为 tf32 的神奇数据类型。它具有与 fp32 相同的数值范围（8 位），但精度只有 10 位（与 fp16 相同），总共使用 19 位。它之所以被称为“神奇”，是因为您可以使用正常的 fp32 训练和/或推理代码，通过启用 tf32 支持，可以将吞吐量提高多达 3 倍。您只需要在代码中添加以下内容：
import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
CUDA 将自动切换到使用 tf32 而不是 fp32，前提是使用的 GPU 是 Ampere 系列。
根据 NVIDIA 研究，大多数机器学习训练工作负载显示，使用 tf32 训练与使用 fp32 具有相同的困惑度和收敛性。如果您已经在使用 fp16 或 bf16 混合精度，它也可能有助于吞吐量。
您可以在 🤗 Trainer 中启用此模式：
TrainingArguments(tf32=True, **default_args)
tf32 不能直接通过 tensor.to(dtype=torch.tf32) 访问，因为它是一种内部 CUDA 数据类型。需要 torch>=1.7 才能使用 tf32 数据类型。
有关 tf32 与其他精度的更多基准测试，请参阅：RTX-3090 和 A100。
Flash Attention 2¶
您可以使用 transformers 中的 Flash Attention 2 集成来加速训练吞吐量。有关如何加载带有 Flash Attention 2 模块的模型的详细信息，请参阅 单 GPU 部分。
优化器选择¶
最常用的优化器是 Adam 或 AdamW（带权重衰减的 Adam）。Adam 通过存储先前梯度的滚动平均值来实现良好的收敛，但会增加与模型参数数量相当的额外内存开销。为了补救这一点，您可以用替代优化器。例如，如果您安装了 NVIDIA GPUs 的 NVIDIA/apex 或 AMD GPUs 的 ROCmSoftwarePlatform/apex，adamw_apex_fused 将在所有支持的 AdamW 优化器中提供最快的训练体验。
Trainer 集成了多种优化器，可以直接使用：adamw_hf、adamw_torch、adamw_torch_fused、adamw_apex_fused、adamw_anyprecision、adafactor 或 adamw_bnb_8bit。还可以通过第三方实现添加更多优化器。
让我们更详细地看一下两种替代 AdamW 的优化器：
adafactor可在 Trainer 中使用adamw_bnb_8bit也在 Trainer 中可用，但下面提供了一个第三方集成示例以供演示。
以一个 3B 参数模型（如 “google-t5/t5-3b”）为例：
- 标准 AdamW 优化器需要 24GB 的 GPU 内存，因为它每个参数使用 8 字节（8 * 3 = 24GB）
- Adafactor 优化器需要超过 12GB 的内存。它每个参数使用略多于 4 字节，所以 4 * 3 加上一些额外的内存。
- 8 位 BNB 量化优化器仅需 (2 * 3) 6GB 的内存，前提是所有优化器状态都被量化。
Adafactor¶
Adafactor 不为权重矩阵中的每个元素存储滚动平均值，而是保留聚合信息（按行和列汇总的滚动平均值），显著减少其内存占用。然而，与 Adam 相比，Adafactor 在某些情况下可能收敛更慢。
您可以通过在 TrainingArguments 中设置 optim="adafactor" 来切换到 Adafactor：
training_args = TrainingArguments(per_device_train_batch_size=4, optim="adafactor", **default_args)
结合其他方法（梯度累加、梯度检查点和混合精度训练），可以看到高达 3 倍的改进，同时保持吞吐量！然而，如前所述，Adafactor 的收敛可能不如 Adam。
8 位 Adam¶
与 Adafactor 不同，8 位 Adam 保留完整的状态并对其进行量化。量化意味着它以较低的精度存储状态，仅在优化时反量化。这类似于混合精度训练的思想。
要使用 adamw_bnb_8bit，只需在 TrainingArguments 中设置 optim="adamw_bnb_8bit"：
training_args = TrainingArguments(per_device_train_batch_size=4, optim="adamw_bnb_8bit", **default_args)
然而，我们也可以使用第三方实现的 8 位优化器进行演示，看看如何集成。
首先，按照 GitHub 仓库 中的安装指南安装实现 8 位 Adam 优化器的 bitsandbytes 库。
接下来，需要初始化优化器。这涉及两个步骤：
- 首先，将模型的参数分为两组 - 一组应用权重衰减，另一组不应用。通常，偏差和层归一化参数不进行权重衰减。
- 然后进行一些参数整理，以使用与之前使用的 AdamW 优化器相同的参数。
import bitsandbytes as bnb
from torch import nn
from transformers.trainer_pt_utils import get_parameter_names
training_args = TrainingArguments(per_device_train_batch_size=4, **default_args)
decay_parameters = get_parameter_names(model, [nn.LayerNorm])
decay_parameters = [name for name in decay_parameters if "bias" not in name]
optimizer_grouped_parameters = [
{
"params": [p for n, p in model.named_parameters() if n in decay_parameters],
"weight_decay": training_args.weight_decay,
},
{
"params": [p for n, p in model.named_parameters() if n not in decay_parameters],
"weight_decay": 0.0,
},
]
optimizer_kwargs = {
"betas": (training_args.adam_beta1, training_args.adam_beta2),
"eps": training_args.adam_epsilon,
}
optimizer_kwargs["lr"] = training_args.learning_rate
adam_bnb_optim = bnb.optim.Adam8bit(
optimizer_grouped_parameters,
betas=(training_args.adam_beta1, training_args.adam_beta2),
eps=training_args.adam_epsilon,
lr=training_args.learning_rate,
)
最后，将自定义优化器作为参数传递给 Trainer：
trainer = Trainer(model=model, args=training_args, train_dataset=ds, optimizers=(adam_bnb_optim, None))
结合其他方法（梯度累加、梯度检查点和混合精度训练），可以预期获得大约 3 倍的内存改进，甚至稍微更高的吞吐量，就像使用 Adafactor 一样。
multi_tensor¶
pytorch-nightly 引入了 torch.optim._multi_tensor，这应该在处理大量小特征张量的情况下显著加快优化器的速度。它最终将成为默认设置，但如果您想提前尝试，可以查看这个 GitHub 问题。
数据预加载¶
实现快速训练速度的一个重要要求是能够以 GPU 能够处理的最大速度喂数据。默认情况下，所有操作都在主进程中进行，可能无法从磁盘读取数据快到足以创建瓶颈，从而导致 GPU 利用率不足。配置以下参数可以减少瓶颈：
DataLoader(pin_memory=True, ...)- 确保数据预加载到 CPU 的固定内存中，通常可以更快地从 CPU 转移到 GPU 内存。DataLoader(num_workers=4, ...)- 启动多个工作线程以更快地预加载数据。在训练过程中，观察 GPU 利用率统计信息；如果远未达到 100%，可以尝试增加工作线程的数量。当然，问题可能出在其他地方，因此增加工作线程数量并不一定会提高性能。
当使用 Trainer 时，对应的 TrainingArguments 是：dataloader_pin_memory（默认为 True），和 dataloader_num_workers（默认为 0）。
DeepSpeed ZeRO¶
DeepSpeed 是一个开源深度学习优化库，已集成到 🤗 Transformers 和 🤗 Accelerate 中。它提供了广泛的功能和优化，旨在提高大规模深度学习训练的效率和可扩展性。
如果您的模型能够适应单个 GPU 并有足够的空间容纳小批量，就不需要使用 DeepSpeed，因为它只会减慢速度。然而，如果模型无法适应单个 GPU，或者无法容纳小批量，您可以利用 DeepSpeed ZeRO + CPU 卸载，或 NVMe 卸载来处理更大的模型。在这种情况下，需要单独 安装库，然后按照其中一个指南创建配置文件并启动 DeepSpeed：
- 对于 Trainer 的 DeepSpeed 集成深入指南，请参阅 相关文档，特别是 [单 GPU 部分](main_classes/deepspeed#部署到单 GPU)。在笔记本中使用 DeepSpeed 需要进行一些调整，请参阅 相应指南。
- 如果您更喜欢使用 🤗 Accelerate，请参阅 🤗 Accelerate DeepSpeed 指南。
使用 torch.compile¶
PyTorch 2.0 引入了一个新的编译函数，无需修改现有 PyTorch 代码，只需添加一行代码即可优化代码：model = torch.compile(model)。
如果使用 Trainer，只需在 TrainingArguments 中传递 torch_compile 选项：
training_args = TrainingArguments(torch_compile=True, **default_args)
torch.compile 使用 Python 的帧评估 API 从现有的 PyTorch 程序中自动创建图。捕获图后，可以部署不同的后端以将图降低到优化引擎。有关更多细节和基准测试，请参阅 PyTorch 文档。
torch.compile 有一个不断增长的后端列表，可以通过调用 torchdynamo.list_backends() 查看，每个后端都有其可选依赖项。
通过在 TrainingArguments 中指定 torch_compile_backend 来选择要使用的后端。一些最常用的后端包括：
调试后端：
dynamo.optimize("eager")- 使用 PyTorch 运行提取的 GraphModule。这在调试 TorchDynamo 问题时非常有用。dynamo.optimize("aot_eager")- 使用 AotAutograd 但不使用编译器，即仅使用 PyTorch 快速模式运行 AotAutograd 提取的前向和后向图。这在调试时很有用，但不太可能提高速度。
训练和推理后端：
dynamo.optimize("inductor")- 使用 TorchInductor 后端和 AotAutograd 以及代码生成的 Triton 内核。了解更多dynamo.optimize("nvfuser")- 使用 nvFuser 与 TorchScript。了解更多dynamo.optimize("aot_nvfuser")- 使用 nvFuser 与 AotAutograd。了解更多dynamo.optimize("aot_cudagraphs")- 使用 cudagraphs 与 AotAutograd。了解更多
仅推理后端：
dynamo.optimize("ofi")- 使用 TorchScript optimize_for_inference。了解更多dynamo.optimize("fx2trt")- 使用 NVIDIA TensorRT 进行推理优化。了解更多dynamo.optimize("onnxrt")- 使用 ONNXRT 进行 CPU/GPU 推理。了解更多dynamo.optimize("ipex")- 使用 IPEX 进行 CPU 推理。了解更多
有关使用 torch.compile 与 🤗 Transformers 的示例，请参阅这篇 博客文章。
使用 🤗 PEFT¶
参数高效微调 (PEFT) 方法在微调过程中冻结预训练模型参数，并在其顶部添加少量可训练参数（适配器）。
结果是 与优化器状态和梯度相关的内存 大幅减少。
例如，使用标准 AdamW 时，优化器状态的内存需求为：
- fp32 参数副本：4 字节/参数
- 动量：4 字节/参数
- 方差：4 字节/参数
假设一个 7B 参数模型，并使用 低秩适配器 注入 2000 万参数。
普通模型的优化器状态内存需求为 12 * 7 = 84 GB（假设 7B 可训练参数）。
添加 Lora 会略微增加与模型权重相关的内存，但会大幅减少优化器状态的内存需求至 12 * 0.2 = 2.4 GB。
有关 PEFT 及其详细使用的更多信息，请参阅 PEFT 文档 或 PEFT 仓库。
使用 🤗 Accelerate¶
使用 🤗 Accelerate，您可以在获得对训练循环的完全控制的同时使用上述方法，基本上可以编写纯 PyTorch 循环，只需进行一些小的修改。
假设您已经组合了 TrainingArguments 中的方法，如下所示：
training_args = TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
gradient_checkpointing=True,
fp16=True,
**default_args,
)
使用 🤗 Accelerate 的完整训练循环示例只有几行代码：
from accelerate import Accelerator
from torch.utils.data.dataloader import DataLoader
dataloader = DataLoader(ds, batch_size=training_args.per_device_train_batch_size)
if training_args.gradient_checkpointing:
model.gradient_checkpointing_enable()
accelerator = Accelerator(fp16=training_args.fp16)
model, optimizer, dataloader = accelerator.prepare(model, adam_bnb_optim, dataloader)
model.train()
for step, batch in enumerate(dataloader, start=1):
loss = model(**batch).loss
loss = loss / training_args.gradient_accumulation_steps
accelerator.backward(loss)
if step % training_args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
首先，我们将数据集包装在一个 DataLoader 中。然后，可以通过调用模型的 gradient_checkpointing_enable() 方法启用梯度检查点。当初始化 Accelerator 时，可以指定是否使用混合精度训练，它将在 prepare 调用中为您处理。在 prepare 调用中，如果使用多个 GPU，数据加载器也将分布在各个工作线程上。我们使用前面示例中的相同 8 位优化器。
最后，可以添加主要的训练循环。注意，backward 调用由 🤗 Accelerate 处理。我们还可以看到梯度累加是如何工作的：我们对损失进行归一化，以便在累积结束时得到平均值，一旦累积了足够的步数，就执行优化。
使用 🤗 Accelerate 实现这些优化技术只需几行代码，而且具有更大的训练循环灵活性。有关所有功能的完整文档，请参阅 Accelerate 文档。
高效软件预构建¶
PyTorch 的 pip 和 conda 构建 预构建了 cuda 工具包，这足以运行 PyTorch，但不足以构建 cuda 扩展。
有时，可能需要额外的努力来预构建某些组件。例如，如果您使用像 apex 这样的库，它们没有预编译。在其他情况下，弄清楚如何系统范围内安装正确的 cuda 工具包可能会很复杂。为了解决这些问题，PyTorch 和 NVIDIA 发布了一个新的 NGC Docker 容器版本，已经预构建了所有组件。您只需在其中安装您的程序，就可以直接运行。
如果要调整 PyTorch 源代码和/或制作新的自定义构建，这种方法也非常有用。要找到所需的 Docker 镜像版本，可以从 PyTorch 发行说明 开始，选择最新的月度发行版之一。进入所需发行版的发行说明，检查环境组件是否符合您的需求（包括 NVIDIA 驱动程序要求！），然后在文档顶部转到相应的 NGC 页面。如果迷失方向，这里有一个 所有 PyTorch NGC 镜像的索引。
接下来，按照说明下载并部署 Docker 镜像。
专家混合(Mixture of Experts)¶
最近的一些论文报告称，通过在 Transformer 模型中集成专家混合（MoE），训练速度提高了 4-5 倍，推理也更快。
自从发现更多的参数可以提高性能以来，这项技术允许在不增加训练成本的情况下将参数数量增加一个数量级。
在这种方法中，每隔一个 FFN 层被替换为 MoE 层，MoE 层由许多专家组成，通过门控函数根据输入令牌在序列中的位置平衡每个专家的训练。
（来源：GLAM）
您可以在本部分末尾列出的论文中找到详尽的细节和比较表。
这种方法的主要缺点是需要巨大的 GPU 内存 - 几乎是其密集型等效模型的十倍。提出了各种蒸馏方法来克服较高的内存要求。
然而，存在直接的权衡。您可以使用少数专家和 2-3 倍更小的基础模型，而不是几十个或几百个专家，从而将模型大小缩小 5 倍，适度提高训练速度，同时适度增加内存需求。
大多数相关论文和实现都围绕 TensorFlow/TPU 构建：
对于 PyTorch，DeepSpeed 也构建了一个：DeepSpeed-MoE: 推动下一代 AI 规模的专家混合推理和训练，专家混合 - 博客文章：1，2 以及特定于大型基于变压器的自然语言生成模型的部署：博客文章，Megatron-DeepSpeed 分支。
使用 PyTorch 原生注意力机制和 Flash Attention¶
PyTorch 的 torch.nn.functional.scaled_dot_product_attention（SDPA） 在底层也可以调用 Flash Attention 和内存高效注意力内核。SDPA 支持目前正被本地添加到 Transformers 中，默认情况下在 torch>=2.1.1 时使用，前提是有可用的实现。有关支持的模型列表和更多详细信息，请参阅 PyTorch 缩放点积注意力。
有关使用 SDPA 加速和节省内存的更多信息，请参阅这篇 博客文章。