PyTorch 推理工程（05）：混合精度、AMP 与吞吐、延迟

1. 本节定位

在 device 与同步问题之外，推理侧常通过降低计算精度以换取显存与吞吐；涉及 FP32、FP16、BF16 的数值范围与尾数位宽差异、是否使用 autocast 而非整网 .half()，以及 batch、吞吐（如 samples/s）与延迟（如 p50/p99）的度量口径。本篇围绕上述选择与约束展开。

2. 要点

混合精度的核心思想不是：

把所有计算都强行变成低精度。

而是：

让更适合低精度的部分用低精度，让更需要数值稳定性的部分继续用更高精度。

这也是为什么叫 AMP（Automatic Mixed Precision），而不是"Automatic Half Precision"。

3. 把 3 个最重要的 dtype 彻底分清

先用一张表看清楚它们的数值属性

1
2
3
4
5
import torch

for dt in [torch.float32, torch.float16, torch.bfloat16]:
    info = torch.finfo(dt)
    print(f"{str(dt):20s}  bits={info.bits:2d}  max={info.max:.2e}  min_pos={info.tiny:.2e}  eps={info.eps:.2e}")

输出：

1
2
3
torch.float32         bits=32  max=3.40e+38  min_pos=1.18e-38  eps=1.19e-07
torch.float16         bits=16  max=6.55e+04  min_pos=6.10e-05  eps=9.77e-04
torch.bfloat16        bits=16  max=3.39e+38  min_pos=1.18e-38  eps=7.81e-03

几个关键发现：

属性 含义 float32 float16 bfloat16
bits 每个数占多少位 32 16 16
max 能表示的最大正数 3.4×10³⁸ 6.55×10⁴ 3.4×10³⁸
eps 1.0 附近能分辨的最小差距 1.2×10⁻⁷ 9.8×10⁻⁴ 7.8×10⁻³

最关键的发现：

  • float16 的 max 只有 65504——一旦中间计算结果超过这个值，就会变成 inf（溢出）
  • bfloat16 的 max 和 float32 一样大——不容易溢出
  • bfloat16 的 eps 比 float16 大——小数点后的精度更粗

为什么 float16 容易数值出问题

最常见的例子就是 softmax：

1
softmax(x) = exp(x_i) / sum(exp(x_j))

如果 x 里有一个值是 300，那 exp(300) ≈ 10¹³⁰，远超 float16 的最大值 65504，直接变成 inf，softmax 输出全是 nan。

用 float32 就没问题，因为 float32 的 max 是 3.4×10³⁸。

为什么 bfloat16 是更安全的低精度方案

float16 和 bfloat16 都是 16 位，但内部 bit 分配不同：

1
2
3
float32：   1 位符号 | 8 位指数 | 23 位尾数
float16：   1 位符号 | 5 位指数 | 10 位尾数  ← 指数位少，范围小，容易溢出
bfloat16：  1 位符号 | 8 位指数 | 7 位尾数   ← 和 float32 一样的指数位，范围大

bfloat16 的指数位和 float32 完全相同，所以能表示的数值范围和 float32 一样，不容易溢出。代价是尾数位更少（7 位 vs 23 位），精度更粗。

float16 与 bfloat16 的差异可归纳为指数位与尾数位分配不同，从而在可表示范围与尾数精度之间形成不同折中；表述时宜对比位宽与典型溢出/舍入行为，而非仅概括“更稳”。

三个 dtype 的工程选择逻辑

dtype 显存 速度 数值稳定性 主要使用场景
float32 4 字节/元素 慢 最强 基线，数值敏感层
float16 2 字节/元素 快（尤其 matmul） 弱，容易溢出 推理加速，需谨慎
bfloat16 2 字节/元素 快 较强，范围大 LLM 推理/训练主流选择

4. 低精度为什么常常更快

从工程角度理解，低精度的加速来自三个层面：

层面 1：每个元素更小，显存和带宽压力减半

一个 shape 为 (B, T, D) = (32, 512, 4096) 的 Tensor：

1
2
float32: 32 × 512 × 4096 × 4 字节 = 268 MB
float16: 32 × 512 × 4096 × 2 字节 = 134 MB  ← 省了一半

显存省了一半，意味着：

  • 同样显存下可以放 2 倍的 batch
  • 模型参数和激活值占用更少
  • 数据在内存和 GPU 之间搬运更快（带宽压力更低）

层面 2：GPU 有专门的低精度计算单元，吞吐更高

现代 GPU（如 NVIDIA A100、H100）有专门的Tensor Core，对 float16 和 bfloat16 的矩阵乘法吞吐是 float32 的几倍：

1
2
3
4
A100 理论峰值：
  float32：  312 TFLOPS
  float16：  312 TFLOPS（TF32）或 624 TFLOPS（FP16）  ← 更高
  bfloat16： 312 TFLOPS（TF32）或 624 TFLOPS（BF16）

矩阵乘法（matmul、linear、conv）是深度学习里最密集的计算，低精度在这里的加速最明显。

层面 3：但"更快"是有条件的

"常常更快"≠ 总是更快。

如果瓶颈不在 matmul 上（比如在数据搬运、同步、Python 控制逻辑），那改精度不会有明显收益。

5. 低精度的数值风险

不是所有 op 都适合低精度。风险主要来自：

风险 1：数值溢出（float16 尤其严重）

如第 3 节说的，float16 最大值只有 65504，稍大的中间值就会变成 inf。

风险 2：精度不足导致梯度消失（更多是训练问题，推理也要注意）

某些操作（如 LayerNorm、归约、log_softmax）在低精度下容易丢失小数值，输出偏移。

风险 3：op 之间的精度积累误差

一个 op 的输出是下一个的输入，如果每一步都有微小精度损失，多层之后可能积累成明显偏差。

哪些 op 适合低精度，哪些要保留高精度

根据 PyTorch 的 AMP 设计规则：

常用 float16 常保留 float32
matmul, mm, bmm softmax, log_softmax
linear layer_norm, group_norm
conv1d/2d/3d cross_entropy
batch_norm（前向） sum, norm（归约）

这个设计逻辑非常合理：

  • 重算力密集的 op（matmul/conv）→ 低精度，吃到 Tensor Core 加速
  • 数值敏感的 op（softmax/norm/loss）→ 保留高精度，避免溢出和不稳定

6. AMP：为什么是"自动混合"，不是"全部半精度"

PyTorch 的解决方案叫 AMP（Automatic Mixed Precision），核心工具是 torch.autocast。

AMP 通过 autocast 按官方规则为各算子选择精度，无需在逐层上手工指定 dtype。

1
2
3
4
5
6
调用形式仍为 `model(x)`：
autocast 做的：Linear → float16
               ReLU   → 不变（非浮点运算）
               LayerNorm → float32  ← 保留！
               Linear → float16
               Softmax → float32  ← 保留！

即按算子类型在 FP16/BF16 与 FP32 之间切换，而非整网统一为 half。

7. torch.autocast 的基本用法

推理标准写法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import torch
import torch.nn as nn

device = "cuda" if torch.cuda.is_available() else "cpu"

