LLM 大语言模型 训练的时候 batchsize 调整大导致梯度爆炸问题解决
优化器AdamW 确实比 SGD 更容易在大 batch 下梯度爆炸,因为自适应学习率会放大稀疏梯度的更新步长。
针对 AdamW + 大 batch,给你几个立竿见影的修复方案:
1. 优化器参数调整(最快见效)
python
from transformers import TrainingArguments
training_args = TrainingArguments(
per_device_train_batch_size=128,
learning_rate=1e-4, # 线性缩放:5e-5 * (128/64)
# AdamW 关键参数
optim="adamw_torch", # 或 "adamw_hf" 更稳定
adam_beta1=0.9,
adam_beta2=0.95, # 从 0.999 降到 0.95,减少历史梯度记忆
adam_epsilon=1e-6, # 从 1e-8 提高,防止除零爆炸
weight_decay=0.01, # 保持,但注意和 lr 的平衡
max_grad_norm=0.5, # 必须收紧!
warmup_ratio=0.15, # 15% 步数 warmup,给 AdamW 缓冲期
lr_scheduler_type="cosine",
)
2. 如果还炸:换 8-bit AdamW 或 Lion
python
# 方案 A:8-bit AdamW(显存省 + 更稳)
optim="adamw_8bit", # bitsandbytes
adam_beta2=0.95,
# 方案 B:Lion 优化器(对大 batch 天生更稳)
# 需要 pip install lion-pytorch,然后在 LLaMA-Factory 或自定义 Trainer 里用
# Lion 的 lr 通常是 AdamW 的 1/3~1/10,比如 1e-4 → 3e-5
3. 分层学习率(Embedding/LM_head 降速)
DeepSeek 的 Embedding 和 LM_head 参数量大,最容易炸:
python
from transformers import Trainer
import torch
class LayerwiseTrainer(Trainer):
def create_optimizer(self):
# 分层参数组
embed_params = []
other_params = []
for name, param in self.model.named_parameters():
if not param.requires_grad:
continue
if "embed" in name or "lm_head" in name:
embed_params.append(param)
else:
other_params.append(param)
# Embedding 学习率 ×0.1
optimizer_grouped_parameters = [
{"params": embed_params, "lr": self.args.learning_rate * 0.1},
{"params": other_params, "lr": self.args.learning_rate},
]
self.optimizer = torch.optim.AdamW(
optimizer_grouped_parameters,
betas=(self.args.adam_beta1, self.args.adam_beta2),
eps=self.args.adam_epsilon,
weight_decay=self.args.weight_decay,
)
return self.optimizer
# 使用:trainer = LayerwiseTrainer(...)
4. 渐进式训练(最稳但慢)
python
# 先 bs=64 训 200 steps 稳定权重分布
# 再切到 bs=128,但前 100 steps 用 max_grad_norm=0.1 极度保守
# 最后恢复正常裁剪
training_args_first = TrainingArguments(
per_device_train_batch_size=64,
learning_rate=5e-5,
max_grad_norm=1.0,
max_steps=200, # 预稳定
)
training_args_second = TrainingArguments(
per_device_train_batch_size=128,
learning_rate=1e-4,
max_grad_norm=0.1, # 前100步保守
warmup_steps=100,
)
推荐配置(AdamW 专用)
python
training_args = TrainingArguments(
per_device_train_batch_size=128,
learning_rate=1e-4, # 线性缩放
# AdamW 防爆炸三件套
adam_beta2=0.95, # 降记忆
adam_epsilon=1e-6, # 防除零
max_grad_norm=0.5, # 硬裁剪
warmup_ratio=0.15, # 足够长的预热
weight_decay=0.01,
bf16=True,
gradient_checkpointing=True,
# 监控
logging_steps=10,
seed=42, # 复现爆炸/稳定
)
先试方案 1(调 beta2 + epsilon + 收紧 grad_norm),如果 50 步内还炸,再上方案 3 分层学习率。
你现在训练到多少步爆炸的?(前 100 步就炸还是几百步后?)