一、混合精度训练（AMP）：FP16/BF16原理、AMP的3种策略（O0/O1/O2/O3）、Loss Scaling技巧

各位同学好，我是你们的GPU工程师老K。今天咱们聊聊混合精度训练，也就是AMP。

说实话，我刚入行那会儿，训练大模型全靠FP32硬扛。显存不够？那就砍batch size呗。后来发现，这简直是暴殄天物。你想想看，一张A100 80G，跑个175B的模型，FP32连参数都塞不下。但换成FP16/BF16混合精度，嘿，不仅能塞下，还能跑得更快。

为什么会这样？因为FP16和BF16的每个数只占2个字节，而FP32占4个字节。显存占用直接砍半。再加上Tensor Core对低精度计算有硬件加速，训练速度能提升2-4倍。我在项目中遇到过，一个原本要跑3天的BERT-large，切到AMP后，1天半就跑完了。嗯，真香。

1.1 FP16和BF16的原理

先说说FP16。它也叫半精度浮点数，用1位符号位、5位指数位、10位尾数位。表示范围大约是±65504，精度大概3位有效数字。

BF16呢？它是Google搞出来的，1位符号位、8位指数位、7位尾数位。表示范围和FP32一样大（±3.4×10³⁸），但精度只有2位有效数字。

我个人的习惯是：如果模型对精度不敏感，比如CV领域的ResNet、ViT，用FP16就够了。但如果是NLP模型，尤其是训练过程中梯度分布很广的那种，我建议用BF16。为什么？因为BF16的动态范围大，不容易溢出。

类型	符号位	指数位	尾数位	表示范围	精度
FP32	1	8	23	±3.4×10³⁸	~7位
FP16	1	5	10	±65504	~3位
BF16	1	8	7	±3.4×10³⁸	~2位

1.2 AMP的4种策略：O0/O1/O2/O3

AMP（Automatic Mixed Precision）是NVIDIA官方推出的混合精度训练工具。它提供了4种策略，从O0到O3，精度逐渐降低，速度逐渐提升。

我刚开始用AMP时，直接上了O3，结果模型不收敛。后来才明白，这玩意儿得根据模型特性来选。

O0：纯FP32

说白了就是不用混合精度。所有计算和存储都用FP32。适合对精度要求极高的场景，或者模型本身很小，显存够用。

O1：保守混合精度

这是我最常用的策略。它会自动识别哪些操作可以用FP16加速，哪些必须保留FP32。比如卷积、矩阵乘法这些计算密集型操作，自动切到FP16。而BatchNorm、Softmax这些对精度敏感的操作，保留FP32。

我在项目中遇到过，O1策略下，模型精度和FP32几乎一样，但速度提升了1.5倍左右。

O2：激进混合精度

除了BatchNorm和少数特殊操作，其他全部用FP16。显存占用更低，速度更快。但风险也大，容易出现梯度下溢或上溢。

O3：纯FP16

所有操作都用FP16。速度最快，但精度损失也最大。我建议只在推理阶段使用，或者对精度要求不高的场景。

策略	计算精度	存储精度	速度提升	风险等级
O0	FP32	FP32	1x	无
O1	混合	混合	1.5-2x	低
O2	FP16为主	FP16为主	2-3x	中
O3	FP16	FP16	3-4x	高

1.3 Loss Scaling技巧

Loss Scaling，说白了就是给损失函数乘一个大数，防止梯度下溢。

为什么会下溢？FP16能表示的最小正数是2⁻²⁴ ≈ 5.96×10⁻⁸。如果梯度比这个还小，就会变成0。梯度为0，参数就不更新了，模型就废了。

我个人的习惯是：初始缩放因子设为128.0，每N步检查一次梯度。如果梯度没有溢出，就把缩放因子翻倍。如果溢出了，就减半。这叫动态Loss Scaling。

# PyTorch AMP 示例
import torch
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler(init_scale=128.0)

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    # Loss Scaling
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

我曾经在一个NLP项目里，用O2策略训练GPT-2，结果跑了2万步后loss突然炸了。排查了半天，发现是某个层的梯度特别大，超过了FP16的表示范围。后来加了Loss Scaling，并把缩放因子从128调到512，问题就解决了。

1.4 知识体系总览

下面这张图，是我自己画的AMP知识体系。你看一眼，就能明白整个流程。

嗯，这张图把AMP的核心要素都串起来了。从数据类型选择，到策略配置，再到Loss Scaling，每一步都有讲究。

想要更多资料，联系可可