一、混合精度训练（AMP）：FP16/BF16原理、AMP实现机制、PyTorch实践与通信影响
大家好，我是你们这节课程的讲师。今天咱们聊聊混合精度训练，也就是AMP。说实话，这玩意儿刚出来的时候，我第一反应是「这不就是偷工减料吗？」——把32位浮点数砍成16位，精度能行吗？后来真在项目里用上了，才发现自己格局小了。
混合精度训练，说白了就是让模型在训练时「该省省，该花花」。关键数据用高精度，次要数据用低精度。这样既省显存，又跑得快。我最早在训练一个8B参数的LLM时，单卡根本放不下，用了AMP之后，显存直接砍半，训练速度还快了将近一倍。嗯，真香。
核心观点：混合精度训练不是简单的精度降低，而是通过「关键路径高精度、非关键路径低精度」的策略，在保证模型收敛质量的前提下，最大化利用硬件算力。
1.1 FP16与BF16：两种半精度格式的对比
先说说FP16和BF16。这两个都是16位浮点数，但内部结构完全不同。
| 格式 | 符号位 | 指数位 | 尾数位 | 表示范围 | 精度 |
|---|---|---|---|---|---|
| FP32 | 1 | 8 | 23 | ~3.4e38 | 高 |
| FP16 | 1 | 5 | 10 | ~65504 | 中等 |
| BF16 | 1 | 8 | 7 | ~3.4e38 | 较低 |
看到区别了吗？FP16的指数位只有5位，所以它能表示的最大值只有65504。而BF16的指数位和FP32一样是8位，所以表示范围跟FP32一样大，只是尾数位少了，精度低一些。
我在项目中遇到过一个问题：用FP16训练时，梯度动不动就溢出变成NaN。后来排查发现，是某个层的激活值太大了，直接超过了65504。换成BF16之后，这个问题就再没出现过。所以我的建议是：如果你的硬件支持BF16（比如A100、H100），优先用BF16。它的动态范围大，不容易溢出。
个人经验：FP16适合小模型或激活值范围可控的场景；BF16适合大模型或训练不稳定时。我一般默认用BF16，除非硬件不支持。
1.2 AMP的实现机制：Master Weights与Loss Scaling
AMP的核心机制有两个：Master Weights和Loss Scaling。咱们一个一个说。
Master Weights（主权重）
训练过程中，权重会不断更新。如果权重本身是FP16的，每次更新时都会累积舍入误差。时间一长，模型就偏了。
解决办法是：保留一份FP32的「主权重」。每次更新时，用FP32的主权重计算梯度更新，然后再把更新后的FP32权重转成FP16用于前向传播。这样既享受了FP16的计算速度，又保证了权重的精度。
我习惯把主权重叫做「保险柜版本」——平时干活用FP16，但真正的「家底」存在FP32里。每次更新完，再同步一下。
Loss Scaling（损失缩放）
这个机制是为了解决梯度下溢的问题。FP16能表示的最小正数是约6e-8。如果梯度比这个还小，就会直接变成0。这在训练后期很常见，尤其是那些小梯度，一缩就没了。
Loss Scaling的做法是：在前向传播时，把loss乘以一个缩放因子（比如1024），让梯度也跟着变大。反向传播完成后，再把梯度除以这个缩放因子，恢复原样。
我曾经在训练一个BERT模型时，发现某些层的梯度全是0。排查了半天，最后发现是Loss Scaling没开。开了之后，那些层的梯度就正常了。嗯，这个坑我踩过，你们别踩。
注意：Loss Scaling的缩放因子不是越大越好。太大可能导致梯度溢出（变成Inf），太小又解决不了下溢。PyTorch的AMP会自动调整这个值，但如果你手动设置，建议从2^16开始，观察梯度是否溢出，再逐步调整。
1.3 AMP在PyTorch中的实践
PyTorch从1.6开始原生支持AMP，用起来非常简单。我直接上代码。
import torch
from torch.cuda.amp import autocast, GradScaler
# 初始化模型、优化器、损失函数
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = torch.nn.CrossEntropyLoss()
# 初始化GradScaler
scaler = GradScaler()
for data, target in dataloader:
data, target = data.cuda(), target.cuda()
# 清零梯度
optimizer.zero_grad()
# 使用autocast自动管理精度
with autocast():
output = model(data)
loss = criterion(output, target)
# 反向传播时使用scaler
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
这段代码里，autocast会自动把前向传播中的操作转成FP16或BF16，而GradScaler负责Loss Scaling。你什么都不用操心，PyTorch帮你搞定。
我个人习惯把scaler.step(optimizer)和scaler.update()放在一起。有一次我忘了调scaler.update()，结果缩放因子一直不变，训练到一半梯度全炸了。所以记住：step和update是成对出现的。
避坑指南：我曾经在分布式训练中用了AMP，结果发现每个GPU的缩放因子不一样。后来才知道，GradScaler需要每个进程独立维护，不能共享。所以分布式训练时，每个进程都要自己创建一个scaler。
1.4 AMP对通信的影响
AMP不仅影响计算，还影响通信。这一点很多人容易忽略。
在多卡训练中，梯度同步是通信的主要开销。如果梯度是FP32的，每次同步需要传输4字节/参数。如果梯度是FP16的，只需要2字节/参数。通信量直接减半。
但这里有个细节：AMP在反向传播时，梯度是FP16的。但在梯度同步之前，PyTorch会把梯度转成FP32再进行AllReduce。为什么？因为FP16的精度太低，多次AllReduce后误差会累积。
我建议的做法是：在分布式训练中，使用torch.distributed的allreduce_hook，把梯度压缩成FP16再通信。这样通信量减半，而且精度损失可以忽略。
import torch.distributed as dist
from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
# 注册FP16压缩的通信钩子
model.register_comm_hook(
state=None,
hook=default_hooks.fp16_compress_hook
)
这段代码会在梯度同步时自动把FP32压缩成FP16，通信完再解压回FP32。我在一个8卡训练任务中试过，通信时间减少了40%，而且模型精度完全没受影响。
核心结论：AMP对通信的影响是双面的。一方面，它减少了通信量（梯度从FP32变成FP16）；另一方面，它增加了额外的精度转换开销。但总体来看，收益远大于成本。尤其是在大模型训练中，通信瓶颈更明显，AMP的加速效果也更显著。
1.5 本章小结
混合精度训练不是魔法，而是一套精密的工程策略。FP16和BF16各有优劣，Master Weights保证了权重精度，Loss Scaling解决了梯度下溢。PyTorch的AMP让这一切变得简单，但背后的原理你得懂，否则出了问题都不知道怎么排查。
最后，我建议你在自己的项目里试试AMP。先跑一个小实验，对比一下开启和关闭AMP的显存占用、训练速度和模型精度。你会发现，AMP带来的收益远超你的预期。
想要更多资料，联系可可