学习率调度策略：Warmup与Cosine退火的工程实践
在现代大规模语言模型（LLM）训练中，一个常被忽视却至关重要的细节是——为什么刚启动训练时 loss 会剧烈震荡？尤其是在 Qwen、LLaMA 等超大模型上，前几百步的梯度更新常常像“醉酒般”不稳定。这并非数据或初始化的问题，而是学习率调度没踩准节奏的结果。
如果你曾手动调过 learning rate，可能有过这样的经历：设高了，loss 直接炸掉；设低了，收敛慢得像蜗牛。而如今几乎所有主流框架——从 Hugging Face Transformers 到 ms-swift ——都在默认使用一种“组合拳”策略：先 Warmup 预热，再 Cosine 退火。这不是偶然，而是经过大量实验验证的有效范式。
那这套机制到底好在哪？它又是如何在真实训练系统中落地的？
我们不妨从一个典型的微调任务说起。假设你正在用 ms-swift 对一个 7B 参数的 LLM 做监督微调（SFT），配置如下：
learning_rate: 2e-5 num_train_epochs: 3 warmup_ratio: 0.1 lr_scheduler_type: cosine看起来很简单，但背后其实藏着一套精密的调度逻辑。整个训练过程的学习率变化曲线大致长这样：
↑ lr │ ↗━━━━━━━━━━━━━━━━━━━━━━━━━┓ │ ↗ ┃ Cosine 退火阶段 │ ↗ ┃（平滑下降至接近0） │ ↗ ┃ │ ↗ ┃ ├──────────────────────────────────→ step ↑ ↑ 第0步 Warmup结束（约10%总步数）这条曲线的前半段靠Warmup打底，后半段由Cosine 退火主导，两者配合，形成了一种“先稳后精”的训练节奏。
为什么需要 Warmup？
Transformer 类模型对初始学习率极为敏感，尤其是注意力层中的残差连接和 LayerNorm 结构。参数随机初始化状态下，某些 head 可能在第一步就输出极大值，若此时使用全量学习率更新，极易引发梯度爆炸。
Warmup 的本质就是“缓起步”。它让学习率从 0 开始，在前warmup_steps步内线性增长到基础值，给模型一个“适应期”。这个过程就像汽车冷启动时不能猛踩油门，得先预热发动机。
具体来说，如果总训练步数为 20,000，warmup_ratio 设为 0.1，则前 2,000 步执行线性升温：
$$
\text{lr}_t = \text{base_lr} \times \frac{t}{\text{warmup_steps}}
$$
这种设计简单却极其有效。尤其与 AdamW 这类自适应优化器结合时，能显著降低前期梯度方差，避免模型“还没学会走路就想跑”。
实践建议：虽然可以硬编码
warmup_steps=2000，但在跨任务迁移时更推荐使用比例形式（如warmup_ratio=0.1），这样无论数据集大小如何变化，预热周期都能自动适配。
当然，也不是越长越好。Warmup 时间过长会导致收敛延迟；太短又起不到稳定作用。经验上看，5%~10% 的训练周期是最优区间，BERT 和 RoBERTa 的原始论文也验证了这一点。
Cosine 退火：不只是衰减，更是搜索
Warmup 解决的是“开头难”，而 Cosine 退火解决的是“收尾糙”。
传统 Step Decay 或指数衰减往往在某个固定节点直接降学习率，容易造成损失曲面跳跃，甚至把模型踢出潜在的最优区域。相比之下，Cosine 提供了一条平滑且非线性的衰减路径：
$$
\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min}) \left(1 + \cos\left(\frac{T_{cur}}{T_{max}} \pi\right)\right)
$$
它的下降趋势是“两头缓、中间快”：初期缓慢释放，中期加速探索，后期细腻微调。这种节奏恰好契合了深度模型的训练动态——前期需要一定幅度跳出初始局部极小，后期则需谨慎逼近更优解。
更重要的是，余弦函数天然具备重启潜力。SGDR（Stochastic Gradient Descent with Warm Restarts）正是基于此思想，通过周期性重置 $ T_{cur} $ 实现多次“冷启动”，帮助模型逃离平坦盆地。
不过在大多数全参数微调或预训练场景中，人们更倾向使用单周期 Cosine，即从 Warmup 结束处开始，一路平滑衰减至训练结束。例如：
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=total_steps - warmup_steps, eta_min=1e-8 )这里最小学习率通常设为 $1 \times 10^{-8}$ 而非 0，防止参数彻底停滞更新。毕竟，哪怕只是“轻轻一推”，也可能让模型跨过一道微妙的泛化边界。
工程实现：如何无缝衔接 Warmup 与 Cosine？
理想情况下，我们希望 Warmup 和 Cosine 是一个连贯的整体，而不是两个割裂的阶段。遗憾的是，PyTorch 原生调度器并不直接支持这种复合逻辑，因此需要封装一层自定义策略。
以下是 ms-swift 框架内部常用的一种实现方式：
from torch.optim.lr_scheduler import LambdaLR def get_warmup_cosine_schedule(optimizer, warmup_steps, total_steps, eta_min=1e-8): def lr_lambda(current_step): if current_step < warmup_steps: # Warmup: 线性上升 return float(current_step) / float(max(1, warmup_steps)) else: # Cosine: 从 base_lr 平滑下降至 eta_min progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps)) return eta_min + 0.5 * (1 - eta_min) * (1 + math.cos(math.pi * progress)) return LambdaLR(optimizer, lr_lambda)这种方式将两种策略融合在一个调度函数中，确保学习率在整个训练过程中连续可导，避免跳变。同时，由于完全基于current_step控制，也能很好地兼容梯度累积、分布式训练等复杂场景。
⚠️ 注意事项：
- 在 DDP 或 FSDP 多卡训练中，务必保证所有 rank 使用相同的全局 step 计数；
- 若启用梯度累积，应按参数更新次数（而非 forward 次数）递增 step；
- 多阶段训练（如 SFT → DPO）建议重新初始化 scheduler，否则后期学习率可能已趋近于零，无法有效更新。
在 ms-swift 的实际架构中，这类调度逻辑已被深度集成至训练引擎的核心流水线：
[用户输入 YAML] ↓ [Config Parser] → 提取 lr_scheduler_type, warmup_ratio 等字段 ↓ [Optimizer Builder] → 构建 AdamW / GaLore / Q-Galore ↓ [Scheduler Injector] → 自动注入 Warmup+Cosine 组合策略 ↓ [Training Loop] ← 每步调用 scheduler.step() ↓ [TensorBoard / EvalScope] ← 实时可视化 lr 曲线这意味着开发者无需关心底层调度代码，只需声明意图即可获得高质量的训练行为。这种“声明式 API + 隐式调度”的设计理念，大大降低了大模型调优门槛。
更重要的是，该组合策略展现出极强的任务普适性。无论是纯文本生成、多模态理解，还是人类偏好对齐（DPO/KTO），其目标函数虽复杂多变，但 Warmup+Cosine 依然表现出良好的鲁棒性。实验表明，在相同条件下，相比固定学习率，该策略平均可提升下游任务指标 2~5 个点，且训练稳定性明显增强。
当然，没有银弹。任何策略都有其适用边界。比如在极短周期微调（<1 epoch）中，过长的 Warmup 反而会压缩有效训练时间；而在持续预训练场景下，可能更适合采用带重启的 Cyclical LR。
但从整体来看，Warmup + Cosine 退火已成为当前大模型训练的事实标准。它不仅是一种技术选择，更是一种工程哲学：不让模型在起点摔倒，也不让它在终点草率收场。
当你下次看到训练日志里那条优雅平滑的学习率曲线时，不妨想想——正是这些看似细微的调度设计，支撑着万亿参数模型一步步走向收敛。
