☰
训练初期loss爆炸?详解warmup预热策略的原理与工程实践
2026/10/17 4:01:41 网站建设 项目流程

去年我在做一个大 batch 的视觉模型训练任务时,遇到过一张非常典型的 loss 曲线:前 50 步 loss 稳稳下降,到 80 步左右突然开始往上窜,再然后就干脆一个 NaN 把我整个实验送走。当时我把问题归咎于数据里有异常样本,排查了一整天才发现,罪魁祸首根本不是数据,而是我压根没给训练初期加 warmup(预热)策略。

如果你也正在跑 Transformer 类结构、大 batch、或者用 Adam 系优化器,那么 warmup 大概率不是“可选项”,而是“必选项”。这篇文章我会从原理、变体、代码实现、步数设定,再到一次完整的排错实录,把自己这几年用 warmup 的经验一次性讲透。不管你是刚接触深度学习的新手,还是已经在调参炼狱里翻滚过一阵的炼丹师,都能从这里拿走一些可以直接用的东西。

1. 先讲我踩过的坑:为什么训练初期那么容易炸

1.1 一个让我怀疑人生的实验现象

那次任务的具体配置我现在还记得:几十张卡做同步训练、batch size 加起来差不多 4096、优化器用的 AdamW、初始学习率 3e-4。我按照以前小 batch 训练的习惯,直接从头到尾用同一个学习率,没有任何预热阶段。

前几十个 step 的输出看起来完全正常,交叉熵损失从 7.2 一路降到 5.1,我还觉得这次训练挺稳。结果到了差不多第 90 个 step,loss 像吃了火药一样反弹到 8.9,再过了两三个 step,直接变成 NaN。我把训练停掉,找了很久的数据问题,又把梯度裁剪加上,情况稍微好了一点,但 loss 依然会在一百步附近出现一个诡异的平台甚至回跳。

后来我才意识到,问题出在训练最开始的“状态”上。模型刚初始化出来的参数分布、优化器内部的动量统计量、以及 BN 层或者 LayerNorm 的统计数据,都是处于一种“还没站稳”的状态。这时候直接上大学习率,就像一辆刚从车库启动的跑车,你一脚油门踩到底,它能不翻吗?

1.2 第一个原因:Adam 早期统计量的不可靠

先说优化器内部的原因。以 Adam 为例,它维护一阶动量 m 和二阶动量 v,每次更新用的是 m 除以 sqrt(v)。这个设计本身是为了给每个参数一个自适应的步长。但问题在于,训练刚开始时 v 的值非常小,因为梯度本来就只是在初始参数附近算出来的那几步,统计意义并不强。

当你除以一个很小的 sqrt(v) 时,相当于把有效步长放大了好多倍。有人可能会说,Adam 不是自带自适应学习率吗?为什么还要 warmup?答案很简单:自适应学习率解决的是“不同参数之间”的量纲问题,它并没有解决“训练初期整体步长偏大且抖动剧烈”的问题。尤其是在稀疏梯度或者深层网络里,早期某个参数的梯度突然变得很大,而它的二阶动量还没跟上,这一下更新就可能把参数推出正常区域。

我后来做过一个对照实验:把 warmup 去掉,但把初始学习率降到原来的十分之一,整体损失曲线依然会出现一波不正常的上扬。这说明问题不只是峰值学习率太高,更是因为优化器状态和模型参数在初期不匹配。

1.3 第二个原因:深层网络参数处于“压导致易碎”的状态

除了优化器的原因,模型本身的参数初始化也是关键因素。Transformer 这类结构里,残差分支上堆叠了很多层,每一层的输出经过残差连接之后会叠加在一起。虽然各种初始化方法会尝试控制每层的方差,但是在训练的最开始,各层之间的梯度和参数更新方向并没有形成稳定的配合。

你可以把整个深层网络想象成一个很多人接力搬运货物的团队。一开始大家方向不一致,每个人都在微调自己的姿态,结果前一个人递给后一个人的东西稍微歪一点,后面就会被逐级放大。warmup 存在的意义,就是给这个团队一个“对齐动作”的缓冲期:先让学习率从很小慢慢升上来,让大家在低学习率下适应彼此的步调,然后再真正开始大步冲刺。

