Transformers 文档

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

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

本指南展示了您可以使用的实用技术,通过优化内存利用率、加快训练速度或两者兼而有之来提高模型训练的效率。如果您想了解训练期间GPU的利用情况,请先参考模型训练解剖概念指南。本指南侧重于实用技术。

如果您可以访问具有多个GPU的机器,这些方法仍然有效,此外您还可以利用多GPU部分中概述的其他方法。

在训练大型模型时,有两个方面需要同时考虑:

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

最大化吞吐量(样本/秒)可以降低训练成本。这通常通过尽可能利用GPU来实现,从而将GPU内存填充到极限。如果所需的批量大小超过GPU内存的限制,可以使用内存优化技术,如梯度累积,来帮助解决。

然而,如果首选的批量大小适合内存,就没有必要应用内存优化技术,因为它们可能会减慢训练速度。仅仅因为可以使用大批量大小,并不意味着应该这样做。作为超参数调优的一部分,您应该确定哪个批量大小能产生最佳结果,然后相应地优化资源。

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

方法/工具 提高训练速度 优化内存利用率
批量大小选择
Gradient accumulation
Gradient checkpointing
混合精度训练 可能*
torch_empty_cache_steps
优化器选择
数据预加载
DeepSpeed Zero
torch.compile
Parameter-Efficient Fine Tuning (PEFT)

*注意:当使用混合精度时,对于小模型和大批量大小,会有一些内存节省,但对于大模型和小批量大小,内存使用将会更大。

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

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

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

批量大小选择

为了实现最佳性能,首先确定适当的批量大小。建议使用大小为2^N的批量大小和输入/输出神经元数量。通常它是8的倍数,但根据所使用的硬件和模型的dtype,它可能会更高。

作为参考,请查看NVIDIA关于输入/输出神经元数量批量大小的建议,这些建议适用于全连接层(涉及GEMMs(通用矩阵乘法))。

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

对于较小的参数,还可以考虑维度量化效应。这是分块发生的地方,正确的乘数可以显著加速。

梯度累积

梯度累积方法旨在通过较小的增量计算梯度,而不是一次性计算整个批次的梯度。这种方法涉及通过模型的前向和后向传递迭代计算较小批次的梯度,并在过程中累积梯度。一旦累积了足够数量的梯度,就执行模型的优化步骤。通过使用梯度累积,可以增加有效批次大小,超越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-3090A100的批量大小和梯度累积基准测试。

梯度检查点

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

为了在反向传播过程中计算梯度而保存前向传播中的所有激活值可能会导致显著的内存开销。另一种方法是在反向传播过程中丢弃激活值并在需要时重新计算它们,这将引入相当大的计算开销并减慢训练过程。

梯度检查点 提供了这两种方法之间的折衷方案,并在整个计算图中策略性地保存选定的激活值,因此只需要重新计算一小部分激活值来获取梯度。关于梯度检查点的深入解释,请参考 这篇精彩的文章

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

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

或者,使用🤗 Accelerate - 在本指南的后面找到🤗 Accelerate示例。

虽然梯度检查点可以提高内存效率,但它会使训练速度减慢约20%。

混合精度训练

混合精度训练是一种旨在通过使用较低精度的数值格式来优化训练模型计算效率的技术。传统上,大多数模型使用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位),但精度只有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-3090A100

Flash Attention 2

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

优化器选择

用于训练Transformer模型的最常见优化器是Adam或AdamW(带权重衰减的Adam)。Adam通过存储先前梯度的滚动平均值来实现良好的收敛性;然而,它增加了与模型参数数量相当的内存占用。为了解决这个问题,你可以使用替代优化器。例如,如果你为NVIDIA GPU安装了NVIDIA/apex,或为AMD GPU安装了ROCmSoftwarePlatform/apexadamw_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 中也可用,但下面提供了一个第三方集成以供演示。

作为比较,对于一个3B参数的模型,如“google-t5/t5-3b”:

  • 标准的AdamW优化器将需要24GB的GPU内存,因为它为每个参数使用8字节(8*3 => 24GB)
  • Adafactor优化器将需要超过12GB的内存。它为每个参数使用略多于4字节的内存,因此4*3再加上一些额外的内存。
  • 8bit 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-bit 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 repo中的安装指南安装bitsandbytes库,该库实现了8位Adam优化器。

接下来你需要初始化优化器。这包括两个步骤:

  • 首先,将模型的参数分为两组 - 一组应该应用权重衰减,另一组不应该应用。通常,偏置和层归一化参数不进行权重衰减。
  • 然后进行一些参数整理,以使用与之前使用的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))

结合其他方法(梯度累积、梯度检查点和混合精度训练),你可以预期在使用Adafactor时获得大约3倍的内存改进,甚至更高的吞吐量。

multi_tensor

pytorch-nightly 引入了 torch.optim._multi_tensor,这应该会显著加快在有许多小特征张量的情况下的优化器速度。它最终应该会成为默认设置,但如果你想更早地尝试它,可以查看这个 GitHub issue

数据预加载

达到高训练速度的重要要求之一是能够以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:

使用 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的eager模式来处理AotAutograd提取的前向和后向图。这对于调试很有用,但不太可能带来速度提升。

训练与推理后端:

  • dynamo.optimize("inductor") - 使用 TorchInductor 后端,结合 AotAutograd 和 cudagraphs,通过利用代码生成的 Triton 内核 阅读更多
  • 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 上进行推理。Read more
  • dynamo.optimize("ipex") - 使用IPEX在CPU上进行推理。Read more

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

使用 🤗 PEFT

Parameter-Efficient Fine Tuning (PEFT) 方法在微调期间冻结预训练模型参数,并在其之上添加少量可训练参数(适配器)。

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

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

  • 参数的fp32副本:每个参数4字节
  • 动量:每个参数4字节
  • 方差:4字节/参数

假设一个模型有7B参数,并注入了2亿个参数,使用了低秩适配器

普通模型的优化器状态的内存需求将是 12 * 7 = 84 GB(假设有 7B 可训练参数)。

添加Lora略微增加了与模型权重相关的内存,并显著减少了优化器状态的内存需求至12 * 0.2 = 2.4GB。

了解更多关于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-bit优化器

最后,我们可以添加主训练循环。请注意,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层,这些专家通过一个门控函数进行训练,该函数根据输入标记在序列中的位置以平衡的方式训练每个专家。

MoE Transformer 2x 模块

(来源: GLAM)

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

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

不过,这里有一个直接的权衡,你可以使用较少的专家,搭配一个2-3倍较小的基础模型,而不是使用数十或数百个专家,这样可以减少5倍的模型大小,从而适度提高训练速度,同时也会适度增加内存需求。

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

对于Pytorch,DeepSpeed也构建了一个:DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale, Mixture of Experts - 博客文章: 1, 2 以及基于大型Transformer的自然语言生成模型的特定部署:博客文章, Megatron-Deepspeed分支.

使用 PyTorch 原生注意力和 Flash Attention

PyTorch 的 torch.nn.functional.scaled_dot_product_attention (SDPA) 也可以在底层调用 FlashAttention 和内存高效的注意力内核。目前正在 Transformers 中本地添加 SDPA 支持,并且在有实现的情况下,默认用于 torch>=2.1.1。请参考 PyTorch 缩放点积注意力 以获取支持的模型列表和更多详细信息。

查看这篇博客文章,了解更多关于使用SDPA进行加速和节省内存的信息。

< > Update on GitHub