class_docstr: Strategy for training multiple XLA devices using the
  :func:`torch_xla.distributed.xla_fully_sharded_data_parallel.XlaFullyShardedDataParallel` method.
  
  .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
  
  For more information check out https://github.com/pytorch/xla/blob/master/docs/fsdp.md
  
  Args:
      auto_wrap_policy: Same as ``auto_wrap_policy`` parameter in
          :class:`torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`.
          For convenience, this also accepts a set of the layer classes to wrap.
      activation_checkpointing_policy: Used when selecting the modules for
          which you want to enable activation checkpointing. Enabling this can free up a significant amount of memory
          at the cost of speed since activations in these layers need to be recomputed during backpropagation.
          This accepts a set of the layer classes to wrap.
  
      state_dict_type: The format in which the state of the model and optimizers gets saved into the checkpoint.
  
          - ``"full"``: The full weights and optimizer states get assembled on rank 0 and saved to a single file.
          - ``"sharded"``: Each rank saves its shard of weights and optimizer states to a file. The checkpoint is
            a folder with files for each shard in the host. Note that TPU VM multihost does not have a shared
            filesystem.
  
      sequential_save: With this enabled, individual ranks consecutively save their state dictionary shards, reducing
          peak system RAM usage, although it elongates the saving process.
      \**kwargs: See available parameters in :class:`torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`.
class_docstr_tok: ["Strategy","for","training","multiple","XLA","devices","using","the",":","func",":","`","torch_xla.distributed.xla_fully_sharded_data_parallel.XlaFullyShardedDataParallel","`","method",".","..","warning",":",":","This","is","an",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","feature",".","For","more","information","check","out","https",":","\/\/github.com\/pytorch\/xla\/blob\/master\/docs\/fsdp.md","Args",":","auto_wrap_policy",":","Same","as","``","auto_wrap_policy","``","parameter","in",":","class",":","`","torch_xla.distributed.fsdp.XlaFullyShardedDataParallel","`",".","For","convenience",",","this","also","accepts","a","set","of","the","layer","classes","to","wrap",".","activation_checkpointing_policy",":","Used","when","selecting","the","modules","for","which","you","want","to","enable","activation","checkpointing",".","Enabling","this","can","free","up","a","significant","amount","of","memory","at","the","cost","of","speed","since","activations","in","these","layers","need","to","be","recomputed","during","backpropagation",".","This","accepts","a","set","of","the","layer","classes","to","wrap",".","state_dict_type",":","The","format","in","which","the","state","of","the","model","and","optimizers","gets","saved","into","the","checkpoint",".","-","``","``","full","''","``",":","The","full","weights","and","optimizer","states","get","assembled","on","rank","0","and","saved","to","a","single","file",".","-","``","``","sharded","''","``",":","Each","rank","saves","its","shard","of","weights","and","optimizer","states","to","a","file",".","The","checkpoint","is","a","folder","with","files","for","each","shard","in","the","host",".","Note","that","TPU","VM","multihost","does","not","have","a","shared","filesystem",".","sequential_save",":","With","this","enabled",",","individual","ranks","consecutively","save","their","state","dictionary","shards",",","reducing","peak","system","RAM","usage",",","although","it","elongates","the","saving","process",".","\\","*","*","kwargs",":","See","available","parameters","in",":","class",":","`","torch_xla.distributed.fsdp.XlaFullyShardedDataParallel","`","."]
docstr: Given a folder, load the contents from a checkpoint and restore the state of the given objects.
  
  The strategy currently only supports saving and loading sharded checkpoints which are stored in form of a
  directory of multiple files rather than a single file.
docstr_tok: ["Given","a","folder",",","load","the","contents","from","a","checkpoint","and","restore","the","state","of","the","given","objects",".","The","strategy","currently","only","supports","saving","and","loading","sharded","checkpoints","which","are","stored","in","form","of","a","directory","of","multiple","files","rather","than","a","single","file","."]
code: def load_checkpoint(
      self,
      path: _PATH,
      state: Optional[
          Union[
              Module,
              Optimizer,
              Dict[str, Union[Module, Optimizer, Any]],
          ]
      ] = None,
      strict: bool = True,
  ) -> Dict[str, Any]:
      """Given a folder, load the contents from a checkpoint and restore the state of the given objects.
  
      The strategy currently only supports saving and loading sharded checkpoints which are stored in form of a
      directory of multiple files rather than a single file.
  
      """
      if not _TORCH_GREATER_EQUAL_2_0:
          raise NotImplementedError(
              "Saving and loading checkpoints with the `FSDPStrategy` is not supported in PyTorch < 2.0."
              " Please upgrade `torch` or file an issue: `https://github.com/Lightning-AI/lightning/issues`."
          )
      if not state:
          raise ValueError(
              f"Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})` but a state with at least "
              " 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, Optimizer)):
          raise NotImplementedError(
              "Loading a single module or optimizer object from a checkpoint"
              " is not supported yet with the XLAFSDP strategy."
          )
  
      from torch_xla.distributed.fsdp import (
          XlaFullyShardedDataParallel as XLAFSDP,
      )
  
      modules = {
          key: module
          for key, module in state.items()
          if isinstance(module, XLAFSDP)
      }
      optimizers = {
          key: optim
          for key, optim in state.items()
          if isinstance(optim, Optimizer)
      }
      if self._state_dict_type == "sharded":
          file = (
              path
              / f"checkpoint_rank-{self.global_rank:08d}-of-{self.world_size:08d}.pth"
          )
          if not file.is_file():
              raise ValueError(
                  f"The path {str(file)!r} does not point to valid sharded checkpoints. Make sure the path points to"
                  " a directory with XLAFSDP checkpoint shards."
              )
          if len(modules) == 0:
              raise ValueError(
                  "Could not find a XLAFSDP 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(modules) > 1:
              raise ValueError(
                  "Found multiple XLAFSDP modules 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 = list(modules.items())[0]
          sharded_ckpt = torch.load(file)
  
          module.load_state_dict(sharded_ckpt["model"], strict=strict)
          for opt_key, opt in optimizers.items():
              opt.load_state_dict(sharded_ckpt[opt_key])
  
          # Load anything leftover from sharded_ckpt
          loaded_metadata_keys = (
              sharded_ckpt.keys() - modules.keys() - optimizers.keys()
          )
          requested_metadata_keys = (
              state.keys() - modules.keys() - optimizers.keys()
          )
          _validate_keys_for_strict_loading(
              requested_metadata_keys,
              loaded_metadata_keys,
              strict=strict,
          )
          for key in requested_metadata_keys:
              if key in loaded_metadata_keys:
                  state[key] = sharded_ckpt[key]
                  loaded_metadata_keys.remove(key)
  
          metadata = {}
          if len(loaded_metadata_keys):
              for key in loaded_metadata_keys:
                  metadata[key] = sharded_ckpt[key]
  
          # remove "shard_metadata" that is loaded in
          if "shard_metadata" in metadata:
              metadata.pop("shard_metadata")
  
          return metadata
  
      if self._state_dict_type == "full":
          if not path.is_file():
              raise ValueError(
                  f"The path {str(path)!r} does not point to a valid full checkpoint. Make sure the path points to a"
                  " directory with a full XLAFSDP checkpoint."
              )
          if (
              len(optimizers) > 0
              or len(state.keys() - modules.keys() - optimizers.keys())
              > 0
          ):
              rank_zero_warn(
                  "Loading a full checkpoint will only load the full model."
                  " The optimizer and any additional metadata are not included."
              )
          if len(modules) > 0:
              raise ValueError(
                  "Found a XLAFSDP model in the provided checkpoint state."
                  " Please provide the model without any XLAFSDP wrapper."
              )
          if "model" not in state or not isinstance(
              model := state["model"], torch.nn.Module
          ):
              raise NotImplementedError(
                  "XLAFSDP only supports a single model instance with 'model' as the key."
              )
          full_ckpt = torch.load(path)
          model.load_state_dict(full_ckpt.pop("model"), strict=strict)
          return full_ckpt
  
      raise ValueError(
          f"Unknown state_dict_type: {self._state_dict_type}"
      )
  
imports: import io
  from contextlib import ExitStack, nullcontext
  from functools import partial
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, Callable, ContextManager, Dict, List, Literal, Optional, Set, Tuple, Type, Union
  import torch
  from torch import Tensor
  from torch.nn import Module
  from torch.optim import Optimizer
  from torch.utils.data import DataLoader
  from typing_extensions import override
  from lightning.fabric.accelerators import Accelerator
  from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
  from lightning.fabric.plugins import XLAPrecision
  from lightning.fabric.plugins.environments import XLAEnvironment
  from lightning.fabric.plugins.io.xla import XLACheckpointIO
  from lightning.fabric.strategies import ParallelStrategy, _StrategyRegistry
  from lightning.fabric.strategies.fsdp import _apply_filter
  from lightning.fabric.strategies.launchers.xla import _XLALauncher
  from lightning.fabric.strategies.strategy import TBroadcast, _BackwardSyncControl, _Sharded, _validate_keys_for_strict_loading
  from lightning.fabric.utilities.cloud_io import get_filesystem
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_1_13, _TORCH_GREATER_EQUAL_2_0
  from lightning.fabric.utilities.init import _EmptyInit
  from lightning.fabric.utilities.rank_zero import rank_zero_only, rank_zero_warn
  from lightning.fabric.utilities.types import _PATH, Optimizable, ReduceOp
label_desc: Use image node_id 30 for calling the XLAFSDPStrategy obj's underlying member method code with example usage: obj.load_checkpoint(path, state, strict) and returns: metadata, full_ckpt