github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

feat(ckpt): save ckpt when reach total step count (#486)

pull/493/head
Guoteng 2023-11-09 21:07:16 +08:00 committed by GitHub
parent 5b67db33d0
commit b7ecdba617
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internlm/utils/model_checkpoint.py
View File

@ -776,7 +776,7 @@ now step_count is {train_state.step_count}",
save_ckpts, save_type, now_break = False, CheckpointSaveType.NORMAL_CHECKPOINT, False save_ckpts, save_type, now_break = False, CheckpointSaveType.NORMAL_CHECKPOINT, False
if self.oss_snapshot_freq > 1 and train_state.step_count % self.oss_snapshot_freq == 0: if self.oss_snapshot_freq > 1 and train_state.step_count % self.oss_snapshot_freq == 0:
save_ckpts, save_type = True, CheckpointSaveType.SNAPSHOT_CHECKPOINT save_ckpts, save_type = True, CheckpointSaveType.SNAPSHOT_CHECKPOINT
if train_state.step_count % self.checkpoint_every == 0: if train_state.step_count % self.checkpoint_every == 0 or train_state.step_count == train_state.total_steps:
save_ckpts, save_type = True, CheckpointSaveType.NORMAL_CHECKPOINT save_ckpts, save_type = True, CheckpointSaveType.NORMAL_CHECKPOINT
now_break, singal_save_ckpts, singal_save_type = self.quit_signal_handler(train_state) now_break, singal_save_ckpts, singal_save_type = self.quit_signal_handler(train_state)
if save_ckpts is False: if save_ckpts is False: