Transformers 文档
GPU
并获得增强的文档体验
开始
显卡
GPU 通常用于训练深度学习模型，因为它们具有高内存带宽和并行处理能力。根据您的 GPU 和模型大小，甚至可以训练具有数十亿参数的模型。关键是在 GPU 内存利用率（数据吞吐量/训练时间）和训练速度之间找到正确的平衡点。
本指南将向您展示 Transformers 和 PyTorch 中可用于在 GPU 上高效训练模型的功能。在许多情况下，您需要结合使用这些功能来优化训练。
请参考下表，以快速帮助您确定与您的训练场景相关的功能。
| 功能 | 训练速度 | 内存使用量 |
|---|---|---|
| 批大小 | 是 | 是 |
| 梯度累积 | 否 | 是 |
| 梯度检查点 | 否 | 是 |
| 混合精度 | 是 | 取决于 |
| 优化器 | 是 | 是 |
| 数据预加载 | 是 | 否 |
| torch_empty_cache_steps | 否 | 是 |
| torch.compile | 是 | 否 |
| PEFT | 否 | 是 |
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 Performance 指南，以了解有关输入特征、输出神经元计数和批大小如何影响性能的更多信息。这些都涉及到 GPU 执行的通用矩阵乘法 (GEMM)。较大的参数更利于并行化和效率。
Tensor Core Requirements 部分对于根据数据类型和 GPU 选择最大化张量乘法速度的批大小也很有用。例如，对于 fp16，建议使用 8 的倍数，除非是 A100 GPU，在这种情况下，请使用 64 的倍数。
最后，对于较小的参数，请考虑 Dimension Quantization Effects。当矩阵维度不可被 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 | 使用 PyTorch eager 模式用于 AOTAutograd 的提取的正向和反向图 | 调试 |
| 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 模型中的原始注意力机制更高效和优化。它支持三种类型的缩放点积注意力。
- FlashAttention2 会自动为具有 fp16 或 bf16 torch 类型的模型启用。请确保首先将您的模型转换为适当的类型。
- xFormers 或 Memory-Efficient Attention 支持具有 fp32 torch 类型的模型。
- 缩放点积注意力的 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")