Prhub

#1775 [Fix] Fix duplicate Megatron LR scheduler resume when optimizer state is not loaded

原始 PR 作者 kaysonyu 合并时间 2026-03-29 13:45 文件变更 1 提交数 2 评论 0 代码增减 +0 / -2

执行摘要

修复 Megatron LR scheduler 重复恢复导致学习率错误的 bug。

根据PR body,当使用--no-load-optim参数时,模型权重从检查点加载,但优化器/scheduler状态故意不恢复,然而新的scheduler仍被快速前进到检查点迭代,导致学习率错误地直接达到min_lr,匹配了观察到的lr=0行为。这是冗余调用引发的bug,需要移除以确保LR调度正确。

该PR变更简单直接,值得快速审阅以确认修复逻辑正确。对于负责megatron_utils模块或涉及LR调度的工程师,这是必要的bugfix,可关注Megatron恢复机制的依赖关系。

讨论亮点

没有review评论或讨论,因此无讨论内容可供分析。

实现拆解

实现非常简单,只修改了slime/backends/megatron_utils/model.py文件中的initialize_model_and_optimizer()函数,删除了opt_param_scheduler.step(increment=iteration * args.global_batch_size)这一行代码,从而消除了重复的scheduler恢复步骤。没有新增代码或其他变动。

文件 模块 状态 重要度
slime/backends/megatron_utils/model.py backends/megatron_utils modified 5.0

关键符号

initialize_model_and_optimizer

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

风险较低,主要在于移除代码可能依赖Megatron的load_checkpoint()是否正确恢复了scheduler状态。如果外部恢复逻辑有误,可能引入回归问题,但根据PR描述,Megatron已负责恢复,因此风险较小。没有新增测试覆盖,但变更简单,回归风险有限。

修复了在特定配置(如--no-load-optim)下学习率错误的问题,避免了训练初期LR直接跳到最小值,提升训练稳定性和效果。影响范围限于使用Megatron后端的用户,特别是那些在不加载优化器状态时恢复训练的场景,对系统其他部分无直接影响。

依赖外部逻辑正确性

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论