Method try_load_model

internlm/utils/model_checkpoint.py:495–531

(self, current_time="")

        return latest_checkpoint

    def try_load_model(self, current_time=""):
        model_load_path = None

        if self.load_ckpt_folder and self.load_model_only_folder:
            raise ValueError(
                "Error, try to use both load_ckpt_folder and load_model_only_folder paths, \
if you only need to load model weights (for example starting an SFT task for the first time), \
set load_model_only_folder path, if you need to resume training from ckpt, \
set load_ckpt_folder or use default value \
(if is the default value, internlm will try to load the latest ckpt from save_ckpt_folder)"
            )

        if self.load_ckpt_folder:
            if gpc.is_rank_for_log():
                logger.info(
                    f"===========Resume training from `{self.load_ckpt_folder}` {current_time} on host:"
                    f"{socket.gethostname()}==========="
                )
            model_load_path = self.load_ckpt_folder
        elif self.load_model_only_folder:
            if gpc.is_rank_for_log():
                logger.info(
                    f"===========Load Model from `{self.load_model_only_folder}` {current_time} on host:"
                    f"{socket.gethostname()}==========="
                )
            model_load_path = self.load_model_only_folder
        else:
            if gpc.is_rank_for_log():
                logger.info(
                    f"===========New Run {current_time} on host:{socket.gethostname()},rank={gpc.get_global_rank()},"
                    f"tp={gpc.get_local_rank(ParallelMode.TENSOR)},pp={gpc.get_local_rank(ParallelMode.PIPELINE)},"
                    f"dp={gpc.get_local_rank(ParallelMode.DATA)}==========="
                )

        # Loading model weights must be done before zero is initialized.
        if model_load_path is not None:
            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.