从 AMP 到 Mixed Precision：训练和推理为什么不使用同一种精度

起因

#

以前在商汤做 CV 的时候，有一段时间在做目标检测模型的 INT8 量化。目标很直接：降低 inference latency，减少 model memory 和 bandwidth，同时利用 INT8 hardware throughput。

把整个模型一把量化成 INT8 以后，检测精度掉得比较明显。

后来排查发现，模型中有一些部分对 quantization error 非常敏感。最后我们没有坚持 entire graph INT8，而是让大部分适合量化的计算走 INT8，同时把少数 precision-sensitive layer 恢复到 FP32。我印象里最后采用的是类似 backbone INT8、部分 detection head 保留 FP32 的方案——但这段经历已经比较久，具体 layer 划分不一定准确。

重点不在当年的具体 layer 名称，而在这个现象：

整个网络使用最低 precision，并不一定是最优解。

自然的问题是：如果有些 layer 对 precision 很敏感，而有些 layer 根本不在乎，为什么一定让整个模型使用同一种 dtype？

这就是 Mixed Precision。

上一篇讨论了一个 number 内部如何分配 bit budget。这一篇把视角从"一个 number"提升到"整个 neural network computational graph"：既然不同 precision 有不同的性能和数值特性，为什么要让整个模型使用同一种 precision？

Mixed Precision 不是某一种 dtype

#

先解决一个概念问题。

FP16、BF16、INT8 是 numerical format。

Mixed Precision 是一种 precision allocation strategy。

不同 tensor、不同 operator、不同 layer、不同 computation stage，可以使用不同 precision。例如：

1Conv / GEMM      → BF16
2Reduction        → FP32
3Accumulator      → FP32

或者：

1Backbone         → INT8
2Sensitive Head   → FP16 / FP32

Mixed Precision 不是一种 datatype。它回答的问题是：

在哪里值得支付高精度的成本？

Training 和 Inference 的 Mixed Precision 必须分开

#

这是第一个重要分叉。训练和推理都叫 Mixed Precision，但面对的 numerical problem 并不一样。

Training 中 numerical error 会进入一个反馈闭环：

flowchart LR
    F["Forward"] --> L["Loss"]
    L --> B["Backward"]
    B --> G["Gradient"]
    G --> U["Parameter Update"]
    U --> F

误差在每一轮 iteration 里不断积累、放大。

Inference 就简单很多：

1Input → Forward → Output

没有 optimizer update，没有 gradient，误差不会反馈回来。因此 inference 往往可以采用比 training 更激进的 precision reduction。

为什么不能直接全部 FP16 训练

#

承接上一篇的 bit layout：

1FP32: 1 | 8  | 23
2FP16: 1 | 5  | 10
3BF16: 1 | 8  | 7

如果只是为了快，很自然会想到：把整个 training graph 从 FP32 替换成 FP16。

但不同 computation 对 numerical error 的敏感程度不同。Conv / GEMM 通常对低精度容忍度比较高，但 reduction、normalization、loss computation、gradient accumulation、optimizer update 可能对精度很敏感。

核心思想：

矩阵乘法非常适合低精度，但并不意味着所有 arithmetic 都适合低精度。

所以 AMP 并不是 model.half() 的同义词。AMP 更接近：

1适合 FP16/BF16 的 op → low precision
2数值敏感的 op         → FP32

AMP 到底自动了什么

#

PyTorch 中的典型 AMP 用法：

1with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
2    output = model(input)
3    loss = criterion(output, target)

这里 autocast 的语义不是"这个 scope 内所有东西全部 BF16"。而是：framework 根据 operator policy 和 backend support，自动选择合适的 execution dtype。

有些 operation（如 GEMM）使用 BF16 执行。有些 operation（如 softmax、layer norm、loss）会保持或提升到 FP32。

AMP = automatic precision placement。

它在 operator 粒度上自动决定：这个 op 用低精度跑性能更好，还是数值上必须用 FP32。不需要手工逐个 op 标注 dtype。

