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."""
        if not state:
            raise ValueError(
                f"Got FSDPStrategy.load_checkpoint(..., state={state!r}) but a state with at least "
                f" a model instance to reload is required. Pass it in like so:"
                " FSDPStrategy.load_checkpoint(..., state={'model': model, ...})"
            )
        # broadcast the path from rank 0 to ensure all the states are loaded from a common path
        path = Path(self.broadcast(path))

        if isinstance(state, Module):
            from lightning.fabric.strategies.model_parallel import _load_raw_module_state_from_path

            _load_raw_module_state_from_path(path, module=state, world_size=self.world_size, strict=strict)
            return {}

        if isinstance(state, Optimizer):
            raise NotImplementedError(
                "Loading a single optimizer object from a checkpoint is not supported yet with the FSDP strategy."
            )

        from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

        modules = {key: module for key, module in state.items() if _has_fsdp_modules(module)}
        if len(modules) == 0:
            raise ValueError(
                "Could not find a FSDP 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."
            )
        optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}
        if len(modules) > 1:
            raise ValueError(
                "Found multiple FSDP models in the given state. Loading checkpoints with FSDP is"
                " currently limited to a single model per checkpoint. To load multiple models, call the"
                " load method for each model separately with a different path."
            )
        module_key, module = list(modules.items())[0]

        if _is_sharded_checkpoint(path):
            state_dict_ctx = _get_sharded_state_dict_context(module)

            with state_dict_ctx:
                module_state = {module_key: module.state_dict()}
                _distributed_checkpoint_load(module_state, path)
                module.load_state_dict(module_state[module_key], strict=strict)

                if optimizers:
                    from torch.distributed.checkpoint import FileSystemReader

                    # TODO: replace with newer APIs
                    # https://github.com/pytorch/pytorch/issues/119800#issuecomment-1942156271
                    reader = FileSystemReader(path=path)
                    # the optimizer states must be loaded separately
                    for optim_key, optim in optimizers.items():
                        optim_state = load_sharded_optimizer_state_dict(
                            model_state_dict=module_state[module_key],
                            optimizer_key=optim_key,
                            storage_reader=reader,
                        )
                        flattened_osd = FSDP.optim_state_dict_to_load(
                            optim_state_dict=optim_state[optim_key],
                            model=module,
                            optim=optim,
                        )
                        optim.load_state_dict(flattened_osd)

            # Load metadata (anything not a module or optimizer)
            metadata = torch.load(path / _METADATA_FILENAME)
            requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
            _validate_keys_for_strict_loading(requested_metadata_keys, metadata.keys(), strict=strict)
            for key in requested_metadata_keys:
                if key not in metadata:
                    continue
                state[key] = metadata.pop(key)

            # return the remaining metadata that wasn't requested as part of `state`
            return metadata

        if _is_full_checkpoint(path):
            checkpoint = _lazy_load(path)

            from lightning.fabric.strategies.model_parallel import (
                _load_raw_module_state,
                _rekey_optimizer_state_if_needed,
            )

            _load_raw_module_state(checkpoint.pop(module_key), module=module, world_size=self.world_size, strict=strict)

            if isinstance(state, Module):
                return {}

            # Materialize lazy tensors if there are any left in the checkpoint
            # The `torch.Optimizer.load_state_dict` method can't load lazy tensors because of deepcopy pickle issues
            checkpoint = _materialize_tensors(checkpoint)

            # Load optimizer states
            for optim_key, optim in optimizers.items():
                # rank0_only should be false because we need to load the optimizer state on all ranks
                with _get_full_state_dict_context(module, world_size=self.world_size, rank0_only=False):
                    temp_state_dict = _rekey_optimizer_state_if_needed(checkpoint.pop(optim_key), module)
                    optim_state_dict = FSDP.optim_state_dict_to_load(
                        optim_state_dict=temp_state_dict,
                        model=module,
                        optim=optim,
                    )
                    optim.load_state_dict(optim_state_dict)

            requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
            _validate_keys_for_strict_loading(requested_metadata_keys, checkpoint.keys(), strict=strict)

            # Load metadata (anything not a module or optimizer)
            _move_state_into(source=checkpoint, destination=state, keys=requested_metadata_keys)

            # return the remaining metadata that wasn't requested as part of `state`
            return checkpoint

        raise ValueError(
            f"The path {str(path)!r} does not point to a valid checkpoint. Make sure the path points to either a"
            " directory with FSDP checkpoint shards, or a single file with a full checkpoint."
        )
Load the contents from a checkpoint and restore the state of the given objects.
src/lightning/fabric/strategies/fsdp.py