单GPU高效训练的方法和工具

本指南演示了您可以使用的实用技术，通过优化内存利用率、加速训练或同时实现两者来提高模型训练效率。如果您想了解 GPU 在训练期间是如何使用的，请先参考模型训练结构概念指南。本指南重点介绍实用技术。

在训练大型模型时，应同时考虑以下两个方面

• 数据吞吐量/训练时间
• 模型性能

最大化吞吐量（样本/秒）可以降低训练成本。这通常是通过尽可能多地利用 GPU 并将 GPU 内存填充到其极限来实现的。如果所需的批大小超过了 GPU 内存的限制，则内存优化技术（例如梯度累积）可以提供帮助。

但是，如果首选的批大小适合内存，则没有理由应用内存优化技术，因为它们会减慢训练速度。仅仅因为可以使用较大的批大小，并不一定意味着应该使用。作为超参数调整的一部分，您应该确定哪个批大小能产生最佳结果，然后相应地优化资源。

本指南中介绍的方法和工具可以根据它们对训练过程的影响进行分类

您可以组合以上方法以获得累积效果。无论您是使用Trainer训练模型还是编写纯 PyTorch 循环，这些技术都可供您使用，在这种情况下，您可以使用 🤗 Accelerate 配置这些优化。

如果这些方法没有带来足够的收益，您可以探索以下选项

• 研究构建自己的自定义 Docker 容器，其中包含高效的软件预构建
• 考虑使用专家混合 (MoE) 的模型
• 将您的模型转换为 BetterTransformer 以利用 PyTorch 原生注意力

最后，如果即使在切换到 A100 等服务器级 GPU 后，以上所有方法仍然不足，请考虑转向多GPU设置。所有这些方法在多GPU设置中仍然有效，此外，您还可以利用多GPU部分中概述的其他并行技术。

批大小选择

为了获得最佳性能，请先确定合适的批大小。建议使用大小为 2^N 的批大小和输入/输出神经元数量。通常它是 8 的倍数，但根据使用的硬件和模型的 dtype，它可以更高。

作为参考，请查看 NVIDIA 对输入/输出神经元数量和批大小的建议，这些建议适用于全连接层（参与 GEMM（通用矩阵乘法））。

Tensor Core 要求根据 dtype 和硬件定义乘数。例如，对于 fp16 数据类型，建议使用 8 的倍数，除非它是 A100 GPU，在这种情况下，请使用 64 的倍数。

对于参数较小的模型，请考虑维度量化效应。这是平铺发生的地方，正确的乘数可以显着提高速度。

梯度累积

梯度累积方法旨在以较小的增量计算梯度，而不是一次计算整个批次的梯度。这种方法涉及通过迭代计算较小批次的梯度来执行此操作，方法是对模型执行前向和后向传递，并在过程中累积梯度。一旦累积了足够数量的梯度，就会执行模型的优化步骤。通过使用梯度累积，可以将有效批大小增加到超出 GPU 内存容量限制的范围。但是，需要注意的是，梯度累积引入的额外前向和后向传递可能会减慢训练过程。

您可以通过将gradient_accumulation_steps参数添加到TrainingArguments来启用梯度累积

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的批大小和梯度累积基准。

梯度检查点

即使将批大小设置为 1 并使用梯度累积，某些大型模型仍然可能面临内存问题。这是因为还有其他组件也需要内存存储。

为了在反向传递期间计算梯度，保存前向传递中的所有激活会导致大量的内存开销。另一种方法是丢弃激活并在反向传递期间需要时重新计算它们，但这会带来相当大的计算开销并减慢训练过程。

梯度检查点在这两种方法之间提供了一种折衷方案，并在整个计算图中保存了经过策略性选择的激活，因此只需要重新计算一小部分激活即可获得梯度。有关梯度检查点的深入解释，请参阅这篇优秀的文章。

要在Trainer中启用梯度检查点，请将相应的标志传递给TrainingArguments

