DeepSpeed LRRT（Learning Rate Range Test）指南：用 PyTorch 快速探测安全学习率边界

• 推理引擎
• 大模型

导读

本文基于 DeepSpeed 的 LRRT（Learning Rate Range Test，学习率范围测试）教程，系统讲解如何在线性递增学习率的短程试训中，快速找到"既能加速收敛、又不导致发散"的最大学习率边界。文章覆盖 LRRT 的四个核心配置参数、PyTorch JSON 配置与命令行两种接入方式、底层LRRangeTest调度器实现原理，并以一个从单卡（batch size 512）扩展到四卡（batch size 2048）的大 batch 调优实战为例，演示如何将 LRRT 探测到的边界直接用于配置 1Cycle（OneCycle）学习率调度。读完本文，你可以独立完成一次完整的 LRRT 实验，并为循环学习率调度器设置可靠的 LR 上下限。

什么是 Learning Rate Range Test（LRRT）

学习率范围测试（LRRT，源自论文A Disciplined Approach to Neural Network Hyper-parameters: Part 1）是一种用于发现"在不导致模型发散的前提下可使用的最大学习率"的方法。数据科学家通常对此很感兴趣，因为：

• 大学习率收敛更快：相比小学习率，大学习率通常能让模型更快地达到目标收敛精度；
• 大学习率是循环调度器的基础：诸如 CLR（Cyclical Learning Rate）和 1Cycle 这类学习率调度策略，正是利用"在上下界之间循环摆动"的方式来配合大 batch 训练，而 LRRT 恰好能为它们提供可信的上下界取值。

DeepSpeed 在 PyTorch 框架内原生提供了 LRRT 能力。从仓库源码看，LRRT 被实现为一个学习率调度器（LR scheduler）：DeepSpeed 的训练引擎会在每次 batch 训练后调用lr_scheduler.step()，从而驱动学习率按预设节奏递增（见 engine.py）。

使用前提

要使用 DeepSpeed 的 LRRT，需要满足两个条件：

1. 集成 DeepSpeed 到训练脚本中：按照 Getting Started 教程 完成 DeepSpeed 与训练脚本的集成（即用deepspeed.initialize包裹模型与优化器）；
2. 为模型配置 LRRT 参数：在 DeepSpeed 的模型配置 JSON 中增加scheduler条目，并把 LRRT 的调度参数写入其params中。

两个条件缺一不可：LRRT 调度器由 DeepSpeed 引擎在_configure_lr_scheduler阶段根据 JSON 配置自动实例化（见 engine.py），因此必须先将 DeepSpeed 接入训练流程。

LRRT 的工作原理与四大参数

LRRT 的工作方式可以概括为：按预定步长（step_size）间隔、以预定倍率（step_rate）线性放大学习率。从这个角度看，LRRT 本身也是一种学习率调度——它定义了训练过程中学习率"何时变、怎么变"。

要配置 LRRT，需要设置以下四个参数（对应的源码常量定义在 lr_schedules.py）：

从源码理解递增公式

在LRRangeTest类的实现中（lr_schedules.py），学习率的计算分为"连续模式"与"阶梯模式"两种：

• 连续模式（staircase=false）：间隔因子为(last_batch_iteration + 1) / step_size；
• 阶梯模式（staircase=true）：间隔因子为floor((last_batch_iteration + 1) / step_size)，即每个step_size区间内学习率保持不变，到区间边界才跳变。

两者统一由以下公式计算当前学习率（源码中的_get_increase与get_lr，见 lr_schedules.py）：

increase = 1 + step_rate × interval_factor current_lr = lr_range_test_min_lr × increase

也就是说，每经过step_size步，学习率会翻约(1 + step_rate)倍。以step_rate=5、step_size=200为例，训练 200 步后学习率约为初始值的 6 倍，400 步后约为 11 倍，依此类推。

另外两点值得注意的实现细节：

• lr_range_test_min_lr既可以传单个浮点数（自动广播到所有param_group），也可以传与优化器param_groups等长的列表/元组（长度不一致会抛出ValueError），用于对不同参数组设置不同起点（lr_schedules.py）；
• LRRangeTest.step()应当在每个 batch 之后被调用（而非每个 epoch），它会把计算出的学习率写回优化器各param_group['lr']，并记录last_batch_iteration以便断点续训（state_dict/load_state_dict仅保存该迭代计数，见 lr_schedules.py）。

