Method load

olmo/train/checkpointer.py:204–247

Load model, optim, and other training state from a local or remote checkpoint directory created via :meth:`save()` or :meth:`save_async()`.

(
        self,
        dir: PathOrStr,
        model: nn.Module,
        optim: Optimizer = None,
        *,
        load_optimizer_state: Optional[bool] = None,
        load_trainer_state: Optional[bool] = None,
        key_mapping: Optional[Dict[str, str]] = None,
    )

        return write_file(dir, fname, contents, self.save_overwrite)

    def load(
        self,
        dir: PathOrStr,
        model: nn.Module,
        optim: Optimizer = None,
        *,
        load_optimizer_state: Optional[bool] = None,
        load_trainer_state: Optional[bool] = None,
        key_mapping: Optional[Dict[str, str]] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Load model, optim, and other training state from a local or remote checkpoint directory
        created via :meth:`save()` or :meth:`save_async()`.
        """
        dir = normalize_path(dir)

        # Maybe load trainer state.
        trainer_state: Optional[Dict[str, Any]] = None
        if load_trainer_state is not False:
            # Try loading the given rank's state first, then fall back to rank 0 train state if it
            # doesn't exist, which can happen when we're restoring a checkpoint with a different world size.
            for path in (f"{dir}/train/rank{get_global_rank()}.pt", f"{dir}/train/rank0.pt"):
                try:
                    trainer_state = torch.load(cached_path(path, quiet=True), weights_only=False)
                    break
                except FileNotFoundError:
                    pass

            if load_trainer_state is True and trainer_state is None:
                raise FileNotFoundError(f"Missing trainer state in checkpoint dir '{dir}'")

        # Load model and optimizer state.
        model_and_optim_dir: str = f"{dir}/model_and_optim"
        load_model_and_optim_state(
            model_and_optim_dir,
            model,
            optim if load_optimizer_state else None,
            process_group=None,
            key_mapping=key_mapping,
            pre_download=is_url(dir) and self.pre_download,
            work_dir=self.work_dir,
            thread_count=self.load_thread_count,
        )
        return trainer_state

    def _save_train_state(self, dir: PathOrStr, wd: Path, train_state: Dict[str, Any]):
        train_dir = wd / "train"