model = nn.Sequential(
    nn.Linear(128, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
).to(device)
model.eval()

x = torch.randn(32, 128).to(device)

with torch.inference_mode():
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        y = model(x)

print(y.dtype)    # 通常是 torch.float16（最后一层 Linear 适合低精度）
print(y.shape)    # torch.Size([32, 10])

bfloat16 版本

1
2
3
with torch.inference_mode():
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        y = model(x)

代码层次结构

1
2
3
4
model.eval()                           # 第 1 层：模块行为模式（Dropout/BN 正确运行）
with torch.inference_mode():           # 第 2 层：关闭 autograd（不构建计算图）
    with torch.autocast(...):          # 第 3 层：自动选择精度
        y = model(x)                   # 实际推理

三层嵌套，每层负责不同的事情，缺一不可。

8. 用代码直接验证 autocast 做了什么

想看清楚 autocast 在每一层选了什么精度，可以加一个 hook：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(128, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# 注册 hook，打印每层输入输出的 dtype
def dtype_hook(module, input, output):
    in_dtype  = input[0].dtype  if isinstance(input, tuple) and len(input) > 0 else "N/A"
    out_dtype = output.dtype    if hasattr(output, 'dtype') else "N/A"
    print(f"{type(module).__name__:12s}  input={in_dtype}  output={out_dtype}")

for layer in model:
    layer.register_forward_hook(dtype_hook)

model.eval()
x = torch.randn(4, 128)

# 先看 float32 下的情况
print("=== float32 ===")
with torch.inference_mode():
    y = model(x)

# 再看 autocast float16 下的情况（需要 CUDA）
if torch.cuda.is_available():
    model = model.cuda()
    x = x.cuda()
    print("\n=== autocast float16 ===")
    with torch.inference_mode():
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            y = model(x)

在 autocast 上下文中，典型 Linear 的参与张量常为 float16，体现自动混合精度的实际效果。

9. 为什么官方不推荐"用 autocast 又 .half()"

官方明确说：使用 autocast 时，不要手动对模型或输入调用 .half() 或 .bfloat16()。

原因是：

autocast 的价值就是"混合"。它会让 softmax、LayerNorm 这些敏感 op 保留 float32。

若事先对整模型调用 .half()，权重已全局降为 float16，autocast 对部分算子保留 FP32 的语义被削弱，数值敏感算子可能被强行置于低精度。

1
2
3
4
5
6
7
8
9
# 反例: 错误用法：先 half，再 autocast
model = model.half()
with torch.autocast(device_type="cuda", dtype=torch.float16):
    y = model(x)   # LayerNorm 的权重是 half，autocast 的保护失效了

# 示例: 正确用法：模型保持 float32，让 autocast 管精度
model = model.to(device)   # float32 权重
with torch.autocast(device_type="cuda", dtype=torch.float16):
    y = model(x)   # autocast 按规则自动决定每层精度

10. “全模型 .half()” 什么时候会出现

除 autocast 外，工程中仍可见：

1
2
model = model.half()
x = x.half()

这种写法不是绝对错误，但它更像是一种：

  • 更手动、更激进的方案
  • 通常出现在"我完全清楚这个模型每一层的数值特性"的场景
  • 某些部署框架（TensorRT、ONNX Runtime）的导出前处理

对于大多数常规推理优化任务：

优先掌握 autocast，比优先掌握"全模型手动 half"更合理、更安全。

11. GradScaler：推理时一般不需要

GradScaler 的适用场景如下：

  • GradScaler 是解决 训练 中 float16 backward 梯度下溢的工具
  • 推理时没有 backward，不需要 GradScaler
场景 autocast GradScaler
推理 需要 不需要
训练 需要 需要（float16 训练时）

12. 吞吐（Throughput）和延迟（Latency）：必须分清

吞吐与延迟是评价混合精度等改动的基本度量；二者定义不同，不可互换。

直觉定义

延迟（Latency）：一次请求从进来到出结果耗时多久。
吞吐（Throughput）：单位时间内能处理多少请求/token/样本。

用餐厅类比：

1
2
延迟 = 一道菜从点单到上桌的时间（单次响应时间）
吞吐 = 餐厅每小时能服务多少桌客人（总处理能力）

用代码建立数字感

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import torch
import torch.nn as nn
import time

model = nn.Linear(1024, 1024).cuda().eval()

def measure(batch_size, n_runs=100):
    x = torch.randn(batch_size, 1024, device="cuda")

    # warm-up
    with torch.inference_mode():
        for _ in range(10):
            _ = model(x)

    torch.cuda.synchronize()
    start = time.time()

    with torch.inference_mode():
        for _ in range(n_runs):
            _ = model(x)

    torch.cuda.synchronize()
    end = time.time()

    total_time = end - start
    latency_ms = total_time / n_runs * 1000          # 单次耗时（毫秒）
    throughput  = batch_size * n_runs / total_time   # 每秒处理样本数

    print(f"batch={batch_size:4d}  latency={latency_ms:.2f} ms  throughput={throughput:.0f} samples/s")

measure(1)
measure(8)
measure(64)
measure(512)

在 A100 等设备上常观察到下列量级关系（示例）：

1
2
3
4
batch=   1  latency=0.08 ms  throughput= 12,500 samples/s
batch=   8  latency=0.09 ms  throughput= 88,000 samples/s
batch=  64  latency=0.25 ms  throughput=256,000 samples/s
batch= 512  latency=1.20 ms  throughput=426,000 samples/s

关键现象：

  • batch 从 1→512：延迟增加了 15 倍，但吞吐增加了 34 倍
  • 对应「提高 batch 可抬吞吐、单请求延迟亦上升」的量化描述

优化目标决定策略

场景 主要指标 策略
在线服务（用户等待结果） 延迟 小 batch，低延迟算子
离线批处理（跑数据集） 吞吐 大 batch，混合精度
推理服务 SLA 延迟 P99 需要同时考虑

13. 混合精度和 batch size：联合优化视角

低精度带来的一个直接好处是显存减半，这意味着：

1
2
float32 下：显存 = 模型参数 × 4字节 + 激活值 × 4字节 + ...
float16 下：显存 = 模型参数 × 2字节 + 激活值 × 2字节 + ...（约少一半）

同样 40GB 显存：

1
2
float32 能跑 batch=32
float16 能跑 batch=64（多了将近一倍）

更大的 batch 对推理的意义：

  • GPU 利用率更高（Tensor Core 利用率提升）
  • 整体吞吐上升
  • 单位成本降低

但代价是：

  • 单请求等待拼 batch 的时间更长
  • 动态长度输入时 padding 浪费更多

所以真实推理优化里，往往是这几个参数联合调：

1
dtype（float32/float16/bfloat16）× batch_size × 硬件型号 → 找最优工作点

14. "半精度没加速"时，怎么系统排查

这在实际工作中非常常见，是工程能力的体现。

现象 可能原因
切了低精度，延迟几乎没变 瓶颈在数据搬运、同步或 Python 控制，不在 matmul
batch 很小时几乎没收益 GPU 没吃满，launch overhead 占比过高
某些层仍然是 float32 正常！autocast 本来就是混合策略，数值敏感层会保留
结果明显飘了 某处触发了 float16 溢出，考虑改用 bfloat16
bfloat16 比 float16 慢 某些老 GPU 对 bfloat16 支持不好，Tensor Core 优化不完整

15. 做精度对比实验的正确方法

比较 FP32、FP16 autocast、BF16 autocast 时，建议至少包含：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import torch
import torch.nn as nn
import time

def benchmark(model, x, dtype_ctx=None, n_runs=200, device="cuda"):
    model.eval()

    # warm-up
    with torch.inference_mode():
        ctx = dtype_ctx if dtype_ctx else torch.inference_mode()
        for _ in range(20):
            if dtype_ctx:
                with torch.inference_mode():
                    with dtype_ctx:
                        _ = model(x)
            else:
                with torch.inference_mode():
                    _ = model(x)

    torch.cuda.synchronize()
    start = time.time()

    with torch.inference_mode():
        for _ in range(n_runs):
            if dtype_ctx:
                with dtype_ctx:
                    _ = model(x)
            else:
                _ = model(x)

    torch.cuda.synchronize()
    elapsed = (time.time() - start) / n_runs * 1000

    mem_mb = torch.cuda.memory_allocated() / 1024**2
    return elapsed, mem_mb

if torch.cuda.is_available():
    device = "cuda"
    model = nn.Sequential(
        nn.Linear(512, 2048), nn.ReLU(),
        nn.Linear(2048, 2048), nn.ReLU(),
        nn.Linear(2048, 512)
    ).to(device)

    x_fp32 = torch.randn(64, 512, device=device)

    # FP32 基线
    lat, mem = benchmark(model, x_fp32)
    print(f"FP32:        latency={lat:.2f} ms  memory={mem:.0f} MB")

    # FP16 autocast
    ctx16 = torch.autocast(device_type="cuda", dtype=torch.float16)
    lat, mem = benchmark(model, x_fp32, ctx16)
    print(f"FP16 autocast: latency={lat:.2f} ms  memory={mem:.0f} MB")

    # BF16 autocast
    ctx_bf = torch.autocast(device_type="cuda", dtype=torch.bfloat16)
    lat, mem = benchmark(model, x_fp32, ctx_bf)
    print(f"BF16 autocast: latency={lat:.2f} ms  memory={mem:.0f} MB")

实验必须有：

  1. warm-up 轮次：首次运行有 kernel 编译和初始化延迟，不代表稳定性能
  2. cuda.synchronize()：确保计时包含 GPU 实际完成时间
  3. 记录显存：不只看速度，也要看显存代价
  4. 保留 FP32 基线：否则难以量化相对加速比

16. 一个最常见的坑：只看速度，忘了验证输出

提速了，但输出对不对？

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 10)).cuda()
model.eval()
x = torch.randn(1, 128, device="cuda")