Gradient 和 Loss Scaling

#

Forward inference 即使用 FP16，很多时候还能工作。Training 更困难的重要原因之一是 gradient 可能非常小。

上一篇讨论过，FP16 exponent 只有 5 bits，最小 positive normal 大约 \(6.1 \times 10^{-5}\)。如果 gradient 更小：

Underflow 到零以后：

参数失去更新。

Loss Scaling 的机制很简单。先把 loss 放大：

那么 backward 得到的 gradient 也同步放大 \(S\) 倍，进入 FP16 更安全的数值范围。Optimizer step 前再除回来：

PyTorch 的 GradScaler 处理这件事：

1scaler = torch.amp.GradScaler("cuda")
2
3with torch.autocast(device_type="cuda", dtype=torch.float16):
4    output = model(input)
5    loss = criterion(output, target)
6
7scaler.scale(loss).backward()
8scaler.step(optimizer)
9scaler.update()

关键行为：

• scale factor 动态调整（训练稳定时逐步增大）
• 如果 scaled gradient 出现 inf/NaN（overflow），本次 update 被 skip，scale 缩小
• 下一轮重试

BF16 为什么让 Mixed Precision Training 简单很多

#

重新比较 bit layout：

1FP16: 1 | 5  | 10
2BF16: 1 | 8  | 7

BF16 保留了和 FP32 完全相同的 8-bit exponent。

BF16 最大的优势不是 precision 更高（事实上 BF16 fraction 只有 7 bits，比 FP16 的 10 bits 更少），而是 dynamic range 足够大。

FP16 的 dynamic range 只覆盖到约 65504，最小 normal 约 \(6 \times 10^{-5}\)。BF16 和 FP32 一样覆盖到约 \(10^{38}\)，最小 normal 约 \(10^{-38}\)。

直接后果：BF16 不容易遇到 gradient underflow 或 activation overflow。

BF16 mixed precision training 通常不需要 FP16 中那种为了弥补 narrow dynamic range 而使用的 loss scaling。

BF16 训练时代码通常直接是：

1with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
2    output = model(input)
3    loss = criterion(output, target)
4
5loss.backward()
6optimizer.step()

不需要 GradScaler。这是现代大模型训练普遍倾向 BF16 的重要原因之一。

训练中到底哪些东西是低精度

#

不要把 AMP 理解成"整个 training state 变成 16 bit"。一个训练流程里存在很多东西，它们不一定使用相同 precision。

一个典型 mixed precision training 的概念化：

 1Parameter storage      FP32（master weight）
 2         │
 3         ▼
 4Cast to BF16 input     BF16
 5         │
 6         ▼
 7GEMM multiply          BF16
 8         │
 9         ▼
10Accumulate             FP32
11         │
12         ▼
13Activation storage     BF16
14         │
15         ▼
16Backward gradient      BF16
17         │
18         ▼
19Optimizer states       FP32（Adam m / v）
20         │
21         ▼
22Master weight update   FP32

讨论 mixed precision 时，必须把 storage dtype、compute dtype 和 accumulator dtype 分开。

这和上一篇 TF32 部分的结论完全一致：面对一个 GPU kernel，不应该只问"这是 BF16 计算吗？"——要分别问 input storage、multiply precision、accumulator dtype、output dtype。

AMP 的收益

#

不要把 AMP 只理解成"显存减半"。收益来自几个方面：

Tensor Core throughput。 BF16 / FP16 GEMM 在 Tensor Core 上通常有比 FP32 更高的 hardware throughput。

Activation memory。 很多 forward intermediate tensor 可以用 BF16 存储，减少约一半 storage。

Memory bandwidth。 读写 activation / tensor 的 bytes 降低。对 memory-bound workload 尤其有帮助。

Cache efficiency。 相同 cache 容纳更多 element。

但要注意：

这和 上一篇显存文章 中 12~16 bytes/parameter 的分析一致。

从 Training 转向 Inference

#

