Repository: allenai/OLMo
File: olmo/checkpoint.py (Python)

Code before:
        load_path: PathOrStr,
        fsdp_model: FSDP,
        optim: Optimizer,
        *,
        local_cache: Optional[PathOrStr] = None,
        load_optimizer_state: bool = True,
    ) -> Dict[str, Any]:
        with FSDP.state_dict_type(
            fsdp_model,
            state_dict_type=StateDictType.FULL_STATE_DICT,
            state_dict_config=FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
            optim_state_dict_config=FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
        ):
            # Load model state.
            log.info("Loading model state...")
            state_dict_to_load, og_keys_to_new = fsdp_model._fsdp_wrapped_module._make_state_dict_compatible(
                load_state_dict(load_path, "model.pt", local_cache=local_cache, map_location="cpu")
            )
            fsdp_model.load_state_dict(state_dict_to_load)
            del state_dict_to_load

            # Load optimizer state.
            gc.collect()
            if load_optimizer_state:
                log.info("Loading optimizer state...")
                if get_global_rank() == 0:
                    optim_state_dict_to_load = self._make_optim_state_dict_compatible(
                        load_state_dict(load_path, "optim.pt", local_cache=local_cache, map_location="cpu"),
                        og_keys_to_new,
                    )
                else:
                    optim_state_dict_to_load = {}
                load_fsdp_optim_state(fsdp_model, optim, optim_state_dict_to_load)

            # Load other state.
            try:
                trainer_state = load_state_dict(load_path, "train.pt", local_cache=local_cache)
            except FileNotFoundError:
                # for backwards compatibility
                trainer_state = load_state_dict(load_path, "other.pt", local_cache=local_cache)
        barrier()
        return trainer_state

    def _make_optim_state_dict_compatible(
        self, optim_state_dict: Dict[str, Any], og_keys_to_new: Dict[str, Set[str]]
    ) -> Dict[str, Any]:
        # This state dict comes in two forms: one where the state keys are integers and one where the
        # keys are fully qualified parameter names. The latter case is easier to deal with here so we
        # first transform the integer key form into the FQN key form.
        if isinstance(next(iter(optim_state_dict["state"].keys())), int):
            id_to_fqn: Dict[int, str] = {}

Review comment (performance):
I think an alternative to this is to set the checkpoint type to not be rank0 only (in the context manager above), and then load optimizer state one local process at a time.

Code after:
    def restore_checkpoint(
        self,
        load_path: PathOrStr,
        fsdp_model: FSDP,
        optim: Optimizer,
        *,
        local_cache: Optional[PathOrStr] = None,
        load_optimizer_state: bool = True,
    ) -> Dict[str, Any]:
        with FSDP.state_dict_type(
            fsdp_model,
            state_dict_type=StateDictType.FULL_STATE_DICT,
            state_dict_config=FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
            optim_state_dict_config=FullOptimStateDictConfig(rank0_only=False, offload_to_cpu=True),
        ):
            # Load model state.
            log.info("Loading model state...")
            state_dict_to_load, og_keys_to_new = fsdp_model._fsdp_wrapped_module._make_state_dict_compatible(
                load_state_dict(load_path, "model.pt", local_cache=local_cache, map_location="cpu")
            )
            fsdp_model.load_state_dict(state_dict_to_load)
            del state_dict_to_load

            # Load optimizer state.
            if load_optimizer_state:
                gc.collect()
                for turn in range(get_local_world_size()):
                    log.info("Loading optimizer state turn %d ...", turn)
                    if turn == get_local_rank():
                        optim_state_dict_to_load = self._make_optim_state_dict_compatible(
                            load_state_dict(load_path, "optim.pt", local_cache=local_cache, map_location="cpu"),
                            og_keys_to_new,
                        )
                        load_fsdp_optim_state(fsdp_model, optim, optim_state_dict_to_load)
                        del optim_state_dict_to_load
                        gc.collect()
                        torch.cuda.empty_cache()
                    barrier()

            # Load other state.
            try:
                trainer_state = load_state_dict(load_path, "train.pt", local_cache=local_cache)
            except FileNotFoundError:
                # for backwards compatibility
                trainer_state = load_state_dict(load_path, "other.pt", local_cache=local_cache)
        barrier()
        return trainer_state

    def _make_optim_state_dict_compatible(
        self, optim_state_dict: Dict[str, Any], og_keys_to_new: Dict[str, Set[str]]
    ) -> Dict[str, Any]: