自动混合精度包 - torch.amp#
创建日期：2025 年 6 月 12 日 | 最后更新日期：2025 年 6 月 12 日
torch.amp 提供了混合精度的便捷方法。在该模式下，部分操作使用 torch.float32 (float) 数据类型，而其他操作则使用低精度浮点数据类型 (lower_precision_fp)：即 torch.float16 (half) 或 torch.bfloat16。某些算子（如线性层和卷积）在 lower_precision_fp 下运行速度明显更快；而另一些算子（如归约操作）通常需要 float32 的动态范围。混合精度旨在为每个算子匹配最合适的数据类型。
通常情况下，数据类型为 torch.float16 的“自动混合精度训练”会同时使用 torch.autocast 和 torch.amp.GradScaler，如 自动混合精度示例 和 自动混合精度指南 所示。然而，torch.autocast 和 torch.GradScaler 是模块化的，如果需要，也可以分别使用。如 torch.autocast 的 CPU 示例部分所示，在 CPU 上使用 torch.bfloat16 数据类型进行“自动混合精度训练/推理”时，仅需使用 torch.autocast。
警告
torch.cuda.amp.autocast(args...) 和 torch.cpu.amp.autocast(args...) 已弃用。请改用 torch.amp.autocast("cuda", args...) 或 torch.amp.autocast("cpu", args...)。同样，torch.cuda.amp.GradScaler(args...) 和 torch.cpu.amp.GradScaler(args...) 也已弃用，请改用 torch.amp.GradScaler("cuda", args...) 或 torch.amp.GradScaler("cpu", args...)。
torch.autocast 和 torch.cpu.amp.autocast 是在 1.10 版本中新增的。
自动混合精度（Autocasting）#
- torch.amp.autocast_mode.is_autocast_available(device_type)[source]#
返回一个布尔值，指示
device_type是否支持 autocast。- 参数:
device_type (str) – 要使用的设备类型。可能的值包括：‘cuda’、‘cpu’、‘mtia’、‘maia’、‘xpu’ 等。该类型与
torch.device的 type 属性相同。因此，您可以通过 Tensor.device.type 获取张量的设备类型。- 返回类型:
- class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)[source]#
autocast的实例可用作上下文管理器或装饰器，允许脚本的特定区域在混合精度下运行。在这些区域中，算子将以 autocast 选择的特定数据类型运行，以在保持精度的同时提高性能。详细信息请参阅 Autocast 算子参考。
进入 autocast 启用的区域时，张量可以是任何类型。使用 autocast 时，不应在模型或输入上显式调用
half()或bfloat16()。autocast应仅包裹网络的前向传递（包括损失计算）。不建议在 autocast 区域内进行反向传播。反向传播算子将以与相应前向算子相同的类型运行。CUDA 设备示例
# Creates model and optimizer in default precision model = Net().cuda() optimizer = optim.SGD(model.parameters(), ...) for input, target in data: optimizer.zero_grad() # Enables autocasting for the forward pass (model + loss) with torch.autocast(device_type="cuda"): output = model(input) loss = loss_fn(output, target) # Exits the context manager before backward() loss.backward() optimizer.step()
有关在更复杂场景（例如梯度惩罚、多个模型/损失、自定义 autograd 函数）中的用法（连同梯度缩放），请参阅 自动混合精度示例。
autocast也可以用作装饰器，例如用于模型的forward方法。class AutocastModel(nn.Module): ... @torch.autocast(device_type="cuda") def forward(self, input): ...
在 autocast 区域内产生的浮点张量可能是
float16类型。返回到 autocast 禁用区域后，如果将其与不同数据类型的浮点张量一起使用，可能会导致类型不匹配错误。如果发生这种情况，请将 autocast 区域产生的张量转换回float32（或所需的其他数据类型）。如果来自 autocast 区域的张量已经是float32，则转换操作不会执行任何操作，也不会产生额外开销。CUDA 示例：# Creates some tensors in default dtype (here assumed to be float32) a_float32 = torch.rand((8, 8), device="cuda") b_float32 = torch.rand((8, 8), device="cuda") c_float32 = torch.rand((8, 8), device="cuda") d_float32 = torch.rand((8, 8), device="cuda") with torch.autocast(device_type="cuda"): # torch.mm is on autocast's list of ops that should run in float16. # Inputs are float32, but the op runs in float16 and produces float16 output. # No manual casts are required. e_float16 = torch.mm(a_float32, b_float32) # Also handles mixed input types f_float16 = torch.mm(d_float32, e_float16) # After exiting autocast, calls f_float16.float() to use with d_float32 g_float32 = torch.mm(d_float32, f_float16.float())
CPU 训练示例
# Creates model and optimizer in default precision model = Net() optimizer = optim.SGD(model.parameters(), ...) for epoch in epochs: for input, target in data: optimizer.zero_grad() # Runs the forward pass with autocasting. with torch.autocast(device_type="cpu", dtype=torch.bfloat16): output = model(input) loss = loss_fn(output, target) loss.backward() optimizer.step()
CPU 推理示例
# Creates model in default precision model = Net().eval() with torch.autocast(device_type="cpu", dtype=torch.bfloat16): for input in data: # Runs the forward pass with autocasting. output = model(input)
带 Jit Trace 的 CPU 推理示例
class TestModel(nn.Module): def __init__(self, input_size, num_classes): super().__init__() self.fc1 = nn.Linear(input_size, num_classes) def forward(self, x): return self.fc1(x) input_size = 2 num_classes = 2 model = TestModel(input_size, num_classes).eval() # For now, we suggest to disable the Jit Autocast Pass, # As the issue: https://github.com/pytorch/pytorch/issues/75956 torch._C._jit_set_autocast_mode(False) with torch.cpu.amp.autocast(cache_enabled=False): model = torch.jit.trace(model, torch.randn(1, input_size)) model = torch.jit.freeze(model) # Models Run for _ in range(3): model(torch.randn(1, input_size))
在 autocast 区域内出现的类型不匹配错误属于 bug；如果您观察到这种情况，请提交 issue。
可以在启用了 autocast 的区域内嵌套
autocast(enabled=False)子区域。局部禁用 autocast 在您希望强制某个子区域以特定dtype运行时很有用。禁用 autocast 可以让您显式控制执行类型。在子区域中，来自周围区域的输入应在使用前转换为目标dtype。# Creates some tensors in default dtype (here assumed to be float32) a_float32 = torch.rand((8, 8), device="cuda") b_float32 = torch.rand((8, 8), device="cuda") c_float32 = torch.rand((8, 8), device="cuda") d_float32 = torch.rand((8, 8), device="cuda") with torch.autocast(device_type="cuda"): e_float16 = torch.mm(a_float32, b_float32) with torch.autocast(device_type="cuda", enabled=False): # Calls e_float16.float() to ensure float32 execution # (necessary because e_float16 was created in an autocasted region) f_float32 = torch.mm(c_float32, e_float16.float()) # No manual casts are required when re-entering the autocast-enabled region. # torch.mm again runs in float16 and produces float16 output, regardless of input types. g_float16 = torch.mm(d_float32, f_float32)
autocast 状态是线程局部的。如果您希望在新的线程中启用它，则必须在该线程中调用上下文管理器或装饰器。当每个进程使用超过一个 GPU 时，这会影响
torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel（请参阅 多 GPU 工作指南）。- 参数:
device_type (str, required) – 要使用的设备类型。可能的值包括：‘cuda’、‘cpu’、‘mtia’、‘maia’、‘xpu’ 和 ‘hpu’。该类型与
torch.device的 type 属性相同。enabled (bool, optional) – 是否在区域中启用自动混合精度。默认值：
True。dtype (torch_dtype, optional) – autocast 运行算子时的数据类型。如果
dtype为None，则使用get_autocast_dtype()给出的默认值（CUDA 为torch.float16，CPU 为torch.bfloat16）。默认值：None。cache_enabled (bool, optional) – 是否启用 autocast 内部的权重缓存。默认值：
True。
- torch.amp.custom_fwd(fwd=None, *, device_type, cast_inputs=None)[source]#
为自定义 autograd 函数的
forward方法创建辅助装饰器。Autograd 函数是
torch.autograd.Function的子类。有关更多详细信息，请参阅 示例页面。- 参数:
device_type (str) – 要使用的设备类型。
cast_inputs (
torch.dtypeor None, optional, default=None) – 如果不为None，当forward在启用了 autocast 的区域中运行时，将传入的浮点张量转换为目标 dtype（非浮点张量不受影响），然后在禁用 autocast 的情况下执行forward。如果为None，则forward的内部算子将按照当前的 autocast 状态执行。
注意
如果修饰后的
forward在 autocast 区域之外被调用，custom_fwd不会执行任何操作，cast_inputs也无效。
- torch.amp.custom_bwd(bwd=None, *, device_type)[source]#
为自定义 autograd 函数的反向传播方法创建辅助装饰器。
确保
backward使用与forward相同的 autocast 状态执行。- 参数:
device_type (str) – 要使用的设备类型。
- class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)[source]#
请参阅
torch.autocast。torch.cuda.amp.autocast(args...)已弃用。请改用torch.amp.autocast("cuda", args...)。
- torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)[source]#
torch.cuda.amp.custom_fwd(args...)已弃用。请改用torch.amp.custom_fwd(args..., device_type='cuda')。
- torch.cuda.amp.custom_bwd(bwd)[source]#
torch.cuda.amp.custom_bwd(args...)已弃用。请改用torch.amp.custom_bwd(args..., device_type='cuda')。
- class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)[source]#
请参阅
torch.autocast。已弃用，请改用torch.amp.autocast("cpu", args...)。
梯度缩放（Gradient Scaling）#
如果某个算子的前向传递具有 float16 输入，其反向传递将产生 float16 梯度。幅度过小的梯度值可能无法在 float16 中表示。这些值会刷新为零（“下溢”），导致相应的参数更新丢失。
为了防止下溢，“梯度缩放”将网络的损失乘以一个缩放因子，并对缩放后的损失进行反向传播。在网络中反向流动的梯度也会按相同因子缩放。这样梯度值的幅度会变大，从而不会刷新为零。
在优化器更新参数之前，应先将每个参数的梯度（.grad 属性）进行反缩放，以避免缩放因子干扰学习率。
注意
AMP/fp16 可能并不适用于所有模型！例如，大多数 bf16 预训练模型无法在 fp16 的最大数值范围（65504）内运行，会导致梯度溢出而非下溢。在这种情况下，缩放因子可能会降至 1 以下，以尝试将梯度带入 fp16 的动态范围。虽然人们通常预期缩放比例始终大于 1，但我们的 GradScaler 为了保持性能并不做此保证。如果您在运行 AMP/fp16 时在损失或梯度中遇到 NaN，请验证您的模型是否兼容。
Autocast 算子参考#
算子适用性#
以 float64 或非浮点数据类型运行的算子不适用，无论是否启用 autocast，它们都将按其原始类型运行。
只有非原地（out-of-place）算子和张量方法适用。原地变体以及显式提供 out=... 张量的调用可以在 autocast 区域内使用，但不会经历 autocasting。例如，在 autocast 区域中，a.addmm(b, c) 可以自动转换，但 a.addmm_(b, c) 和 a.addmm(b, c, out=d) 则不行。为了获得最佳性能和稳定性，请在 autocast 区域中优先使用非原地算子。
通过显式 dtype=... 参数调用的算子不适用，并将产生符合 dtype 参数的输出。
CUDA 算子特定行为#
以下列表描述了适用算子在 autocast 区域中的行为。这些算子无论作为 torch.nn.Module 的一部分、函数还是 torch.Tensor 方法调用，都会经过 autocasting。
未列出的算子不会进行 autocasting。它们按输入定义的类型运行。但是，如果未列出的算子位于经过 autocasting 的算子下游，autocasting 仍可能会改变其运行类型。
如果某个算子未列出，我们假设它在 float16 下数值稳定。如果您认为某个未列出的算子在 float16 下数值不稳定，请提交 issue。
可自动转换为 float16 的 CUDA 算子#
__matmul__, addbmm, addmm, addmv, addr, baddbmm, bmm, chain_matmul, multi_dot, conv1d, conv2d, conv3d, conv_transpose1d, conv_transpose2d, conv_transpose3d, GRUCell, linear, LSTMCell, matmul, mm, mv, prelu, RNNCell
可自动转换为 float32 的 CUDA 算子#
__pow__, __rdiv__, __rpow__, __rtruediv__, acos, asin, binary_cross_entropy_with_logits, cosh, cosine_embedding_loss, cdist, cosine_similarity, cross_entropy, cumprod, cumsum, dist, erfinv, exp, expm1, group_norm, hinge_embedding_loss, kl_div, l1_loss, layer_norm, log, log_softmax, log10, log1p, log2, margin_ranking_loss, mse_loss, multilabel_margin_loss, multi_margin_loss, nll_loss, norm, normalize, pdist, poisson_nll_loss, pow, prod, reciprocal, rsqrt, sinh, smooth_l1_loss, soft_margin_loss, softmax, softmin, softplus, sum, renorm, tan, triplet_margin_loss
提升至最宽输入类型的 CUDA 算子#
这些算子不需要特定的 dtype 来保持稳定性，但如果接收多个输入，则要求输入类型一致。如果所有输入均为 float16，则算子以 float16 运行。如果任一输入为 float32，autocast 会将所有输入转换为 float32 并以该类型运行算子。
addcdiv, addcmul, atan2, bilinear, cross, dot, grid_sample, index_put, scatter_add, tensordot
此处未列出的某些算子（例如 add 等二元算子）无需自动混合精度（autocast）干预即可原生提升输入类型。如果输入是 float16 和 float32 的混合，无论是否启用了 autocast，这些算子都会在 float32 下运行并产生 float32 输出。
优先使用 binary_cross_entropy_with_logits 而非 binary_cross_entropy#
torch.nn.functional.binary_cross_entropy()（以及封装它的 torch.nn.BCELoss）的反向传播过程可能会产生在 float16 中无法表示的梯度。在启用了 autocast 的区域中，前向传播输入可能是 float16，这意味着反向传播梯度必须能在 float16 中表示（将 float16 前向输入自动提升为 float32 并无帮助，因为该转换必须在反向传播中逆转）。因此，binary_cross_entropy 和 BCELoss 在启用 autocast 的区域中会抛出错误。
许多模型在二元交叉熵层之前使用 sigmoid 层。在这种情况下，请使用 torch.nn.functional.binary_cross_entropy_with_logits() 或 torch.nn.BCEWithLogitsLoss 将这两层合并。使用 binary_cross_entropy_with_logits 和 BCEWithLogits 进行 autocast 是安全的。
XPU 算子特定行为（实验性）#
以下列表描述了适用算子在 autocast 区域中的行为。这些算子无论作为 torch.nn.Module 的一部分、函数还是 torch.Tensor 方法调用，都会经过 autocasting。
未列出的算子不会进行 autocasting。它们按输入定义的类型运行。但是，如果未列出的算子位于经过 autocasting 的算子下游，autocasting 仍可能会改变其运行类型。
如果某个算子未列出，我们假设它在 float16 下数值稳定。如果您认为某个未列出的算子在 float16 下数值不稳定，请提交 issue。
可自动转换为 float16 的 XPU 算子#
addbmm, addmm, addmv, addr, baddbmm, bmm, chain_matmul, multi_dot, conv1d, conv2d, conv3d, conv_transpose1d, conv_transpose2d, conv_transpose3d, GRUCell, linear, LSTMCell, matmul, mm, mv, RNNCell
可自动转换为 float32 的 XPU 算子#
__pow__, __rdiv__, __rpow__, __rtruediv__, binary_cross_entropy_with_logits, cosine_embedding_loss, cosine_similarity, cumsum, dist, exp, group_norm, hinge_embedding_loss, kl_div, l1_loss, layer_norm, log, log_softmax, margin_ranking_loss, nll_loss, normalize, poisson_nll_loss, pow, reciprocal, rsqrt, soft_margin_loss, softmax, softmin, sum, triplet_margin_loss
提升至最宽输入类型的 XPU 算子#
这些算子不需要特定的 dtype 来保持稳定性，但如果接收多个输入，则要求输入类型一致。如果所有输入均为 float16，则算子以 float16 运行。如果任一输入为 float32，autocast 会将所有输入转换为 float32 并以该类型运行算子。
bilinear, cross, grid_sample, index_put, scatter_add, tensordot
此处未列出的某些算子（例如 add 等二元算子）无需自动混合精度（autocast）干预即可原生提升输入类型。如果输入是 float16 和 float32 的混合，无论是否启用了 autocast，这些算子都会在 float32 下运行并产生 float32 输出。
CPU 算子特定行为#
以下列表描述了适用算子在 autocast 区域中的行为。这些算子无论作为 torch.nn.Module 的一部分、函数还是 torch.Tensor 方法调用，都会经过 autocasting。
未列出的算子不会进行 autocasting。它们按输入定义的类型运行。但是，如果未列出的算子位于经过 autocasting 的算子下游，autocasting 仍可能会改变其运行类型。
如果算子未在此列出，我们假设它在 bfloat16 下是数值稳定的。如果您认为某个未列出的算子在 bfloat16 下数值不稳定，请提交问题。 float16 与 bfloat16 共享相同的算子列表。
可自动转换为 bfloat16 的 CPU 算子#
conv1d, conv2d, conv3d, bmm, mm, linalg_vecdot, baddbmm, addmm, addbmm, linear, matmul, _convolution, conv_tbc, mkldnn_rnn_layer, conv_transpose1d, conv_transpose2d, conv_transpose3d, prelu, scaled_dot_product_attention, _native_multi_head_attention
可自动转换为 float32 的 CPU 算子#
avg_pool3d, binary_cross_entropy, grid_sampler, grid_sampler_2d, _grid_sampler_2d_cpu_fallback, grid_sampler_3d, polar, prod, quantile, nanquantile, stft, cdist, trace, view_as_complex, cholesky, cholesky_inverse, cholesky_solve, inverse, lu_solve, orgqr, inverse, ormqr, pinverse, max_pool3d, max_unpool2d, max_unpool3d, adaptive_avg_pool3d, reflection_pad1d, reflection_pad2d, replication_pad1d, replication_pad2d, replication_pad3d, mse_loss, cosine_embedding_loss, nll_loss, nll_loss2d, hinge_embedding_loss, poisson_nll_loss, cross_entropy_loss, l1_loss, huber_loss, margin_ranking_loss, soft_margin_loss, triplet_margin_loss, multi_margin_loss, ctc_loss, kl_div, multilabel_margin_loss, binary_cross_entropy_with_logits, fft_fft, fft_ifft, fft_fft2, fft_ifft2, fft_fftn, fft_ifftn, fft_rfft, fft_irfft, fft_rfft2, fft_irfft2, fft_rfftn, fft_irfftn, fft_hfft, fft_ihfft, linalg_cond, linalg_matrix_rank, linalg_solve, linalg_cholesky, linalg_svdvals, linalg_eigvals, linalg_eigvalsh, linalg_inv, linalg_householder_product, linalg_tensorinv, linalg_tensorsolve, fake_quantize_per_tensor_affine, geqrf, _lu_with_info, qr, svd, triangular_solve, fractional_max_pool2d, fractional_max_pool3d, adaptive_max_pool3d, multilabel_margin_loss_forward, linalg_qr, linalg_cholesky_ex, linalg_svd, linalg_eig, linalg_eigh, linalg_lstsq, linalg_inv_ex
提升至最宽输入类型的 CPU 算子#
这些算子不需要特定的数据类型（dtype）来保证稳定性，但它们接收多个输入，并要求这些输入的 dtype 匹配。如果所有输入均为 bfloat16，则算子以 bfloat16 运行。如果任何输入为 float32，autocast 会将所有输入转换为 float32 并在 float32 下运行算子。
cat, stack, index_copy
此处未列出的某些算子（例如 add 等二元算子）无需自动混合精度（autocast）干预即可原生提升输入类型。如果输入是 bfloat16 和 float32 的混合，无论是否启用了 autocast，这些算子都会在 float32 下运行并产生 float32 输出。
