大模型训练优化实战：提升效率的 10 个技巧

训练大模型，效率就是金钱。本文分享 10 个实战优化技巧。

---

技巧 1：混合精度训练

效果： 显存减少 50%，速度提升 2-3 倍

实现：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    loss = model(inputs)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

---

技巧 2：梯度累积

效果： 等效于更大的 batch size，无需增加显存

实现：

accumulation_steps = 4

for i, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

---

技巧 3：激活值重计算

效果： 显存减少 60-70%，速度略降

适用： 显存紧张时的训练

实现：

import torch.utils.checkpoint as checkpoint

output = checkpoint.checkpoint(model_layer, input)

---

技巧 4：优化数据加载

问题： GPU 等待数据，利用率低

解决：

dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=8,
    pin_memory=True,
    prefetch_factor=2
)

---

技巧 5：ZeRO 优化

效果： 显存大幅降低，支持更大模型

DeepSpeed ZeRO：

• ZeRO-1：优化器分片
• ZeRO-2：梯度分片
• ZeRO-3：参数分片

---

技巧 6：混合并行策略

组合拳：

• 数据并行（DP）
• 张量并行（TP）
• 流水线并行（PP）

推荐配置：

• 8 卡机器：DP=8 或 DP=4+TP=2
• 多机：DP+TP+PP 组合

---

技巧 7：学习率调度

Warmup + Decay：

from transformers import get_cosine_schedule_with_warmup

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=1000,
    num_training_steps=total_steps
)

---

技巧 8：梯度裁剪

防止梯度爆炸：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

---

技巧 9：检查点策略

智能保存：

• 每 N 个 epoch 保存一次
• 只保留最近的 K 个检查点
• 遇到最优指标时保存

---

技巧 10：监控与调试

关键指标：

• GPU 利用率（目标 > 80%）
• 显存使用率
• 训练 loss 曲线
• 梯度范数

工具推荐：

• WandB / TensorBoard
• NVIDIA-smi
• PyTorch Profiler

---

总结

技巧	显存优化	速度提升	实现难度
混合精度	⭐⭐⭐	⭐⭐⭐	⭐
梯度累积	⭐⭐	⭐	⭐
激活重计算	⭐⭐⭐	⭐	⭐⭐
数据加载优化	⭐	⭐⭐	⭐
ZeRO	⭐⭐⭐⭐	⭐⭐	⭐⭐⭐

建议： 从简单的开始，逐步叠加优化。

---

实践出真知，欢迎分享你的优化经验！