AdamW与学习率调度实战：Warmup、Cosine、Linear在大模型训练中的用法

2026-03-19 Ai大模型训练教程,AdamW,Warmup,Cosine学习率调度,Linear学习率衰减,大模型优化器,Transformers训练,PyTorch学习率调度

原始问题：

本文为《Ai大模型训练教程》系列实战篇，详解大模型训练中AdamW与学习率调度的组合用法，重点讲清Warmup设置方法、Cosine与Linear调度的选择与参数建议，并提供Transformers/PyTorch可直接复用的代码示例与常见排障清单，帮助训练更稳定、更可复现。

重新提问

这篇在系列中的位置与目标

在《Ai大模型训练教程：从入门到实战落地的系统课程》里，优化器与学习率（Learning Rate, LR）调度是“训练能不能稳定收敛、能不能把算力花在刀刃上”的关键一环。本篇聚焦 AdamW + 学习率调度 的实战组合，给出你在大模型训练中最常用的三种调度策略：Warmup、Cosine、Linear 的具体用法、推荐参数、常见坑和可直接套用的配置与代码。

你会看到：

为什么大模型训练几乎标配 AdamW（而不是 Adam）
为什么 Warmup 常常“必须有”，以及 warmup 步数怎么估
Cosine vs Linear 怎么选，分别更适合什么训练阶段
以 HuggingFace Transformers / PyTorch 为例，如何把它们接到一起并验证是否生效

AdamW：大模型训练的默认优化器选择

Adam 和 AdamW 的关键差异：权重衰减是否“解耦”

很多人会把 L2 正则（weight decay）当成“在 loss 上加一个 λ||w||²”，但 Adam 的自适应矩估计会让这种衰减与梯度更新耦合在一起，导致等效行为与“真正的权重衰减”不一致。

AdamW 的核心是：

先按 Adam 的方式算一套梯度更新（带动量与二阶矩）
再额外对参数做一次与梯度无关的衰减：w = w - lr * weight_decay * w

这在大模型训练里更稳定、可控，实践中更常见于：

语言模型预训练 / 指令微调
大规模视觉/多模态模型训练

AdamW 的常用参数范围（可作为起点）

以下是经验起点，实际仍要结合 batch size、序列长度、数据噪声等：

betas=(0.9, 0.95) 或 (0.9, 0.999)
- 大模型（尤其 Transformer）里 (0.9, 0.95) 很常见，收敛更“紧凑”
eps=1e-8（默认通常够用；混合精度不稳时可尝试 1e-6）
weight_decay=0.01（常见默认）
- 微调场景有时用 0.0 ~ 0.1，看是否过拟合

参数分组：哪些参数不做 weight decay（非常重要）

Transformer 结构中通常对以下参数不做 weight decay：

LayerNorm / RMSNorm 的 weight（或 gamma）
bias

原因：对归一化层做衰减可能破坏尺度稳定性，对 bias 做衰减意义也不大。

推荐做法：参数分组。

学习率调度为什么在大模型训练里更关键

大模型训练不只是“设一个 lr 就跑”。你通常会遇到：

初期梯度不稳定，loss 抖动甚至 NaN
中后期需要更小 lr 才能把验证指标压下去
训练步数很长，不做调度会浪费大量算力

因此业界常见组合是：

Warmup（预热）：前 N 步把 lr 从 0 拉到目标 lr
再接一种衰减策略：Cosine 或 Linear（最常见）

Warmup：从“能跑起来”到“稳定收敛”的第一道保险

Warmup 在解决什么问题

大模型训练初期，参数还在随机附近，梯度分布与激活尺度可能很不稳定。如果一上来就用较大学习率：

可能出现梯度爆炸、loss 震荡
混合精度下更容易溢出（GradScaler 不一定救得回来）

Warmup 的作用是：

让优化器与网络逐渐进入稳定区
在较小 lr 下先“找方向”，再加速

Warmup 步数/比例怎么设（给你可执行的估算方法）

常见设置方式有两种：

1) 按比例：warmup_ratio = 0.01 ~ 0.1

预训练步数超长时常用 0.01 ~ 0.03
微调数据少、训练步数少时可用 0.05 ~ 0.1

2) 按步数：warmup_steps = 100 ~ 2000（视总步数）

若总步数只有几千步，warmup 设太大反而浪费训练

实操建议（可直接套用）：

先算 total_steps = epochs * steps_per_epoch
再设 warmup_steps = min(2000, int(total_steps * 0.03))
如果你发现前 100~300 step loss 仍然剧烈抖动，可把 warmup 提到 5% 甚至 10%

Warmup 曲线：Linear Warmup 最常用

最常见是线性预热：

第 0 步 lr = 0
第 warmup_steps 步 lr = base_lr

也有人用 cosine warmup，但线性已经足够稳定且易解释。