模型配置变更：在 PyTorch 中启用 LRRangeTest

对 PyTorch 模型而言，LRRT 被实现为一个学习率调度器（PyTorch 1.0.1 及以上版本均支持该特性），因此在模型配置中加入一个"type": "LRRangeTest"的"scheduler"条目即可。

示例调度配置

以下示例展示了一个典型的 LRRT 调度配置：

1. 初始学习率为0.0001；
2. 缩放倍率step_rate为5；
3. 缩放间隔step_size为 200 个训练步；
4. 每个训练步都缩放学习率（即不使用 staircase 阶梯模式）。

"scheduler": { "type": "LRRangeTest", "params": { "lr_range_test_min_lr": 0.0001, "lr_range_test_step_size": 200, "lr_range_test_step_rate": 5, "lr_range_test_staircase": false } }

将该 JSON 片段合并进 DeepSpeed 模型配置（与train_batch_size、optimizer等并列），再用deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())初始化即可。引擎会读取"scheduler"条目，在lr_schedules模块中查找同名类并实例化（见 engine.py）。

命令行参数方式（无需修改 JSON）

除 JSON 配置外，DeepSpeed 还提供了一组等价的命令行参数（定义在 lr_schedules.py 的add_tuning_arguments中）。当你通过命令行传入--lr_schedule LRRangeTest时，引擎会调用override_lr_range_test_params用命令行值覆盖配置：

--lr_schedule LRRangeTest --lr_range_test_min_lr 0.0001 --lr_range_test_step_size 200 --lr_range_test_step_rate 5 --lr_range_test_staircase false

命令行的默认值与类的默认值略有差异：命令行中lr_range_test_step_size默认1000、lr_range_test_min_lr默认0.001、step_rate默认1.0、staircase默认False（lr_schedules.py）。建议显式设置全部四个参数，避免依赖隐式默认值。

断点续训支持

LRRangeTest实现了state_dict()/load_state_dict()（仅保存last_batch_iteration），而 DeepSpeed 引擎在保存 checkpoint 时会同步写入lr_scheduler状态、加载时按load_lr_scheduler_states恢复（见 engine.py）。这意味着 LRRT 实验中断后可精确恢复到中断时刻的学习率，不会因重启而丢失进度。

实战案例：用 LRRT 为大 batch 训练调参

官方教程用一个真实案例展示了 LRRT 的价值：将某个内部生产模型从单卡（batch size 512）扩展到四卡（batch size 2048）时，目标是在相同数据量下用大 batch 达到与小 batch 相当的性能。大 batch 训练的经典难题是收敛变慢，官方团队的做法是使用 DeepSpeed 的 1Cycle 调度教程 配合 LRRT 来完成调参。

实验设计与发现

上图展示了 batch size 2048 下前 9000 个 batch 的实验结果：左图是不同学习率策略下的验证损失，右图是对应时段的学习率取值。通过网格搜索可知，batch size 2048 下最优的固定学习率为0.0002（图中蓝色lr=0.0002线）；两条 LRRT 曲线（橙色lr_range_test_step_rate=5、灰色lr_range_test_step_rate=50）从相同的初始学习率出发，但灰色曲线的学习率增长速度约为橙色的 10 倍，分别被称为"快增长"与"慢增长" LRRT 调度。

由此得到两条关键观察：

1. 较大的学习率在到达某个临界点前明显有利于性能：快增长 LRRT 在 3000 个 batch 时就取得 0.46 的验证损失，而固定学习率在 9000 个 batch 时仍未达到；慢增长 LRRT 直到 6000 个 batch 后才追平该指标，但其性能优势随训练持续扩大。
2. 学习率存在"有效上界"，超出即发散：快增长 LRRT 很快触及上界并发散，慢增长 LRRT 也因同一原因在更晚时刻发散。LRRT 帮助团队仅用不到 2% 的训练数据就快速定位了这些边界，而这些边界信息正是构建学习率调度的关键输入。

用 LRRT 结果配置 1Cycle 调度

基于上述观察，官方团队为 1Cycle 调度配置了学习率边界与周期跨度，最终解决了大 batch 收敛问题，其配置如下：

"OneCycle": { "cycle_min_lr": 0.002, "cycle_max_lr": 0.005, "cycle_first_step_size": 2000, "cycle_second_step_size": 2000, ... }

在官方经验中，1Cycle 调度最关键的四个参数是cycle_min_lr、cycle_max_lr、cycle_first_step_size和cycle_second_step_size，其设定依据如下：

1. 选择慢增长 LRRT（lr_range_test_step_rate=5）的结果来设定cycle_min_lr：因为它取得了最佳损失，而快增长调度发散得相当快，其数据不可靠；
2. cycle_max_lr取 0.005 而非图中更优的更高学习率：虽然图中性能在略高的学习率下仍在改善，但若等到逼近最大学习率再停止，模型可能已处于发散边缘、无法恢复；
3. cycle_first_step_size/cycle_second_step_size均设为 2000：因为学习率增长到 0.005 需要 8000 个 batch，而四卡并行下处理 8000 个 batch 恰好需要 2000 步。

从源码看 1Cycle 如何消费这些参数

1Cycle 调度器（OneCycle类，lr_schedules.py）按cycle_first_step_size + cycle_second_step_size构成一个完整周期：前半周期学习率从cycle_min_lr线性升至cycle_max_lr，后半周期再线性降回，之后进入可选的衰减阶段（decay_lr_rate）。它还会以与学习率相反的相位循环动量（cycle_min_mom/cycle_max_mom，默认 0.8 / 0.9）。对应的单元测试在 test_lr_schedulers.py 中验证了"起始学习率等于cycle_min_lr、在第cycle_first_step_size步达到cycle_max_lr"等关键性质。

测试与验证：如何确认 LRRT 行为符合预期

仓库中的单元测试 test_lr_schedulers.py 覆盖了 LRRT 的多组关键行为，可作为理解与验证 LRRT 行为的参考：

• TestLrRange（test_lr_schedulers.py）：对(min_lr, step_rate, step_size, staircase)的 6 组组合做参数化测试，断言：
  • 起始学习率等于lr_range_test_min_lr；
  • staircase=true时，学习率在每个step_size区间内保持不变（阶梯式递增）；
  • staircase=false时，学习率随步数严格连续递增。
• TestSchedulerOptimizerParity（test_lr_schedulers.py）：验证lr_scheduler.get_lr()与model.get_lr()在每个训练步后完全一致，即调度器确实把学习率写回了优化器。
• TestGetLrBeforeTrain（test_lr_schedulers.py）：验证在训练开始前调用get_lr()不会报错，且训练循环全程可用。

这些测试使用的配置骨架（train_batch_size、optimizer、scheduler、gradient_clipping的 JSON 结构）与本文前面的配置示例一致，可以直接照搬作为你本地验证的最小样例。

实践建议与注意事项

• 把 LRRT 当作一次性的"探测实验"：LRRT 的目的不是完成训练，而是在极少量数据（官方经验中不到 2%）内扫描学习率边界，因此应当用少量步数尽早停止，然后根据 loss 曲线确定安全区间；
• 观察发散点要留有余量：如官方案例所示，最优学习率往往略低于曲线看起来"还能继续提升"的点——因为一旦越过发散边界，模型可能无法恢复。取cycle_max_lr时应保守一些；
• step_rate的选择决定探测精度与速度：较大的step_rate（如 50）能更快扫到上界但容易错过细微变化，较小的step_rate（如 5）更平稳但耗时更长，可先用快扫描粗定位、再用慢扫描精读；
• 结合 DeepSpeed 引擎使用：LRRT 依赖 DeepSpeed 引擎在训练循环内自动调用step()（见 engine.py），因此不要手动重复调用调度器的step()，避免学习率跳变失序；
• 调度器类型清单：DeepSpeed 支持的调度器包括LRRangeTest、OneCycle、WarmupLR、WarmupDecayLR（见 lr_schedules.py），LRRT 探测出的边界可直接服务于OneCycle的上下界配置。

小结

LRRT 以极低的实验成本（少量训练数据）换来了对模型"安全学习率区间"的可靠估计，是与大 batch 训练、循环学习率调度（CLR / 1Cycle）配合使用的关键调参工具。借助 DeepSpeed 的LRRangeTest调度器，你只需在 JSON 配置中增加一个"scheduler"条目，即可在 PyTorch 训练流程中自动完成学习率扫描，并将探测结果无缝转化为 1Cycle 调度的cycle_min_lr/cycle_max_lr与周期跨度，从而解决大 batch 场景下的收敛难题。相关源码入口为 lr_schedules.py 与 engine.py，配套测试见 test_lr_schedulers.py。

• 推理引擎
• 大模型