Megatron-LM训练中学习率调度器的覆盖问题解析

在大型语言模型训练过程中，学习率调度是一个关键的超参数配置。本文针对Megatron-LM框架中一个典型的学习率调度问题展开分析，该问题出现在分阶段训练时尝试覆盖学习率参数的情况下。

问题现象

当用户进行分阶段训练时，可能会遇到以下情况：

1. 第一阶段训练使用较高学习率（如3e-4）和最小学习率（3e-5）
2. 第二阶段希望降低学习率（如3e-5）和最小学习率（2e-5）
3. 虽然设置了--override-opt_param-scheduler等参数，但实际训练时学习率仍保持第一阶段的值

技术原理

Megatron-LM的训练恢复机制包含多个组件：

• 模型参数（默认加载）
• 优化器状态（默认加载）
• 数据加载器状态（可通过--reset-dataloader重置）
• 迭代计数（可通过--reset-iteration重置）
• 学习率调度器（可通过--override-opt_param-scheduler覆盖）

关键发现

问题的根源在于优化器状态的加载机制。即使覆盖了学习率调度器参数，优化器内部保存的动量等信息仍会保持之前的状态，这可能导致学习率表现不符合预期。

解决方案

针对不同场景，可采取以下方法：

1. 完全重置优化器
   使用--no-load-optim参数跳过优化器状态的加载，这将使优化器完全重新初始化，包括学习率在内的所有参数都将使用新配置。

2. 部分覆盖优化器参数
   若需要保留优化器的动量等状态，但仅修改学习率，可考虑：

   • 修改源码中优化器加载逻辑
   • 在训练脚本中添加学习率强制覆盖代码

3. 分阶段训练最佳实践

   • 明确每个训练阶段的目标
   • 合理规划学习率衰减策略
   • 在阶段转换时进行充分验证

实施建议

对于大多数用户，推荐采用第一种方案，即在第二阶段训练时添加--no-load-optim参数。这种方法简单可靠，能确保学习率参数被正确应用。