明白了这两个原因,你就能理解为什么 warmup 并不是一个“锦上添花”的技巧,而是很多模型训练里真正兜底的东西。

2. warmup 的几种玩法与选型思路

2.1 四种常用 warmup 曲线的对比

warmup 本身不是一个固定的东西,它有多种实现曲线,最常见的有这么几种:

类型学习率变化方式特点适合场景
常量 warmup固定一个小学习率,跑 N 步后直接跳到峰值实现简单,但跳变容易造成 loss 波动实验对比、快速验证
线性 warmup从 0(或一个小值)线性增加到峰值最常见,稳定可控大多数 CNN、Transformer 任务
指数/非线性 warmup先快后慢或先慢后快地增长到峰值可精细控制,需多调一组指数参数大 batch、batch 特别大的预训练任务
逐步 warmup每过一定步数学习率翻倍,分段爬升在分布式训练里方便同步多机多卡、用 checkpoint 断点续训时

我自己的经验是:如果是第一次用 warmup,直接选线性 warmup 最稳妥。它没有额外参数,行为也最容易预期。等你把实验跑到一定程度,确认了峰值学习率在什么范围合适,再考虑换指数或者余弦变体。

2.2 与学习率调度器组合的真实搭配

实际训练里 warmup 很少单独出现,它通常是和一个衰减调度器拼在一起。最常见的组合有三个:

第一种是 warmup + 余弦退火。这是我在大部分视觉任务里的标配:前 5% 步数做线性 warmup 到峰值学习率,之后按照余弦曲线把学习率降到接近零。这组合的好处是衰减平滑,后期不容易在局部最小值附近来回震荡。

第二种是 warmup + 阶梯式衰减。比如跑到总步数的 50% 和 80% 时学习率各乘 0.1。这种搭配比较适合某些对后期探索要求不高的任务,但我个人觉得阶梯跳变会让损失曲线出现可见的抖动,不如余弦顺畅。

第三种是 warmup + 恒定学习率。这常见于微调任务,我一般只会在数据量特别小、训练步数很短的任务里这样干,因为衰减还没开始训练就结束了,恒定学习率反而省心。

2.3 选型口诀

我自己总结了一条很朴素的选型路径:不确定怎么选,就用线性 warmup + 余弦退火,warmup 占比设 5%。这个组合在大 batch 和小 batch 下都表现稳定,基本不会出大问题。等你有明确的证据表明需要更激进的行为,再换别的曲线,别一开始就在 warmup 类型上纠结太多时间。

3. 手写一个 warmup 调度器并集成到训练循环

3.1 用 LambdaLR 实现线性 warmup + cosine 退火

最常见的做法是直接用 PyTorch 的 LambdaLR。这个类支持传入一个计算缩放因子的函数,训练循环里每个 step 它都会帮你算出当前学习率。

import math import torch from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR total_steps = 10000 warmup_steps = 500 def lr_lambda(step): if step < warmup_steps: # 线性从 0 爬到 1,step 从0开始所以要+1 return (step + 1) / warmup_steps # 余弦退火到 0.05 * 峰值 progress = (step - warmup_steps) / (total_steps - warmup_steps) return 0.05 + 0.95 * 0.5 * (1.0 + math.cos(math.pi * progress)) optimizer = AdamW(model.parameters(), lr=1e-3) scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)

注意这里lr_lambda返回的是一个缩放倍数,不是绝对学习率。所以你在AdamW里配置的lr=1e-3是峰值学习率,warmup 阶段会从这个峰值的(step+1)/warmup_steps开始爬升。

一个容易被忽略的细节:PyTorch 的LambdaLR在第一个 step 时传进来的step是 0,如果你直接用step / warmup_steps,第一次会得到学习率为 0,虽然不倒腾问题不大,但有些实现里 optimizer 对 0 学习率会直接不动,为了保险我习惯写成(step + 1)。

3.2 自定义 Scheduler 的完整实现