# FP32 基线输出
with torch.inference_mode():
    y_fp32 = model(x).float()   # 保证是 float32

# FP16 autocast 输出
with torch.inference_mode():
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        y_fp16 = model(x).float()

# 计算差异
max_err = (y_fp32 - y_fp16).abs().max().item()
mean_err = (y_fp32 - y_fp16).abs().mean().item()
print(f"最大绝对误差: {max_err:.6f}")
print(f"平均绝对误差: {mean_err:.6f}")

如果误差在小数点后 4 位以内，通常可以接受。如果出现 inf 或 nan，或者误差很大，就需要检查是否有数值溢出。

17. 混合精度推理的标准模板（汇总）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import torch
import torch.nn as nn

device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. 模型保持 float32 权重（不要手动 half）
model = model.to(device)
model.eval()

# 2. 输入保持 float32（autocast 会处理）
x = x.to(device)

# 3. 三层嵌套推理
with torch.inference_mode():                              # 关闭 autograd
    with torch.autocast(device_type="cuda",              # 自动混合精度
                        dtype=torch.float16):             # 或 bfloat16
        y = model(x)

# 4. 如果需要后处理，注意 dtype 可能是 float16，必要时转回 float32
y = y.float()

18. 面试常见问题

Q：float16 和 bfloat16 的区别是什么？
A：都是 16 位浮点，但 bit 分配不同。float16 有 5 位指数+10 位尾数，max ≈ 65504，容易溢出；bfloat16 有 8 位指数+7 位尾数，和 float32 指数位相同，max ≈ 3.4×10³⁸，不容易溢出，但小数精度更粗。

