11.27课堂作业：显存优化与混合精度训练

核心思想

• 保持FP32精度的主权重副本用于精度保证
• 使用FP16进行前向和反向传播节省显存
• 累加梯度时使用FP32精度防止精度损失

PyTorch实现

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

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

for data, target in train_loader:
    optimizer.zero_grad()
    
    # 自动混合精度上下文
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    # 缩放损失以防止梯度下溢
    scaler.scale(loss).backward()
    
    # 优化器步进（自动缩放回原始精度）
    scaler.step(optimizer)
    scaler.update()

原理

• 在反向传播前放大损失值，防止小梯度下溢
• 梯度更新时缩小回原始尺度，保持数值稳定

实现示例

def manual_gradient_scaling(model, loss, scale_factor=1024.0):
    """手动梯度缩放实现"""
    scaled_loss = loss * scale_factor
    scaled_loss.backward()
    
    # 梯度缩放回原始值
    for param in model.parameters():
        if param.grad is not None:
            param.grad = param.grad / scale_factor

自适应策略

• 自动调整损失缩放因子
• 检测梯度溢出并动态调整
• 平衡数值稳定性和精度

算法伪代码

class DynamicLossScaler:
    def __init__(self, init_scale=65536.0, scale_factor=2.0, scale_window=2000):
        self.scale = init_scale
        self.scale_factor = scale_factor
        self.scale_window = scale_window
        self.counter = 0
    
    def step(self, model, optimizer):
        # 检查梯度溢出
        has_inf = torch.isinf(torch.cat([p.grad.flatten()
                                        for p in model.parameters()
                                        if p.grad is not None])).any()
        
        if has_inf:
            # 梯度溢出，减小缩放因子
            self.scale /= self.scale_factor
            optimizer.zero_grad()
            return False  # 需要重新前向传播
        else:
            # 梯度正常，更新缩放因子
            self.counter += 1
            if self.counter >= self.scale_window:
                self.scale *= self.scale_factor
                self.counter = 0
            return True  # 可以继续优化

BF16（BFloat16）格式

• 保持FP32的动态范围
• 精度略低于FP16但比FP16更适合训练
• 数值范围：10⁻³⁸到10³⁸（与FP32相同）

实现示例

# 使用BF16格式（需要硬件支持）
model = model.to(dtype=torch.bfloat16)

# 在支持BF16的硬件上训练
for data, target in train_loader:
    data = data.to(dtype=torch.bfloat16)
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

数值稳定的网络设计

class NumericalStableNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 使用数值稳定的激活函数
        self.activation = nn.ReLU()  # 而非nn.Tanh()
        self.norm = nn.LayerNorm(eps=1e-6)  # 增加数值稳定性
    
    def forward(self, x):
        x = self.activation(x)
        x = self.norm(x)
        return x

# 数值稳定的softmax实现
def stable_softmax(x, dim=1):
    """数值稳定的softmax，防止上溢下溢"""
    x_max = torch.max(x, dim=dim, keepdim=True)[0]
    x_exp = torch.exp(x - x_max)  # 减去最大值防止上溢
    return x_exp / torch.sum(x_exp, dim=dim, keepdim=True)

Tensor Core优化

• 利用现代GPU的Tensor Core硬件加速
• 选择合适的计算精度配置
• 优化内存访问模式

针对不同硬件的精度选择

def get_optimal_dtype():
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name()
        if 'A100' in gpu_name or 'V100' in gpu_name:
            return torch.bfloat16  # 新一代GPU支持BF16
        else:
            return torch.float16   # 老一代GPU使用FP16
    return torch.float32