torch.cuda 与混合精度

torch.cuda 负责 GPU 设备的发现、选择与显存治理；在其之上，autocast + GradScaler 构成的 AMP 机制用 fp16/bf16 执行前向与反向，以一半的显存与带宽换取接近无损的训练精度。本篇覆盖设备管理、显存诊断、TF32 与 AMP 全链路。

设备选择与迁移

torch.cuda 是 PyTorch 对 CUDA 运行时的封装，处于 PyTorch 架构总览中 ATen 调度层之下的设备后端：每个 Tensor 与 Storage 都携带 device 属性，Storage 中的数据实际驻留在对应设备的显存里。

import torch
 
torch.cuda.is_available()       # CUDA 运行时是否可用
torch.cuda.device_count()       # 可见 GPU 数量
torch.cuda.get_device_name(0)   # "NVIDIA A100-SXM4-80GB"
 
device = "cuda" if torch.cuda.is_available() else "cpu"

"cuda" 不带索引时绑定当前设备（默认 cuda:0）；torch.cuda.set_device(1) 修改当前设备，影响所有无索引的 "cuda" 引用。
多进程场景惯用环境变量 CUDA_VISIBLE_DEVICES 在进程外部做卡隔离，进程内永远只看到重编号后的设备。
torch.cuda.device(1) 是上下文管理器，块内临时切换默认设备，离开后恢复。

.to(device) 的迁移惯例：

model = model.to(device)   # nn.Module 原地修改（返回 self）
x = x.to(device)           # Tensor 返回新对象，必须重新赋值
y = y.to(device)           # 目标张量同设备，否则 kernel 报错

torch.nn 的 Module.to() 是 in-place，Tensor 的 .to() 是纯函数——漏写 x = x.to(device) 是最常见的新手错误。模型、输入、目标三者必须同设备：跨设备运算没有隐式搬运，直接抛出 RuntimeError。训练开始前迁移一次模型即可，数据在每个 batch 到达时迁移。

显存统计与释放

PyTorch 使用缓存分配器（caching allocator）：释放的 Tensor 显存进入进程内缓存池而非归还 CUDA driver，供后续分配复用。因此存在两层口径：

API	含义
`memory_allocated()`	活跃 Tensor 实际占用
`max_memory_allocated()`	峰值占用（配合 `reset_peak_memory_stats()`）
`memory_reserved()`	向 CUDA 累计申请的缓存池总量
`mem_get_info()`	设备级（空闲，总量），不区分进程

reserved − allocated 即碎片与空闲块。训练 OOM 时先看 max_memory_allocated：若峰值贴近 allocated，是真容量不足；若 reserved 远大于 allocated，是碎片化，可设 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 缓解。torch.cuda.empty_cache() 把缓存池中未使用的块归还 driver——只影响其他进程的可用显存与 nvidia-smi 读数，不会释放任何活跃 Tensor，训练循环内调用纯属浪费。

显存占用构成（训练态）：

参数：权重本体（fp32 下 1 参数 = 4 B）
梯度：与参数等大，由 torch.autograd 在 backward 时写入 param.grad
优化器状态：Adam 变体持有 momentum + variance，再乘 2；fp32 训练的 Adam 总账约 16 B/参数
激活：前向中间结果，随 batch 与序列长度增长，是训练区别于推理的主要开销
碎片：缓存池内的不可用空洞

TF32

Ampere 及之后的 GPU 上，FP32 的 matmul/卷积可改用 TF32（TensorFloat-32）：19 bit 总宽（8 位指数、10 位尾数），动态范围与 FP32 相同，吞吐接近 FP16。精度高于 FP16、低于 FP32，对绝大多数训练任务无感。

torch.backends.cudnn.allow_tf32 = True   # 卷积，默认开启
torch.backends.cuda.matmul.allow_tf32 = False  # matmul，默认关闭
 
# 推荐的新式开关，一次控制 matmul 精度档位
torch.set_float32_matmul_precision("high")  # "highest"=FP32, "high"=TF32, "medium"=bf16

TF32 只作用于存储仍为 FP32 的 Tensor，与 AMP 是两条独立的加速路径；两者叠加时 autocast 通常已把算子降到 fp16/bf16，TF32 不再介入。

AMP：autocast + GradScaler

动机：fp16 将显存与带宽减半、在 Tensor Core 上吞吐翻倍以上，但训练直接全量 fp16 会崩——损失尺度异常、梯度下溢。AMP（Automatic Mixed Precision）的策略是按算子选择精度而非按层手工标注。

autocast：前向类型自动转换

with torch.autocast(device_type="cuda", dtype=torch.float16):
    out = model(x)          # conv / matmul 自动降为 fp16
    loss = criterion(out, y)  # softmax / logsumexp / loss 保持 fp32

autocast 区域内，算子按内置清单分流：

降精度：matmul、conv*、linear——对带宽敏感、对精度不敏感，是收益来源
保 fp32：softmax、exp、log、reduction、loss——数值敏感，fp16 下易溢出或丢精度

权重本身保持 fp32，kernel 在执行时按需读取低精度副本；autocast 只改变算子执行 dtype，不改存储。区域外的运算不受影响；前向中被缓存的张量若要在区域外复用，用 torch.autocast("cuda", enabled=False) 显式回到 fp32，避免精度漂移。

GradScaler：防 fp16 梯度下溢

fp16 最小正规数约 $6 \times 1 0^{- 5}$ ，深层网络中小梯度乘积极易落入 denormal 区被截断为 0。GradScaler 将 loss 乘以放大系数 $s$ （初值 $2^{16}$ ），梯度整体放大后再 backward；step 前 unscale 回真实值，若检测到 inf/nan 则跳过本步更新。

scaler = torch.amp.GradScaler("cuda")  # 旧写法 torch.cuda.amp.GradScaler 已弃用
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
 
for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()
 
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = criterion(model(x), y)
 
    scaler.scale(loss).backward()       # 放大后的梯度写入 param.grad
    scaler.unscale_(optimizer)          # 还原真实梯度，梯度裁剪前必须调用
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)              # 内部检查 inf/nan，异常则跳过 [[torch.optim]] 的 step
    scaler.update()                     # 动态调整放大系数：异常则减半，平稳则回增

链路要点：scale → backward → unscale_ → clip → step → update。放大系数按指数间隔试探增长、遇溢出立即减半，使有效梯度始终落在 fp16 可表示区间。step 内部已做 unscale，不要再手动调用第二次。

构造时可传 enabled 条件——GradScaler("cuda", enabled=(use_amp and amp_dtype == torch.float16))：条件不满足时 scaler 退化为 noop，scale/step/update 全部直通，一套训练代码无需分支即可同时兼容 fp16 / bf16 / fp32。

bf16：无需 scaler

格式	符号	指数	尾数	动态范围
fp32	1	8	23	与 fp32 同
tf32	1	8	10	与 fp32 同
fp16	1	5	10	窄，梯度易下溢
bf16	1	8	7	与 fp32 同

bf16 指数位宽与 fp32 相同，动态范围不缩水，梯度既不下溢也不上溢，因此 GradScaler 可整体省略，只保留 autocast；代价是 7 位尾数带来的更高单步舍入误差，实践中由随机性抵消。Ampere 之后的 GPU 对 bf16 原生支持，当前社区默认已是 dtype=torch.bfloat16；fp16 + scaler 主要见于 V100 及更早的硬件。注意 CPU 不支持 bf16 autocast，个别老算子对 bf16 支持差时会出现精度漂移，排查时先降 fp16 验证。多卡训练下 AMP 与 NCCL 集合通信兼容，衔接见分布式训练总览。

知识花园

探索