到这里 Training Mixed Precision 的核心就讲完了。目标是让 training 稳定收敛，同时获得 throughput 和 memory 的收益。

Inference Mixed Precision 的目标不同：让最终 prediction 的 accuracy 不明显下降。

Inference 没有 gradient、没有 optimizer、没有 parameter update。所以 precision space 大很多——可以使用 FP32、FP16、BF16、FP8、INT8、INT4，甚至在一个 graph 内混合使用。

Inference Mixed Precision：以 INT8 为例

#

回到开头的目标检测经历。假设一个 detection model：

flowchart LR
    I["Input"] --> B["Backbone"]
    B --> N["Neck"]
    N --> H["Detection Head"]
    H --> O["NMS / Output"]

如果整个模型全部 INT8，accuracy 下降。那么一个非常自然的策略是：

flowchart LR
    I["Input"] --> B["INT8 Backbone"]
    B --> N["INT8 Neck"]
    N --> H["FP16/FP32 Head"]
    H --> O["Output"]

这就是 inference mixed precision。

目标不是让 INT8 coverage 达到 100%，而是在满足 accuracy constraint 的前提下，让尽可能多的 expensive computation 使用低精度。

为什么 Backbone 适合 INT8 而 Head 可能更敏感

#

这不是 universal rule，但在很多 CV workload 中确实存在这样的现象。

Backbone 中的 large Conv 贡献了大量 FLOPs、memory traffic 和 latency，所以把 backbone 量化成 INT8 有很高的 performance ROI。同时这些层通常有很多冗余和 redundancy，对 quantization noise 容忍度比较高。

而 detection head 中的一些 computation 可能更 numerical-sensitive。例如 confidence score：

1score: 0.501 → 经过 quantization noise → 0.492

在中间 activation 看起来只是很小的误差。但如果后面有：

1if score > 0.5:  keep
2else:            drop

一个连续的 numerical error 就被 downstream 的 discrete decision 放大了。

Bounding box regression 也类似——coordinate 的微小偏移可能改变 IoU，进而影响 NMS 结果和最终 mAP。

不要把这理解成"detection head 一定不能量化"。正确表述是：某些 detection model 的 head 对 quantization error 更敏感，因此实际部署中经常需要通过 sensitivity analysis 决定哪些 layer 保持高精度。

INT8 为什么比 BF16 更复杂

#

BF16 仍然是 floating point——有 sign、exponent、fraction，天然覆盖很大的 dynamic range。

INT8 只有 256 个 discrete integer levels。把一个 floating-point tensor 映射到 INT8 通常需要：

反过来：

其中 \(s\) 是 scale factor。

两类 error：

Rounding Error。 两个很接近的浮点值可能被映射到同一个 integer bucket，无法区分。

Clipping / Saturation Error。 如果真实 tensor 分布超出 calibration range：

1outlier → clip → 信息丢失

INT8 的误差很大程度取决于 tensor distribution，而不仅仅取决于 operator 类型。

一个 tensor 如果 distribution 很集中、没有离群值，INT8 的 256 个 level 用起来很高效。如果 distribution 有 heavy tail 或 outlier，INT8 会非常吃力。

这是为什么 inference quantization 比 training autocast 更难自动化——autocast 主要根据 operator 类型做 policy，而 INT8 quantization 还需要考虑数据分布。

Calibration 的真正作用

#

PTQ（Post-Training Quantization）中 calibration 经常被简单描述成"找 scale"。但本质上它在做一个更根本的决定：

INT8 一共只有 256 个 level。我们必须决定这 256 个 level 覆盖哪一段 floating-point range。

如果 range 选太大：

1[-1000, 1000] → 每个 level 间距 ≈ 7.8

小数字之间完全无法区分。

如果 range 选太小：

1[-1, 1] → outlier 全部 clip

大量信息丢失。

Calibration 本质也是一个 trade-off：

1quantization resolution ←→ clipping loss

