deepspeed存在一個(gè)bug，即在訓(xùn)練時(shí)不保存調(diào)度器狀態(tài)，因此如果訓(xùn)練中斷后再重新開(kāi)始訓(xùn)練，調(diào)度器還是會(huì)從頭開(kāi)始而不是接著上一個(gè)checkpoint的調(diào)度器狀態(tài)來(lái)訓(xùn)練。這個(gè)bug在deepspeed的github中也有其他人提出：https://github.com/microsoft/DeepSpeed/issues/3875
因此我們需要寫(xiě)一個(gè)保存調(diào)度器狀態(tài)的代碼，才可以解決這個(gè)問(wèn)題。
具體方法是加一個(gè)callback類(lèi)，專(zhuān)門(mén)負(fù)責(zé)保存調(diào)度器的狀態(tài)以及在訓(xùn)練重新開(kāi)始時(shí)加載調(diào)度器的狀態(tài)：
先在訓(xùn)練文件中給trainer加一個(gè)callback文章來(lái)源：http://www.zghlxwxcb.cn/news/detail-704636.html
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 創(chuàng)建檢查點(diǎn)目錄名
checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
# 完整的檢查點(diǎn)目錄路徑
checkpoint_path = os.path.join(args.output_dir, checkpoint_folder)
# 如果目錄不存在，則創(chuàng)建它
if not os.path.exists(checkpoint_path):
os.makedirs(checkpoint_path)
# 完整的保存路徑
save_path = os.path.join(checkpoint_path, SCHEDULER_NAME)
# 保存scheduler狀態(tài)
torch.save(scheduler_state, save_path)
def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
# 如果resume_from_checkpoint設(shè)置了有效路徑
if args.resume_from_checkpoint is not None:
load_path = os.path.join(args.resume_from_checkpoint, SCHEDULER_NAME)
# 如果該路徑下有保存的調(diào)度器狀態(tài)，則加載它
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重新開(kāi)始訓(xùn)練的時(shí)候，學(xué)習(xí)率是接著之前的學(xué)習(xí)率開(kāi)始的（5.5e-7)，而不是從頭開(kāi)始(0.5e-7)：
文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-704636.html
到了這里，關(guān)于解決deepspeed框架的bug：不保存調(diào)度器狀態(tài)，模型訓(xùn)練重啟時(shí)學(xué)習(xí)率從頭開(kāi)始的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
