← 返回文章列表

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

大模型训练优化实战:提升效率的 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⭐⭐⭐⭐⭐⭐⭐⭐⭐
建议: 从简单的开始,逐步叠加优化。


*实践出真知,欢迎分享你的优化经验!*