Method try_resume_training

internlm/utils/model_checkpoint.py:533–561

Attempt to restore the training state of the last ckpt. Args: lr_scheduler (_LRScheduler): lr_scheduler object. optimizer (Optimizer): optimizer object. lr (float): learning rate. train_state (dict): traing states. train_dl (DataLo

(self, lr_scheduler, optimizer, lr, train_state, train_dl)

            load_model_checkpoint(folder=model_load_path, model=self.model)

    def try_resume_training(self, lr_scheduler, optimizer, lr, train_state, train_dl):
        """Attempt to restore the training state of the last ckpt.

        Args:
            lr_scheduler (_LRScheduler): lr_scheduler object.
            optimizer (Optimizer): optimizer object.
            lr (float): learning rate.
            train_state (dict): traing states.
            train_dl (DataLoader): traning dataloader object
        """
        if self.load_ckpt_folder is not None:
            # load optimzier states.
            if self.load_optimizer:
                load_optimizer_checkpoint(self.load_ckpt_folder, optimizer)
            # load lr scheduler states.
            load_scheduler(self.load_ckpt_folder, lr_scheduler, optimizer, lr, train_state)
            # load training states.
            load_context(self.load_ckpt_folder, train_dl, train_state)
            # load dataloader sampler states.
            if hasattr(train_state, "batch_sampler") and not isinstance(
                train_state.batch_sampler, torch.utils.data.sampler.BatchSampler
            ):
                load_sampler(self.load_ckpt_folder, train_dl.batch_sampler)
            if hasattr(train_state, "data_state_dict"):
                train_dl.dataset.load_state_dict(
                    llm_load(os.path.join(self.load_ckpt_folder, "sampler_0.pt")), ckpt_path=self.load_ckpt_folder
                )
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler

    def save_checkpoint(
        self,