从 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 划分不一定准确。

这个现象是清楚的：

整个网络使用最低 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 会进入一个反馈闭环：

误差在每一轮 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。

但上面这张图里还有一个问题值得深挖：第一行和最后一行都标了 FP32。BF16 不是有和 FP32 一样大的 dynamic range 吗？为什么参数存储和更新还需要回到 FP32？

为什么需要 FP32 Master Weights

#

问题不在 range，在 precision。

BF16 和 FP32 共享 8 位 exponent，dynamic range 确实一样大。但 BF16 只有 7 位 fraction，FP32 有 23 位。数值范围足够大，不等于数值精度足够高。

举个极简例子。假设参数 \( w = 1.25 \)，每次更新量 \( \Delta w = -0.001 \)。BF16 在 1.25 附近相邻可表示数的间距是 \( 2^0 \times 2^{-7} = 0.0078125 \)，也就是说 1.25 的下一个 BF16 可表示数是 1.2421875。

\( \Delta w = -0.001 \) 远小于这个间距：

更新被舍入吞掉了。如果直接用 BF16 存参数并逐步更新，参数永远不动。

但如果维护一份 FP32 master weight，小更新可以持续累积。累积变化跨过 BF16 相邻可表示数的边界时，cast 回来的 BF16 副本才发生变化：

Step 4 时 FP32 master 累积了 4 次 \( -0.001 \)，跨过了舍入边界，cast 出来的 BF16 从 1.25 跳到了 1.2422（精确值 1.2421875）。而 BF16 直接更新那列，始终卡在 1.25。

需要注意两个区分：

这和 FP16 的 gradient underflow / Loss Scaling 不是同一个问题。 Loss Scaling 解决的是梯度本身太小、在 FP16 narrow dynamic range 下直接变成零。Master Weights 解决的是参数更新步长太小、在低精度下被舍入吞掉。BF16 的 dynamic range 和 FP32 一样大，梯度不容易 underflow，但 7 位 fraction 的精度仍然不够保留微小的参数变化。

FP32 Master Weights 和 Adam Optimizer States（m / v）不是同一个东西。 Adam 的 first moment m 和 second moment v 是优化器维护的统计量，用于计算自适应学习率。Master weight 是参数本身的高精度副本，用于接收 optimizer 的更新结果。两者都是 FP32，但功能和数据内容完全不同。

FP32 Master Weights 是否抵消了 BF16 的显存收益

#

想通了 master weight 的必要性以后，紧接着就是下一个疑问：BF16 原本是为了省显存，额外加一份 FP32 master weight，省下的不是又花回去了吗？

用 1B 参数模型 + Adam 优化器做一个简单的显存账本（1B × bytes/param = GB）：

BF16 在参数和梯度上省下的 4 GB，被额外的 FP32 master weight 原样补回来了。Model States 总量完全一样。

需要强调这是一个特定配置下的估算：BF16 梯度 2 bytes/param，FP32 Adam states，独立的 FP32 master copy。不同框架和配置可能有差异，不能把 16 bytes/parameter 说成所有 BF16 混合精度训练的普遍结论。

那 BF16 省下的显存在哪里？Activations。

Forward 过程中大量中间 tensor 可以用 BF16 存储。假设 FP32 训练时 Activations 占 12 GB，BF16 时降为 6 GB：

这些数字是示意估算，不代表实际训练的激活值一定减半，也不能视为真实 Peak Memory 测量。但它展示了一个关键模式：

Mixed Precision 不是把所有数据都变成低精度，而是将不同精度分配给适合它们的计算和状态。降低计算精度，并不意味着训练的 Model States 显存一定减少。

此外 BF16 的收益远不只是显存。Tensor Core 上 BF16 GEMM 的 hardware throughput 通常显著高于 FP32，读写 activation 的 memory bandwidth 也因 bytes 减半而获益——把 BF16 仅仅描述为"显存优化技术"会低估它的价值。

Activation Checkpointing 会削弱 BF16 的显存收益吗

#

前面看到 BF16 在 Model States 上没有显存优势，主要靠 Activations 省空间。那如果同时开了 Activation Checkpointing，Activations 大幅减少以后，BF16 还能省多少？

延续前面的示意估算：

Checkpointing 已经把 Activations 从 12 GB 砍到 2 GB，BF16 在此基础上只再省 1 GB。

但这不意味着 BF16 在这个场景下没有意义。Checkpointing 并不消灭所有 Activations——重计算时仍需要临时显存；BF16 的 Tensor Core 计算吞吐优势不因此消失；实际 Peak Memory 还取决于 tensor lifetime 和 workspace 的重叠方式。

这里有一个更一般的工程观察：显存优化技术的收益往往不能简单相加——它们可能同时作用于同一部分内存。 Checkpointing 减少 Activations，BF16 也减少 Activations，两者同时使用时，收益不是各自效果的加总。

PyTorch AMP 与 FP32 Master Weights

#

前面的 ASCII 图写了 Parameter storage FP32（master weight），容易让人以为 PyTorch AMP 一定会额外创建一份 FP32 master copy。实际上有两种常见模式：

模式 A：FP32 参数 + BF16 Autocast。 最常见的 torch.autocast 用法。模型参数本身就是 FP32，autocast 让适合的 op 在执行时 cast 到 BF16。参数直接接收 FP32 精度的 optimizer 更新——FP32 参数本身就承担了 master weight 的角色，不需要额外副本：

1model = MyModel().cuda()  # 参数默认 FP32
2optimizer = torch.optim.Adam(model.parameters())
3
4with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
5    loss = model(input)
6loss.backward()           # 梯度通常是 FP32
7optimizer.step()          # 直接更新 FP32 参数

这种模式下，常规 FP32 参数的梯度通常也是 FP32。

模式 B：BF16 参数 + 独立 FP32 Master Weights。 模型持久化存储 BF16 权重，优化器内部维护一份独立的 FP32 master copy，更新后再 cast 回 BF16。这种模式需要额外的 FP32 显存开销。

两者的区别在于参数本身的 dtype。模式 A 的参数就是 FP32，不存在额外 master copy 的显存开销；模式 B 的参数是 BF16，需要 master copy 来保留更新精度。

注意：直接 model.to(torch.bfloat16) 让参数变成 BF16 后，并不会自动获得 FP32 master weights。没有 master 的纯 BF16 参数更新，就会遇到前面数值例子中的舍入问题。

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：

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

这就是 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。

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。

回到当年的量化方案

#

当年把部分 layer 恢复成 FP32，是因为整图 INT8 达不到检测精度要求。具体 layer 划分记不准确了，但这个选择的依据很明确：先找到误差敏感的位置，再检查恢复精度后的 latency 和内存收益。

训练里的 AMP 也会按计算选择精度，不过它还要处理 backward、梯度范围和参数更新。把训练 AMP 的经验直接套到 inference quantization 上，会漏掉 calibration、tensor 分布和目标硬件这些条件。

对我来说，选 mixed precision 方案时最有用的记录仍然是：哪些计算改了 dtype，精度指标怎么变，端到端开销又怎么变。低精度覆盖了多少层，只是其中一个中间量。

补充：精度分配的抽象与延伸

#

约束优化、数值格式、accumulator 与更低精度的延伸

写成 Optimization Problem

#

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

优化的目标就是：

$$ \min \quad \text{Latency}(\text{或 Memory / Bandwidth / Cost}) $$

subject to:

$$ \text{Accuracy Loss} \lt \epsilon $$

更 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。

理解这个抽象以后，无论未来出现什么新 dtype，思考方式都不变：在 accuracy constraint 下，把最便宜的 precision 分配给最不敏感的 computation。

Training AMP 和 Inference Mixed Precision 的统一

#

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。

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

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

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

不要忘记 Accumulator Precision

#

即使 input 是 BF16，也不意味着所有 arithmetic 都是 BF16。

GEMM：

$$ C_{ij} = \sum_k A_{ik} B_{kj} $$

典型执行：

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。

参考资料

#

• 从 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.