解决deepspeed框架的bug:不保存调度器状态,模型训练重启时学习率从头开始

这篇具有很好参考价值的文章主要介绍了解决deepspeed框架的bug:不保存调度器状态,模型训练重启时学习率从头开始。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

deepspeed存在一个bug,即在训练时不保存调度器状态,因此如果训练中断后再重新开始训练,调度器还是会从头开始而不是接着上一个checkpoint的调度器状态来训练。这个bug在deepspeed的github中也有其他人提出:https://github.com/microsoft/DeepSpeed/issues/3875
因此我们需要写一个保存调度器状态的代码,才可以解决这个问题。
具体方法是加一个callback类,专门负责保存调度器的状态以及在训练重新开始时加载调度器的状态:
先在训练文件中给trainer加一个callback

from smoe.callbacks.save_model import SchedulerStateCallback
trainer.add_callback(SchedulerStateCallback)
class SchedulerStateCallback(TrainerCallback):
    def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        if os.environ.get("RANK", "0") == "0":
            #scheduler = kwargs['lr_scheduler']
            scheduler = kwargs.get("lr_scheduler")
            if scheduler is None:
                return 
            scheduler_state = scheduler.state_dict()
            #save_path = os.path.join(args.output_dir, SCHEDULER_NAME)
             # 使用 PREFIX_CHECKPOINT_DIR 和 global_step 创建检查点目录名
            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
            # 完整的检查点目录路径
            checkpoint_path = os.path.join(args.output_dir, checkpoint_folder)
            # 如果目录不存在,则创建它
            if not os.path.exists(checkpoint_path):
                os.makedirs(checkpoint_path)
            # 完整的保存路径
            save_path = os.path.join(checkpoint_path, SCHEDULER_NAME)
            # 保存scheduler状态
            torch.save(scheduler_state, save_path)

    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        # 如果resume_from_checkpoint设置了有效路径
        if args.resume_from_checkpoint is not None:
            load_path = os.path.join(args.resume_from_checkpoint, SCHEDULER_NAME)
            # 如果该路径下有保存的调度器状态,则加载它
            if os.path.exists(load_path):
                #scheduler = kwargs['lr_scheduler']
                scheduler = kwargs.get("lr_scheduler")
                if scheduler is None:
                    return 
                scheduler_state = torch.load(load_path)
                scheduler.load_state_dict(scheduler_state)

解决效果如下,我们可以看到,在chaeckpoint10重新开始训练的时候,学习率是接着之前的学习率开始的(5.5e-7),而不是从头开始(0.5e-7):
解决deepspeed框架的bug:不保存调度器状态,模型训练重启时学习率从头开始,deepspeed解决deepspeed框架的bug:不保存调度器状态,模型训练重启时学习率从头开始,deepspeed文章来源地址https://www.toymoban.com/news/detail-704636.html

到了这里,关于解决deepspeed框架的bug:不保存调度器状态,模型训练重启时学习率从头开始的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包