Cosine 学习率与预热

Cosine 学习率与预热

本节摘要:学习率调度是损失函数之后第二重要的决策。带权重衰减分组的 AdamW + 余弦退火 + 线性预热是语言模型训练的现代默认,因为它让模型在脆弱的前一千更新里看到小有效步长、爬到配置的峰值、再平滑降回接近零。本节单独把这个调度抽出来:精确实现任意步数的值(无跨运行浮点漂移)、把梯度 L2 范数与学习率并排日志、把调度渲染成肉眼可读的文本图与 CSV。调度有三区:预热线性爬、余弦退火降、超出总步数钉在地板。建侧的问题是调度易差一:差一会在训练六小时后于模型开始过拟合的瞬间表现为学习率偏高或偏低 1%,除非在边界穷尽测试,否则不可见。

对应原课程:Phase 19 · Lesson 44 · cosine-lr-warmup(原英文 phases/19-capstone-projects/44-cosine-lr-warmup/docs/en.md)。本节属「预训练/分布式」赛道第三节。

学习目标

阅读完本节,你应当能够:

1. 实现接到余弦学习率调度(带线性预热)的 AdamW 优化器。
2. 在任意步数计算调度的精确值,跨运行无浮点漂移。
3. 把梯度 L2 范数与学习率并排日志,使训练健康可观察。
4. 把调度渲染成肉眼可读的文本图与任意工具可消费的 CSV。

一、问题与直觉

前一千训练更新最吵。模型权重仍近初始化,优化器的二阶矩估计未稳,梯度范数大且噪。若此时学习率在峰值,模型要么直接发散,要么陷进永远逃不出的损失平台。两个知名修复是梯度裁剪(第 43 节)与一个从小起步、爬上去的学习率调度。

预热公式

step 在 [0, warmup_steps](warmup_steps > 0)时,学习率是 lr_max * step / warmup_steps。退化的 warmup_steps = 0 当作「无预热」:调度直接从步零的 lr_max 起步、立即进余弦衰减。

余弦公式

step 在 (warmup_steps, total_steps] 时,学习率是 lr_min + 0.5*(lr_max - lr_min)*(1 + cos(pi*progress)),其中 progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)。step = warmup_steps 时余弦取 cos(0)=1,给 lr_max,精确匹配预热端点;step = total_steps 时取 cos(pi)=-1,给 lr_min,精确匹配衰减端点。

两端连续性不是偶然——这是调度实现成单一 step 函数(而非三个函数粘起来)的原因。粘起来的调度一改 lr_max 就丢一个边界。

总步数后钉地板

step > total_steps 时学习率停在 lr_min。契约明确:调度不报错、不外推,钉地板让训练器打警告。要延长训练的训练器改调度的 total_steps,而非改循环。

梯度范数与学习率并排

调度只是训练健康的一半。另一半是梯度范数。若梯度范数在学习率峰值附近尖峰,模型在发散边缘;若范数随训练单调增,过拟合在逼近。把两者并排日志,眼睛能同时看步长与噪声。

二、从零实现

code/main.py 实现:

• cosine_with_warmup(step, warmup_steps, total_steps, max_lr, min_lr):单一函数覆盖三区,返回该步学习率。
• AdamWScheduled(model, ...):AdamW + 每 step 前按调度设 param_groups 的 lr。
• compute_grad_norm(model):全局 L2 范数(所有参数梯度拼成的向量的欧氏范数)。
• render_text_plot(values, width):把调度渲染成 ASCII 文本图。
• dump_csv(rows, path):把步日志(步、学习率、梯度范数、损失)落 CSV。

调度核心:

import math

def cosine_with_warmup(step, warmup_steps, total_steps, max_lr, min_lr):
    if step <= warmup_steps and warmup_steps > 0:
        return max_lr * step / warmup_steps                       # 线性预热
    if step <= total_steps:
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))
    return min_lr                                                  # 钉地板
​

并排日志:

def train(model, loader, total, warmup, max_lr, min_lr):
    opt = torch.optim.AdamW(model.parameters(), lr=max_lr,
                            betas=(0.9, 0.95), weight_decay=0.1)
    rows = []
    for step in range(total):
        x, y = next(loader)
        lr = cosine_with_warmup(step, warmup, total, max_lr, min_lr)
        for g in opt.param_groups: g["lr"] = lr
        loss = calc_loss(model, x, y); loss.backward()
        grad_norm = compute_grad_norm(model)                      # 并排日志
        opt.step(); opt.zero_grad()
        rows.append({"step": step, "lr": lr, "grad_norm": grad_norm,
                     "loss": loss.item()})
    dump_csv(rows, "outputs/schedule.csv")
    render_text_plot([r["lr"] for r in rows], 60)
​

设计要点:单一函数覆盖三区是连续性的保证——预热端点的 lr_max 与余弦起点 cos(0)=1 给的 lr_max 必须相同,粘三个函数会在改 lr_max 时漏一个边界。max(1, total_steps - warmup_steps) 防 total == warmup 的除零。文本图与 CSV 让调度可肉眼读、可工具消费,二者都不依赖运行中的画图。

三、框架对比

HuggingFace transformers 的 get_cosine_schedule_with_warmup 是同调度的库实现,逻辑与本节一致。PyTorch Lightning 的 CosineAnnealingWarmRestarts 支持周期重启(SGDR)。timm 的 create_scheduler 把多种调度配进一个工厂。本节手写让你看清三区切换、端点连续、地板契约。生产上调度常配 WSD(warmup-stable-decay,稳定段后衰减,Chinchilla 后流行)或 WSR(warmup-stable-restart),但底层都是「预热 + 形状 + 地板」三段,本节是它们的共同骨架。

四、可复用产物

code/main.py:cosine_with_warmup、compute_grad_norm、render_text_plot、dump_csv 均可复用。demo 在 CPU 上几秒跑完,把调度渲染成文本图、把步日志落 CSV、在边界步(0、warmup、total、total+1)断言值符合公式。CSV 是可复现的产物——同一调度同一种子,跨运行的学习率与梯度范数逐位相同。

五、练习

1. 边界断言:给 step 取 0、warmup、warmup+1、total、total+1,断言返回值与手算一致。
2. 零预热:设 warmup_steps=0,确认调度从步零的 lr_max 起步、立即进余弦。
3. 地板抬高:设 min_lr = 0.1*max_lr,重绘曲线,讨论对末期微调的影响。
4. WSD 变体:在预热与衰减间加一段稳定段(恒 lr_max),对比纯余弦的损失曲线。
5. 梯度范数告警:训练时若梯度范数超阈值打告警,讨论它作为发散早期信号的有效性。

本节要点回顾

1. 调度是第二重要决策:损失之后,学习率调度决定训练成败。
2. 三区:预热线性爬、余弦退火降、超总步钉地板。
3. 端点连续:单一函数覆盖三区,改 lr_max 不漏边界。
4. 预热给优化器时间:前一千步步长小,让二阶矩填满。
5. 地板契约:超总步不报错不外推,改 total_steps 延长训练。
6. 梯度范数并排:步长与噪声同看,发散早期可见。

下一节,我们做「梯度裁剪与混合精度」——全局 L2 裁剪加 autocast + GradScaler,扛住 FP16 溢出与坏批尖峰。