Q：为什么 PyTorch 推荐 AMP / autocast，而不是全模型 .half()？
A：不是所有 op 都适合低精度，autocast 按照 op-specific 规则自动为 matmul 等重计算 op 使用低精度，为 softmax、LayerNorm 等数值敏感 op 保留 float32。手动全模型 half 会跳过这道保护。

Q：吞吐和延迟的区别是什么？
A：延迟是单次请求的响应时间，吞吐是单位时间的总处理量。大 batch 通常提高吞吐，但可能增加单请求等待时间。在线服务关注延迟，离线批处理关注吞吐。

Q：推理里为什么一般不需要 GradScaler？
A：GradScaler 解决的是 float16 训练时反向传播梯度下溢的问题，推理没有 backward，不需要。

Q：切了低精度但没有加速，原因是什么？
A：低精度的收益主要在 matmul/conv 等重计算上。如果瓶颈在数据搬运、同步点、Python 控制逻辑或者 batch 太小，切精度不会有明显收益。

19. 思考题

练习 1：读懂 dtype 数值属性

1
2
3
4
5
import torch

for dt in [torch.float32, torch.float16, torch.bfloat16]:
    info = torch.finfo(dt)
    print(f"{str(dt):20s}  max={info.max:.2e}  eps={info.eps:.2e}")

思考：为什么 float16 的 max 只有 6.55×10⁴？softmax 里如果 logit 很大，会发生什么？

练习 2：验证溢出

1
2
3
4
5
6
7
8
9
10
import torch

x = torch.tensor([100.0], dtype=torch.float32)
print(torch.exp(x))   # 正常

x16 = torch.tensor([100.0], dtype=torch.float16)
print(torch.exp(x16))   # inf！

x_bf = torch.tensor([100.0], dtype=torch.bfloat16)
print(torch.exp(x_bf))   # 正常，bfloat16 不溢出

softmax 等在 float16 下易数值不稳；autocast 常将其保留在 float32 以降低溢出与舍入风险。

练习 3：写一个标准 AMP 推理脚本

自己写一个包含 Linear + ReLU + Linear 的模型：

  1. 保持 float32 权重
  2. 用 inference_mode + autocast(float16) 做推理
  3. 打印输出的 dtype 和 shape
  4. 和不加 autocast 的 float32 输出做对比，计算最大误差

练习 4：batch size vs 延迟/吞吐实验

对同一个模型，测试 batch_size = 1, 8, 64, 512 时的：

  • 单次推理时间（毫秒）
  • 每秒处理样本数

思考：哪个 batch size 让 GPU 利用率最高？哪个让单请求等待最短？

练习 5：对比 FP32 / FP16 / BF16 的速度和精度

可用第 15 节示例在本机复现并记录：

  • 延迟（ms）
  • 最大输出误差

思考：bfloat16 的误差比 float16 大还是小？为什么？（提示：尾数位不同）

20. 本节要点与自检

  • 能解释 float32、float16、bfloat16 的 bit 分配差异和对应的数值特性
  • 理解低精度为什么更快（显存、带宽、Tensor Core）
  • 知道哪些 op 适合低精度，哪些要保留高精度
  • 正确使用 torch.autocast(...) 做推理混合精度
  • 知道为什么不推荐"用了 autocast 还手动全模型 half"
  • 清晰区分吞吐和延迟，理解 batch size 对两者的影响
  • 做精度实验时能保留 FP32 基线并验证输出质量

21. 小结

混合精度不是"把模型全变小数位"，而是"让 matmul 这类重计算吃到低精度加速，让 softmax 这类数值敏感 op 保留稳定性，并在吞吐、延迟、显存和输出质量之间做工程权衡"。

系列导航