混合精度训练如何降低显存

混合精度训练（Mixed Precision Training）主要通过减少数据占用的内存空间和利用硬件加速特性来降低显存使用。以下是详细解释：

一、核心原理

1. 使用更低精度的数据类型

FP32（单精度）：32位浮点数，每个参数占 4字节
FP16（半精度）：16位浮点数，每个参数占 2字节
BF16（脑浮点）：16位，每个参数占 2字节

将模型参数、激活值从 FP32 转为 FP16/BF16，理论上显存直接减半。

二、具体降低显存的机制

1. 参数存储减半

FP32 参数: 1亿参数 × 4字节 = 400MB
FP16 参数: 1亿参数 × 2字节 = 200MB

2. 梯度存储减半

反向传播时梯度同样以 FP16 存储，显存占用再减半。

3. 激活值（Activations）减半

前向传播的激活值若用 FP16 存储，同样减少显存。

4. 优化器状态优化（关键！）

以 Adam 优化器为例：

存储内容	FP32训练	混合精度训练
参数	4字节	2字节（FP16）+ 4字节（FP32副本）
梯度	4字节	2字节（FP16）
一阶动量m	4字节	4字节（FP32）
二阶动量v	4字节	4字节（FP32）

虽然参数有 FP32 副本，但梯度 + 激活值的大幅减少仍使总显存显著降低。

三、实际显存节省效果

组件	FP32	混合精度	节省
模型参数	4x	2x + 4x（副本）	~33%
梯度	4x	2x	50%
优化器状态	8x	8x	0%
激活值	4x	2x	50%

综合来看，混合精度训练通常可节省 30%~50% 显存，具体取决于模型结构。

四、为什么不会严重损失精度？

1. 使用 FP32 主权重副本（Master Weights）

前向/反向: FP16（快速、省显存）
参数更新: 转回 FP32（保持精度）

更新后再将 FP32 参数 cast 回 FP16。

2. Loss Scaling（损失缩放）

FP16 存在下溢问题（接近0的小梯度会变为0）
对 loss 乘以一个缩放因子（如 1024、2048）
反向传播后再除回来，保留小梯度信息

五、代码示例（PyTorch）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data, target in dataloader:
optimizer.zero_grad()
with autocast():  # 自动混合精度
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()  # 缩放loss，防止下溢
scaler.step(optimizer)
scaler.update()

六、硬件支持

硬件	支持情况
NVIDIA V100	支持 FP16 TensorCore
NVIDIA A100	支持 FP16 + BF16 + TF32
NVIDIA H100	支持 FP8 + FP16 + BF16
AMD MI250	支持 FP16

TensorCore 在 FP16 下计算吞吐量通常是 FP32 的 2~8倍，既省显存又提速。

七、总结

混合精度训练降低显存的路径：
FP32 → FP16/BF16
├── 参数显存 ↓ 50%
├── 梯度显存 ↓ 50%
├── 激活值显存 ↓ 50%
└── 配合优化技术保持精度

一句话：混合精度训练通过用 16 位浮点数替代 32 位浮点数，在保证训练精度的前提下，直接大幅减少参数、梯度和激活值的显存占用。