从满足需求的基准开始
选择一个简短的数据集,包含普通输入、边界情况和决策边界。固定模型、权重、预处理以及 train 或 eval 模式。两个模型或两个 batch 之间的比较,无法将其差异归因于精度。
记录对应用有用的输出,而不仅仅是损失。对于分类器,这可能包括分数和决策;对于回归,则包括误差和极值。先检查基准中是否已出现 NaN 或 inf。一次错误的 FP32 运行不会因为它有更多位数就成为可靠的基础。
在试运行前就确定容差、最低质量以及无非有限值。PyTorch 提醒,浮点计算并不保证在不同设备或执行路径之间得到完全相同的结果。
区分 autocast、数值格式和 GradScaler
autocast 会根据各操作的计算策略选择其类型,而不会把整个程序转换为单一格式。使用它时,应避免手动用 half() 转换整个模型。当前文档推荐使用 torch.autocast 或 torch.amp.autocast;旧的 torch.cuda.amp 接口已弃用。
GradScaler 作用于训练过程中损失和梯度的缩放。它不能用作推理加速器,因为推理不执行 backward。FP16 的数值范围比 BF16 更窄;为 BF16 设计的模型在 FP16 下可能溢出。因此,缩放系数反复下降并不能说明问题已经解决。
应根据模型的约束条件和实际使用的操作来选择格式,然后验证目标平台是否支持。显卡的商业名称或 PyTorch 的某种准备偏好,并不能证明你的自定义算子拥有所需的核函数。
左右滚动表格即可查看所有列。| 选择 | 角色 | 需要检查的内容 |
|---|---|---|
| FP32 基准 | 项目的对比基准 | 有限输出与预期质量 |
| FP16 autocast | 部分操作以降低精度执行 | 数值范围与梯度 |
| BF16 autocast | 范围/精度的另一种权衡 | 可用算子与质量 |
| GradScaler | 梯度缩放的处理 | 实际执行的更新 |
将训练步骤安排在正确的顺序
所给出的代码片段假定模型和优化器已构建完成,输入和目标位于同一块 GPU 上,且损失为标量。该片段未经过实际运行,也不构成对任何方案的验证。autocast 上下文包围前向传播和损失计算;反向传播在其关闭之后进行。scaler 在训练会话中只创建一次,而不是每个 batch 都创建。
若要检查或裁剪梯度,应先用 unscale_ 去除其缩放系数。官方 AMP 示例指出,每个优化器只应执行一次,且在累积完用于其更新的梯度之后进行。下面的裁剪阈值 1.0 只是一个示例值,应根据你的项目选择,并非通用建议。
此处的保护性检查会在损失、梯度或总范数非有限时中断诊断。这些读取会传输到 CPU,具有侵入性:不要用此代码片段做计时。梯度累积、多个优化器以及调度器都需要各自定义更新方式。
import torch
# 前提:model、optimizer、loss_fn、inputs 和 targets 均已存在。
# 模型和输入位于同一个 CUDA/HIP 设备上。
dtype = torch.float16 # 需要验证的选择;BF16 是另一种可选方案。
scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))
# 放入你的循环中,并在各 batch 之间保留 scaler。
optimizer.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=dtype):
prediction = model(inputs)
loss = loss_fn(prediction, targets)
if not bool(torch.isfinite(loss).item()):
raise FloatingPointError("损失非有限:中断诊断")
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
if any(p.grad is not None and
not bool(torch.isfinite(p.grad).all().item())
for p in model.parameters()):
raise FloatingPointError("梯度非有限:中断诊断")
torch.nn.utils.clip_grad_norm_(
model.parameters(), max_norm=1.0, error_if_nonfinite=True,
)
scaler.step(optimizer)
scaler.update()实例讲解:两个相近的决策并不等价
假设有一个根据最高分选择类别的服务。在一份教学用输入上,基准产生两个非常接近的分数:1.0000 和 1.0003。另一条数值路径可能会改变它们的顺序或造成相等。这些数字只是用来说明决策边界,并非 FP16 或 BF16 实测输出。
正确的验证包含两个层面。先用明确的容差比较分数,再比较决策本身以及处理相等情况的规则。绝对值上的微小差异可能改变所选择的操作。反之,一个肉眼可见的数值差异,对于阈值远离所观测分数的任务来说,可能并无影响。
记录标识符、基准输出、AMP 试验及对决策的影响。在读取结果之前先确定接受规则。不要为了让某个棘手案例消失而放宽容差;不同规模的输出可能需要不同的判据。
左右滚动表格即可查看所有列。| 判定标准 | 参考编号 | AMP 试验 | 决定 |
|---|---|---|---|
| 已完成的输出 | 待检查 | 待检查 | 拒绝无法解释的未完成项 |
| 数值差异 | 保留的数值 | 待计算的差异 | 试验前已定义的容差 |
| 应用决策 | 类别或操作 | 类别或操作 | 检查变更 |
| 固定数据集上的质量 | 待测量 | 待测量 | 遵守项目阈值 |
解读 NaN 与被跳过的更新
当出现未完成项时,找出产生它们的最早步骤:输入、中间输出、损失或梯度。用基准重放同一案例,然后在可疑运算周围局部禁用 autocast,同时也要检查其输入的类型。将整个训练改回 FP32 可作为对比,但不会自动定位问题。
当梯度包含 inf 或 NaN 时,scaler 可能会跳过某次更新。因此一个仍在继续的循环未必执行的更新次数与迭代次数相同。在诊断期间记录这一行为。不要盲目推进一个被认为会跟随实际更新的学习策略。
损失有限并不保证梯度有限。反之,一次偶发事件也不足以判定训练不可用:要检查其频率、进展和质量。AMP 配方提供了一种方法,当 autocast 和 scaling 二者之一可疑时,分别隔离它们。
准备可复现的回退方案
在试验之前,保留基准配置、权重、优化器状态以及一致的 checkpoint。如果你的训练使用 scaler,其状态也属于恢复的一部分。记录 dtype 以及任何保持 FP32 的区域。使用不同策略进行恢复是一项需要明确的实验性变更,而不是隐含等价的延续。
如果输出变为非有限值、质量超出既定判据,或更新停止以可用的方式推进,就回到之前的配置。保留促成此次回退的案例。修改后,先在相同数据集上重新进行对比,再延长时长。
Kernodeck 的 checkpoint 练习验证的是不带 AMP 的 CPU 恢复。复用其对比方法,并加入你的循环实际消耗的状态。
只有在数值验证之后才测量收益
验证之后,在不进行详细诊断的情况下测量内存和时间。保持形状、batch、模型和质量不变。标量读取、同步和分析器可能会改变耗时;在最终测量中移除侵入式检查。
在 ROCm 上,PyTorch 的设备名仍为 cuda,相应接口被复用。这既不能保证相同的内核,也不能保证与 NVIDIA 相同的结果。请在目标上验证项目的后端和算子。本指南不承诺任何固定的内存缩减、速度倍增或 Kernodeck 准备兼容性。