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: Save model, optimizer, and other state in the provided checkpoint directory.
  
  If the user specifies sharded checkpointing, the directory will contain one file per process, with model- and
  optimizer shards stored per file. If the user specifies full checkpointing, the directory will contain a
  consolidated checkpoint combining all of the sharded checkpoints.
docstr_tok: ["Save","model",",","optimizer",",","and","other","state","in","the","provided","checkpoint","directory",".","If","the","user","specifies","sharded","checkpointing",",","the","directory","will","contain","one","file","per","process",",","with","model-","and","optimizer","shards","stored","per","file",".","If","the","user","specifies","full","checkpointing",",","the","directory","will","contain","a","consolidated","checkpoint","combining","all","of","the","sharded","checkpoints","."]
code: def save_checkpoint(
      self,
      path: _PATH,
      state: Dict[str, Union[Module, Optimizer, Any]],
      storage_options: Optional[Any] = None,
      filter: Optional[Dict[str, Callable[[str, Any], bool]]] = None,
  ) -> None:
      """Save model, optimizer, and other state in the provided checkpoint directory.
  
      If the user specifies sharded checkpointing, the directory will contain one file per process, with model- and
      optimizer shards stored per file. If the user specifies full checkpointing, the directory will contain a
      consolidated checkpoint combining all of the sharded checkpoints.
  
      """
      if not _TORCH_GREATER_EQUAL_2_0:
          raise NotImplementedError(
              "Saving and loading checkpoints with the `XLAFSDPStrategy` is not supported in PyTorch < 2.0."
              " Please upgrade `torch`."
          )
      # broadcast the path from rank 0 to ensure all the states are saved in a common path
      path = Path(self.broadcast(path))
      if path.is_dir() and any(path.iterdir()):
          raise FileExistsError(
              f"The checkpoint directory already exists and is not empty: {path}"
          )
      from torch_xla.distributed.fsdp import (
          XlaFullyShardedDataParallel as XLAFSDP,
      )
  
      modules = [
          module
          for module in state.values()
          if isinstance(module, XLAFSDP)
      ]
      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: `save_checkpoint(..., state={'model': model, ...})`. Make sure"
              " you set up the model (and optimizers if any) through the strategy before saving the checkpoint."
          )
      if len(modules) > 1:
          raise ValueError(
              "Found multiple XLAFSDP modules in the given state. Saving checkpoints with FSDP is"
              " currently limited to a single model per checkpoint. To save multiple models, call the"
              " save method for each model separately with a different path."
          )
      import torch_xla.core.xla_model as xm
  
      # ensure model parameters are updated
      xm.mark_step()
  
      parallel_devices = self.parallel_devices
      assert parallel_devices is not None
      if self._sequential_save:
          # each host runs this in parallel, but the ranks in the host run it sequentially
          for rank in range(len(parallel_devices)):
              if rank == self.local_rank:
                  self._save_checkpoint_shard(
                      path, state, storage_options, filter
                  )
              self.barrier(f"wait-for-{rank}-save")
      else:
          self._save_checkpoint_shard(
              path, state, storage_options, filter
          )
  
      if self._state_dict_type == "full":
          ckpt_prefix = str(path / "checkpoint")
          ckpt_suffix = "_rank-*-of-*.pth"
          if len(parallel_devices) != self.world_size:  # multihost
              raise OSError(
                  "Multihost setups do not have a shared filesystem, so the checkpoint shards cannot be consolidated"
                  " into a single checkpoint after saving them. Please switch to"
                  " `XLAFSDPStrategy(state_dict_type='sharded')`. TIP: You can consolidate them manually by getting"
                  " them together into a single directory and running `python -m"
                  f" torch_xla.distributed.fsdp.consolidate_sharded_ckpts --ckpt_prefix {ckpt_prefix!r} --ckpt_suffix"
                  f" {ckpt_suffix!r} --save_path 'path/to/consolidated.ckpt'`."
              )
  
          from torch_xla.distributed.fsdp import (
              consolidate_sharded_model_checkpoints,
          )
  
          self.barrier("before_ckpt_consolidation")
          if self.is_global_zero:
              save_path = path.parent / "consolidated.ckpt"
              # save consolidated checkpoint separate to the shards
              consolidate_sharded_model_checkpoints(
                  ckpt_prefix, ckpt_suffix, str(save_path)
              )
              # remove the shards directory
              self.checkpoint_io.remove_checkpoint(path)
              # mv the consolidated checkpoint where the user would expect it
              get_filesystem(save_path).mv(str(save_path), str(path))
          self.barrier("after_ckpt_consolidation")
  
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 28 for calling the XLAFSDPStrategy obj's underlying member method code with example usage: obj.save_checkpoint(path, state, storage_options, filter) without return types