Repository: microsoft/DeepSpeed
File: deepspeed/runtime/engine.py (Python)

Code before:

        load_zero_checkpoint = load_optimizer_states and load_path is not None and (self.zero_optimization()
                                                                                    or self.bfloat16_enabled())
        if load_zero_checkpoint:
            success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states)
            if not success:
                self.optimizer._restore_from_bit16_weights()

        if self.zero_has_nvme_offload():
            from shutil import copytree, disk_usage
            offload_dir = self.optimizer.optimizer_swapper.swap_folder
            offload_ckpt_dir = os.path.join(load_dir, tag, "offloaded_tensors")
            _, _, free = disk_usage(offload_dir)
            logger.info(
                f"Copying NVMe offload checkpoint from {offload_ckpt_dir} to {offload_dir}, {free / 1e9:,.2f} GB free on target filesystem..."
            )
            copytree(offload_ckpt_dir, offload_dir, dirs_exist_ok=True)
            _, _, free = disk_usage(offload_dir)
            logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem")
            self.optimizer.reset_swap_buffers()

        if self._optimizer_has_ckpt_event_epilogue():
            self.optimizer.checkpoint_event_epilogue()

        if self.load_universal_checkpoint():
            self.optimizer.update_lp_params()
            if load_zero_checkpoint:
                self.update_optimizer_step(step=client_states['iteration'] + 1)

        return load_path, client_states

    def _load_checkpoint(self,
                         load_dir,
                         tag,
                         load_module_strict=True,
                         load_optimizer_states=True,
                         load_lr_scheduler_states=True,
                         load_module_only=False,
                         custom_load_fn=None):

        from deepspeed.runtime.state_dict_factory import SDLoaderFactory

        ckpt_list = self._get_all_ckpt_names(load_dir, tag)
        sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine)

        is_pipe_parallel = isinstance(self.module, PipelineModule)

        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
        load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel)

        if checkpoint is None:

Review comment (suggestion):
Replace "offloaded_tensors" with symbolic constant, such as `OFFLOADED_TENSORS_DIR`, which should be defined in deepspeed/checkpoints/constants.py

Code after:
                    )
                    return None, None

        if self._optimizer_has_ckpt_event_prologue():
            # Prepare for checkpoint load by ensuring all parameters are partitioned
            self.optimizer.checkpoint_event_prologue()

        load_path, client_states = self._load_checkpoint(load_dir,
                                                         tag,
                                                         load_module_strict=load_module_strict,
                                                         load_optimizer_states=load_optimizer_states,
                                                         load_lr_scheduler_states=load_lr_scheduler_states,
                                                         load_module_only=load_module_only,
                                                         custom_load_fn=custom_load_fn)

        load_zero_checkpoint = load_optimizer_states and load_path is not None and (self.zero_optimization()
                                                                                    or self.bfloat16_enabled())
        if load_zero_checkpoint:
            success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states)
            if not success:
                self.optimizer._restore_from_bit16_weights()

        if self.zero_has_nvme_offload():
            from shutil import copytree, disk_usage
            offload_dir = self.optimizer.optimizer_swapper.swap_folder
            offload_ckpt_dir = os.path.join(load_dir, tag, "offloaded_tensors")
            _, _, free = disk_usage(offload_dir)
            logger.info(
                f"Copying NVMe offload checkpoint from {offload_ckpt_dir} to {offload_dir}, {free / 1e9:,.2f} GB free on target filesystem..."
            )
            copytree(offload_ckpt_dir, offload_dir, dirs_exist_ok=True)
            _, _, free = disk_usage(offload_dir)
            logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem")
            self.optimizer.reset_swap_buffers()

        if self._optimizer_has_ckpt_event_epilogue():
            self.optimizer.checkpoint_event_epilogue()

        if self.load_universal_checkpoint():
            self.optimizer.update_lp_params()
            if load_zero_checkpoint:
                self.update_optimizer_step(step=client_states['iteration'] + 1)

        return load_path, client_states

    def _load_checkpoint(self,
                         load_dir,
                         tag,
                         load_module_strict=True,
                         load_optimizer_states=True,
                         load_lr_scheduler_states=True,