如果你想要一个更可控、能同时响应进度百分比和 warmup 的调度器,也可以继承LRScheduler写一个自定义类。下面这个是我常用的简化版:

import math from torch.optim.lr_scheduler import LRScheduler class WarmupCosineLR(LRScheduler): def __init__(self, optimizer, warmup_steps, total_steps, min_factor=0.05, last_epoch=-1): self.warmup_steps = warmup_steps self.total_steps = total_steps self.min_factor = min_factor super().__init__(optimizer, last_epoch) def get_lr(self): step = max(0, self.last_epoch) if step < self.warmup_steps: return [base_lr * (step + 1) / self.warmup_steps for base_lr in self.base_lrs] progress = (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps) cosine_factor = 0.5 * (1.0 + math.cos(math.pi * progress)) factor = self.min_factor + (1.0 - self.min_factor) * cosine_factor return [base_lr * factor for base_lr in self.base_lrs]

这个类和前面LambdaLR的效果是等价的,但好处是你可以把warmup_steps和total_steps直接封装成参数,后续做配置化调参更方便。如果你的训练脚本里有很多实验配置,我建议优先采用这种类封装的方式。

3.3 集成时的几个关键注意点

调度器的集成看起来简单,但我在不同框架里都踩过不少小坑,这里列三个重点:

第一个是计数单位问题。LRScheduler里的last_epoch默认按 optimizer step 计算,不是按 batch 计算。如果你开了梯度累积,一个 optimizer step 可能对应了好几个 batch,那么调度器里的 step 数应该基于 optimizer step 来设计,否则 warmup 的实际长度会比你想的要短一大截。

第二个是恢复断点。如果你的训练中断了,用 checkpoint 恢复时一定要把 scheduler 的last_epoch或者全局 step 同时恢复。否则 scheduler 会从头再来,而模型参数已经是训练中后期的状态,这时一个低学习率爬升阶段会莫名其妙地打断正常训练节奏,造成 loss 出现“W”形反转。

第三个是打印确认。正式开跑之前,我建议前几十个 step 每个 step 都打印一次当前学习率,确认它的爬升速度和你预想的一致。很多时候你以为设了 warmup,结果因为scheduler.step()调用时机的问题,或者LambdaLR的 step 参数传错了,整个 warmup 其实根本没生效。

4. 热身步数怎么定:不同任务规模下的经验法则

4.1 warmup steps 与总步数、批大小的关系

warmup 步数的设定,理论上并没有一个严格的数学公式,但业界有几个很通用的经验法则。最常见的是按总训练步数的一定比例来设定,通常取 3% 到 10%。比如总步数是 10 万,warmup 步数可以取 3000 到 10000。

还有一个更细的观察是:batch size 越大,warmup 步数往往要越充分。原因是大 batch 下每个 optimizer step 看到的梯度更接近全量梯度,单个 step 的更新质量更高,但也更容易在一开始就沿着错误方向大步跑偏。小 batch 反而因为本身噪声大、有效步长被稀释,对 warmup 的需求没那么强烈。我见过一些大规模预训练任务,总步数几十万,warmup 甚至会占到 10% 以上,这不是因为模型更脆弱,而是因为它们太贵了,不允许第一次就炸掉。

4.2 不同任务形态的推荐区间

按我的经验,可以把任务大致分成几类,分别给个参考区间:

任务类型典型总步数推荐 warmup 占比备注
小规模分类(几十万样本)1000-50005%-8%可快速实验,不要把步数耗太多
中规模视觉任务1万-5万3%-5%线性 warmup + cosine 好用
中大规模语言模型训练10万-50万1%-3%总步数长,占比可以少一点,但步数绝对值不小
超大 batch 预训练50万以上5%-10%宁可多热身,别冒险

这些数值不是拍脑袋,而是我根据多次实验观察到的“损失曲线何时开始明显下降”去反推的。你可以在自己的任务里做个小实验:固定其他所有超参,只改 warmup 步数,观察前 1000 步的 loss 下降速度。如果 warmup 结束后 loss 立即拉到峰值学习率时还能继续稳步下降,说明步数偏短也不错;如果 warmup 结束附近出现 loss 平台或回跳,就说明热身不够充分。

4.3 我见过的两种错误用法

第一种是把 warmup 设为 0 或者小到几乎可以忽略。很多人觉得“反正我学习率也不高”,结果在训练前期照样遇到 NaN。我之前做过一个对比测试,同样配置下,warmup 从 0 调到 100 步后,NaN 出现的概率大幅降低,这让我很长一段时间内都习惯性地保留 100 步以上的 warmup。

第二种是反过来,warmup 太长,占了总步数的 50%。这种配置会让训练在最开始很长一段时间内都用很小的学习率原地磨蹭,整体收敛速度被严重拖慢,而且后面的衰减区间被压缩得太短,模型来不及收敛到满意的状态。我之前在小规模数据上试过一次 40% 的 warmup,最终效果明显比 5% 的版本要差,因为后半段有效学习时间太短了。

5. 深入一点:warmup 与优化器、模型结构的关系

5.1 为什么 Adam 系优化器对 warmup 更敏感

前面提到 Adam 早期二阶动量很小,这里我想再展开一点。Adam 的更新公式里,每一步的有效学习率大致等于lr / sqrt(v + epsilon)。当 v 非常小时,这个值可能比预设的 lr 大得多。虽然 Adam 论文里用了 bias correction,但在前几十步仍然可能出现步长偏大的情况。

所以你会发现一个有意思的现象:SGD 训练里 warmup 没那么重要,甚至不加 warmup 也能跑;但 Adam 系优化器,尤其是 AdamW,几乎都会被建议配上 warmup。PyTorch 里 AdamW 默认的eps=1e-8,在有些半精度训练场景下这个值偏小,会进一步放大早期的数值风险。我个人现在做大规模训练时习惯把eps设到1e-6或1e-5,配合 warmup 一起用,效果比单开 warmup 更稳。

5.2 LAMB 大 batch 训练:warmup 从“可选”变“必须”

LAMB 优化器在大 batch 训练里非常出名,比如超大批次预训练任务里用得很多。LAMB 在 Adam 的基础上加了一个逐层归一化的操作,使得每层更新幅度都相对可控。但正因为如此,它在大 batch 下的有效步长可以做得很大,这时候如果没有 warmup,训练初期非常容易直接失稳。

我记得有一次试过用 LAMB 在 batch size 很大的任务上跑,不开 warmup,loss 大概一百步内必崩。当时我把峰值学习率从1e-3一路降到1e-4,虽然没崩,但收敛速度慢得让人抓狂。后来把学习率调回升到1e-3,只加了 2% 的 warmup,居然又稳又快。这个对比让我非常明确:对于大 batch 训练,warmup 不是避免训练崩坏的技巧,而是支撑高学习率的前提条件。

5.3 什么场景下 warmup 几乎没用

也不是所有场景都必须 warmup。比如你加载了一个预训练好的模型,只做少量步数的微调,且学习率本来就设得很低,那么 warmup 的收益很小。再比如某些用固定初始化方式且网络非常浅的任务,或者强化学习里非平稳分布变化极快、每一步数据都在变化的场景,warmup 的作用也可能被其他机制覆盖。

我在做某些迁移学习任务时,会把 warmup 步数压到极低甚至为 0,因为预训练特征本身已经很稳定,低学习率微调足够安全。这时候强行加很长的 warmup 反而浪费训练时间。所以一个健康的心态是:warmup 是默认配置,但永远要根据你能接受的实验风险去调整。

6. 排错实录:warmup 导致训练再次崩坏的完整链路

6.1 现象:loss 先降后炸,记录仪表盘

有一次我在复现一个文本模型训练任务时,用的是 1000 步 warmup + cosine 退火。日志记录显示:前 120 步 loss 从 6.8 顺利降到 5.2,可视化曲线非常漂亮;但到 160 步左右,loss 出现了明显的回升;到 200 步附近,直接跳到 NaN。

我当时的第一反应是学习率峰值设置太高,于是从3e-4降到1e-4。奇怪的是,这次跑到了 1000 步也没炸,但最终收敛 effect 明显变差,模型在验证集上的指标一直没有提升。这说明降低学习率确实“治好了”NaN,但本质上是用牺牲收敛速度换来的,问题根源并没有找到。

6.2 排查:逐步排除数据、梯度、调度器

我按下面这套流程去排查,你也可以参考:

  1. 先固定数据:把每个 batch 换成随机生成的白噪声数据,跑 300 步。如果依然 NaN,基本可以排除数据问题。
  2. 打印梯度范数:在每个 step 记录total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1e9),然后观察是不是某一步梯度范数突然暴涨。
  3. 单独验证调度器:写一个小脚本,只创建 optimizer 和 scheduler,不跑模型,打印出每一步的学习率。

排查到第三步时我发现了问题:我的调度器在 warmup 结束那一刻虽然到了峰值学习率,但我在step()之前额外做了一次optimizer.zero_grad()和scaler.scale(loss),半精度梯度在 warmup 结束瞬间因为学习率已经是峰值,又叠加了混合精度计算带来的数值误差,导致那一两步的梯度范数直接爆掉。

6.3 修复:warmup + grad clip + 恢复断点的细节

最终修复方案其实很简单:把梯度裁剪的max_norm从 1.0 调到了 0.5,同时把AdamW的eps从默认的1e-8调到了1e-6。这两个修改让训练一下子就稳定了。更关键的是,我把检查点里保存的scheduler.last_epoch恢复逻辑补上,下次恢复训练时不会再出现 warmup 重复跑一遍的问题。

当时还有一个很隐蔽的坑:我在恢复 checkpoint 时,只恢复了 model 和 optimizer,没有恢复 scheduler 的状态,于是 warmup 从第 0 步重新开始。虽然看起来模型还在继续训练,但学习率突然变成了一个很低的值,整个训练节奏被打乱,loss 曲线出现了一个明显的“W”形。这个问题不看学习率曲线根本发现不了,我也因此养成了一个习惯:不管用什么框架,恢复断点时一定把 scheduler 的 state dict 一并恢复,并且在日志里记录当前 global step 和当前学习率。

7. 最后分享一些真正让我收益的小细节

这些细节算不上什么高大上的理论,都是我在实际训练里一点点试出来的经验,希望能帮你少走点弯路。

第一,warmup 步数最好按 optimizer step 计算,不是按 epoch 或 batch 数。如果你的 batch size 是 256,梯度累积 4 步,那么一个 optimizer step 其实对应了 1024 个样本。不同实验之间的 warmup 步数若要可比,必须先统一到这个维度上。我之前在对比实验时因为没注意这一点,差点把两组结果的差异归错原因。

第二,混合精度训练里 warmup 更值得保留。半精度浮点数的表示范围有限,训练前期梯度和更新幅度都很小时,很容易发生下溢;而 warmup 恰好把原始很小的步长变得更为温和,给了模型和优化器一个缓冲。配合torch.cuda.amp.GradScaler使用时,我会把梯度裁剪放在scaler.unscale_之后再执行,顺序错的话裁剪效果会打折扣。

第三,把学习率和梯度范数都记录到可视化面板里。很多人只盯 loss 曲线,但其实学习率曲线和梯度范数曲线能提供更早的异常信号。比如你看到梯度范数在 warmup 结束瞬间出现尖峰,就知道峰值学习率可能偏高了,在 loss 真正炸掉之前还有机会调整。

最后说一个有点反直觉的体会:warmup 不是越短越好,也不是越长越好,它只是一个让你的模型和优化器“对齐节奏”的缓冲区。我现在的默认做法是:第一次跑新任务时,先固定用 5% 比例的线性 warmup 配 cosine 退火,然后只调峰值学习率,其他都先不动。这个套路虽然朴素,但在我经历过的大多数任务里都稳定得让我意外。等有一天你发现不管怎么调 warmup 都救不回来训练时,那大概率不是 warmup 本身的问题,而要去看看你的数据管线和损失函数设计是不是出了更大的问题。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询