Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/plugins/training_type/deepspeed.py (Python)

Code before:
        if zero_optimization:
            cfg = {
                "zero_allow_untested_optimizer": zero_allow_untested_optimizer,
                "zero_optimization": zero_kwargs,
                **cfg
            }
        return cfg

    def _filepath_to_dir(self, filepath: str) -> str:
        return os.path.dirname(filepath)

    @property
    def deepspeed_engine(self):
        return self.model

    def save_checkpoint(self, checkpoint: Dict, filepath: str) -> None:
        """Save model/training states as a checkpoint file through state-dump and file-write.

        Args:
            checkpoint: The checkpoint state dictionary
            filepath: write-target file's path
        """
        if self.world_size > 1 and self.zero_stage_3:
            save_dir = self._filepath_to_dir(filepath)
            if self.save_full_weights:
                # todo: expose this as a unprotected function in deepspeed
                state_dict = self.deepspeed_engine._zero3_consolidated_fp16_state_dict()
                if self.is_global_zero:
                    # Remove module portion of the checkpoint before saving
                    state_dict = {k.partition('module.')[2]: state_dict[k] for k in state_dict.keys()}
                    checkpoint['state_dict'] = state_dict
                    return super().save_checkpoint(checkpoint, filepath)
                return

            # Use deepspeed's internal checkpointing function to handle partitioned weights across processes
            # dump states as a checkpoint dictionary object
            _exclude_keys = ['state_dict', 'optimizer_states', 'lr_schedulers']
            checkpoint = {k: v for k, v in checkpoint.items() if k not in _exclude_keys}
            self.deepspeed_engine.save_checkpoint(save_dir, client_state=checkpoint)
        else:
            super().save_checkpoint(checkpoint, filepath)

    def restore_model_state_from_ckpt_path(
        self,
        ckpt_path: str,
        map_location: Callable = lambda storage, loc: storage,
    ) -> Tuple[Dict, bool]:
        if not self.save_full_weights and self.world_size > 1:
            # Rely on deepspeed to load the checkpoint and necessary information
            from pytorch_lightning.trainer.states import TrainerState
            stage_is_fit = self.lightning_module.trainer.state == TrainerState.FITTING

Review comment (suggestion):
```suggestion
                # todo: expose this as an unprotected function in deepspeed
```

Code after:
        if zero_optimization:
            cfg = {
                "zero_allow_untested_optimizer": zero_allow_untested_optimizer,
                "zero_optimization": zero_kwargs,
                **cfg
            }
        return cfg

    def _filepath_to_dir(self, filepath: str) -> str:
        return os.path.dirname(filepath)

    @property
    def deepspeed_engine(self):
        return self.model

    def save_checkpoint(self, checkpoint: Dict, filepath: str) -> None:
        """Save model/training states as a checkpoint file through state-dump and file-write.

        Args:
            checkpoint: The checkpoint state dictionary
            filepath: write-target file's path
        """
        if self.world_size > 1 and self.zero_stage_3:
            if self.save_full_weights:
                # todo: expose this as general function in deepspeed
                state_dict = self.deepspeed_engine._zero3_consolidated_fp16_state_dict()
                if self.is_global_zero:
                    # State dict keys will include reference to wrapper LightningDeepSpeedModule
                    # Delete `module` prefix before saving.
                    state_dict = {k.partition('module.')[2]: state_dict[k] for k in state_dict.keys()}
                    checkpoint['state_dict'] = state_dict
                    return super().save_checkpoint(checkpoint, filepath)
                return

            # Use deepspeed's internal checkpointing function to handle partitioned weights across processes
            # dump states as a checkpoint dictionary object
            save_dir = self._filepath_to_dir(filepath)
            _exclude_keys = ['state_dict', 'optimizer_states', 'lr_schedulers']
            checkpoint = {k: v for k, v in checkpoint.items() if k not in _exclude_keys}
            self.deepspeed_engine.save_checkpoint(save_dir, client_state=checkpoint)
        else:
            super().save_checkpoint(checkpoint, filepath)

    def restore_model_state_from_ckpt_path(
        self,
        ckpt_path: str,
        map_location: Callable = lambda storage, loc: storage,
    ) -> Tuple[Dict, bool]:
        if not self.save_full_weights and self.world_size > 1:
            # Rely on deepspeed to load the checkpoint and necessary information
            from pytorch_lightning.trainer.states import TrainerState