AI 原生开放教材

外观

AI 工程从零到一（中文）源码研读·中文原文 ↗原件 ↓仓库 ↗

Cosine 学习率 + 线性 Warmup ​

学习率 schedule 是仅次于 loss 函数的第二重要决策。AdamW 配 cosine decay + linear warmup 是语言模型训练的现代默认配置——前面那一千步脆弱期给模型一个小的有效步长，然后升到配置好的峰值，最后平滑衰减回接近零。这节课实现这个 schedule、画出训练步上的曲线、在 schedule 旁边记录 gradient norm、并证明 schedule 在 warmup、peak 和 decay 边界上的行为正确。

类型： Build 语言： Python 前置要求： 第 19 阶段第 30-37 课 预计时间： ~90 分钟

学习目标 ​

  • 实现一个连接了 cosine 学习率 schedule + linear warmup 的 AdamW optimizer。
  • 在任意 step 精确计算 schedule 的值，跨 run 无浮点漂移。
  • 把 gradient L2 norm 和学习率并排记录，让训练健康状态可观测。
  • 把 schedule 渲染成肉眼可读的文本图和任意工具可消费的 CSV。

问题 ​

训练的前一千步是最吵的。模型权重还贴近初始化。Optimizer 的二阶矩估计还没稳定。梯度 norm 又大又噪声。如果学习率在这些 step 就已经到峰值，模型要么直接发散，要么掉进一个再也爬不出来的 loss 平台。两个已知的修复方法是 gradient clipping（第 19 阶段第 45 课的主题）和一个从小到大慢慢升的学习率 schedule。

Cosine-with-warmup schedule 有三个区间。从 step 0 到 step warmup_steps，学习率从零线性升到配置的峰值 lr_max。从 step warmup_steps 到 step total_steps，学习率沿 cosine 曲线的上半段从 lr_max 衰减到 lr_min。total_steps 之后学习率钉在 lr_min，这样配置错了的 trainer 超了步也不会悄悄跑出 schedule。

构建上的难点在于 schedule 很容易 off-by-one。这个 off-by-one 会在训练跑了六个小时后才暴露为一个高了或低了 1% 的学习率——正好赶上模型开始过拟合的节骨眼——除非你在边界上做了穷举测试，否则根本看不见。

概念 ​

mermaid
flowchart TD
  Step[训练步] --> Branch{步数状态}
  Branch -- step <= warmup --> Linear[从 0 到 lr_max 线性爬升]
  Branch -- warmup < step <= total --> Cosine[从 lr_max 到 lr_min 余弦衰减]
  Branch -- step > total --> Floor[钉在 lr_min]
  Linear --> Apply[AdamW.step]
  Cosine --> Apply
  Floor --> Apply
  Apply --> GradNorm[计算梯度 L2 范数]
  GradNorm --> Log[步日志行]
  Log --> Plot[文本图 + CSV]

Warmup 公式 ​

对于 step 在 [0, warmup_steps] 且 warmup_steps > 0 的情况，学习率是 lr_max * step / warmup_steps。退化情况 warmup_steps = 0 视为"无 warmup"：step 0 直接从 lr_max 开始，立即进入 cosine decay。有些测试 harness 会传 warmup_steps = 0 来检查 schedule 是否仍能产出可用曲线。

Cosine 公式 ​

对于 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 时 cosine 算出 cos(0) = 1，得到 lr_max，精确匹配 warmup 终点。在 step = total_steps 时 cosine 算出 cos(pi) = -1，得到 lr_min，精确匹配 decay 终点。

两个端点的连续性不是巧合。这就是为什么 schedule 要实现为关于 step 的单一函数，而不是三个不同函数粘在一起。粘起来的 schedule 在第一次改 lr_max 时就会丢掉一个边界。

total_steps 之后的地板 ​

step > total_steps 时学习率保持 lr_min。契约是明确的：schedule 不会报错也不会外推；它钉在地板上，让 trainer 记一条 warning。需要延长训练就改 schedule 的 total_steps，别改循环。

Gradient norm 和学习率并排记录 ​

Schedule 是训练健康的一半。Gradient norm 是另一半。训练循环每步都记录两者。发散的训练在 gradient norm spike 时就会暴露，比 loss 更早；调得好的 warmup 会让 norm 随学习率线性上升；peak 太激进会表现为 warmup 结束后 norm 居高不下。磁盘上的数据集是 step, lr, grad_l2_norm, loss。CSV 是唯一的持久记录。

