自动混合精度¶

PyTorch/XLA 的 AMP 扩展了 PyTorch 的 AMP 包，支持在 XLA:GPU 和 XLA:TPU 设备上进行自动混合精度。AMP 通过执行某些操作的 float32 和其他操作的较低精度数据类型（float16 或 bfloat16，具体取决于硬件支持）来加速训练和推理。本文档介绍了如何在 XLA 设备上使用 AMP 以及最佳实践。

AMP for XLA:TPU¶

TPU 上的 AMP 会自动将操作转换为以 float32 或 bfloat16 运行，因为 TPU 原生支持 bfloat16。下面是一个简单的 TPU AMP 示例

# Creates model and optimizer in default precision
model = Net().to(xm.xla_device())
# Pytorch/XLA provides sync-free optimizers for improved performance
optimizer = syncfree.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass
    with autocast(xm.xla_device()):
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward()
    loss.backward()
    xm.optimizer_step.(optimizer)

autocast(xm.xla_device()) 是 torch.autocast('xla') 的别名，当 XLA 设备是 TPU 时。或者，如果一个脚本仅用于 TPU，那么可以直接使用 torch.autocast('xla', dtype=torch.bfloat16)。

如果某个操作本应进行自动类型转换但未包含在内，请提交一个 issue 或 pull request。

AMP for XLA:TPU 最佳实践¶

autocast 应该只包装网络的正向传播和损失计算。反向传播操作将以与相应正向传播操作相同的类型运行。
由于 TPU 使用 bfloat16 混合精度，因此不需要进行梯度缩放。
PyTorch/XLA 提供了优化器的修改版本，它们避免了设备和主机之间的额外同步。

支持的操作¶

TPU 上的 AMP 的工作方式类似于 PyTorch 的 AMP。自动类型转换规则总结如下

只有非原地操作（out-of-place ops）和 Tensor 方法才有资格进行自动类型转换。在启用自动类型转换的区域中，允许使用原地变体和显式提供 out=… Tensor 的调用，但它们不会经过自动类型转换。例如，在启用自动类型转换的区域中，a.addmm(b, c) 可以进行自动类型转换，但 a.addmm_(b, c) 和 a.addmm(b, c, out=d) 则不能。为了获得最佳性能和稳定性，建议在启用自动类型转换的区域中使用非原地操作。

以 float64 或非浮点数据类型运行的操作不符合条件，无论是否启用自动类型转换，它们都将以这些类型运行。此外，显式使用 dtype=… 参数调用的操作也不符合条件，并且将生成符合 dtype 参数的输出。

下面未列出的操作不会经过自动类型转换。它们将以其输入的类型运行。如果未列出的操作是从经过自动类型转换的操作的下游，自动类型转换仍可能更改其运行的类型。

自动类型转换为 ``bfloat16`` 的操作

__matmul__, addbmm, addmm, addmv, addr, baddbmm,bmm, conv1d, conv2d, conv3d, conv_transpose1d, conv_transpose2d, conv_transpose3d, linear, matmul, mm, relu, prelu, max_pool2d

自动类型转换为 ``float32`` 的操作

batch_norm, log_softmax, binary_cross_entropy, binary_cross_entropy_with_logits, prod, cdist, trace, chloesky ,inverse, reflection_pad, replication_pad, mse_loss, cosine_embbeding_loss, nll_loss, multilabel_margin_loss, qr, svd, triangular_solve, linalg_svd, linalg_inv_ex

自动类型转换为最宽输入类型的操作

stack, cat, index_copy

AMP for XLA:GPU¶

XLA:GPU 设备上的 AMP 重用了 PyTorch 的 AMP 规则。有关 CUDA 特定行为，请参阅 PyTorch 的 AMP 文档。下面是一个简单的 CUDA AMP 示例

# Creates model and optimizer in default precision
model = Net().to(xm.xla_device())
# Pytorch/XLA provides sync-free optimizers for improved performance
optimizer = syncfree.SGD(model.parameters(), ...)
scaler = GradScaler()

for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass
    with autocast(xm.xla_device()):
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward pass
    scaler.scale(loss).backward()
    gradients = xm._fetch_gradients(optimizer)
    xm.all_reduce('sum', gradients, scale=1.0 / xr.world_size())
    scaler.step(optimizer)
    scaler.update()

当 XLA 设备是 CUDA 设备（XLA:GPU）时，autocast(xm.xla_device()) 是 torch.cuda.amp.autocast() 的别名。或者，如果一个脚本仅用于 CUDA 设备，那么可以直接使用 torch.cuda.amp.autocast，但要求 torch 是用 cuda 支持编译的，以支持 torch.bfloat16 数据类型。我们建议在 XLA:GPU 上使用 autocast(xm.xla_device())，因为它不需要 torch.cuda 支持任何数据类型，包括 torch.bfloat16。

AMP for XLA:GPU 最佳实践¶

autocast 应该只包装网络的正向传播和损失计算。反向传播操作将以与相应正向传播操作相同的类型运行。
在使用 Cuda 设备上的 AMP 时，请勿设置 XLA_USE_F16 标志。这将覆盖 AMP 提供的每个操作的精度设置，并导致所有操作以 float16 执行。
使用梯度缩放来防止 float16 梯度发生下溢。
PyTorch/XLA 提供了优化器的修改版本，它们避免了设备和主机之间的额外同步。

示例¶

我们的 mnist 训练脚本和 imagenet 训练脚本展示了 AMP 在 TPU 和 GPU 上的使用情况。