Cosine 调度：中后期“又快又稳”的常用选择

Cosine 的形状与直觉

Cosine（余弦退火）会让学习率从 base_lr 平滑下降到 min_lr（或 0），下降前期慢、中段快、后期更慢，适合：

长训练
希望后期仍保留一点探索空间（不会像线性那样“直线砍到底”）

什么时候优先选 Cosine

训练步数较长（例如几万到几十万 step）
你更看重最终指标（后期更细腻）
你打算在一个 schedule 内跑完整个训练，不做复杂分段

min_lr 要不要设为 0

在不少实现里 cosine 会退火到 0。但大模型训练里也常见设一个 min_lr（例如 base_lr 的 1%~10%），避免 lr 过小导致：

后期更新几乎停滞
训练时间浪费

若你的框架只支持退火到 0：

也可以接受，尤其是训练后期本来就是精修阶段

Linear 调度：最简单的“稳妥衰减”

Linear 的形状与直觉

Linear decay 是把学习率从 base_lr 线性下降到 0（或某个最小值）。它的优点是：

行为可预测
超参更少
在微调里非常常见（例如 warmup + linear）

什么时候优先选 Linear

指令微调 / 领域微调：训练步数往往不长
你希望 schedule 简单、可控、复现性强
你在做对比实验，需要减少调度因素带来的变量

三种调度在大模型训练中的典型搭配

下面给出三个“最常用、最少踩坑”的搭配，你可以按场景直接选。

搭配 A：Warmup + Linear（微调首选）

场景：SFT、LoRA 微调、领域适配
特点：稳定、简单、好复现
建议：warmup_ratio 0.03~0.1；weight_decay 0~0.1

搭配 B：Warmup + Cosine（长训练首选）

场景：较长预训练、较大数据的持续训练
特点：后期更细腻，最终指标常更好
建议：warmup_ratio 0.01~0.03；可考虑 min_lr

搭配 C：Warmup + Cosine（带 restarts，不太建议新手）

场景：需要多次“重启探索”的训练
风险：restarts 的周期、振幅等参数不当会导致指标波动

本篇重点放在 A/B，因为它们覆盖了 90% 的大模型训练任务。

关键超参如何一起定：给你一套可落地的决策流程

Step 1：先定“有效 batch size”和 base_lr

大模型训练里常用概念：

global_batch = per_device_batch * num_devices * grad_accum

经验上（非严格公式），你可以：

先用社区/论文/开源配方的 base_lr 当起点（同类型模型最可靠）
如果要按 batch 做缩放：常见是线性缩放（batch 翻倍，lr 近似翻倍），但大模型不一定严格成立，建议保守一点

Step 2：选择调度：短训练 Linear，长训练 Cosine

total_steps < 10k：优先 warmup + linear
total_steps >= 10k：优先 warmup + cosine

Step 3：定 warmup

默认：warmup_steps = int(total_steps * 0.03)
若训练非常短：保证 warmup 至少 100 步，但不超过总步数的 10%

Step 4：权重衰减与参数分组

默认 weight_decay=0.01
对 norm 和 bias 不衰减
LoRA 微调时：很多人会对 base model 参数冻结，只训练 adapter；weight_decay 可更小甚至 0（取决于是否过拟合）

实战：PyTorch + Transformers 里正确使用 AdamW + 调度器

下面示例强调三点：
1) AdamW 参数分组（no_decay）
2) 正确计算 total_steps / warmup_steps
3) 调度器与 optimizer 的 step 顺序（用 Trainer 会自动处理；手写循环要注意）

参数分组与优化器

import torch
from torch.optim import AdamW

no_decay = ["bias", "LayerNorm.weight", "layer_norm.weight", "ln_f.weight", "norm.weight"]

def build_param_groups(model, weight_decay=0.01):
    decay_params = []
    nodecay_params = []

    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if any(nd in n for nd in no_decay):
            nodecay_params.append(p)
        else:
            decay_params.append(p)

    return [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ]

optimizer = AdamW(
    build_param_groups(model, weight_decay=0.01),
    lr=2e-5,
    betas=(0.9, 0.95),
    eps=1e-8,
)

计算训练步数与 warmup 步数

from math import ceil

def compute_steps(num_samples, per_device_batch, grad_accum, num_devices, epochs):
    steps_per_epoch = ceil(num_samples / (per_device_batch * num_devices * 1.0))
    # 注意：如果你用 grad_accum，优化器更新步数会除以 grad_accum
    optim_steps_per_epoch = ceil(steps_per_epoch / grad_accum)
    total_steps = optim_steps_per_epoch * epochs
    return total_steps

total_steps = compute_steps(
    num_samples=len(train_dataset),
    per_device_batch=per_device_batch_size,
    grad_accum=gradient_accumulation_steps,
    num_devices=torch.cuda.device_count(),
    epochs=num_train_epochs,
)