training_args = TrainingArguments(
    per_device_train_batch_size=1, gradient_accumulation_steps=4, gradient_checkpointing=True, **default_args
)

或者，使用 🤗 Accelerate - 请参阅本指南后面部分的 🤗 Accelerate 示例。

混合精度训练

混合精度训练是一种旨在通过利用较低精度的数值格式来表示某些变量，从而优化模型训练计算效率的技术。传统上，大多数模型使用 32 位浮点数精度 (fp32 或 float32) 来表示和处理变量。但是，并非所有变量都需要这种高精度级别才能获得准确的结果。通过将某些变量的精度降低到较低的数值格式（如 16 位浮点数 (fp16 或 float16)），我们可以加快计算速度。由于在这种方法中，一些计算以半精度执行，而另一些计算仍然以全精度执行，因此该方法称为混合精度训练。

混合精度训练最常见的是使用 fp16 (float16) 数据类型来实现，但是，一些 GPU 架构（例如 Ampere 架构）提供了 bf16 和 tf32（CUDA 内部数据类型）数据类型。查看NVIDIA 博客，了解更多关于这些数据类型之间差异的信息。

fp16

混合精度训练的主要优势在于以半精度 (fp16) 保存激活值。虽然梯度也以半精度计算，但它们会转换回全精度以进行优化步骤，因此此处不会节省内存。虽然混合精度训练会导致更快的计算，但它也可能导致使用更多的 GPU 内存，尤其是在小批量大小的情况下。这是因为模型现在以 16 位和 32 位精度同时存在于 GPU 上（GPU 上原始模型的 1.5 倍）。

要启用混合精度训练，请将fp16标志设置为True

training_args = TrainingArguments(per_device_train_batch_size=4, fp16=True, **default_args)

如果您更喜欢使用 🤗 Accelerate，请在本指南的后面部分查找 🤗 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 位），但精度不是 23 位，而是只有 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 与其他精度的更多信息，请参阅以下基准测试：RTX-3090 和 A100。

Flash Attention 2

您可以通过在 transformers 中使用 Flash Attention 2 集成来加快训练吞吐量。查看单个 GPU 部分中的相应部分，了解有关如何加载具有 Flash Attention 2 模块的模型的更多信息。

优化器选择

用于训练 Transformer 模型的最常用优化器是 Adam 或 AdamW（带权重衰减的 Adam）。Adam 通过存储先前梯度的滚动平均值来实现良好的收敛；但是，它增加了与模型参数数量同阶的额外内存占用。为了解决这个问题，您可以使用替代优化器。例如，如果您为 NVIDIA GPU 安装了NVIDIA/apex，或为 AMD GPU 安装了ROCmSoftwarePlatform/apex，则adamw_apex_fused将为您提供所有受支持的 AdamW 优化器中最快的训练体验。

Trainer集成了多种开箱即用的优化器：adamw_hf、adamw_torch、adamw_torch_fused、adamw_apex_fused、adamw_anyprecision、adafactor或adamw_bnb_8bit。更多优化器可以通过第三方实现进行插入。

让我们仔细看看 AdamW 优化器的两种替代方案

1. adafactor在Trainer中可用
2. adamw_bnb_8bit也在 Trainer 中可用，但下面提供了第三方集成以进行演示。

为了进行比较，对于一个 30 亿参数的模型，例如“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

8 位 Adam 并没有像 Adafactor 那样聚合优化器状态，而是保留了完整的状态并对其进行量化。量化意味着它以较低的精度存储状态，并且仅在优化时对其进行反量化。这类似于混合精度训练背后的思想。

要使用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 获得略高的吞吐量。

多张量

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 Offload 或 NVMe Offload 来处理更大的模型。在这种情况下，您需要单独 安装库，然后按照其中一个指南创建配置文件并启动 DeepSpeed

• 有关 DeepSpeed 与 Trainer 集成的深入指南，请查看 相应的文档，特别是 单个 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 且没有编译器，即仅对 AotAutograd 提取的前向和后向图使用 PyTorch eager。这对于调试很有用，并且不太可能提高速度。