选 range 小一点，resolution 更高（representable numbers 更密集），但 clip 更多 outlier。选 range 大一点，outlier 不会被 clip，但所有数字的分辨率都变差。

这和全文更大的"precision allocation"主题完全一致——即使在单个 tensor 的量化中，也在做 resource allocation。

PTQ 和 QAT

#

简单讲清楚两者和 mixed precision 的关系。

PTQ — Post-Training Quantization

#

1Trained FP32/BF16 model
2        ↓
3Calibration（跑一批 representative data）
4        ↓
5Quantized inference model

优点：简单，不需要重新训练。缺点：sensitive model 可能掉 accuracy。

QAT — Quantization-Aware Training

#

Training 中模拟 quantization 的效果：

1forward 中插入 fake quantize：
2  x → quantize → dequantize → x_approx

让模型学会适应 quantization noise。最终 inference 再真正用 INT8。

QAT 解决的是"模型如何适应低精度误差"；mixed precision 解决的是"哪些地方值得使用低精度"。

两者不互斥。完全可以：

1QAT 训练
2  +
3Mixed INT8 / FP16 graph 部署

一起使用。例如对整个模型做 QAT，但 sensitivity analysis 后仍然让少数 layer fallback 到 FP16。

Sensitivity Analysis：决定 Precision Placement

#

一种简单但有效的工程做法：

1先全部 INT8 → accuracy 掉很多
2         ↓
3逐层恢复 FP32 → 观察 accuracy recovery

或者反过来：

1逐层量化 → 测 accuracy delta

可以得到类似：

之后做 precision allocation：

1低 sensitivity  → INT8
2中等 sensitivity → INT8 / FP16（看 accuracy budget）
3高 sensitivity  → FP16 / FP32

这个过程的本质：

Precision 本身就是一个需要优化的 resource。

写成 Optimization Problem

#

把全文最重要的 abstraction 提出来。

假设模型 \(f = f_n \circ f_{n-1} \circ \cdots \circ f_1\)，每一个 layer \(f_i\) 可以选择一种 precision（FP32、BF16、FP16、INT8 等）。

我们真正希望优化的是：

subject to:

更 ML Infra 一点：

1minimize:
2    latency / memory / bandwidth / compute cost
3
4subject to:
5    mAP drop < 0.1%
6    accuracy drop < threshold
7    numerical stability OK

Mixed Precision 是一个 constrained optimization problem。

这比单纯教 torch.autocast() 重要得多。理解了这个抽象以后，无论未来出现什么新 dtype，思考方式都不变：在 accuracy constraint 下，把最便宜的 precision 分配给最不敏感的 computation。

Training AMP 和 Inference Mixed Precision 的统一

#

