Transformers 文档
GPU
并获得增强的文档体验
开始使用
GPU
GPU 因其高内存带宽和并行处理能力，常用于深度学习模型训练。根据你的 GPU 和模型大小，甚至可以训练数十亿参数的模型。关键在于找到 GPU 内存利用率（数据吞吐量/训练时间）和训练速度之间的平衡。
本指南将向你展示 Transformers 和 PyTorch 中可用于高效训练 GPU 模型的特性。在许多情况下，你需要结合使用这些特性来优化训练。
请参考下表，以便快速找到与你的训练场景相关的特性。
| 特性 | 训练速度 | 内存使用 |
|---|---|---|
| 批量大小 | 是的 | 是的 |
| 梯度累积 | 否 | 是的 |
| 梯度检查点 | 否 | 是的 |
| 混合精度 | 是的 | 取决于 |
| 优化器 | 是的 | 是的 |
| 数据预加载 | 是的 | 否 |
| torch_empty_cache_steps | 否 | 是的 |
| torch.compile | 是的 | 否 |
| 缩放点积注意力 (SDPA) | 是的 | 是的 |
Trainer
Trainer 支持许多有用的训练特性，可以通过 TrainingArguments 进行配置。本节重点介绍一些对优化训练更重要的特性。
批量大小
批量大小是高效 GPU 训练最重要的超参数之一，因为它会影响内存使用和训练速度。较大的批量大小可以利用 GPU 的并行处理能力，从而加快训练速度。建议使用 2 的幂次方作为批量大小，例如 8、64、128、256、512 等。批量大小取决于你的 GPU 和模型的数据类型。
在 TrainingArguments 中配置 per_device_train_batch_size()。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=256,
per_device_eval_batch_size=256,
)请参阅 NVIDIA 的 性能 指南，了解输入特征、输出神经元数量和批量大小如何影响性能。这些都与 GPU 执行的通用矩阵乘法 (GEMM) 有关。参数越大，并行化和效率越好。
Tensor Core 要求 一节对于根据数据类型和 GPU 选择最大化张量乘法速度的批量大小也很有用。例如，对于 fp16，建议使用 8 的倍数，除非是 A100 GPU，这种情况下请使用 64 的倍数。
最后，考虑 维度量化效应 对于较小参数的影响。当矩阵维度不能被 GPU 的线程块瓦片大小整除时，会导致瓦片量化，从而使 GPU 资源利用不足。选择正确的批量大小乘数，使矩阵可以被瓦片大小整除，可以显著加快训练速度。
梯度累积
梯度累积通过在多次小批量数据上累积梯度来克服内存限制（对于在单个 GPU 上无法容纳的超大型模型非常有用），然后在更新参数之前进行一次更新。这通过存储更少的梯度来减少内存使用，并能够使用更大的*有效批量大小*进行训练，因为通常参数是从单个批次数据中更新的。然而，由于梯度累积引入了额外的正向和反向传播，训练速度可能会变慢。
在 TrainingArguments 中配置 per_device_train_batch_size() 以启用梯度累积。
from transformers import TrainingArguments
# effective batch size of 64
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
)尽量避免过多的梯度累积步骤，因为它会大大降低训练速度。考虑以下示例，你的 GPU 上能容纳的最大批量大小是 4。你应该将批量大小保持在 4 以更好地利用 GPU。
| 批量大小 | 梯度累积步数 | 有效批量大小 | |
|---|---|---|---|
| 1 | 64 | 64 | 👎 |
| 4 | 16 | 64 | 👍 |
梯度检查点
梯度检查点通过在反向传播过程中仅存储部分中间激活并重新计算剩余激活来减少内存使用。这避免了存储正向传播中的*所有*中间激活，后者可能需要大量内存开销。然而，其代价是训练速度变慢（约 20%）。
在 TrainingArguments 中配置 gradient_checkpointing() 以启用梯度检查点。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
)混合精度
混合精度通过在半精度（fp16）和全精度（fp32）下执行部分计算来加速训练速度。半精度计算能提升训练速度，因为它不像全精度计算那样需要大量计算。同时，保留部分全精度计算能保持准确性。
有几种数据类型可用于混合精度训练。
混合精度训练的主要优点是将激活存储为 fp16。
在 TrainingArguments 中配置 fp16() 以启用 fp16 数据类型的混合精度训练。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
fp16=True.
)fp16 不是内存优化的，因为以 fp16 计算的梯度在优化步骤中会转换回 fp32。你最终可能会使用更多的 GPU 内存，特别是对于小批量大小，因为现在 GPU 上有模型的两个版本（fp16 和 fp32）。
优化器
Transformers 默认实现 PyTorch 的 AdamW (adamw_torch) 优化器。但由于它存储了过去梯度的加权平均值，因此需要额外的内存，其大小与模型参数数量成比例，以存储过去梯度。这在训练非常大的模型时可能会成为问题，在这种情况下，你应该考虑选择不同的优化器。例如，如果你在 NVIDIA 或 AMD 上安装了 Apex，那么使用 adamw_apex_fused 优化器可以为所有 AdamW 优化器提供最快的训练。
在 TrainingArguments 中配置 optim() 以选择优化器。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit"
)根据你的训练场景，有许多优化器可供选择（请参阅 OptimizerNames 获取完整支持列表）。例如，Adafactor 可以通过存储行或列的加权平均值而不是矩阵中的每个元素来显著减少内存需求，代价是收敛速度较慢。另一个例子是使用来自 bitsandbytes 的 8 位 AdamW 优化器 来量化优化器状态。优化器状态以较低精度存储，并在优化器步骤中使用之前进行反量化。
请参阅 优化器 指南以了解更多专用优化器。
数据预加载
数据预加载可在 CPU 上提前加载并准备批处理数据，以确保 GPU 持续工作，从而减少 GPU 空闲时间和提高利用率。有两种方法可以预加载数据，以确保 GPU 始终处于工作状态。
- 在 CPU 上分配固定内存以存储数据，并将其直接传输到 GPU。
- 增加 CPU 线程或工作线程的数量以更快地预加载数据。
在 TrainingArguments 中配置 dataloader_pin_memory() 和 dataloader_num_workers() 以分配固定内存和增加工作线程数量。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
)PyTorch
PyTorch 提供了多种功能来减少内存需求并提高训练速度。通过仅添加少量代码，通常可以在 Transformers 中启用这些功能。
torch.empty_cache_steps
torch.cuda.empty_cache 函数释放未使用的缓存内存，这有助于避免内存不足 (OOM) 错误，但代价是训练速度会慢约 10%。
在 TrainingArguments 中使用 torch_empty_cache_steps()，在一定数量的训练步骤后启用它。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
torch_empty_cache_steps=4,
)torch.compile
torch.compile 将 PyTorch 代码编译成优化的内核，显著加快了训练速度。此功能依赖于 TorchDynamo 通过 Frame Evaluation API 捕获 PyTorch 图。该图可以进一步编译为针对不同后端的优化内核。
在 TrainingArguments 中配置 torch_compile() 以启用它，并配置 torch_compile_backend() 以选择要使用的后端。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
torch_empty_cache_steps=4,
torch_compile=True,
torch_compile_backend="inductor"
)请参考下表，以帮助您为训练场景选择合适的后端。
| 后端 | 描述 | 目标 |
|---|---|---|
| eager | 使用 PyTorch 运行提取的 GraphModule | 调试 |
| aot_eager | 对 AOTAutograd 提取的正向和反向图使用 PyTorch eager 模式 | 调试 |
| inductor | 通过利用 Triton 内核，将 TorchInductor 与 AOTAutograd 和 CUDA Graphs 结合使用 | 训练和推理 |
| nvfuser | 将 nvFuser 与 TorchScript 结合使用 | 训练和推理 |
| aot_nvfuser | 将 nvFuser 与 AOTAutograd 结合使用 | 训练和推理 |
| aot_cudagraphs | 将 CUDA Graphs 与 AOTAutograd 结合使用 | 训练和推理 |
| ofi | 使用 TorchScripts optimize_for_inference | 推理 |
| fx2trt | 使用 Torch-TensorRT | 推理 |
| onnxrt | 将 ONNX-RT 用于 CPU 和 GPU 推理 | 推理 |
| ipex | 将 IPEX 用于 CPU 推理 | 推理 |
缩放点积注意力
torch.nn.functional.scaled_dot_product_attention (SDPA) 是缩放点积注意力机制的 PyTorch 原生实现。SDPA 比 Transformer 模型中原始的注意力机制更高效且更优化。它支持三种类型的缩放点积注意力：
- 对于 torch 类型为 fp16 或 bf16 的模型，将自动启用 FlashAttention2。请务必先将模型转换为相应的类型。
- xFormers 或内存高效注意力支持 torch 类型为 fp32 的模型。
- 缩放点积注意力的 C++ 实现。
SDPA 默认在 PyTorch 2.1.1+ 中启用，但可以通过在 from_pretrained() 中设置 attn_implementation="sdpa" 来显式启用。
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", attn_implementation="sdpa")