warmup_steps = min(2000, int(total_steps * 0.03))

常见坑：把 dataloader step 当成 optimizer step。开启梯度累积后，真正的参数更新次数会变少，total_steps 需要按优化器 step 来算。

Linear / Cosine 调度器（Transformers 提供）

from transformers import get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup

# 方案1：Warmup + Linear
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps,
)

# 方案2：Warmup + Cosine
# lr_scheduler = get_cosine_schedule_with_warmup(
#     optimizer,
#     num_warmup_steps=warmup_steps,
#     num_training_steps=total_steps,
# )

手写训练循环时的 step 顺序

scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

optimizer.zero_grad(set_to_none=True)

for step, batch in enumerate(train_loader):
    with torch.cuda.amp.autocast(enabled=use_amp):
        loss = model(**batch).loss
        loss = loss / gradient_accumulation_steps

    scaler.scale(loss).backward()

    if (step + 1) % gradient_accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

        # 一般在 optimizer.step() 之后调用 scheduler.step()
        lr_scheduler.step()

常见坑：把 lr_scheduler.step() 放在 optimizer.step() 前；或者每个 micro-batch 都 step scheduler，导致 lr 下降过快。

如何验证你的 Warmup / 调度是否真正生效

训练时请务必记录并检查以下信号（很多训练“玄学问题”其实是调度没按预期跑）：

1) 打印或记录 lr 曲线

每隔 N 个 optimizer step 打印当前 lr
你应该看到：
- warmup 阶段 lr 单调上升
- 之后按 linear 或 cosine 下降

示例：

current_lr = optimizer.param_groups[0]["lr"]
if global_step % 50 == 0:
    print(global_step, current_lr)

2) 关注 warmup 结束点的 loss 变化

正常现象：warmup 结束后 loss 可能短暂波动一下，但整体趋于更快下降
异常现象：warmup 结束后直接发散或 NaN，说明 base_lr 可能过大，或需要更长 warmup

3) 对比实验要固定 schedule

做不同模型/数据的对比时：

尽量固定 total_steps / warmup_ratio / schedule
否则你很难判断收益来自哪里

常见问题与排障清单（大模型训练高频踩坑）

问题 1：loss 在前几百步疯狂抖动甚至 NaN

优先检查：

是否没有 warmup 或 warmup 太短
base_lr 是否过大（先降 2~4 倍试试）
混合精度下 eps 是否太小（可试 1e-6）
是否忘了 gradient clipping（例如 clip_norm=1.0）

问题 2：训练到后期不怎么动了，指标停滞

可能原因：

linear 衰减到 0 太早（总步数算错，或 scheduler step 次数过多）
cosine 退火到 0 后学习停止（可考虑 min_lr 版本或缩短训练）
weight_decay 过大导致欠拟合

问题 3：开了梯度累积后，lr 曲线“跑得飞快”

原因：

你把 total_steps 当成 dataloader steps
或者每个 micro-batch 都调用了 scheduler.step()

修复：

total_steps 用“optimizer 更新步数”
scheduler.step() 放到 optimizer.step() 同频

问题 4：LoRA 微调时用 AdamW + weight_decay 反而变差

排查：

你是否只训练 LoRA 参数？如果是，weight_decay 可以更小甚至 0
数据量很小时，衰减可能加重欠拟合

可直接复用的推荐配置（起步配方）

配方 1：指令微调/SFT（Warmup + Linear）

optimizer: AdamW
lr: 1e-5 ~ 5e-5（按模型大小与 batch 调）
betas: (0.9, 0.95)
weight_decay: 0.0 ~ 0.01
warmup_ratio: 0.05
scheduler: linear decay to 0
gradient clipping: 1.0

配方 2：较长训练（Warmup + Cosine）

optimizer: AdamW
lr: 1e-4 ~ 2e-4（预训练常见范围之一；按任务调整）
betas: (0.9, 0.95)
weight_decay: 0.01
warmup_ratio: 0.01 ~ 0.03
scheduler: cosine decay（可考虑保留 min_lr）
gradient clipping: 1.0

小结：把 AdamW + Warmup + 调度用“对”的三个关键

1) AdamW + 参数分组：norm/bias 不做 weight decay
2) Warmup 基本必备：按总 optimizer steps 的 1%~10% 设，默认 3% 很好用
3) Scheduler step 频率要对：梯度累积时按 optimizer step 调度，不要按 micro-batch

掌握这套组合，你在大多数大模型训练（预训练、继续预训、SFT、LoRA）里都能得到稳定、可复现的收敛曲线。下一步再进阶，通常就是：更细的分段 schedule、不同层的 lr（layer-wise decay）、以及更复杂的优化器（Adafactor、Lion 等）与更大规模并行训练策略的配合。