上海 网站备案,平面设计,中国建设银行的网站特色,一个网站如何做cdn加速器deepspeed存在一个bug#xff0c;即在训练时不保存调度器状态#xff0c;因此如果训练中断后再重新开始训练#xff0c;调度器还是会从头开始而不是接着上一个checkpoint的调度器状态来训练。这个bug在deepspeed的github中也有其他人提出#xff1a;https://github.com/mic…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)
