def load_checkpoint(
        self,
        path: _PATH,
        state: Optional[Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]]]] = None,
        strict: bool = True,
    ) -> dict[str, Any]:
        """Load the contents from a checkpoint and restore the state of the given objects.

        Args:
            path: A path to where the file is located
            state: A dictionary of objects whose state will be restored in-place from the checkpoint path.
                This should contain exactly one model, and the model must already be set up by DeepSpeed.
            strict: Whether to enforce that the keys in `state` match the keys in the checkpoint.

        Returns:
            Dictionary with the state inside DeepSpeed's engine

        Raises:
            ValueError:
                If no state is provided, when no :class:`deepspeed.DeepSpeedEngine` objects were found in the
                state, or when multiple :class:`deepspeed.DeepSpeedEngine` objects were found.
            RuntimeError:
                If DeepSpeed was unable to load the checkpoint due to missing files or because the checkpoint is
                not in the expected DeepSpeed format.

        """
        if isinstance(state, (Module, Optimizer)) or self.load_full_weights and self.zero_stage_3:
            # This code path to enables loading a checkpoint from a non-deepspeed checkpoint or from
            # a consolidated checkpoint
            path = self.broadcast(path)
            return super().load_checkpoint(path=path, state=state, strict=strict)

        if not state:
            raise ValueError(
                f"Got DeepSpeedStrategy.load_checkpoint(..., state={state!r}) but a state with at least "
                f" a model instance to reload is required. Pass it in like so:"
                " DeepSpeedStrategy.load_checkpoint(..., state={'model': model, ...})"
            )
        _validate_checkpoint_directory(path)

        engines = _get_deepspeed_engines_from_state(state)
        if len(engines) == 0:
            raise ValueError(
                "Could not find a DeepSpeed model in the provided checkpoint state. Please provide the model as"
                " part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure"
                " you set up the model (and optimizers if any) through the strategy before loading the checkpoint."
            )
        if len(engines) > 1:
            raise ValueError(
                "Found multiple DeepSpeed engine modules in the given state. Saving and loading checkpoints"
                " with DeepSpeed is currently limited to a single model per checkpoint. To load multiple model"
                " states, call the load method for each model checkpoint separately."
            )
        engine = engines[0]

        if _DEEPSPEED_GREATER_EQUAL_0_14_1:
            from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
        else:
            from deepspeed.runtime import DeepSpeedOptimizer

        optimzer_state_requested = any(isinstance(item, (Optimizer, DeepSpeedOptimizer)) for item in state.values())

        torch.cuda.empty_cache()
        _, client_state = engine.load_checkpoint(
            path,
            tag="checkpoint",
            load_optimizer_states=optimzer_state_requested,
            load_lr_scheduler_states=False,
            load_module_strict=strict,
        )

        if client_state is None:
            raise RuntimeError(
                "DeepSpeed was unable to load the checkpoint. Ensure you passed in a DeepSpeed compatible checkpoint"
                " or a single checkpoint file by setting `DeepSpeedStrategy(..., load_full_weights=True)`."
            )

        # `Engine.load_checkpoint` adds useless keys 'optimizer' and 'lr_scheduler' to the client state; remove
        # them to avoid name collision with user state
        keys = set(client_state) & set(state) - {"optimizer", "lr_scheduler"}
        _move_state_into(source=client_state, destination=state, keys=keys)
        return client_state
Load the contents from a checkpoint and restore the state of the given objects.

        Args:
            path: A path to where the file is located
            state: A dictionary of objects whose state will be restored in-place from the checkpoint path.
                This should contain exactly one model, and the model must already be set up by DeepSpeed.
            strict: Whether to enforce that the keys in `state` match the keys in the checkpoint.

        Returns:
            Dictionary with the state inside DeepSpeed's engine

        Raises:
            ValueError:
                If no state is provided, when no :class:`deepspeed.DeepSpeedEngine` objects were found in the
                state, or when multiple :class:`deepspeed.DeepSpeedEngine` objects were found.
            RuntimeError:
                If DeepSpeed was unable to load the checkpoint due to missing files or because the checkpoint is
                not in the expected DeepSpeed format.

        
src/lightning/fabric/strategies/deepspeed.py