开源大模型训练中optimizer配置错误记录
在进行开源大模型微调过程中，optimizer配置错误是常见但容易被忽视的问题。本文记录一次典型的optimizer配置错误及其排查过程。
问题现象
使用Hugging Face Transformers库进行LLaMA模型微调时，训练loss不下降，甚至出现NaN值。通过日志分析发现optimizer状态异常。
错误配置示例
from transformers import AutoTokenizer, AutoModelForCausalLM
from torch.optim import AdamW
model = AutoModelForCausalLM.from_pretrained("huggyllama/llama-7b")
optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
问题分析
主要问题在于：
- 学习率设置不当：对于大模型，1e-4的学习率过高导致训练不稳定
- 权重衰减配置错误：在某些优化器中，weight_decay的默认行为可能与预期不符
- 梯度裁剪缺失：未设置梯度裁剪防止梯度爆炸
正确配置方案
from transformers import AdamW, get_linear_schedule_with_warmup
import torch.nn.utils.clip_grad_norm_
# 1. 使用更小的学习率
optimizer = AdamW(model.parameters(), lr=5e-6, weight_decay=0.01)
# 2. 添加学习率调度器
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=1000,
num_training_steps=total_steps
)
# 3. 添加梯度裁剪
for step, batch in enumerate(dataloader):
outputs = model(**batch)
loss = outputs.loss
loss.backward()
# 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
验证方法
通过监控训练过程中的loss变化、梯度范数以及模型参数更新幅度来验证配置正确性。建议使用tensorboard或wandb进行可视化监控。
总结
在大模型微调中，optimizer配置需要根据模型规模、数据集大小和任务类型进行精细调整。建议遵循以下原则：
- 使用较小的学习率（5e-6到1e-5）
- 启用学习率预热
- 添加梯度裁剪
- 定期验证训练稳定性
讨论