Transformers documentation
混合精度训练
You are viewing main version, which requires installation from source. If you'd like
regular pip install, checkout the latest stable version (v5.22.2).
混合精度训练
全精度（fp32）训练以 32 位存储和计算所有数据。混合精度使用 fp16 或 bf16 执行计算密集的前向和反向传播，同时保留权重的 fp32 副本用于优化器更新。计算速度更快，权重和激活的内存占用更低，且训练稳定性得以保持。
┌─────────────────────────────────────────────────────┐ │ MIXED PRECISION TRAINING LOOP │ │ │ │ fp32 master weights ──cast──▶ fp16/bf16 │ │ ▲ │ │ │ │ FORWARD (autocast) │ │ │ matmuls in fp16/bf16 │ │ │ reductions stay fp32 │ │ │ │ loss │ │ │ LOSS SCALE ×S ──fp16 │ │ │ │ │ │ │ BACKWARD │ │ │ grads in fp16/bf16 │ │ │ │ │ │ │ UNSCALE ÷S ──fp16 │ │ │ check inf/nan │ │ │ cast grads → fp32 │ │ └────────────────────────── optimizer.step │ └─────────────────────────────────────────────────────┘
将 bf16() 或 fp16() 设置为 True 以启用混合精度训练。两者均为 16 位类型，但 bf16 的指数范围与 fp32 相同，因此几乎不会发生溢出。在 Ampere 或更新的 GPU（如 A100、H100）上使用 bf16，在 V100 或 T4 等较旧的硬件上则回退到 fp16。
以 fp32 加载模型，否则 autocast 将不起作用。若以 bf16 或 fp16 加载，优化器将没有 fp32 主副本可供更新。
from transformers import TrainingArguments
args = TrainingArguments(..., bf16=True)
args = TrainingArguments(..., fp16=True)如果您的模型在 bf16/fp16 下数值稳定，您可以完全跳过混合精度，直接以 bf16/fp16 加载并训练模型。这样可以避免在内存中保留权重的 fp32 副本。
tf32
tf32 是 Ampere GPU 上的一种计算模式，它在矩阵乘法中使用 10 位尾数，而非 23 位。这可以带来速度提升，尤其是在与 bf16/fp16 配合使用时。PyTorch 在 Ampere 及更新版本的 GPU 上默认对矩阵乘法启用 tf32，但在 TrainingArguments 中显式设置它，可以确保无论 PyTorch 版本或环境默认值如何，该模式都处于激活状态。
from transformers import TrainingArguments
args = TrainingArguments(..., bf16=True, tf32=True)后续步骤
- 请参阅 Kernels 指南，了解如何使用自定义融合内核加速训练。
- 请参阅 torch.compile 指南，了解如何编译前向和反向传播以提升吞吐量。