figure
cap-cosine-warmup

构建 ​

code/main.py 实现：

  • CosineWithWarmup - 一个无状态函数 lr(step) -> float，覆盖配置好的 schedule。
  • TrainState - 把模型、AdamW optimizer 和 schedule 包装成一个 step 函数。
  • TrainState.step - 跑一次 forward pass、一次 backward pass、记录 gradient L2 norm、然后把 lr(step) 应用到 optimizer。
  • plot_schedule_ascii - 把 schedule 渲染为肉眼可读的文本图。
  • write_schedule_csv - 每步一行，输出学习率。

文件底部的 demo 构建一个 tiny nn.Linear 模型，在固定输入 batch 上训练 20 步，打印每步的学习率、gradient norm 和 loss。Schedule 也渲染成文本图做视觉 sanity check。

运行：

bash
python3 code/main.py

脚本 exit 0 并打印 per-step 训练日志加 schedule 图。

生产模式 ​

四个模式把 schedule 升级为生产制品。

Schedule 放在配置里，不在代码里。 Trainer 从 git 里提交的 YAML 或 JSON 配置中读取 warmup_steps、total_steps、lr_max、lr_min。Schedule 可复现因为配置是内容寻址的；schedule 可审计因为配置是 PR diff 的一部分。

Step counter 单调递增且与 epoch 解耦。 有些框架在数据分片或 dataloader 重启时会搞混 step 和 epoch。Schedule 从 trainer 的 checkpoint 读 global_step，不是从本地计数器读。恢复训练时 schedule 位置正确，因为 step counter 是持久化的轴。

Schedule 图放在 run 目录里。 每次训练运行都在 run 目录下写一个 outputs/lr_schedule.png（本课是文本图）。Reviewer 扫一眼目录就能 sanity-check schedule，不需要重跑。这在 PR 时就能抓住配错 schedule 这类 bug。

Log row schema 固定。 step, lr, grad_l2_norm, loss，按这个顺序。下游 notebook 或 dashboard 按这个 schema 读；不 bump version 就改列名会废掉所有现有 dashboard。

使用方式 ​

生产模式：

  • 先 sweep peak，再 sweep 其他任何东西。 lr_max 是最敏感的旋钮。先在小模型上 sweep；最优 lr_max 随模型大小弱缩放，所以小模型 sweep 是个很强的先验。
  • Warmup 是 total steps 的比例，不是绝对步数。 一个 2 亿步的运行配 2,000 步 warmup 几乎一开始就到峰值；一个 20,000 步的运行配同样的数字却是 10% 的 warmup。把 warmup 配成比例（典型：1-3%），schedule 就能随训练时长自动缩放。
  • lr_min 故意非零。 一个 10% lr_max 的地板让 optimizer 在长尾阶段还能学。lr_min = 0 的 schedule 画出来的曲线很漂亮，但模型其实还没训完。

拿去用 ​

outputs/skill-cosine-warmup.md 在真实项目中会描述哪个配置文件承载 schedule、trainer step 的全局计数器从哪里读、部署值的 lr_max sweep 产出了什么。本课交付的是引擎。

练习 ​

  1. 加一个 inverse-square-root 变体 schedule，在 200 步 toy 训练运行上对比。哪条曲线的最终 loss 更低？
  2. 加一个 --restart flag，在 total_steps / 2 处做第二次 warmup。论证 warm restart 在 toy 运行上是改善还是损害。
  3. 加一个单元测试检查 schedule 连续性：对 [0, total_steps] 中的每个 step，差值 |lr(step+1) - lr(step)| 必须被 lr_max / warmup_steps 上界。
  4. 把 schedule 接入 torch.optim.lr_scheduler.LambdaLR，让它能和框架代码组合。本课用的是纯 step 函数；wrapper 改变了什么？
  5. 加一个 --plot-png flag 用 matplotlib 画真图。论证本课的文本图和 PNG 哪个更适合作为 CI 运行的默认值。

关键术语 ​

术语大家嘴上说的实际含义
Warmup"慢启动"在前 warmup_steps 步内从零线性升到 lr_max
Cosine decay"平滑下降"在剩余步数上从 lr_max 到 lr_min 的上半 cosine 曲线
地板"训练之后"total_steps 之后 schedule 钉死的 lr_min 值
Gradient norm"梯度的 L2"拼接后梯度向量的欧几里得范数，每步记录
Global step"Schedule 轴"一个能扛住重启的单调 step 计数器，驱动 schedule

延伸阅读 ​