大模型训练成本控制：混合精度训练（AMP）与梯度累积的实践技巧

在大模型训练中，显存占用和计算效率是成本控制的核心挑战。混合精度训练（Automatic Mixed Precision, AMP）和梯度累积是两种关键优化技术，以下从原理到实践逐步解析：

一、混合精度训练（AMP）

原理：
通过组合$fp16$（半精度）和$fp32$（单精度），在保持数值稳定性的同时加速计算：

• 前向传播和梯度计算使用$fp16$，减少50%显存占用
• 权重更新和累加器使用$fp32$，避免下溢误差
• 动态损失缩放（Loss Scaling）补偿$fp16$的精度损失

实践技巧：

1. 框架选择：
   PyTorch使用torch.cuda.amp，TensorFlow使用tf.train.experimental.enable_mixed_precision_graph_rewrite

   # PyTorch示例
   from torch.cuda.amp import autocast, GradScaler
   scaler = GradScaler()
   
   with autocast():
       outputs = model(inputs)
       loss = criterion(outputs, labels)
   scaler.scale(loss).backward()
   scaler.step(optimizer)
   scaler.update()
   

2. 超参调优：

   • 初始缩放因子设为$2^{16}$，根据梯度溢出动态调整
   • 监控NaN出现频率，若>5%需增大缩放因子

3. 稳定性保障：

   • 对Softmax、LayerNorm等敏感操作保留$fp32$
   • 使用amp.register_float_function自定义精度规则

二、梯度累积（Gradient Accumulation）

原理：
将大批量拆分为小批量，多次前向后累积梯度再更新：
$$ \nabla W_{\text{accum}} = \sum_{i=1}^{N} \nabla W_i $$
其中$N$为累积步数，等效批量大小$=N \times \text{micro_batch_size}$

实践技巧：

1. 显存与吞吐平衡：

   • 累积步数$N$满足：$ \text{GPU显存} \propto \frac{1}{N} $
   • 典型配置：$ \text{micro_batch_size}=8, N=4 \to \text{等效batch_size}=32 $

2. 代码实现：

   optimizer.zero_grad()
   for step in range(accum_steps):
       inputs = next_batch()
       with autocast():
           loss = model(inputs) / accum_steps  # 损失归一化
       scaler.scale(loss).backward()  # 梯度累积
   scaler.step(optimizer)  # 累积后更新
   scaler.update()
   

3. 学习率调整：

   • 等效批量增大$k$倍时，学习率需缩放$\sqrt{k}$倍
   • 使用线性预热：$ \eta_t = \eta_{\max} \times \frac{t}{T_{\text{warmup}}} $

三、联合优化策略

1. AMP+梯度累积协同：

   • AMP降低单步显存，梯度累积突破单卡批量限制
   • 总显存节省可达$4\times$（以GPT-3 175B为例）

2. 收敛性保障：

   • 每$N$步同步一次BatchNorm统计量
   • 梯度裁剪阈值按累积步数缩放：$ \text{clip_threshold} \times \sqrt{N} $

3. 通信优化：

   • 在累积结束后调用optimizer.step()，减少分布式通信次数

四、实验效果对比

注：实际收益与模型结构相关，Transformer类模型加速比可达$2.8\times$

五、避坑指南

1. 数值不稳定：

   • AMP下出现NaN：检查损失缩放，添加torch.isnan(grad).any()监控
   • 梯度累积时使用loss.mean()而非loss.sum()

2. 超参陷阱：

   • 学习率需随等效批量二次根缩放，而非线性缩放
   • 避免在累积中途调用optimizer.zero_grad()

3. 硬件适配：

   • Volta+架构GPU（如V100/A100）才支持Tensor Core加速
   • 使用nvidia-smi dmon监控显存与利用率波动

通过合理组合AMP与梯度累积，可在显存减少$4\times$的条件下保持90%+训练速度，显著降低大模型训练成本。