flowchart TD
    A["Neural Network Graph"] --> B{"Numerically Sensitive?"}
    B -->|Yes| C["Higher Precision
FP32 / BF16"]
    B -->|No| D["Lower Precision
BF16 / FP16 / INT8"]
    C --> E["Accuracy / Stability"]
    D --> F["Memory / Bandwidth / Throughput"]
    E --> G["Precision Placement Decision"]
    F --> G

Training AMP 和 inference mixed precision 的实现方式不同，但底层思想完全一致：不要平均地给所有计算相同的 precision，而是把 precision budget 花在最敏感的位置。

Training 中，framework 根据 operator 的 numerical property 做 autocast。

Inference quantization 中，系统根据 tensor distribution、layer sensitivity、accuracy metric、hardware support 决定 precision。

机制不同，问题相同。

和上一篇 IEEE 754 文章的对应

#

这一段非常值得想清楚。

上一篇的核心问题是：一个 number 只有固定 bit budget，怎么分配？

1多少 bit 给 exponent？→ dynamic range
2多少 bit 给 fraction？→ precision

于是产生了 FP32、FP16、BF16、FP8 E4M3、FP8 E5M2。

这是 number-level precision allocation。

这一篇的核心问题是：一个 neural network 有很多 operator，怎么分配 precision？

1哪些 op 用 FP32？
2哪些用 BF16？
3哪些用 INT8？

这是 graph-level precision allocation。

flowchart TD
    subgraph NumberLevel["Number-Level
上一篇"]
        NB["固定 bit budget"] --> NE["Exponent ↔ Fraction"]
        NE --> NR["Range ↔ Precision"]
    end
    subgraph GraphLevel["Graph-Level
本篇"]
        GB["固定高精度 budget"] --> GL["Layer A ↔ Layer B ↔ Layer C"]
        GL --> GR["Accuracy ↔ Efficiency"]
    end
    NumberLevel -.->|"相同的 trade-off 思想"| GraphLevel

两件事情本质上非常相似：

1IEEE floating-point design:
2有限 bit budget 怎么分？
3
4Mixed precision:
5有限高精度 budget 怎么分？

最终都是：在 correctness constraint 下，尽可能降低 system cost。

不要忘记 Accumulator Precision

#

和上一篇保持一致，强调：即使 input 是 BF16，也不意味着所有 arithmetic 都是 BF16。

GEMM：

典型执行：

1BF16 input → BF16 multiply → FP32 accumulation → BF16 output

因为大量 multiply result 不断 accumulation，误差容易积累。因此硬件普遍采用：

low-precision multiply + higher-precision accumulation

这本身就是 mixed precision。不要只把 mixed precision 理解成"Layer A FP16，Layer B FP32"。它甚至存在于单个 Tensor Core operation 内部。

现代 LLM 的延伸

#

文章主体讨论的是 FP16、BF16、INT8。但今天 LLM training / inference 已经进一步出现 FP8、INT4、FP4。

例如：

1Weight      → INT4 / FP4
2Activation  → BF16 / FP8
3Accumulator → FP32

或者：

1Large GEMM  → FP8 E4M3
2Sensitive op → BF16
3Accumulate   → FP32

基本问题完全没有改变。即使未来 dtype 再变化：

Mixed Precision 的思想不会变化。 仍然是在 precision / memory / bandwidth / throughput / accuracy 之间做 trade-off。

结论

#

1. Mixed Precision 不是某一种 dtype，而是一种 precision allocation strategy。

2. AMP 的核心不是"把模型变成 FP16"，而是自动决定哪些计算值得使用低精度。

3. FP16 training 需要 loss scaling 来弥补 narrow dynamic range；BF16 保留了 FP32 的 exponent range，通常不需要。

4. 训练中 storage dtype、compute dtype 和 accumulator dtype 三者需要分开讨论。

5. Inference 可以比 training 使用更激进的 precision reduction，因为没有梯度反馈闭环。

6. INT8 quantization 的误差取决于 tensor distribution，不仅取决于 operator 类型。

7. Inference mixed precision 的目标不是 100% 低精度 coverage，而是在 accuracy constraint 下最大化系统收益。

8. Mixed Precision 本质上是一个 constrained optimization problem：minimize system cost, subject to accuracy loss < ε。

上一篇解决的是：一个 number 内部怎么分 precision budget。

这一篇解决的是：整个 neural network graph 怎么分 precision budget。

参考资料

#

• 从 IEEE 754 到 BF16：理解 ML Infra 中的浮点精度选择（前篇）
• 从参数量到 Peak Memory：模型显存到底应该怎么算？（姊妹篇）
• 从 FLOPs 到 Latency：GPU 推理性能到底由什么决定？（姊妹篇）
• PyTorch. Automatic Mixed Precision package — torch.amp.
• PyTorch. Automatic Mixed Precision Examples.
• Micikevicius et al. Mixed Precision Training. ICLR 2018.
• NVIDIA. Training with Mixed Precision.
• NVIDIA. TensorRT Developer Guide — Working with Reduced Precision.
• NVIDIA. Achieving FP32 Accuracy for INT8 Inference Using Quantization Aware Training with NVIDIA TensorRT.
• Kalamkar et al. A Study of BFLOAT16 for Deep Learning Training. 2019.