共计 1745 个字符，预计需要花费 5 分钟才能阅读完成。
1. 为什么 4090 的 FP16 算力值得关注
NVIDIA RTX 4090 采用 Ada Lovelace 架构，其 FP16（半精度浮点）计算能力达到 330 TFLOPS，是 FP32（单精度）的 4 倍。这种算力优势特别适合：
- 大规模矩阵运算（如 Transformer 注意力机制）
- 显存密集型任务（可减少 50% 显存占用）
- 实时推理场景（需要低延迟）
实际测试中，ResNet50 训练在启用 FP16 后，4090 的吞吐量比 FP32 模式提升 1.8-2.3 倍。
2. 核心概念解析
2.1 FP16 vs FP32
| 类型 | 位数 | 指数位 | 小数位 | 数值范围 |
|---|---|---|---|---|
| FP32 | 32 | 8 | 23 | ±1.18×10⁻³⁸~3.4×10³⁸ |
| FP16 | 16 | 5 | 10 | ±6.1×10⁻⁵~6.5×10⁴ |
FP16 的缺点是：
– 容易数值溢出（超过 6.5×10⁴会变成 inf）
– 容易下溢（小于 6.1×10⁻⁵会变成 0）
2.2 混合精度训练原理
现代框架采用三管齐下的策略：
- 权重保留 FP32 副本 ：作为主权重（Master Weights）
- 自动梯度缩放 ：通过 Loss Scaling 防止梯度下溢
- 动态精度转换 ：仅在矩阵乘法等操作使用 FP16
3. PyTorch 实现方案
3.1 基础代码框架
import torch
from torch.cuda.amp import GradScaler, autocast
# 初始化（必须！）scaler = GradScaler() # 自动梯度缩放
for epoch in range(epochs):
for inputs, labels in dataloader:
inputs, labels = inputs.cuda(), labels.cuda()
# 前向传播（自动混合精度）with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播（自动处理精度转换）scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()3.2 关键配置说明
- GradScaler：初始 scale=65536，每 2000 步自动调整
- autocast 区域 ：应包含前向计算和 loss 计算
- 优化器 ：推荐使用 AdamW（对精度不敏感）
4. 常见问题与解决方案
4.1 梯度爆炸 / 消失
现象 ：Loss 变成 NaN 或震荡剧烈
解决方法 ：
- 检查 scaler 的 scale 值：
print(scaler.get_scale()) - 减小初始学习率（通常为 FP32 的 0.5-0.8 倍）
- 添加梯度裁剪：
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
4.2 显存不足
现象 ：明明用了 FP16 却报 OOM
根本原因 ：部分操作强制使用 FP32（如 softmax）
检查方法 ：
from torch.cuda.amp import custom_fwd, custom_bwd
@custom_fwd(cast_inputs=torch.float32) # 显式声明需要 FP32 的操作
def danger_op(x):
return x.softmax(dim=1)5. 性能对比实测
在 BERT-base 模型上的测试结果（batch_size=32）：
| 指标 | FP32 模式 | FP16 模式 | 提升幅度 |
|---|---|---|---|
| 单 step 耗时 | 580ms | 320ms | 44.8% |
| 显存占用 | 9.8GB | 5.2GB | 46.9% |
| 最高温度 | 72℃ | 68℃ | -4℃ |
6. 调优建议
- 学习率策略 ：
- 初始值 = FP32 学习率 × 0.6
-
配合线性 warmup（前 10% 训练步）
-
监控指标 ：
# 在 validation 阶段关闭 autocast 获得准确指标 with torch.no_grad(): outputs = model(inputs.float()) # 显式转为 FP32 -
模型特定调整 ：
- LayerNorm 保持 FP32 计算
- 输出层建议保留 FP32
延伸思考
- 为什么某些模型（如 GAN）不适合 FP16 训练？
- 如何通过 NSight 工具分析 FP16 的计算效率？
- 在模型量化和 FP16 加速之间如何选择？
实践建议：先用 FP32 训练 1 个 epoch 作为 baseline，再开启 FP16 对比收敛曲线。遇到问题时，可以逐层禁用 FP16（
torch.nn.Layer.float()）进行问题定位。
正文完
