5090 fp16算力入门指南：从零开始掌握混合精度训练

1次阅读

没有评论

共计 1984 个字符，预计需要花费 5 分钟才能阅读完成。

fp16（半精度浮点数）训练在深度学习领域越来越受欢迎，主要因为它能带来两大核心优势：

显存占用减半：相比传统的 fp32（单精度），fp16 的每个参数仅占用 2 字节，理论上可节省 50% 显存，这对于大模型训练尤为重要
计算速度提升：现代 GPU（如 5090）的 Tensor Core 针对 fp16 计算做了专门优化，吞吐量可达 fp32 的 2 - 8 倍

但新手直接使用 fp16 可能会遇到：

数值下溢：小于 6e- 8 的值会变为 0，导致梯度消失
精度损失：某些计算（如指数运算）在 fp16 下误差会放大
训练不稳定：梯度值范围过大 / 过小都可能引发 NaN

推荐使用 PyTorch 2.0+ 的 AMP（自动混合精度）工具链：

确认驱动版本（需≥545.23）
```
nvidia-smi | grep "Driver Version"
```

安装 CUDA 12.1+ 和对应 PyTorch

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

验证 Tensor Core 可用性

import torch
print(torch.cuda.get_device_capability())  # 应显示 (8,9) 或更高

由于 fp16 的表示范围有限（±65,504），梯度值可能超出这个范围。梯度缩放通过以下步骤解决：

前向计算使用 fp16
损失值乘以缩放因子（scale factor，通常 8 -32k）
反向传播得到放大的梯度
将梯度反缩放后更新参数

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast

# 初始化缩放器
scaler = GradScaler()

model = ResNet18().cuda()
optimizer = optim.SGD(model.parameters(), lr=0.1)
criterion = nn.CrossEntropyLoss()

for epoch in range(100):
    for inputs, targets in train_loader:
        inputs, targets = inputs.cuda(), targets.cuda()

        optimizer.zero_grad()

        # 前向传播（自动选择 fp16/fp32）with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, targets)

        # 反向传播 + 梯度缩放
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

测试环境：
– GPU: NVIDIA 5090 (24GB)
– CUDA: 12.1
– Batch Size: 256

精度	显存占用	吞吐量(imgs/sec)	训练时间 /epoch
FP32	18.7GB	1,200	3m21s
FP16	9.2GB	3,800	1m02s

初始 scale 设为 8192
每 2000 步检查 NaN 情况

遇到 NaN 时：

if torch.isnan(loss):
    scaler.update(new_scale=scaler.get_scale()*0.8)

出现 NaN 的检查步骤：
先关闭 AMP 验证是否模型问题
检查第一个出现 NaN 的 layer
尝试减小 learning rate
不适用 fp16 的算子：
指数运算（如 Softmax）
大型累加（如 BatchNorm）
小数值范围运算（如 Log）

当 batch_size 受显存限制时，可通过梯度累积模拟更大 batch：

accum_steps = 4

for i, (inputs, targets) in enumerate(train_loader):
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets) / accum_steps

    scaler.scale(loss).backward()

    if (i+1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

通过合理使用 fp16 混合精度训练，我在 5090 上实现了：
– 训练速度提升 3.2 倍
– 显存占用减少 51%
– 最终模型准确率保持与 fp32 相当

建议初学者从 CIFAR-10 等小数据集开始实践，逐步掌握 scale 调整和 NaN 处理技巧。遇到问题时，PyTorch 的 torch.autograd.detect_anomaly() 是非常有用的调试工具。

正文完