训练和推理后端:

• dynamo.optimize("inductor") - 通过利用代码生成的 Triton 内核使用 TorchInductor 后端以及 AotAutograd 和 cudagraphs 阅读更多
• dynamo.optimize("nvfuser") - 带有 TorchScript 的 nvFuser。 阅读更多
• dynamo.optimize("aot_nvfuser") - 带有 AotAutograd 的 nvFuser。 阅读更多
• dynamo.optimize("aot_cudagraphs") - 带有 AotAutograd 的 cudagraphs。 阅读更多

仅推理后端

• dynamo.optimize("ofi") - 使用 TorchScript optimize_for_inference。 阅读更多
• dynamo.optimize("fx2trt") - 使用 NVIDIA TensorRT 进行推理优化。 阅读更多
• dynamo.optimize("onnxrt") - 使用 ONNXRT 在 CPU/GPU 上进行推理。 阅读更多
• dynamo.optimize("ipex") - 使用 IPEX 在 CPU 上进行推理。 阅读更多

有关使用 🤗 Transformers 的 torch.compile 示例，请查看此 关于使用最新 PyTorch 2.0 功能微调 BERT 模型进行文本分类的博文

使用 🤗 PEFT

参数高效微调 (PEFT) 方法在微调期间冻结预训练模型参数，并在其之上添加少量可训练参数（适配器）。

因此，与 优化器状态和梯度相关的内存 大大减少了。

例如，使用普通的 AdamW，优化器状态的内存需求将是

• 参数的 fp32 副本：4 字节/参数
• 动量：4 字节/参数
• 方差：4 字节/参数

假设一个具有 70 亿个参数的模型，并使用 低秩适配器注入 2 亿个参数。

对于普通模型的优化器状态，内存需求为 12 * 7 = 84 GB（假设有 70 亿个可训练参数）。

添加 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 镜像。

专家混合

一些最近的论文报告说，通过将专家混合 (MoE) 集成到 Transformer 模型中，可以将训练速度提高 4-5 倍，并加快推理速度。

由于已经发现更多的参数会导致更好的性能，因此此技术允许将参数数量增加一个数量级，而不会增加训练成本。

在这种方法中，每隔一个 FFN 层都会被一个 MoE 层替换，该层由许多专家组成，并带有一个门控函数，该函数根据输入标记在序列中的位置以平衡的方式训练每个专家。

（来源：GLAM）

您可以在本节末尾列出的论文中找到详尽的细节和比较表。

这种方法的主要缺点是它需要惊人的 GPU 内存量——几乎比其密集等效物大一个数量级。提出了各种蒸馏和方法来克服更高的内存需求。

不过，存在直接的权衡，您可以使用少量专家和 2-3 倍更小的基础模型，而不是几十或数百个专家，从而使模型缩小 5 倍，从而适度提高训练速度，同时适度提高内存需求。

大多数相关的论文和实现都是围绕 Tensorflow/TPU 构建的。

• GShard：使用条件计算和自动分片扩展大型模型
• Switch Transformers：使用简单高效的稀疏性扩展到万亿参数模型
• GLaM：通用语言模型 (GLaM)

对于 Pytorch，DeepSpeed 也构建了一个：DeepSpeed-MoE：推进专家混合推理和训练，为下一代 AI 规模提供动力，专家混合 - 博客文章：1，2 以及使用大型基于 Transformer 的自然语言生成模型的特定部署：博客文章，Megatron-Deepspeed 分支。

使用 PyTorch 原生注意力和 Flash Attention

PyTorch 的torch.nn.functional.scaled_dot_product_attention (SDPA) 也可以在后台调用 FlashAttention 和内存高效的注意力内核。SDPA 支持目前正在 Transformers 中原生添加，并且在 torch>=2.1.1 中默认使用，前提是存在实现。请参阅PyTorch 标量点积注意力以获取受支持模型的列表和更多详细信息。

查看此博文，以了解有关使用 SDPA 加速和节省内存的更多信息。