大模型训练优化实战:提升效率的 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 | ⭐⭐⭐⭐ | ⭐⭐ | ⭐⭐⭐ |
*实践出真知,欢迎分享你的优化经验!*