AI混合精度折腾手记

import matplotlib.pyplot as plt
import numpy as np

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 创建对比数据
formats = {
    'FP32': {'exp': 8, 'mant': 23, 'range': 38},
    'BF16': {'exp': 8, 'mant': 7, 'range': 38},
    'FP16': {'exp': 5, 'mant': 10, 'range': 4}
}

# 创建图形
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

# 左图：指数位对比
ax1.bar(formats.keys(), [v['exp'] for v in formats.values()], color=['#1f77b4', '#ff7f0e', '#2ca02c'])
ax1.set_ylabel('指数位 (Bits)')
ax1.set_title('动态范围能力（指数位）')
ax1.set_ylim(0, 10)

# 右图：尾数位对比
ax2.bar(formats.keys(), [v['mant'] for v in formats.values()], color=['#1f77b4', '#ff7f0e', '#2ca02c'])
ax2.set_ylabel('尾数位 (Bits)')
ax2.set_title('数值精度能力（尾数位）')
ax2.set_ylim(0, 25)

plt.tight_layout()
plt.savefig('/home/liqinsi/Documents/project/thinkblog/static/images/posts/274-ai-mixed-precision-fp32-fp16-bf16-practice/precision-comparison.webp', dpi=150, bbox_inches='tight')

import torch
from torch.cuda.amp import autocast, GradScaler

# 模型和优化器
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# GradScaler 仅在 FP16 时需要
use_bf16 = torch.cuda.is_bf16_supported()
scaler = None if use_bf16 else GradScaler()

for batch in dataloader:
    optimizer.zero_grad()

    # 前向传播自动使用混合精度
    with autocast(dtype=torch.bfloat16 if use_bf16 else torch.float16):
        loss = model(batch)

    # 反向传播
    if use_bf16:
        loss.backward()
        optimizer.step()
    else:
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

# 错误示范：不调用 update()
scaler.scale(loss).backward()
scaler.step(optimizer)
# 缺少 scaler.update() 会导致缩放因子一直不更新

# 正确写法
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# 自定义 Loss Scaling 策略
scaler = GradScaler(
    init_scale=2.0,      # 初始缩放因子，根据你的 loss 范围调整
    growth_factor=2.0,   # 每次成功后的增长倍数
    backoff_factor=0.5,  # 失败后的回退倍数
    growth_interval=2000 # 连续成功多少次后才增长
)

# 更安全的 BF16 检测方式
def safe_bf16_supported():
    if not torch.cuda.is_available():
        return False
    device = torch.cuda.current_device()
    device_name = torch.cuda.get_device_name(device)
    # 检查已知支持 BF16 的显卡
    bf16_gpus = ['A100', 'RTX 40', 'H100', 'L40']
    return any(gpu in device_name for gpu in bf16_gpus) and torch.cuda.is_bf16_supported()

use_bf16 = safe_bf16_supported()

# 错误示范：自定义算子没有考虑精度
def custom_attention(q, k, v):
    # 默认用 FP32 计算注意力分数
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
    attn = torch.softmax(scores, dim=-1)
    return torch.matmul(attn, v)

# 正确做法：显式处理精度
def custom_attention(q, k, v):
    with autocast():
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
        attn = torch.softmax(scores, dim=-1)
        return torch.matmul(attn, v)

# 手动清理中间结果
for batch in dataloader:
    optimizer.zero_grad()

    with autocast():
        loss = model(batch)

    loss.backward()

    # 清理不需要的中间结果
    torch.cuda.empty_cache()

    optimizer.step()

class StableLayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-12):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.bias = nn.Parameter(torch.zeros(dim))
        self.eps = eps

    def forward(self, x):
        # LayerNorm 的均值和方差计算强制用 FP32
        mean = x.mean(-1, keepdim=True, dtype=torch.float32)
        var = x.var(-1, keepdim=True, unbiased=False, dtype=torch.float32)
        x = (x - mean) / torch.sqrt(var + self.eps)
        return (self.weight * x + self.bias).to(x.dtype)