ScheduleFree优化器在训练中断恢复时的注意事项

背景介绍

在深度学习模型训练过程中，使用ScheduleFree优化器（如AdamWScheduleFree和RAdamScheduleFree）时，训练中断后的恢复操作需要特别注意。这类优化器与传统优化器不同，它们移除了学习率调度机制，转而采用更稳定的更新策略，这使得它们在恢复训练时表现出独特的行为特性。

问题现象

用户在实际使用kohya sd-scripts进行模型训练时发现，当使用ScheduleFree优化器（特别是AdamWScheduleFree）并启用梯度累积时，训练中断后恢复会出现以下问题：

1. 生成的样本图像立即出现明显质量下降，表现为卡通化和对比度异常
2. 损失函数在第一步出现大幅下降
3. 虽然后续训练能部分恢复，但网格状伪影持续存在
4. 这种现象在使用梯度累积时尤为明显，而batch size为1无梯度累积时则不会出现

原因分析

这种现象的根本原因在于ScheduleFree优化器的特性：

1. 动量状态初始化：ScheduleFree优化器依赖动量机制来稳定训练，中断后重新初始化会导致动量状态丢失
2. 梯度累积影响：梯度累积会放大初始几步的更新幅度，加剧了状态不一致的问题
3. 预热必要性：虽然ScheduleFree设计上不需要学习率预热，但在恢复训练时，优化器状态需要渐进调整

解决方案

针对这一问题，实践中发现两种有效解决方法：

1. 手动添加预热步骤：通过--optimizer_args "warmup_steps=5"参数强制添加5步预热，缓解初始状态冲击
2. 改用RAdamScheduleFree：该优化器变体内置了自动预热机制，能更好地处理训练恢复场景

最佳实践建议

基于这些经验，我们建议在使用ScheduleFree优化器时：

1. 对于关键训练任务，优先考虑使用RAdamScheduleFree而非AdamWScheduleFree
2. 如果必须使用AdamWScheduleFree，确保在恢复训练时添加适当的预热步骤
3. 定期保存优化器状态，而不仅仅是模型权重，以便更完整地恢复训练
4. 对于使用梯度累积的场景，要特别注意初始几步的监控

技术原理深入

ScheduleFree优化器通过消除学习率调度来简化超参数调整，其核心思想是：

1. 使用指数移动平均(EMA)来稳定参数更新
2. 依赖优化器自身的自适应机制而非外部调度
3. 在稳定状态下提供更平滑的收敛特性

正是这些特性使得它在恢复训练时对初始状态更加敏感，特别是当优化器状态未被保存或正确恢复时。理解这一原理有助于用户更好地应对各种训练场景。