class_docstr: Strategy for Fully Sharded Data Parallel provided by torch.distributed.
  
  .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
  
  Fully Sharded Training shards the entire model across all available GPUs, allowing you to scale model
  size, whilst using efficient communication to reduce overhead. In practice, this means we can remain
  at parity with PyTorch DDP, whilst scaling our model sizes dramatically. The technique is similar
  to ZeRO-Stage 3.
  
  For more information check out
  `this blogpost <https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api>`__.
  
  Defaults have been set and options have been exposed, but may require configuration
  based on your level of memory/speed efficiency. We suggest having a look at
  `this tutorial <https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html>`__ for more information.
  
  Arguments:
      cpu_offload: See ``cpu_offload`` parameter in :class:`torch.distributed.fsdp.FullyShardedDataParallel`.
      mixed_precision: See ``mixed_precision`` parameter in :class:`torch.distributed.fsdp.FullyShardedDataParallel`.
      auto_wrap_policy: Same as ``auto_wrap_policy`` parameter in
          :class:`torch.distributed.fsdp.FullyShardedDataParallel`. For convenience, this also accepts a set of the
          layer classes to wrap.
      activation_checkpointing: Deprecated. Use ``activation_checkpointing_policy``.
      activation_checkpointing_policy: Same as ``auto_wrap_policy`` parameter in
          :class:`torch.distributed.fsdp.FullyShardedDataParallel` but 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. For
          convenience, this also accepts a set of the layer classes to wrap.
      sharding_strategy: Select whether to shard model parameters, gradients, optimizer states, or a combination of
          them. Available values are:
  
          - ``"FULL_SHARD"``: Shards model parameters, gradients, and optimizer states (default).
          - ``"SHARD_GRAD_OP"``: Shards gradients and optimizer states only. Model parameters get replicated.
          - ``"NO_SHARD"``: No sharding (identical to regular DDP).
          - ``"HYBRID_SHARD"``: Shards model parameters, gradients, and optimizer states within a single machine, but
            replicates across machines.
  
          Also accepts a :class:`torch.distributed.fsdp.ShardingStrategy` enum value.
  
      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 as many files as the world size.
  
      \**kwargs: See available parameters in :class:`torch.distributed.fsdp.FullyShardedDataParallel`.
class_docstr_tok: ["Strategy","for","Fully","Sharded","Data","Parallel","provided","by","torch.distributed",".","..","warning",":",":","This","is","an",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","feature",".","Fully","Sharded","Training","shards","the","entire","model","across","all","available","GPUs",",","allowing","you","to","scale","model","size",",","whilst","using","efficient","communication","to","reduce","overhead",".","In","practice",",","this","means","we","can","remain","at","parity","with","PyTorch","DDP",",","whilst","scaling","our","model","sizes","dramatically",".","The","technique","is","similar","to","ZeRO-Stage","3",".","For","more","information","check","out","`","this","blogpost","<","https",":","\/\/pytorch.org\/blog\/introducing-pytorch-fully-sharded-data-parallel-api",">","`","__",".","Defaults","have","been","set","and","options","have","been","exposed",",","but","may","require","configuration","based","on","your","level","of","memory\/speed","efficiency",".","We","suggest","having","a","look","at","`","this","tutorial","<","https",":","\/\/pytorch.org\/tutorials\/intermediate\/FSDP_tutorial.html",">","`","__","for","more","information",".","Arguments",":","cpu_offload",":","See","``","cpu_offload","``","parameter","in",":","class",":","`","torch.distributed.fsdp.FullyShardedDataParallel","`",".","mixed_precision",":","See","``","mixed_precision","``","parameter","in",":","class",":","`","torch.distributed.fsdp.FullyShardedDataParallel","`",".","auto_wrap_policy",":","Same","as","``","auto_wrap_policy","``","parameter","in",":","class",":","`","torch.distributed.fsdp.FullyShardedDataParallel","`",".","For","convenience",",","this","also","accepts","a","set","of","the","layer","classes","to","wrap",".","activation_checkpointing",":","Deprecated",".","Use","``","activation_checkpointing_policy","``",".","activation_checkpointing_policy",":","Same","as","``","auto_wrap_policy","``","parameter","in",":","class",":","`","torch.distributed.fsdp.FullyShardedDataParallel","`","but","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",".","For","convenience",",","this","also","accepts","a","set","of","the","layer","classes","to","wrap",".","sharding_strategy",":","Select","whether","to","shard","model","parameters",",","gradients",",","optimizer","states",",","or","a","combination","of","them",".","Available","values","are",":","-","``","``","FULL_SHARD","''","``",":","Shards","model","parameters",",","gradients",",","and","optimizer","states","(","default",")",".","-","``","``","SHARD_GRAD_OP","''","``",":","Shards","gradients","and","optimizer","states","only",".","Model","parameters","get","replicated",".","-","``","``","NO_SHARD","''","``",":","No","sharding","(","identical","to","regular","DDP",")",".","-","``","``","HYBRID_SHARD","''","``",":","Shards","model","parameters",",","gradients",",","and","optimizer","states","within","a","single","machine",",","but","replicates","across","machines",".","Also","accepts","a",":","class",":","`","torch.distributed.fsdp.ShardingStrategy","`","enum","value",".","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","as","many","files","as","the","world","size",".","\\","*","*","kwargs",":","See","available","parameters","in",":","class",":","`","torch.distributed.fsdp.FullyShardedDataParallel","`","."]
code: def root_device(self) -> torch.device:
      assert self.parallel_devices is not None
      return self.parallel_devices[self.local_rank]
  
imports: import logging
  import shutil
  from contextlib import contextmanager, nullcontext
  from datetime import timedelta
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Literal, Mapping, Optional, Set, Type, Union
  import torch
  from lightning_utilities.core.rank_zero import rank_zero_only
  from torch import Tensor
  from torch.nn import Module
  from torch.optim import Optimizer
  from typing_extensions import override
  import lightning.pytorch
  from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
  from lightning.fabric.plugins.collectives.torch_collective import default_pg_timeout
  from lightning.fabric.strategies import _StrategyRegistry
  from lightning.fabric.strategies.fsdp import _METADATA_FILENAME, _activation_checkpointing_kwargs, _auto_wrap_policy_kwargs, _get_full_state_dict_context, _get_sharded_state_dict_context, _has_meta_device_parameters, _init_cpu_offload, _init_sharding_strategy, _is_full_checkpoint, _is_sharded_checkpoint, _load_raw_module_state, _move_torchmetrics_to_device, _optimizer_has_flat_params, _setup_activation_checkpointing
  from lightning.fabric.utilities.distributed import _distributed_is_initialized, _get_default_process_group_backend_for_device, _init_dist_connection, _sync_ddp_if_available
  from lightning.fabric.utilities.distributed import group
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_1_13, _TORCH_GREATER_EQUAL_2_0, _TORCH_GREATER_EQUAL_2_1, _TORCH_GREATER_EQUAL_2_2
  from lightning.fabric.utilities.init import _EmptyInit
  from lightning.fabric.utilities.load import _lazy_load, _materialize_tensors
  from lightning.fabric.utilities.optimizer import _optimizers_to_device
  from lightning.fabric.utilities.seed import reset_seed
  from lightning.fabric.utilities.types import _PATH, ReduceOp
  from lightning.pytorch.core.optimizer import LightningOptimizer
  from lightning.pytorch.plugins.precision import Precision
  from lightning.pytorch.plugins.precision.fsdp import FSDPPrecision
  from lightning.pytorch.strategies.launchers.subprocess_script import _SubprocessScriptLauncher
  from lightning.pytorch.strategies.parallel import ParallelStrategy
  from lightning.pytorch.strategies.strategy import TBroadcast
  from lightning.pytorch.trainer.states import TrainerFn
  from lightning.pytorch.utilities.model_helpers import is_overridden
  from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_only, rank_zero_warn
label_desc: Use image node_id 2 for calling the FSDPStrategy obj's underlying member method code with example usage: obj.root_device() and returns: self