一、混合精度训练（AMP）：FP16/BF16原理、AMP的GradScaler机制、Loss Scaling技巧、硬件支持与Tensor Cores利用
1.1 为什么我们需要混合精度？
先说说我个人的经历。几年前我在训练一个百亿参数的NLP模型，显存32GB，结果模型刚加载完就OOM了。当时我盯着nvidia-smi发呆，心想这还怎么玩？后来切换到混合精度训练，显存占用直接砍半，训练速度还快了2倍多。
说白了，混合精度训练的核心思路就一句话：用更少的内存存数据，用更快的硬件算数据。但这里有个坑——精度降低会导致梯度消失或溢出，所以需要一些技巧来兜底。
1.2 FP16与BF16：两种低精度格式
先看两张图，理解FP16和BF16的位宽分布。
看到区别了吗？BF16保留了和FP32一样的8位指数，所以它的动态范围跟FP32一样大。这意味着什么？你想想看，训练过程中梯度爆炸或消失的概率会小很多。而FP16只有5位指数，范围窄，容易溢出。
我的经验：如果显卡支持BF16（比如A100、H100），我无脑选BF16。FP16在训练大模型时经常需要频繁调Loss Scaling，而BF16基本不用操心。
1.3 AMP的GradScaler机制
AMP（Automatic Mixed Precision）的核心组件就是GradScaler。它的工作流程是这样的：
- 前向传播：用FP16/BF16计算，节省显存和带宽
- Loss计算：用FP32计算，保证精度
- Loss Scaling：把Loss乘以一个缩放因子（比如2^16），防止梯度下溢
- 反向传播：在FP16下计算缩放后的梯度
- 梯度反缩放：把梯度除以缩放因子，恢复真实值
- 优化器更新：用FP32的Master Weights更新参数
代码实现其实很简单，PyTorch已经封装好了：
import torch
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data, target in dataloader:
optimizer.zero_grad()
with autocast():
output = model(data)
loss = loss_fn(output, target)
# 关键步骤：用scaler缩放loss
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()嗯，这里要注意：scaler.update()会在每个step结束后检查是否有梯度溢出。如果发现溢出，它会跳过参数更新，并自动增大缩放因子。我遇到过好几次这种情况，训练刚开始时loss突然变成NaN，其实就是梯度溢出了，scaler会自动处理。
1.4 Loss Scaling技巧
Loss Scaling说白了就是给Loss乘一个大数，让梯度在FP16的表示范围内。但乘多大合适？
| 策略 | 缩放因子 | 适用场景 |
|---|---|---|
| 保守策略 | 2^8 = 256 | 小模型、稳定训练 |
| 激进策略 | 2^16 = 65536 | 大模型、梯度容易消失 |
| 动态调整 | 自动增长/衰减 | AMP默认策略 |
避坑指南：我曾经在训练一个Transformer模型时，手动把缩放因子设成了2^20，结果梯度全部变成0。后来才发现，缩放因子太大导致梯度在反缩放时精度丢失。建议用动态调整，让框架自己决定。
1.5 硬件支持与Tensor Cores利用
Tensor Cores是NVIDIA从Volta架构开始引入的专用硬件单元。它们专门用来做混合精度矩阵乘法，速度比普通CUDA Core快好几倍。
要充分利用Tensor Cores，需要满足几个条件：
- 维度对齐：矩阵维度最好是8的倍数（FP16）或16的倍数（BF16）
- 数据类型：输入必须是FP16或BF16，累加可以用FP32
- 算子支持：卷积、全连接、LayerNorm等常用算子都支持
我习惯在模型定义时加一个维度对齐的检查：
def align_dim(dim, alignment=8):
"""将维度对齐到alignment的倍数"""
return ((dim + alignment - 1) // alignment) * alignment
# 在Linear层中使用
class AlignedLinear(nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
self.in_features = align_dim(in_features)
self.out_features = align_dim(out_features)
self.weight = nn.Parameter(torch.randn(self.out_features, self.in_features))
def forward(self, x):
return F.linear(x, self.weight)注意：不是所有算子都适合用Tensor Cores。比如LayerNorm、Softmax这些逐元素操作，用FP32算反而更快。AMP的autocast会自动处理这些细节，你不需要手动干预。
1.6 知识体系总览
下面这张图总结了混合精度训练的核心逻辑：
你看，这四个组件缺一不可。数据格式决定了你能省多少显存，GradScaler和Loss Scaling保证了训练稳定性，Tensor Cores则把硬件性能榨干。我刚开始用AMP时只开了FP16，结果模型训到一半就炸了。后来加上GradScaler，世界就清净了。
最后说一句：混合精度训练现在已经不是「高级技巧」了，而是大模型训练的标配。如果你还在用纯FP32训练，赶紧试试AMP，你会回来感谢我的。