text: import contextlib
  import copy
  import functools
  import itertools
  import math
  import traceback
  import warnings
  from contextlib import contextmanager
  from dataclasses import dataclass
  from enum import Enum, auto
  from typing import (
      TYPE_CHECKING,
      Any,
      Callable,
      Dict,
      Generator,
      Iterable,
      Iterator,
      List,
      Mapping,
      NamedTuple,
      Optional,
      Set,
      Tuple,
      Union,
      cast,
  )
  
  import torch
  import torch.distributed as dist
  from torch.distributed.utils import _to_kwargs, _sync_params_and_buffers, _replace_by_prefix
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  from torch.distributed import ProcessGroup
  from torch.distributed._shard.sharded_tensor import (
      Shard,
      ShardedTensor,
      init_from_local_shards,
  )
  from torch.distributed.distributed_c10d import _get_default_group
  from torch.nn.parameter import Parameter
  
  from ._optim_utils import (
      _broadcast_pos_dim_tensor_states,
      _broadcast_processed_optim_state_dict,
      _flatten_full_optim_state_dict,
      _get_flat_param_to_fsdp_module,
      _get_param_id_to_param,
      _get_param_to_param_id,
      _process_pos_dim_tensor_state,
      _unflatten_optim_state,
  )
  from ._utils import (
      _apply_to_modules, _apply_to_tensors,
      _override_batchnorm_mixed_precision, _contains_batchnorm
  )
  from .flatten_params_wrapper import (
      FLAT_PARAM,
      FPW_MODULE,
      FlatParameter,
      FlattenParamsWrapper,
  )
  from .wrap import _recursive_wrap, _wrap_batchnorm_individually, _or_policy
  
  if TYPE_CHECKING:
      from collections import OrderedDict  # noqa: F401
  
  _TORCHDISTX_AVAIL = True
  try:
      from torchdistx import deferred_init, fake
  except ImportError:
      _TORCHDISTX_AVAIL = False
  
  
  FSDP_WRAPPED_MODULE = "_fsdp_wrapped_module"
  FSDP_PREFIX = FSDP_WRAPPED_MODULE + "." + FPW_MODULE + "."
  
  _PARAM_BROADCAST_BUCKET_SIZE = int(250 * 1024 * 1024)
  
  
  def _default_meta_device_init_fn(module):
      """
      Default initializer for modules initialized on the meta device.
      """
      # TODO: move module to device_id here once device_id is available.
      module.to_empty(device=torch.cuda.current_device())
      try:
          with torch.no_grad():
              module.reset_parameters()
      except BaseException as e:
          warnings.warn(
              f"Unable to call reset_parameters() for module on meta device with error {str(e)}. "
              "Please ensure your module implements a ``reset_parameters`` function."
          )
          raise e
  
  
  class ShardingStrategy(Enum):
      """
      Specify which sharding strategy will be used for the distributed training.
      FULL_SHARD: Shards parameters, gradients and optimizer states. This algorithm
                  inserts ``all_gather`` before forward and backward computation to gather
                  parameters, also inserts ``reduce_scatter`` after backward computation for
                  synchronizing and sharding gradients. Sharded optimizer states are
                  updated locally.
      SHARD_GRAD_OP: Shard optimizer states and gradients, this algorithm inserts all_gather
                     before forward computation and keeps the full parameters in
                     GPU memory until backward computation is done. It inserts reduce_scater
                     after backward computation for synchronizing and sharding gradients.
                     Sharded optimizer states are updated locally.
      NO_SHARD: This is similar to PyTorch ``DistributedDataParallel`` API. Parameters, gradients
                and optimizer states are replicated among ranks, ``all_reduce`` is inserted after
                backward computation is done for synchronizing gradients. Full optimizer states
                are updated in each rank.
      HYBRID_SHARD(future support): apply FULL_SHARD algorithm in the intra node and
                                    apply NO_SHARD algorithm in the inter nodes.
  
      """
      FULL_SHARD = auto()
      SHARD_GRAD_OP = auto()
      NO_SHARD = auto()
      # TODO
      # HYBRID_SHARD = auto()
  
  
  @dataclass
  class MixedPrecision:
      """
      A config to enable mixed precision training with FullyShardedDataParallel.
      This class can be constructed with three flags:
          ``param_dtype`` controls the precision of model parameters, inputs, and
          therefore the precision under which computation happens. After forward
          and backward passes, FSDP parameters point to full precision shards
          that are kept in memory. Full precision parameters are always
          checkpointed.
          ``reduce_dtype`` controls the precision under which gradient reduction
          would occur, which can potentially be different than ``param_dtype``
          for use cases such as communication efficiency.
          ``buffer_dtype`` controls the precision that buffers are cast to. Note
          that buffers are unsharded and are cast in the first forward pass, and
          remain in their reduced precision state even after forward/backward
          passes. However, when taking checkpoints with ``state_dict``, buffers
          are checkpointed in their full precision (and then restored back to
          to their reduced precision) as expected. Note that this checkpoint
          support is currently limited to ``StateDictType.FULL_STATE_DICT``.
  
      .. note:: In ``summon_full_params``, parameters are summoned in full
          precision but buffers are not.
  
      .. note:: Parameters and buffers are checkpointed in full precision. For
          buffers, this is only guaranteed to work for ``StateDictType.FULL_STATE_DICT``.
  
      .. note:: This API is experimental and subject to change.
  
      .. note:: Specification of reduced precision types must be explicit, in that
          if, for example, ``param_dtype`` is not specified, it will not be cast by
          FSDP. Thus, a config such as ``MixedPrecision(reduce_dtype=torch.float16)``
          will not cast buffers or parameters. Note that if a ``MixedPrecision``
          config is specified without a ``reduce_dtype``, gradient communication
          would occur in the `param_dtype` precision, if given, otherwise, in the
          original parameter precision.
      """
      # maintain a tensor of this dtype that the fp32 param shard will be cast to.
      # Will control the precision of model params, inputs, and thus compute as
      # well.
      param_dtype: Optional[torch.dtype] = None
      # Gradient communication precision.
      reduce_dtype: Optional[torch.dtype] = None
      # Buffer precision.
      # TODO: buffer + param are usually of the same type, if user specifies
      # param but not buffer, should we automatically make buffer be the same?
      buffer_dtype: Optional[torch.dtype] = None
  
  
  @dataclass
  class CPUOffload:
      """
      CPU offloading config. Currently, only parameter and gradient CPU
      offload are supported.
      offload_params: Offloading parameters to CPUs when these parameters are
                      not used for computation on GPUs. This implicitly enables
                      gradient offloading to CPUs in order for parameters and
                      gradients to be on the same device to work with optimizer.
      """
  
      offload_params: bool = False
  
  
  class BackwardPrefetch(Enum):
      """
      Specify where to prefetch next layer's full parameters
      during backward pass.
      BACKWARD_PRE: prefetch right before current layer's backward computation
                    starts, this approach will increase backward communication
                    and computation overalpping and potentialy improve training
                    performance, but it may increase the peak memory usage as
                    the prefetched full parameters will be kept in the GPU memory
                    until next layer's backward computation is done.
      BACKWARD_POST: prefetch right after current layer's backward computation finishes,
                     this approach will not increase peak memory as prefetching happens
                     after current layer's full parameters are freed.
                     It could potentially improve backward communication and computation
                     overlapping as it avoids all_gather and reduce_scatter are blocked
                     each other in the single NCCL stream. However, based on our experiments,
                     for some models, the backward post backward hook fire order is not always
                     the reversed forward computation order, so this
                     approach may prefetch full parameters for layers ahead of next layer,
                     this 'ahead' all_gather could delay next layer's all_gather in the
                     single NCCL stream and cause the next layer's computation delay. So it may
                     cause some performance regession for some models.
      """
  
      BACKWARD_PRE = auto()
      BACKWARD_POST = auto()
      # TODO, BACKWARD_PRE_CPU, prefetch full parameters and keep them in the CPU memory
  
  
  class TrainingState_(Enum):
      """
      Simple enum to indicate what state FSDP is in. Used for asserting
      to make sure APIs are called in the correct state.
      ..note::
          ``BACKWARD_PRE`` and ``BACKWARD_POST`` states are used to ensure we
          receives backward hooks in the correct order. It is used to catch
          unexpected order of hooks being called (likely due to our
          hook registration logic or autograd engine logic changes).
      """
  
      IDLE = auto()
      FORWARD = auto()
      BACKWARD_PRE = auto()
      BACKWARD_POST = auto()
      SUMMON_FULL_PARAMS = auto()
  
  
  class StateDictType(Enum):
      """
      This enum indicates that which type of ``state_dict`` the FSDP module is
      currently processing (returning or loading).
      The default value is FULL_STATE_DICT to comply the PyTorch convention.
      ..note::
          FSDP currently supports two types of ``state_dict``:
              1. ``state_dict/load_state_dict`: this pair of APIs return and load
                 the non-sharded, unflattened parameters. The semantics is the
                 same as using DDP.
              2. ``_local_state_dict/_load_local_state_dict``: this pair of APIs return
                 and load local sharded, flattened parameters. The values returned
                 by ``_local_state_dict`` can be directly used by FSDP and is only
                 meaningful to FSDP (because parameters are flattened). Note that
                 these APIs are meant for use via the :func:`state_dict_type`
                 context manager as follows:
                     >>> with fsdp.state_dict_type(StateDictType.LOCAL_STATE_DICT):
                     >>>     state = fsdp.state_dict()  # loads local state dict
              3. ``_sharded_state_dict/_load_sharded_state_dict``: this pair of APIs
                 return and load sharded, unflattened parameters. The ``state_dict``
                 return by ``sharded_state_dict`` can be used by all other parallel
                 schemes (resharding may be required).
      """
  
      FULL_STATE_DICT = auto()
      LOCAL_STATE_DICT = auto()
      SHARDED_STATE_DICT = auto()
  
  @dataclass
  class StateDictConfig:
      """
      ``StateDictConfig`` is the base class for all state_dict configuration classes.
      Users should instantiate a child version (i.e. ``FullStateDictConfig``) in
      order to configure settings for the particular type of ``state_dict``
      implementation FSDP will use.
      """
      pass
  
  @dataclass
  class FullStateDictConfig(StateDictConfig):
      """
      ``FullStateDictConfig`` is a config class meant to be used with
      ``StateDictType.FULL_STATE_DICT``. Currently, it accepts two parameters,
      ``offload_to_cpu`` and ``rank0_only`` which can be configured to offload
      the full ``state_dict`` to CPU and to materialize the ``state_dict`` on
      rank 0 only. When used, it is recommended to enable both of these flags
      together to optimize memory savings when taking checkpoints. Note that
      this config class is meant for user via the :func:`state_dict_type`
      context manager as follows:
          >>> fsdp = FSDP(model, auto_wrap_policy=...)
          >>> cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
          >>> with FullyShardedDataParallel.state_dict_type(fsdp, StateDictType.FULL_STATE_DICT, cfg):
          >>>     state = fsdp.state_dict()
          >>>     # state will be empty on non rank 0 and contain CPU tensors on rank 0.
          >>> # To reload checkpoint for inference, finetuning, transfer learning, etc:
          >>> model = model_fn() # Initialize model on CPU in preparation for wrapping with FSDP
          >>> if dist.get_rank() == 0:
          >>>     # Load checkpoint only on rank 0 to avoid memory redundancy
          >>>     state_dict = torch.load("my_checkpoint.pt")
          >>>     model.load_state_dict(state_dict)
          >>> # All ranks initialize FSDP module as usual. ``sync_module_states`` argument
          >>> # communicates loaded checkpoint states from rank 0 to rest of the world.
          >>> fsdp = FSDP(model, device_id=torch.cuda.current_device(), auto_wrap_policy=..., sync_module_states=True)
          >>> # After this point, all ranks have FSDP model with loaded checkpoint.
      """
      offload_to_cpu: bool = False
      rank0_only: bool = False
  
  @dataclass
  class LocalStateDictConfig(StateDictConfig):
      pass
  
  @dataclass
  class ShardedStateDictConfig(StateDictConfig):
      pass
  
  _state_dict_type_to_config = {
      StateDictType.FULL_STATE_DICT: FullStateDictConfig,
      StateDictType.LOCAL_STATE_DICT: LocalStateDictConfig,
      StateDictType.SHARDED_STATE_DICT: ShardedStateDictConfig,
  }
  
  class OptimStateKeyType(Enum):
      PARAM_NAME = auto()
      PARAM_ID = auto()
  
  
  class _ExecOrderWarnStatus(Enum):
      """Used internally for execution order validation."""
      NONE = auto()     # no deviation yet
      WARNING = auto()  # deviated this iteration; currently issuing warnings
      WARNED = auto()   # deviated in a previous iteration
  
  
  class _ExecOrderData():
      """
      This contains the data used for validating execution order across ranks.
  
      Attributes:
          _all_flat_params (List[FlatParameter]): A :class:`list` of all
              flattened parameters contained in the FSDP module hierarchy with
              the list index implicitly giving a unique parameter index.
          _param_to_unflat_param_names (Dict[FlatParameter, List[str]]): A
              mapping from flattened parameter to the comprising unflattened
              parameters' names.
          is_first_iter (bool): Whether executing in the first iteration or not.
          param_order (List[int]): Order that parameters participate in the
              forward pass; constructed on the first iteration and validated
              against in subsequent iterations.
          index (int): Index tracking the position in ``param_order``
              when validating the forward pass execution order in subsequent
              iterations.
          warn_status (_ExecOrderWarnStatus): To avoid flooding the console, we
              only issue warnings throughout the first deviating iteration and no
              longer check thereafter; this tracks the warning status.
      """
      def __init__(self) -> None:
          self._all_flat_params: List[FlatParameter] = []
          self._param_to_unflat_param_names: Dict[FlatParameter, List[str]] = []
          # Modified in the first iteration:
          self.is_first_iter: bool = True
          self.param_order: List[int] = []
          # Modified in the subsequent iterations:
          self.index: int = 0
          self.warn_status: _ExecOrderWarnStatus = _ExecOrderWarnStatus.NONE
  
      def init(self, root_module: "FullyShardedDataParallel"):
          assert root_module._is_root, "This data structure should only be " \
              "initialized on an FSDP root module"
          # Save `root_modules.parameters()` to `_all_flat_params` instead of
          # re-materializing each time to avoid the result depending on the
          # calling context (e.g. when some parameters have been rebuilt)
          self._all_flat_params = list(root_module.parameters())
          self._param_to_unflat_param_names = cast(
              Dict[FlatParameter, List[str]],
              _get_param_to_unflat_param_names(root_module)
          )  # `root_module.parameters()` should only contain `FlatParameter`s
  
      def get_param_index(self, param: FlatParameter) -> int:
          """Returns a unique non-negative parameter index for ``param`` if it is
          valid or -1 otherwise. Critically, this index assignment must be the
          same across ranks."""
          assert isinstance(param, FlatParameter), \
              f"Expects `param` is a `FlatParameter` but got {type(param)}"
          for i, p in enumerate(self._all_flat_params):
              if p is param:
                  return i
          return -1
  
      def get_param(self, param_index: int) -> Optional[FlatParameter]:
          """Returns the parameter corresponding to ``param_index`` or ``None``
          if the index is invalid."""
          for i, p in enumerate(self._all_flat_params):
              if i == param_index:
                  return p
          return None
  
      def get_unflat_param_names(self, param_index: int) -> List[str]:
          """Returns a :class:`list` of unflattened parameter names comprising
          the flattened parameter with index ``param_index`` or an empty
          :class:`list` if ``param_index`` is invalid."""
          param = self.get_param(param_index)
          if param is None:
              return []
          assert param in self._param_to_unflat_param_names, \
              "Internal data structures out of sync; check `init()`"
          return self._param_to_unflat_param_names[param]
  
      def reset(self):
          """Called in :meth:`_wait_for_post_backward` to reset data for the next
          iteration."""
          self.is_first_iter = False
          self.index = 0
          # `reset()` marks the end of an iteration, so transition if needed
          if self.warn_status == _ExecOrderWarnStatus.WARNING:
              self.warn_status = _ExecOrderWarnStatus.WARNED
  
  
  class FullyShardedDataParallel(nn.Module):
      """
      A wrapper for sharding Module parameters across data parallel workers. This
      is inspired by `Xu et al.`_ as well as the ZeRO Stage 3 from DeepSpeed_.
      FullyShardedDataParallel is commonly shortened to FSDP.
  
      .. _`Xu et al.`: https://arxiv.org/abs/2004.13336
      .. _DeepSpeed: https://www.deepspeed.ai/
  
      Example::
  
          >>> import torch
          >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
          >>> torch.cuda.set_device(device_id)
          >>> sharded_module = FSDP(my_module)
          >>> optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
          >>> x = sharded_module(x, y=3, z=torch.Tensor([1]))
          >>> loss = x.sum()
          >>> loss.backward()
          >>> optim.step()
  
      .. warning::
          The optimizer must be initialized *after* the module has been wrapped,
          since FSDP will shard parameters in-place and this will break any
          previously initialized optimizers.
  
      .. warning::
          If the destination CUDA device has ID ``dev_id``, either (1)
          ``module`` should already be placed on that device, (2) the device
          should be set using ``torch.cuda.set_device(dev_id)``, or (3)
          ``dev_id`` should be passed into the ``device_id`` constructor
          argument. This FSDP instance's compute device will be that destination
          device. For (1) and (3), the FSDP initialization always occurs on GPU.
          For (2), the FSDP initialization happens on ``module`` 's current
          device, which may be CPU.
  
      .. warning::
          FSDP currently does not support gradient accumulation outside
          ``no_sync()`` when using CPU offloading. Trying to do so yields
          incorrect results since FSDP will use the newly-reduced gradient
          instead of accumulating with any existing gradient.
  
      .. warning::
          Changing the original parameter variable names after construction will
          lead to undefined behavior.
  
      .. warning::
          Passing in `sync_module_states=True` flag requires module to be put
          on GPU, or to use ``device_id`` argument to specify a CUDA device that
          FSDP will move module to. This is because ``sync_module_states=True``
          requires GPU communication.
  
      .. warning::
          As of PyTorch 1.12, FSDP only offers limited support for shared parameters
          (for example, setting one ``Linear`` layer's weight to another's). In
          particular, modules that share parameters must be wrapped as part of the
          same FSDP unit. If enhanced shared parameter support is needed for your
          use case, please ping https://github.com/pytorch/pytorch/issues/77724
  
      .. note::
          Inputs into FSDP ``forward`` function will be moved to compute device
          (same device FSDP module is on) before running ``forward``, so user does
          not have to manually move inputs from CPU -> GPU.
  
      Args:
          module (nn.Module):
              module to be wrapped with FSDP.
          process_group (Optional[ProcessGroup]):
              process group for sharding
          sharding_strategy (Optional[ShardingStrategy]):
              Config sharding algorithm, different sharding algorithm has trade
              off between memory saving and communication overhead. ``FULL_SHARD``
              will be chosen if sharding_strategy is not specified.
          cpu_offload (Optional[CPUOffload]):
              CPU offloading config. Currently, only parameter and gradient CPU
              offload is supported. It can be enabled via passing in
              ``cpu_offload=CPUOffload(offload_params=True)``. Note that this
              currently implicitly enables gradient offloading to CPU in order for
              params and grads to be on same device to work with optimizer. This
              API is subject to change. Default is ``None`` in which case there
              will be no offloading.
          auto_wrap_policy (Optional[Callable]):
              A callable specifying a policy to recursively wrap layers with FSDP.
              Note that this policy currently will only apply to child modules of
              the passed in module. The remainder modules are always wrapped in
              the returned FSDP root instance.
              ``size_based_auto_wrap_policy`` written in ``torch.distributed.fsdp.wrap`` is
              an example of ``auto_wrap_policy`` callable, this policy wraps layers
              with the number of parameters larger than 100M. ``transformer_auto_wrap_policy``
              written in ``torch.distributed.fsdp.wrap`` is an example of ``auto_wrap_policy``
              callable for tranformer-like model architectures. Users can supply the customized
              ``auto_wrap_policy`` callable that should accept following arguments:
              ``module: nn.Module``, ``recurse: bool``, ``unwrapped_params: int``,
              extra customized arguments could be added to the customized
              ``auto_wrap_policy`` callable as well. It is a good practice to print out
              the sharded model and check whether the sharded model is what
              the application wants and then adjust accordingly.
  
              Example::
  
                  >>> def custom_auto_wrap_policy(
                  >>>     module: nn.Module,
                  >>>     recurse: bool,
                  >>>     unwrapped_params: int,
                  >>>     # These are customizable for this policy function.
                  >>>     min_num_params: int = int(1e8),
                  >>> ) -> bool:
                  >>>     return unwrapped_params >= min_num_params
  
          backward_prefetch (Optional[BackwardPrefetch]):
              This is an experimental feature that is subject to change in the
              the near future. It allows users to enable two different backward_prefetch
              algorithms to help backward communication and computation overlapping.
              Pros and cons of each algorithm is explained in the class ``BackwardPrefetch``.
          mixed_precision (Optional[MixedPrecision]): A ``MixedPrecision`` instance
              describing the mixed precision training config to be used. ``MixedPrecision``
              supports configuring parameter, buffer, and gradient communication dtype. Note
              that only floating point data is cast to the reduced precision. This allows
              users potential memory saving and training speedup while trading off
              accuracy during model training. If ``None``, no mixed precision is applied.
              Note that if ``mixed_precision`` is enabled for FSDP model that
              contains ``BatchNorm`` with ``auto_wrap_policy``, FSDP will take
              care to disable mixed precision for ``BatchNorm`` units by wrapping
              them separately in their own FSDP unit with ``mixed_precision=None``.
              This is done because several ``BatchNorm`` kernels do not implement
              reduced type support at the moment. If individually wrapping the model,
              users must take care to set ``mixed_precision=None`` for
              ``BatchNorm`` units.
              (Default: ``None``)
          ignored_modules (Optional[Iterable[torch.nn.Module]]): Modules whose
              own parameters and child modules' parameters and buffers are
              ignored by this instance. None of the modules directly in
              ``ignored_modules`` should be :class:`FullyShardedDataParallel`
              instances, and any child modules that are already-constructed
              :class:`FullyShardedDataParallel` instances will not be ignored if
              they are nested under this instance. This argument may be used to
              avoid sharding specific parameters at module granularity when using an
              ``auto_wrap_policy`` or if parameters' sharding is not managed by
              FSDP. (Default: ``None``)
          param_init_fn (Optional[Callable[[nn.Module], None]]):
              A ``Callable[torch.nn.Module] -> None`` that
              specifies how modules that are currently on the meta device should be initialized
              onto an actual device. Note that as of v1.12, we detect modules on the meta
              device via ``is_meta`` check and apply a default initialization that calls
              ``reset_parameters`` method on the passed in ``nn.Module`` if ``param_init_fn``
              is not specified, otherwise we run ``param_init_fn`` to initialize the passed
              in ``nn.Module``. In particular, this means that if ``is_meta=True`` for any
              module parameters for modules that will be wrapped with FSDP and ``param_init_fn``
              is not specified, we assume your module properly implements a ``reset_paramters()``
              and will throw errors if not. Note that additionally, we offer support for modules
              initialized with torchdistX's (https://github.com/pytorch/torchdistX)
              ``deferred_init`` API. In this case, deferred modules would be initialized
              by a default initialization function that calls torchdistX's
              ``materialize_module``, or the passed in ``param_init_fn``, if it is not
              ``None``. The same ``Callable`` is applied to initialize all meta modules.
              Note that this initialization function is applied before doing any FSDP sharding
              logic.
  
              Example::
  
                  >>> module = MyModule(device="meta")
                  >>> def my_init_fn(module):
                  >>>     # responsible for initializing a module, such as with reset_parameters
                  >>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy)
                  >>> print(next(fsdp_model.parameters()).device) # current CUDA device
                  >>> # With torchdistX
                  >>> module = deferred_init.deferred_init(MyModule, device="cuda")
                  >>> # Will initialize via deferred_init.materialize_module().
                  >>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)
  
          device_id (Optional[Union[int, torch.device]]): An ``int`` or ``torch.device``
              describing the CUDA device the FSDP module should be moved to determining where
              initialization such as sharding takes place. If this argument is not specified
              and ``module`` is on CPU, we will move ``module`` to current CUDA device for faster
              initialization and move ``module`` back to CPU before returning.
              If specified, resulting FSDP instances will reside on this device.
              Note that if ``device_id`` is specified but ``module`` is already
              on a different CUDA device, an error will be thrown. (Default: ``None``)
  
          sync_module_states (bool): If ``True``, each individually wrapped FSDP unit will broadcast
              module parameters from rank 0 to ensure they are the same across all ranks after
              initialization. This helps ensure model parameters are the same across ranks
              before starting training, but adds communication overhead to ``__init__``, as at least
              one broadcast is triggered per individually wrapped FSDP unit.
              This can also help load checkpoints taken by ``state_dict`` and to be loaded by
              ``load_state_dict`` in a memory efficient way. See documentation for
              :class:`FullStateDictConfig` for an example of this. (Default: ``False``)
  
      """
  
      def __init__(
          self,
          module: nn.Module,
          process_group: Optional[ProcessGroup] = None,
          sharding_strategy: Optional[ShardingStrategy] = None,
          cpu_offload: Optional[CPUOffload] = None,
          auto_wrap_policy: Optional[Callable] = None,
          backward_prefetch: Optional[BackwardPrefetch] = None,
          mixed_precision: Optional[MixedPrecision] = None,
          ignored_modules: Optional[Iterable[torch.nn.Module]] = None,
          param_init_fn: Optional[Callable[[nn.Module], None]] = None,
          device_id: Optional[Union[int, torch.device]] = None,
          sync_module_states: bool = False,
      ):
          torch._C._log_api_usage_once("torch.distributed.fsdp")
          super().__init__()
          # Validate the ignored modules and derive the ignored parameters/buffers
          ignored_modules = self._get_ignored_modules(module, ignored_modules)
          self._ignored_modules = ignored_modules
          ignored_params, ignored_param_names = \
              self._get_ignored_params(module, ignored_modules)
          buffer_names = self._get_buffer_names(module)
          # Compute the names to ignore for full state dict cloning (i.e. those
          # of the ignored modules' parameters and of all modules' buffers)
          self._ignored_param_names = ignored_param_names
          self._buffer_names = buffer_names
          # NOTE: Since the names are computed at construction time, if the user
          # changes them later, then FSDP will not properly ignore them. However,
          # the `FlatParameter` implementation already relies on this assumption.
          # We do this at construction time since we want the fully prefixed
          # parameter names matching the keys in the model state dict (namely,
          # including the wrapped module's name in the prefix), which may be done
          # most non-intrusively here before flattening.
  
          # if auto_wrap_policy is specified, submodules should not be
          # already wrapped, otherwise we'd attempt to double wrap them resulting
          # in errors.
          if auto_wrap_policy is not None:
              self._check_wrapped(
                  module,
                  check_fn=lambda mod: not isinstance(mod, FullyShardedDataParallel),
                  err_fn=lambda mod: f"Expected {mod} to NOT be FullyShardedDataParallel if auto_wrap is enabled.",
              )
              if mixed_precision is not None and _contains_batchnorm(module):
                  _override_batchnorm_mixed_precision(module)
                  policy_to_use = functools.partial(
                      _or_policy,
                      policies=[_wrap_batchnorm_individually, auto_wrap_policy]
                  )
                  warnings.warn(
                      "Mixed precision was specified for FSDP module with"
                      " batchnorm submodules wrapped via ``auto_wrap_policy``."
                      " BatchNorm units will be wrapped as a separate FSDP unit,"
                      " with mixed_precision disabled (i.e. set to ``None``)"
                      " as several BatchNorm kernels would raise errors when"
                      " operating on reduced precision inputs."
                  )
              else:
                  policy_to_use = auto_wrap_policy
              _recursive_wrap(
                  module,
                  auto_wrap_policy=policy_to_use,
                  wrapper_cls=FullyShardedDataParallel,
                  ignored_modules=ignored_modules,
                  ignored_params=ignored_params,
                  # Note that we have the recursive_wrap skip wrapping for
                  # the outermost (this) module otherwise it will result in a
                  # double-wrap causing issues.
                  only_wrap_children=True,
                  # FSDP arguments follow.
                  process_group=process_group,
                  sharding_strategy=sharding_strategy,
                  cpu_offload=cpu_offload,
                  backward_prefetch=backward_prefetch,
                  mixed_precision=mixed_precision,
                  param_init_fn=param_init_fn,
                  device_id=device_id,
                  sync_module_states=sync_module_states,
              )
  
          self.process_group = process_group or _get_default_group()
          self.rank = self.process_group.rank()
          self.world_size = self.process_group.size()
          if device_id is not None:
              self.device_id = (
                  device_id if isinstance(device_id, torch.device)
                  else torch.device(device_id)
              )
              # If user passed in something like torch.device("cuda"),
              # device index of current device is unclear, make it explicit.
              if self.device_id == torch.device("cuda"):
                  warnings.warn(
                      f"Passed in {self.device_id} does not have explicit index, "
                      f"setting it to current index: {torch.cuda.current_device()}. "
                      "If this is not correct, please explicitly call torch.cuda.set_device()"
                      "before FSDP initialization or pass in explicit device index as device_id argument."
                  )
                  self.device_id = torch.device("cuda", torch.cuda.current_device())
          else:
              self.device_id = None
  
  
          is_meta_module = any(p.is_meta for p in module.parameters())
          is_torchdistX_deferred_init = (
              not is_meta_module and _TORCHDISTX_AVAIL
              and any(fake.is_fake(p) for p in module.parameters())
          )
  
          def _run_param_init_fn():
              # Call user-specified initialization function.
              if not callable(param_init_fn):
                  raise ValueError(
                      f"Expected {param_init_fn} to be callable, but got {type(param_init_fn)}"
                  )
              param_init_fn(module)
  
          if is_meta_module:
              if param_init_fn is not None:
                  _run_param_init_fn()
              else:
                  # Call default initialization function that is dependent on
                  # reset_parameters.
                  _default_meta_device_init_fn(module)
          elif is_torchdistX_deferred_init:
              assert _TORCHDISTX_AVAIL, "Got torchdistX initialized module but torchdistX lib is not available."
              if param_init_fn is not None:
                  _run_param_init_fn()
              else:
                  # Call default torchdistX initialization function. Omit re-initialization of FSDP submodules
                  # which is unnecessary.
                  check_fn = lambda k: not isinstance(k, FullyShardedDataParallel)  # noqa: E731
                  deferred_init.materialize_module(module, check_fn=check_fn)
  
          # Check that module was placed onto a single device.
          module_devices = set(
              p.device for p in module.parameters() if p not in ignored_params and not isinstance(p, FlatParameter)
          )
  
          if len(module_devices) > 1:
              raise RuntimeError(
                  f"FSDP only supports single device modules, but got params on {module_devices}"
              )
  
          # Move module appropriately depending on device_id and whether module is on CPU.
          self._move_module_if_needed(module)
  
          # device for computation, if module is on GPU, use module.device;
          # if module is on CPU, use current device;
          self.compute_device = _get_default_cuda_device(module)
  
          # if device_id is specified, ensure it is the same
          assert (
              self.device_id is None or self.compute_device == self.device_id
          ), f"Inconsistent compute_device and device_id: {self.compute_device} vs {self.device_id}"
  
          # Enum to indicate if we're in the forward/backward pass, idle, etc.
          self.training_state = TrainingState_.IDLE
  
          # setting two factors to avoid underflow and overflow
          self.gradient_predivide_factor: float = self._get_gradient_predivide_factor(
              self.world_size
          )
          self.gradient_postdivide_factor: float = (
              self.world_size / self.gradient_predivide_factor
          )
  
          self.numel_padded_per_param: List[int] = []
          self.cpu_offload = cpu_offload or CPUOffload()
          self.backward_prefetch = backward_prefetch
          self.sharding_strategy = sharding_strategy or ShardingStrategy.FULL_SHARD
          self.mixed_precision = mixed_precision
          # Original buffer type (mapping since all buffers may not be of same type). In
          # the case of mixed precision training, this is used to restore buffers
          # to their original type (which may not be the same as that of the
          # parameters in the model) when checkpointing.
          self._orig_buffer_dtypes: Dict[str, torch.dtype] = {}
  
          # Only handle params which are not already sharded. This enables
          # sharding individual layers of a Module, with an outer wrapper to
          # shard any leftover parameters.
          params = [
              p for p in module.parameters()
              if p not in ignored_params and not isinstance(p, FlatParameter)
          ]
  
          if sync_module_states:
              if params != [] and params[0].device == torch.device("cpu"):
                  raise ValueError(
                      "Module has CPU parameters, but sync_module_states=True is specified."
                      "This only works for GPU module, please specify `device_id` argument or move"
                      " module to GPU before init."
                  )
              # Collect buffers we have to synchronize, avoiding buffers that have already
              # been synchronized to avoid redundant synchronization.
              bufs_to_sync = []
              for buf in module.buffers():
                  if not getattr(buf, '_fsdp_has_been_sync', False):
                      buf._fsdp_has_been_sync = True
                      bufs_to_sync.append(buf.detach())
  
              states_to_sync = [param.detach() for param in params]
              states_to_sync.extend(bufs_to_sync)
              _sync_params_and_buffers(
                  process_group=self.process_group,
                  module_states=states_to_sync,
                  # Same bucket size as DDP
                  broadcast_bucket_size=_PARAM_BROADCAST_BUCKET_SIZE,
                  src=0,
              )
  
          self._fsdp_wrapped_module: FlattenParamsWrapper = FlattenParamsWrapper(
              module, param_list=params
          )
          assert getattr(self, FSDP_WRAPPED_MODULE) is self._fsdp_wrapped_module
          del module  # free original module in case it helps garbage collection
          if self._fsdp_wrapped_module.flat_param is not None:
              self.params = [self._fsdp_wrapped_module.flat_param]
          else:
              self.params = []
  
          # Shard module parameters in place
          self._shard_parameters()
  
          # Make sure all parameters are sharded.
          for n, p in self.named_parameters():
              if p not in ignored_params and not isinstance(p, FlatParameter):
                  raise RuntimeError(
                      f"found unflattened parameter: {n} ; {p.size()} {p.__class__}"
                  )
          self._reset_lazy_init()
  
          # Flag indicating if we require gradient reduction in the backward
          # pass (set to `False` in the `no_sync()` context manager)
          self._require_backward_grad_sync: bool = True
  
          self._state_dict_type = StateDictType.FULL_STATE_DICT
          self._state_dict_config = FullStateDictConfig()
  
          # FSDP currently provides three different state_dicts. The actual
          # state_dict that will be saved/loaded is decided by
          # self._state_dict_type. And the main logic of each state_dict is
          # implemented in the hook. Therefore, for each hook (post-save and
          # pre-load), there is a dispatcher dictionary to dispatch the execution
          # flow to the correct implementation.
          self._register_state_dict_hook(self._post_state_dict_hook)
          self._post_state_dict_hook_fn = {
              StateDictType.FULL_STATE_DICT: self._full_post_state_dict_hook,
              StateDictType.LOCAL_STATE_DICT: self._local_post_state_dict_hook,
              StateDictType.SHARDED_STATE_DICT: self._sharded_post_state_dict_hook,
          }
          self._register_load_state_dict_pre_hook(
              self._pre_load_state_dict_hook, with_module=True
          )
          self._pre_load_state_dict_hook_fn = {
              StateDictType.FULL_STATE_DICT: self._full_pre_load_state_dict_hook,
              StateDictType.LOCAL_STATE_DICT: self._local_pre_load_state_dict_hook,
              StateDictType.SHARDED_STATE_DICT: self._sharded_pre_load_state_dict_hook,
          }
          self.register_load_state_dict_post_hook(
              self._post_load_state_dict_hook
          )
          self._post_load_state_dict_hook_fn = {
              StateDictType.FULL_STATE_DICT: self._full_post_load_state_dict_hook,
              StateDictType.LOCAL_STATE_DICT: self._local_post_load_state_dict_hook,
              StateDictType.SHARDED_STATE_DICT: self._sharded_post_load_state_dict_hook,
          }
  
          # Flag to guard against preparing gradients multiple times per backward pass.
          self._pre_backward_hook_has_run = False
          # Used for prefetching all gather full params in post backward hook
          self._need_rebuild_full_params = False
  
          # If specified, offload parameter shard to CPU.
          if self.cpu_offload.offload_params:
              for p in self.params:
                  self._offload_to_cpu(p)
  
          # For validating execution order across ranks
          self._exec_order_data = _ExecOrderData()
  
      def _move_module_if_needed(self, module) -> None:
          """
          Moves module appropriately depending on device_id and
          whether module is on CPU. Returns a ``bool`` indicating
          whether the module needs to be moved back to CPU before
          returning to user.
          """
          # Move module to device specified. Note that this is done prior to
          # setting compute_device to ensure that they align.
          if self.device_id is not None:
              param = None
              try:
                  # Get the next unflat param
                  param_gen = module.parameters()
                  while True:
                      param = next(param_gen)
                      if not isinstance(param, FlatParameter):
                          break
  
                  if param.device == torch.device("cpu"):
                      module = module.to(self.device_id)
              except StopIteration:
                  # this FSDP instance manages no parameters.
                  pass
  
              # For GPU modules, module device should match device_id.
              if param is not None and param.device != self.device_id:
                  raise RuntimeError(
                      f"Module on rank {self.rank} is given device_id argument "
                      f"{self.device_id}, but is on {param.device}. "
                      " Either move module before FSDP init or omit device_id argument."
                  )
          else:
              # device_id argument is not specified
              # If module is on CPU, log a warning asking user to use `device_id` for faster
              # GPU init.
              try:
                  # Get the next unflat param
                  param_gen = module.parameters()
                  while True:
                      param = next(param_gen)
                      if not isinstance(param, FlatParameter):
                          break
  
                  if param.device == torch.device("cpu"):
                      warnings.warn(
                          "Module is put on CPU and will thus have flattening and sharding"
                          " run on CPU, which is less efficient than on GPU. We recommend passing in "
                          "`device_id` argument which will enable FSDP to put module on GPU device,"
                          " module must also be on GPU device to work with `sync_module_states=True` flag"
                          " which requires GPU communication."
                      )
              except StopIteration:
                  # this FSDP instance manages no parameters
                  pass
  
      def _init_reshard_after_forward(self):
          if self.sharding_strategy == ShardingStrategy.FULL_SHARD:
              # Free full params and keep shard only after forward
              self.reshard_after_forward = True
          elif self.sharding_strategy == ShardingStrategy.SHARD_GRAD_OP:
              # Keep full params in the GPU memory until backward
              # computation is done
              self.reshard_after_forward = False
          elif self.sharding_strategy == ShardingStrategy.NO_SHARD:
              # self.reshard_after_forward is not used when NO_SHARD
              # is set, just setting it as False here
              self.reshard_after_forward = False
          else:
              raise RuntimeError(
                  "sharding_strategy only supports FULL_SHARD, SHARD_GRAD_OP and NO_SHARD right now."
              )
  
      def _get_ignored_modules(
          self,
          root_module: torch.nn.Module,
          _ignored_modules: Any,
      ) -> Set[torch.nn.Module]:
          """
          Checks that ``_ignored_modules`` (1) is an iterable of
          ``torch.nn.Module`` s without any :class:`FullyShardedDataParallel`
          instances and does not contain the top-level ``module`` itself, and
          then returns them and their children as a :class:`set`, excluding
          nested :class:`FullyShardedDataParallel` instances.
  
          We include the child modules of modules in ``_ignored_modules`` to be
          more intuitive since ignoring a module should ignore its child modules
          as well, and we exclude :class:`FullyShardedDataParallel` instances
          since ``self`` may be the intended root instance that manages them.
          """
          if _ignored_modules is None:
              return set()
          msg_prefix = "`ignored_modules` should be an iterable of " \
              "`torch.nn.Module`s "
          try:
              ignored_root_modules = set(_ignored_modules)
          except TypeError:
              raise TypeError(msg_prefix + f"but got {type(_ignored_modules)}")
          for module in ignored_root_modules:
              if not isinstance(module, torch.nn.Module):
                  raise TypeError(
                      msg_prefix + f"but got an iterable with {type(module)}"
                  )
              if isinstance(module, FullyShardedDataParallel):
                  raise ValueError(
                      "`ignored_modules` should not include FSDP modules"
                  )
          # Include child modules and exclude nested FSDP modules
          ignored_modules = set(
              child for module in ignored_root_modules
              for child in module.modules()
              if not isinstance(child, FullyShardedDataParallel) and
              not isinstance(child, FlattenParamsWrapper)
          )
          if root_module in ignored_modules:
              warnings.warn(
                  "Trying to ignore the top-level module passed into the FSDP "
                  "constructor itself will result in all parameters being "
                  f"ignored and is not supported: {module}"
              )
          for submodule in root_module.modules():
              if isinstance(submodule, FullyShardedDataParallel):
                  assert hasattr(submodule, "_ignored_modules")
                  ignored_modules.update(submodule._ignored_modules)
          return ignored_modules
  
      def _get_ignored_params(
          self,
          root_module: torch.nn.Module,
          ignored_modules: Set[torch.nn.Module],
      ) -> Tuple[Set[torch.nn.Parameter], Set[str]]:
          """
          Returns the parameters of the modules in ``ignored_modules``,
          excluding any :class:`FlatParameter` s and their fully prefixed names,
          both as :class:`set` s.
  
          Args:
              root_module (torch.nn.Module): Top-level module passed into the
                  FSDP constructor from which to derive the fully prefixed names.
              ignored_modules (Set[torch.nn.Module]): Modules to ignore.
          """
          ignored_params = set(
              p for m in ignored_modules for p in m.parameters()
              if not isinstance(p, FlatParameter)
          )
          param_to_unflat_param_names = _get_param_to_unflat_param_names(
              root_module, dedup_shared_params=False,
          )
          ignored_param_names = set()
          for param in ignored_params:
              unflat_param_names = param_to_unflat_param_names[param]
              clean_names = []
              for k in unflat_param_names:
                  clean_names.append(clean_tensor_name(k))
              ignored_param_names.update(clean_names)
          return ignored_params, ignored_param_names
  
      def _get_buffer_names(self, root_module: torch.nn.Module) -> Set[str]:
          """
          Returns the fully prefixed names of all buffers in the module hierarchy
          rooted at ``root_module`` as a class:`set`.
  
          Args:
              root_module (torch.nn.Module): Top-level module passed into the
                  FSDP constructor from which to derive the fully prefixed names.
          """
          def module_fn(module, prefix, buffer_names):
              # For FSDP modules, only add the entry when considering the
              # contained `FlattenParamsWrapper` to avoid duplication
              if not isinstance(module, FullyShardedDataParallel):
                  for buffer_name, _ in module.named_buffers(recurse=False):
                      prefixed_buffer_name = clean_tensor_name(prefix + buffer_name)
                      buffer_names.add(prefixed_buffer_name)
  
          def return_fn(buffer_names, *args):
              return buffer_names
  
          buffer_names: Set[str] = set()
          return _apply_to_modules(
              root_module, module_fn, return_fn, buffer_names,
          )
  
      @classmethod
      def _check_wrapped(cls, begin_module, check_fn, err_fn):
          for _, mod in begin_module.named_modules():
              if not check_fn(mod):
                  raise ValueError(err_fn(mod))
  
      @property
      def module(self) -> FlattenParamsWrapper:
          """make model.module accessible, just like DDP."""
          assert isinstance(self._fsdp_wrapped_module, FlattenParamsWrapper)
          return self._fsdp_wrapped_module
  
      def check_is_root(self) -> bool:
          self._lazy_init()
          assert self._is_root is not None
          return self._is_root
  
      @staticmethod
      def fsdp_modules(
          module: nn.Module,
          root_only: bool = False,
      ) -> List["FullyShardedDataParallel"]:
          """
          Returns all nested FSDP instances, possibly including ``module`` itself
          and only including FSDP root modules if ``root_only=True``.
  
          Args:
              module (torch.nn.Module): Root module, which may or may not be an
                  ``FSDP`` module.
              root_only (bool): Whether to return only FSDP root modules.
                  (Default: ``False``)
  
          Returns:
              List[FullyShardedDataParallel]: FSDP modules that are nested in
              the input ``module``.
          """
          return [
              submodule for submodule in module.modules()
              if isinstance(submodule, FullyShardedDataParallel) and
              (not root_only or submodule.check_is_root())
          ]
  
      def apply(self, fn: Callable[[nn.Module], None]) -> "FullyShardedDataParallel":
          r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
          as well as self. Typical use includes initializing the parameters of a model
          (see also :ref:`nn-init-doc`).
  
          Compared to ``torch.nn.Module.apply``, this version additionally gathers
          the full parameters before applying ``fn``. It should not be called from
          within another ``summon_full_params`` context.
  
          Args:
              fn (:class:`Module` -> None): function to be applied to each submodule
  
          Returns:
              Module: self
          """
          uninitialized = self._is_root is None
          self._assert_state(TrainingState_.IDLE)
          with self._summon_full_params(recurse=False, writeback=True):
              ret = super().apply(fn)
  
          # Reset lazy init that might be called by _summon_full_params, since
          # it could have set is_root incorrectly for non-root FSDP instances.
          if uninitialized and self._is_root:
              for module in self.fsdp_modules(self):
                  module._reset_lazy_init()
  
          return ret
  
      # setting two factors 'self.gradient_predivide_factor'
      # and 'self.gradient_postdivide_factor' to avoid underflow and overflow
      def _get_gradient_predivide_factor(self, world_size: int) -> float:
          factor: int = 1
          while world_size % factor == 0 and world_size / factor > factor:
              factor *= 2
          return float(factor)
  
      def _offload_to_cpu(self, p):
          """
          Offloads parameter to CPU from self.compute_device. If the parameter is
          already on CPU then this is a noop.
          """
          cpu_device = torch.device("cpu")
          if p.device == cpu_device:
              return
          with torch.no_grad():
              p.data = p.to(cpu_device)
  
      def _mixed_precision_enabled_for_params(self) -> bool:
          """
          Whether user explicitly enabled mixed precision for
          parameters or not.
          """
          return (
              self.mixed_precision is not None
              and self.mixed_precision.param_dtype is not None
          )
  
      def _mixed_precision_enabled_for_buffers(self) -> bool:
          """
          Whether user explicitly enabled mixed precision for
          buffers or not.
          """
          return (
              self.mixed_precision is not None
              and self.mixed_precision.buffer_dtype is not None
          )
  
      def _mixed_precision_enabled_for_reduce(self) -> bool:
          """
          Whether user explicitly enabled mixed precision for
          gradient reduction or not.
          """
          return (
              self.mixed_precision is not None
              and self.mixed_precision.reduce_dtype is not None
          )
  
      def _cast_fp_inputs_to_precision(
          self, dtype: torch.dtype, *args: Any, **kwargs: Any
      ) -> Tuple[Any, Any]:
          """
          Casts floating point tensors in args and kwargs to precision given by dtype.
          requires_grad field is respected.
          """
          def cast_fn(x: torch.Tensor) -> torch.Tensor:
              if not torch.is_floating_point(x):
                  return x
              y = x.to(dtype)
              # Explicitly copy over requires_grad context since this is happening
              # within torch.no_grad.
              if x.is_leaf:
                  y.requires_grad = x.requires_grad
              return y
  
          with torch.no_grad():
              return (
                  _apply_to_tensors(cast_fn, args),
                  _apply_to_tensors(cast_fn, kwargs)
              )
  
      @torch.no_grad()
      def _cast_param_shards_to_dtype(self):
          """
          Allocates a mixed precision paramter shard and casts parameter shards to
          reduced precision by copying into this mixed precision shard. Note that
          if we are CPU offloading, this also implicitly loads the parameter shard
          back to GPU.
          """
          assert (
              self._mixed_precision_enabled_for_params()
          ), "Expected to only be called when mixed precision for parameters is enabled."
          with torch.cuda.stream(self._streams["mixed_precision_params"]):
              for p in self.params:
                  assert p._mp_shard is not None
                  _alloc_storage(data=p._mp_shard, size=p._local_shard.size())
                  # Cast is done by copy
                  p._mp_shard.copy_(
                      # no-op if not CPU offloading, otherwise nonblocking because
                      # p._local_shard is pinned in _init_param_attributes.
                      p._local_shard.to(p._mp_shard.device, non_blocking=True)
                  )
                  # Point p to the mp shard
                  p.data = p._mp_shard
          # Block current stream on this copy work.
          torch.cuda.current_stream().wait_stream(self._streams["mixed_precision_params"])
  
      @torch.no_grad()
      def _free_mp_shard(self, params: List[FlatParameter]):
          """
          Deallocate storage for parameter's mixed precision shard.
          """
          assert (
              self._mixed_precision_enabled_for_params()
          ), "Expected to only be called when mixed precision for parameters is enabled."
          current_stream = torch.cuda.current_stream()
          for p in params:
              # mp_shard should always be allocated.
              assert p._mp_shard is not None
              # Shard is allocated in "mixed_precision_stream" and then we block
              # current stream on this stream, so don't free it until work in the
              # current stream is completed.
              p._mp_shard.record_stream(current_stream)
              _free_storage(p._mp_shard)
  
      def _cast_buffers(
          self,
          device: Optional[torch.device] = None,
          dtype: Optional[Dict[str, torch.dtype]] = None,
          memo: Optional[Set] = None,
          recurse: bool = True,
      ) -> None:
          """Move all buffers to the given *device* and *dtype*.
          If *device* is not given, then it will default to
          ``self.compute_device``, otherwise buffer will be moved to ``device``.
          In the case of nested FSDP instances, we will respect the child instance's
          ``compute_device`` configuration.
          If *dtype* is given, it must be a mapping of buffer name to buffer dtype,
              and this argument is currently only given to restore back to original
              buffer types during checkpoint. If *dtype* is not given, and we are
              in mixed precision training, the buffer will be cast to buffer_dtype,
              otherwise the buffer will not be cast.
          Args:
              device (torch.device, Optional):
                  device to cast buffers to (defaults to compute_device)
              dtype: (Dict[str, torch.dtype], Optional):
                  Mapping of buffer name to their dtype to cast to.
              memo (Set, Optional):
                  set of modules that have already been processed
              recurse (bool, Optional):
                  Whether to call _cast_buffers recursively on nested FSDP
                  instances (default is True).
          """
          if memo is None:
              memo = set()
          for module in self.modules():
              if module is not self and isinstance(module, FullyShardedDataParallel) and recurse:
                  # Allow any child FSDP instances to handle their own buffers.
                  module._cast_buffers(device=device, dtype=dtype, memo=memo, recurse=recurse)
              elif module not in memo:
                  memo.add(module)
                  for name, buf in module.named_buffers(recurse=False):
                      if buf is None:
                          continue
                      buf = buf.to(device=device or self.compute_device)
                      if name not in self._orig_buffer_dtypes:
                          self._orig_buffer_dtypes[name] = buf.dtype
                      # If given, cast buffer to the given dtype. This is used to
                      # suppport mixed precision for buffers
                      # (given by self.mixed_precision.buffer_dtype) and also used
                      # to restore the buffer dtype to the original precision for
                      # state_dict() calls.
                      # Note that non-floating point buffers are not casted.
                      if torch.is_floating_point(buf):
                          # We are restoring the original buffer type in
                          # preparation for checkpoint.
                          if dtype:
                              buf = buf.to(dtype=dtype[name])
                          # Note that we don't pass in self.mixed_precision.buffer_dtype
                          # recursively into _cast_buffers, as we want to respect
                          # mp config for child FSDP instances.
                          elif self._mixed_precision_enabled_for_buffers():
                              buf = buf.to(self.mixed_precision.buffer_dtype)
  
                      setattr(module, name, buf)
  
      @torch.no_grad()
      def _shard_parameters(self) -> None:
          """
          At initialization we wrap a module with full parameters and shard the
          parameters in-place. Sharding is implemented by viewing each parameter
          as a 1D Tensor and retaining only a single slice, where the slice size
          is determined by the number of data parallel workers.
          After this initial sharding is complete, the user can initialize a
          ``torch.optim.Optimizer`` in the usual way, i.e.::
          .. code-block:: python
              optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
          The optimizer will see only a single slice of parameters and will thus
          allocate less memory for optimizer state, avoiding redundancy across
          data parallel workers.
          """
          for p in self.params:
              assert not p._is_sharded, "Param should have not been sharded yet."
              assert (
                  p.is_floating_point()
              ), "Autograd does not support operations for integer type."
  
              # Sharding is done only when world_size is larger than 1 and
              # sharding_strategy!=NO_SHARD.
              p._is_sharded = (  # type: ignore[attr-defined]
                  self.world_size > 1
                  and self.sharding_strategy != ShardingStrategy.NO_SHARD
              )
              p._orig_size = p.size()  # type: ignore[attr-defined]
  
              if not p._is_sharded:  # type: ignore[attr-defined]
                  self.numel_padded_per_param.append(0)
                  continue
  
              # Save the original storage and free it later on.
              # Since we're modifying the tensor's storage directly,
              # make sure the tensor is the sole occupant of the storage.
              assert (
                  p.storage_offset() == 0
              ), "The tensor is not the sole occupant of the storage."
              orig_storage = p.storage()
  
              # Replace p with the relevant shard.
              local_shard, num_padded = self._get_shard(p)
              p.set_(local_shard)  # type: ignore[call-overload]
              p.shard_by_offsets(
                  self.rank * local_shard.numel(),
                  (self.rank + 1) * local_shard.numel() - 1,
                  num_padded,
              )
              self.numel_padded_per_param.append(num_padded)
  
              # Free storage that contains the original full data.
              if orig_storage.size() > 0:
                  orig_storage.resize_(0)  # type: ignore[attr-defined]
  
          assert len(self.numel_padded_per_param) == len(
              self.params
          ), "numel_padded_per_param is not populated correctly."
  
      @staticmethod
      def _get_chunk(
          tensor: torch.Tensor,
          rank: int,
          world_size: int,
      ) -> Tuple[torch.Tensor, int]:
          """Returns the unpadded chunk as a view and the number of padding
          elements of a full tensor for the given rank and world size."""
          # Shard using `torch.chunk()` to match all-gather/reduce-scatter.
          chunks = torch.flatten(tensor).chunk(world_size)
          if len(chunks) < (rank + 1):
              # If there are not enough chunks to shard across ranks, create an
              # empty chunk that will just be padded with zeros to be the
              # appropriate size.
              chunk = chunks[0].new_empty(0)
          else:
              chunk = chunks[rank]
          # Determine number of padding elements.
          num_to_pad = chunks[0].numel() - chunk.numel()
          assert num_to_pad >= 0, \
              "Chunk's size should at most the first chunk's size"
          return chunk, num_to_pad
  
      @staticmethod
      def _get_shard_functional(
          tensor: torch.Tensor,
          rank: int,
          world_size: int,
      ) -> Tuple[torch.Tensor, int]:
          """Functional version of :meth:`_get_shard`."""
          chunk, num_to_pad = FullyShardedDataParallel._get_chunk(
              tensor, rank, world_size,
          )
          # We always need to clone here regardless of the padding and even
          # though `chunk` is a view of `tensor` because `tensor` may be
          # deallocated after this method returns
          shard = chunk.clone()
          if num_to_pad > 0:
              shard = F.pad(shard, [0, num_to_pad])
          return shard, num_to_pad
  
      def _get_shard(
          self,
          tensor: torch.Tensor,
          rank: Optional[int] = None,
      ) -> Tuple[torch.Tensor, int]:
          """Returns the local shard and the number of padding elements of a full
          tensor for the calling rank if ``rank=None`` or for the rank ``rank``
          if not ``None``."""
          rank = self.rank if rank is None else rank
          return FullyShardedDataParallel._get_shard_functional(
              tensor, rank, self.world_size,
          )
  
      def __getattr__(self, name: str) -> Any:
          """Forward missing attributes to wrapped module."""
          try:
              return super().__getattr__(name)  # defer to nn.Module's logic
          except AttributeError:
              return getattr(self.module, name)
  
      def __getitem__(self, key: int) -> Any:
          """Forward indexing calls in case the module is a nn.Sequential."""
          return self.module.__getitem__(key)  # type: ignore[operator]
  
      def _reset_lazy_init(self) -> None:
          """
          Reset instance so :func:`_lazy_init` will run on the next forward.
          Currently this is only called in __init__
          """
          self._is_root: Optional[bool] = None
          self._streams: Dict[str, torch.cuda.Stream] = {}
          self._fsdp_graph_order: List[nn.Module] = []
          self._my_fsdp_idx_in_graph: Optional[int] = None
          for p in self.params:
              if hasattr(p, "_local_shard"):
                  # reset attributes that are added in _init_param_attributes, as
                  # part of _lazy_init
                  del p._local_shard  # type: ignore[attr-defined]
          # set 'self.reshard_after_forward' flag based on self.sharding_strategy
          self._init_reshard_after_forward()
  
      def _lazy_init(self) -> None:
          """Initialization steps that should happen lazily, typically right
          before the first forward pass.
          """
          # Initialize param attributes lazily, in case the param's dtype or
          # device changes after __init__.
          for p in self.params:
              self._init_param_attributes(p)
  
          # Initialize _is_root and setup streams. These steps would ideally
          # happen in __init__, but _is_root can only be determined after the
          # entire model hierarchy is setup, thus we run it lazily.
          if self._is_root is None:
              # _is_root means that we are in the outermost module's forward.
              self._set_is_root()
              self._setup_streams()
  
          if self._is_root:
              # Buffers stay on GPU, and don't get sharded. Since _cast_buffers
              # applies recursively, we only call this from the root instance.
              self._cast_buffers(recurse=True)
  
              # Don't free the full params for the outer-most (root) instance,
              # In most cases, root instance contains params in the last layers
              # or has no params. In these cases, those params will be needed
              # immediately after for the backward pass. Note that this only
              # applies currently when freeing parameters at end of layer's
              # forward pass.
              self.reshard_after_forward = False
  
              # Due to the use of streams, we need to make sure the previous
              # ``optim.step()`` is done before we all-gather parameters.
              self._wait_for_previous_optim_step()
  
      @torch.no_grad()
      def _init_param_attributes(self, p: Parameter) -> None:
          """
          We manage several attributes on each Parameter instance. The first two
          are set by :func:`_shard_parameters`:
              ``_is_sharded``: ``True`` if the Parameter is sharded or ``False``
                  if the Parameter is intentionally not sharded (in which case we
                  will all-reduce grads for this param). Currently the way
                  `_is_sharded = False` is if world_size = 1 or sharding strategy
                  is NO_SHARD.
              ``_orig_size``: the size of the original Parameter (before sharding)
          A few attributes are set here:
              ``_local_shard``: a single shard of the parameter. This is needed to
                  recover the shard after rebuilding full parameter in forward
                  and backward.
              ``_full_param_padded``: the full weight (padded to be evenly
                  divisible by ``world_size``), used for computation in the
                  forward and backward pass. It is initialized with the
                  appropriate size and then has its storage freed. This will be
                  resized in place and only materialized (via all-gather) as needed.
          Another attribute is set by :func:`_register_post_backward_hooks`:
              ``_shard_bwd_hook``: it holds the parameter's AccumulateGrad object
                  and the registered post hook handle.
          """
          assert hasattr(p, "_is_sharded") and hasattr(
              p, "_orig_size"
          ), "Parameters should have been sharded during construction."
          # If _local_shard has been set in the first lazy init and
          # current parameter is pointed to _local_shard, no need to
          # set the _local_shard again.
          if hasattr(p, "_local_shard"):
              # If CPU offloading, p._local_shard should have been placed on CPU
              # during its first lazy construction.
              if self.cpu_offload.offload_params:
                  assert p._local_shard.device == torch.device(  # type: ignore[attr-defined]
                      "cpu"
                  ), (
                      "Expected p._local_shard to be on CPU, "  # type: ignore[attr-defined]
                      f"but it's on {p._local_shard.device}"  # type: ignore[attr-defined]
                  )
              return
  
          # A single shard of the parameters. Also makes p._local_shard to be on
          # CPU if we are CPU offloading, since p.data would be on CPU during
          # init.
          if self.cpu_offload.offload_params:
              assert p.device == torch.device("cpu"), (
                  "Expected param to be on CPU when cpu_offloading is enabled. "
                  "If CPU offloading is enabled correctly, you may be "
                  "accidentally moving the model to CUDA after FSDP initialization."
              )
          p._local_shard = p.data  # type: ignore[attr-defined]
          # If CPU offloading, pin the memory to enable faster CPU -> GPU device
          # transfer.
          if self.cpu_offload.offload_params:
              assert p._local_shard.device == torch.device("cpu")  # type: ignore[attr-defined]
              p._local_shard.pin_memory()  # type: ignore[attr-defined]
              # When offloading parameters, also move the grad shard to CPU during
              # backward pass. In this case, it's important to pre-allocate the
              # CPU grad shard in pinned memory so that we can do a non-blocking
              # transfer.
              p._cpu_grad = torch.zeros_like(  # type: ignore[attr-defined]
                  p, device=torch.device("cpu")
              ).pin_memory()
  
          # If mixed_precision, maintain reduced precision param shard on
          # compute_device for computation in fwd/bwd. We resize storage to 0 here
          # and rematerialize before building the full param when needed. After
          # fwd/bwd, it is freed and we only hold on to the full precision shard.
          # As a result, this reduced precision shard is not allocated if we are
          # not in the forward/backward pass.
          if (
              self._mixed_precision_enabled_for_params()
          ):
              p._mp_shard = torch.zeros_like(
                  p._local_shard,
                  device=self.compute_device,
                  dtype=self.mixed_precision.param_dtype
              )
              _free_storage(p._mp_shard)
  
          # We also maintain a full-sized parameter of type self.compute_dtype.
          # We resize the storage to size 0 at init (here) and only materialize
          # as needed. The storage may contain padding elements so that it is
          # evenly divisible by world_size, although these padding elements will
          # be removed before the relevant computation.
          if p._is_sharded:  # type: ignore[attr-defined]
              # We set p._full_param_padded's dtype to the desired parameter dtype
              # in the case of mixed precision. This is so that when we all_gather
              # into full_param_padded it can occur without issues and result in
              # full_param_padded having the expected param_dtype.
              full_param_dtype = (
                  p.dtype if not self._mixed_precision_enabled_for_params()
                  else self.mixed_precision.param_dtype
              )
              p._full_param_padded = torch.zeros(  # type: ignore[attr-defined]
                  p.numel() * self.world_size,
                  device=self.compute_device,
                  dtype=full_param_dtype,
              )
              _free_storage(p._full_param_padded)  # type: ignore[attr-defined]
  
      def _set_is_root(self) -> None:
          """If ``True``, implies that no other :class:`FullyShardedDataParallel`
          instance wraps this one. Called once by :func:`_lazy_init`.
          """
          if self._is_root is not None:
              return
          # No FSDP instance wraps this, else _is_root would be set to False.
          self._is_root = True
          self._exec_order_data.init(self)
          # If final backward callback is never been queued, state should be IDLE.
          # If final backward callback is queued, the callback should be finished
          # and the state was reset to be IDLE.
          # This should be asserted at the beginning of forward pass in the root instance only.
          # For children instances, if they are checkpointed, state will not be reset to
          # IDLE after each inner forward/backward.
          self._assert_state(TrainingState_.IDLE)
          for m in self.modules():
              if m is not self and isinstance(m, FullyShardedDataParallel):
                  # We relax the assert for non-root instance, when the nested initialized module is wrapped
                  # again in FSDP later, for example after training to run inference.
                  assert (
                      m._is_root is None or not m._is_root
                  ), "Non-root instance's _is_root flag should have not been set yet" \
                      "or has already been set as False."
                  if m._is_root is None:
                      m._is_root = False
  
      def _setup_streams(self) -> None:
          """Create streams to overlap data transfer and computation."""
          if len(self._streams) > 0 or not self._is_root:
              return
  
          if torch.cuda.is_available():
              # Stream for all-gathering parameters.
              self._streams["all_gather"] = torch.cuda.Stream()
              # Stream for overlapping grad reduction with the backward pass.
              self._streams["post_backward"] = torch.cuda.Stream()
              # Stream to move main params to self.mixed_precision.param_dtype
              # for forward pass.
              if self._mixed_precision_enabled_for_params():
                  self._streams["mixed_precision_params"] = torch.cuda.Stream()
  
          # We share streams with all children instances, which allows them to
          # overlap transfers across the forward pass without synchronizing with
          # the default stream.
          for m in self.modules():
              if m is not self and isinstance(m, FullyShardedDataParallel):
                  m._streams = self._streams
                  m._fsdp_graph_order = self._fsdp_graph_order
                  # Give each non-root FSDP module an alias to the root's
                  # execution order data structure and the root's ignored
                  # parameters and all buffer names since only the root's names
                  # are fully prefixed like the state dict keys
                  m._exec_order_data = self._exec_order_data
  
      def _wait_for_previous_optim_step(self) -> None:
          """
          The outer-most :class:`FullyShardedDataParallel` instance (i.e., the root
          instance) needs to synchronize with the default stream to ensure the
          previous optimizer step is done.
          """
          if not torch.cuda.is_available():
              return
  
          if self._mixed_precision_enabled_for_params():
              self._streams["mixed_precision_params"].wait_stream(
                  torch.cuda.current_stream()
              )
  
          self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
  
      def _need_prefetch_pre_backward_hook(self) -> bool:
          if (
              self.backward_prefetch == BackwardPrefetch.BACKWARD_PRE
              and self._fsdp_graph_order is not None
              and self._my_fsdp_idx_in_graph is not None
              and self._my_fsdp_idx_in_graph > 0
              and self._fsdp_graph_order[self._my_fsdp_idx_in_graph - 1].training_state
              != TrainingState_.BACKWARD_POST
          ):
              return True
          else:
              return False
  
      def _need_prefetch_post_backward_hook(self) -> bool:
          if (
              self.backward_prefetch == BackwardPrefetch.BACKWARD_POST
              and self._fsdp_graph_order is not None
              and self._my_fsdp_idx_in_graph is not None
              and self._my_fsdp_idx_in_graph > 0
              and self._fsdp_graph_order[self._my_fsdp_idx_in_graph - 1].training_state
              != TrainingState_.BACKWARD_POST
              and self._fsdp_graph_order[
                  self._my_fsdp_idx_in_graph - 1
              ]._need_rebuild_full_params
          ):
              return True
          else:
              return False
  
      @staticmethod
      @contextlib.contextmanager
      def state_dict_type(
          module: nn.Module,
          state_dict_type: StateDictType,
          state_dict_config: Optional[StateDictConfig] = None,
      ) -> Generator:
          """
          A context manager to set the ``state_dict_type`` of all the descendant
          FSDP modules of the target module. The target module does not have to
          be a FSDP module. If the target module is a FSDP module, its
          ``state_dict_type`` will also be changed.
  
          .. note:: This API should be called for only the top-level (root)
              module.
  
          .. note:: This API enables users to transparently use the conventional
              ``state_dict`` API to take model checkpoints in cases where the
              root FSDP module is wrapped by another ``nn.Module``. For example,
              the following will ensure ``state_dict``  is called on all non-FSDP
              instances, while dispatching into `local_state_dict` implementation
              for FSDP:
  
          Example::
  
          >>> model = DDP(FSDP(...))
          >>> with FSDP.state_dict_type(model, StateDictType.LOCAL_STATE_DICT):
          >>>     checkpoint = model.state_dict()
  
          Args:
              module (torch.nn.Module): Root module.
              state_dict_type (StateDictType): the desired ``state_dict_type`` to set.
          """
          prev_state_dict_type = None
          prev_state_dict_config = None
          # Use default config a state_dict config is not set.
          if state_dict_config is None:
              state_dict_config = _state_dict_type_to_config[state_dict_type]()
          for submodule in FullyShardedDataParallel.fsdp_modules(module):
              if prev_state_dict_type is None:
                  prev_state_dict_type = submodule._state_dict_type
              if prev_state_dict_config is None:
                  prev_state_dict_config = submodule._state_dict_config
              if prev_state_dict_type != submodule._state_dict_type:
                  raise RuntimeError("All FSDP module should the same state_dict_type.")
              if type(prev_state_dict_config) != type(submodule._state_dict_config):
                  raise RuntimeError(
                      "All FSDP modules should have the same type of state_dict_config."
                  )
  
              expected_state_dict_config_type = _state_dict_type_to_config[state_dict_type]
              if expected_state_dict_config_type != type(state_dict_config):
                  raise RuntimeError(
                      f"Expected state_dict_config of type {expected_state_dict_config_type} but got {type(state_dict_config)}"
                  )
              submodule._state_dict_type = state_dict_type
              submodule._state_dict_config = state_dict_config
          try:
              yield
          finally:
              assert prev_state_dict_type is not None  # Avoid mypy warning
              assert prev_state_dict_config is not None  # Avoid mypy warning
              for submodule in FullyShardedDataParallel.fsdp_modules(module):
                  submodule._state_dict_type = prev_state_dict_type
                  submodule._state_dict_config = prev_state_dict_config
  
      def _full_post_state_dict_hook(
          self,
          state_dict: Dict[str, Any],
          prefix: str,
      ) -> Dict[str, Any]:
          """
          Hook that runs after model.state_dict() is called before returning result to
          user. For FSDP, we may have to clone the tensors in state_dict as params go
          back to sharded version after _summon_full_params ends, and also remove
          "_fsdp_wrapped_module" prefix.
          """
          self._assert_state([TrainingState_.SUMMON_FULL_PARAMS])
          # state_dict is empty for nonzero ranks if `rank0_only` was enabled.
          if not state_dict:
              return state_dict
  
          offload_to_cpu = self._state_dict_config.offload_to_cpu
          cpu_device = torch.device("cpu")
          for key in state_dict:
              clean_key = clean_tensor_name(key)
              clean_prefix = clean_tensor_name(prefix)
              # Strip prefix out of key if needed as buffer names and param names
              # do not have prefix considered as they are not computed in `state_dict`
              # call.
              if clean_key.startswith(clean_prefix):
                  clean_key = clean_key[len(clean_prefix):]
              # Do not need to clone buffers since they are not sharded
              if clean_key in self._buffer_names:
                  # Offload the buffer to CPU if needed -- we do not do this in
                  # `_summon_full_params()` since without care, that would free
                  # the original buffer's GPU memory and require reallocating
                  # that memory later; this only affects the state dict's buffer
                  # variable and leaves the original buffer's GPU memory intact
                  if offload_to_cpu and state_dict[key].device != cpu_device:
                      state_dict[key] = state_dict[key].to(cpu_device)
                  continue
              # Clone non-ignored parameters before exiting the
              # `_summon_full_params()` context
              if clean_key not in self._ignored_param_names and \
                      not getattr(state_dict[key], "_has_been_cloned", False):
                  try:
                      state_dict[key] = state_dict[key].clone().detach()
                      state_dict[key]._has_been_cloned = True  # type: ignore[attr-defined]
                  except BaseException as e:
                      warnings.warn(
                          f"Failed to clone() tensor with name {key}. This may mean "
                          "that this state_dict entry could point to invalid memory "
                          "regions after returning from state_dict() call if this "
                          "parameter is managed by FSDP. Please check clone "
                          f"implementation of {key}. Error: {str(e)}"
                      )
  
          _replace_by_prefix(state_dict, prefix + f"{FSDP_WRAPPED_MODULE}.", prefix)
          return state_dict
  
      def _local_post_state_dict_hook(
          self,
          state_dict: Dict[str, Any],
          prefix: str,
      ) -> Dict[str, Any]:
          """
          This hook create a ShardedTensor from the local flat_param and replace
          the state_dict[f"{prefix}{FLAT_PARAM}] with the ShardedTensor. No copy
          will happen. The underlying storage is the same.
          """
          _replace_by_prefix(state_dict, f"{prefix}{FSDP_WRAPPED_MODULE}.", prefix)
          if self.module.no_params:
              return state_dict
  
          # state_dict[f"{prefix}{FLAT_PARAM}"] exists and has the same tensor
          # value as the flat_param but it is a pure Tensor because
          # nn.Module.state_dict() will detach the parameter. Therefore, we need
          # to get flat_param from the FlattenParamsWrapper to get the metadata.
          flat_param = getattr(self.module, FLAT_PARAM, None)
          # Construct a ShardedTensor from the flat_param.
          full_numel = flat_param.full_numel
          shard_offset = flat_param.numel() * self.rank
          valid_data_size = flat_param.numel() - flat_param.num_padded
          if valid_data_size > 0 and flat_param.num_padded > 0:
              flat_param = flat_param.narrow(0, 0, valid_data_size)
          local_shards = [
              Shard.from_tensor_and_offsets(flat_param, [shard_offset], self.rank)
          ]
          state_dict[f"{prefix}{FLAT_PARAM}"] = init_from_local_shards(
              local_shards, full_numel, process_group=self.process_group
          )  # type: ignore[assignment]
  
          return state_dict
  
      @torch.no_grad()
      def _sharded_post_state_dict_hook(
          self,
          state_dict: Dict[str, Any],
          prefix: str,
      ) -> Dict[str, Any]:
          """
          The hook replaces the unflattened, unsharded parameter in the state_dict
          with a unflattened, sharded parameter (a ShardedTensor).
          """
          _replace_by_prefix(state_dict, f"{prefix}{FSDP_WRAPPED_MODULE}.", prefix)
          if self.module.no_params:
              return state_dict
  
          for module_name, _, param_name in self.module.orig_flat_param[0].param_info:
              module_name = module_name.replace(f"{FPW_MODULE}.", "")
              module_name = module_name.replace(f"{FPW_MODULE}", "")
              if module_name:
                  module_name = f"{module_name}."
              fqn = f"{prefix}{module_name}{param_name}"
  
              # Create a ShardedTensor for the unflattened, non-sharded parameter.
              param = state_dict[fqn]
              local_shard = param.chunk(self.world_size)[self.rank].clone()
              offsets = [0 for _ in param.size()]
              offsets[0] = math.ceil(param.size()[0] / self.world_size) * self.rank
              local_shards = [
                  Shard.from_tensor_and_offsets(local_shard, offsets, self.rank)
              ]
              state_dict[fqn] = init_from_local_shards(
                  local_shards, param.size(), process_group=self.process_group
              )  # type: ignore[assignment]
          return state_dict
  
      @staticmethod
      def _post_state_dict_hook(
          module: nn.Module,
          state_dict: Dict[str, Any],
          prefix: str,
          *args: Any,
      ) -> Dict[str, Any]:
          """
          _post_state_dict_hook() is called after the state_dict() of this
          FSDP module is executed. ``self._state_dict_type`` is used to decide
          what postprocessing will be done.
          """
          self = cast(FullyShardedDataParallel, module)
          processed_state_dict = self._post_state_dict_hook_fn[self._state_dict_type](state_dict, prefix)
          # Restore buffers, which currently are in their full precision type,
          # back to their mixed precision type. This is because buffers are cast
          # during lazy_init() and stay at their mixed precision type before/after
          # forward/backward. As a result state_dict() should maintain this.
          if (
              self._is_root
              and self._mixed_precision_enabled_for_buffers()
          ):
              self._cast_buffers(recurse=True)
          return processed_state_dict
  
      def state_dict(self, *args, **kwargs):
          """
          This is the entry point of all three FSDP ``state_dict`` APIs: full,
          local, and sharded. For the full state dict
          (``StateDictType.FULL_STATE_DICT``), FSDP attempts to unshard the model
          on all ranks, which may result in an OOM error if the full model cannot
          fit on a single GPU. In that case, users may pass in a
          :class:`FullStateDictConfig` to only save the checkpoint on rank 0 and/
          or to offload it to CPU memory layer by layer, enabling much larger
          checkpoints. If the full model cannot fit in CPU memory, then users may
          instead take a local state dict (``StateDictType.LOCAL_STATE_DICT``)
          that only saves the local shard of the model. The sharded state dict
          (``StateDictType.SHARDED_STATE_DICT``) saves the model parameters as
          ``ShardedTensor`` s. The ``state_dict`` type can be configured using
          the :meth:`state_dict_type` context manager.
  
          Example::
  
          >>> import torch
          >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
          >>> from torch.distributed.fsdp import StateDictType
          >>> torch.cuda.set_device(device_id)
          >>> my_module = nn.Linear(...)
          >>> sharded_module = FSDP(my_module)
          >>> full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
          >>> with FSDP.state_dict_type(sharded_module, StateDictType.FULL_STATE_DICT, full_state_dict_config):
          >>>     full_dict = sharded_module.state_dict()
          >>> full_dict.keys()
          >>> odict_keys(['weight', 'bias'])
          >>> # using local state dict
          >>> with FSDP.state_dict_type(sharded_module, StateDictType.LOCAL_STATE_DICT):
          >>>     local_dict = sharded_module.state_dict()
          >>> local_dict.keys()
          >>> odict_keys(['flat_param', 'inner.flat_param'])
  
          .. warning:: This needs to be called on all ranks, since synchronization
              primitives may be used.
          """
          # TODO (rohan-varma): separate these out once a state_dict pre-hook
          # is available.
          if torch.cuda.is_available():
              torch.cuda.synchronize()
  
          self._lazy_init()
          if self._state_dict_type == StateDictType.FULL_STATE_DICT:
              # Get config args
              full_state_dict_config = (
                  self._state_dict_config if self._state_dict_config is not None
                  else FullStateDictConfig()
              )
              rank0_only = full_state_dict_config.rank0_only
              offload_to_cpu = full_state_dict_config.offload_to_cpu
              summon_ctx = (
                  self._summon_full_params(
                      recurse=False, writeback=False, offload_to_cpu=offload_to_cpu, rank0_only=rank0_only
                  )
                  if self.training_state != TrainingState_.SUMMON_FULL_PARAMS else
                  contextlib.suppress()
              )
              with summon_ctx:
                  # Since buffers are not sharded and stay casted, restore them to their
                  # original user module specified types for checkpoint. We take care to
                  # recast in post_state_dict_hook for consistency with the fact that
                  # buffers stay casted after forward/backward. We must have the
                  # call here instead of above because _summon_full_params itself
                  # calls _lazy_init() which would cast the buffers.
                  if (
                      self._is_root
                      and self._mixed_precision_enabled_for_buffers()
                  ):
                      self._cast_buffers(
                          dtype=self._orig_buffer_dtypes, recurse=False
                      )
                  state_dict = super().state_dict(*args, **kwargs)
  
              # TODO: support offload to CPU in post state dict hook.
              if not rank0_only or self.rank == 0:
                  return state_dict
              else:
                  return {}
  
          elif self._state_dict_type == StateDictType.LOCAL_STATE_DICT:
              if (
                  self.module.flat_param is not None and
                  not self.module.flat_param._is_sharded
              ):
                  raise RuntimeError(
                      "local_state_dict can only be called "
                      "when parameters are flatten and sharded."
                  )
              return super().state_dict(*args, **kwargs)
          elif self._state_dict_type == StateDictType.SHARDED_STATE_DICT:
              summon_ctx = (
                  self._summon_full_params(recurse=False, writeback=False)
                  if self.training_state != TrainingState_.SUMMON_FULL_PARAMS else
                  contextlib.suppress()
              )
              with summon_ctx:
                  return super().state_dict(*args, **kwargs)
          else:
              raise ValueError(f"Unknown StateDictType {self._state_dict_type}.")
  
      def _local_state_dict(self, *args: Any, **kwargs: Any) -> Any:
          """
          Returns the local state of the module. Parameters are flattened and
          sharded, so the resulting state_dict can only be loaded after the module
          has been wrapped with FSDP.
          """
          with self.state_dict_type(self, StateDictType.LOCAL_STATE_DICT):
              return self.state_dict(*args, **kwargs)
  
      def _full_post_load_state_dict_hook(self, *args, **kwargs) -> None:
          # We should exit summon_full_params context.
          self._assert_state([TrainingState_.SUMMON_FULL_PARAMS])
          assert getattr(self, '_full_param_ctx', None) is not None
          self._full_param_ctx.__exit__(None, None, None)
          self._full_param_ctx = None
  
      def _sharded_state_dict(self, *args: Any, **kwargs: Any) -> Any:
          """
          Returns the sharded states of the module. Parameters are unflattened and
          sharded, so the resulting state_dict can be used with any parallelism
          (e.g., DPP, model parallelism, and single trainer) after a valid
          resharding.
          """
          with self.set_state_dict_type(StateDictType.SHARDED_STATE_DICT):
              return self.state_dict(self, *args, **kwargs)
  
      def _full_pre_load_state_dict_hook(
          self,
          state_dict: Dict[str, Any],
          prefix: str,
      ) -> None:
          # We do not expect to be calling pre-hooks twice without post-hook
          # call in between.
          assert getattr(self, '_full_param_ctx', None) is None
          # Note that it needs writeback=True to persist.
          self._full_param_ctx = self._summon_full_params(
              recurse=False, writeback=True
          )
          self._full_param_ctx.__enter__()
          _replace_by_prefix(state_dict, prefix, prefix + f"{FSDP_WRAPPED_MODULE}.")
  
      def _local_post_load_state_dict_hook(self, *args, **kwargs) -> None:
          pass
  
      def _local_pre_load_state_dict_hook(
          self,
          state_dict: Dict[str, Any],
          prefix: str,
      ) -> None:
          """
          This hook finds the local flat_param for this FSDP module from the
          state_dict. The flat_param should be a ShardedTensor. This hook converts
          the ShardedTensor to a tensor. No copy happen unless padding is required.
          """
          _replace_by_prefix(state_dict, prefix, f"{prefix}{FSDP_WRAPPED_MODULE}.")
          fqn = f"{prefix}{FSDP_WRAPPED_MODULE}.{FLAT_PARAM}"
          if fqn not in state_dict:
              assert getattr(self.module, FLAT_PARAM, None) is None, (
                  "No flat parameter in state_dict but self.module.flat_param is not None"
              )
              return
          load_tensor = state_dict[fqn]
          assert isinstance(
              load_tensor, ShardedTensor
          ), "Tensors in local_state_dict should be ShardedTensor."
  
          # Convert the ShardedTensor to a Tensor.
          shards = load_tensor.local_shards()
          assert len(shards), "load_local_state_dict assume one shard per ShardedTensor."
          load_tensor = cast(torch.Tensor, shards[0].tensor)
  
          # Get the metada of the flat_param to decide whether to pad the loaded
          # tensor.
          flat_param = self.module.flat_param
          assert flat_param is not None
          if flat_param.num_padded not in (0, flat_param.numel()):
              assert load_tensor.numel() < flat_param.numel(), (
                  f"Local shard size = {flat_param.numel()} and the tensor in "
                  f"the state_dict is {load_tensor.numel()}."
              )
              load_tensor = F.pad(load_tensor, [0, flat_param.num_padded])
          state_dict[fqn] = load_tensor
  
      def _sharded_post_load_state_dict_hook(self, *args, **kwargs) -> None:
          pass
  
      def _sharded_pre_load_state_dict_hook(
          self,
          state_dict: Dict[str, Any],
          prefix: str,
      ) -> None:
          """
          The hook combines the unflattened, sharded parameters (ShardedTensor) to
          a new FlatParameter and shards the new FlatParameter to the local chunk.
          """
          _replace_by_prefix(state_dict, prefix, prefix + f"{FSDP_WRAPPED_MODULE}.")
          if self.module.no_params:
              return
  
          if not self.module.flat_param._is_sharded:
              raise RuntimeError(
                  "load_sharded_state_dict can only be called when parameters "
                  "are flatten and sharded."
              )
  
          nonsharded_tensors = []
          # TODO: Reduce the communication by using only one _all_gather_base to
          # gather all the parameters in this layer. This can be achieved by
          # concatenated all the local shards and then append the padding.
          # https://github.com/pytorch/pytorch/issues/77461
          for module_name, _, param_name in self.module.flat_param._param_infos:
              module_name = module_name.replace(f"{FPW_MODULE}.", "")
              module_name = module_name.replace(f"{FPW_MODULE}", "")
              if module_name:
                  module_name = f"{module_name}."
              fqn = f"{prefix}{FSDP_WRAPPED_MODULE}.{module_name}{param_name}"
              param = state_dict.pop(fqn)
  
              # All-gather the param (ShardedTensor)
              shards = param.local_shards()
              local_tensor = cast(torch.Tensor, shards[0].tensor).flatten()
              dim_0_size = param.size()[0]
              param_numel = param.size().numel()
              chunk_size = (
                  math.ceil(dim_0_size / self.world_size) * param_numel // dim_0_size
              )
              num_padding = chunk_size - local_tensor.numel()
              if num_padding > 0:
                  local_tensor = F.pad(local_tensor, [0, num_padding])
              tensor = torch.empty(
                  chunk_size * self.world_size, dtype=local_tensor.dtype
              ).cuda()
              dist._all_gather_base(tensor, local_tensor, group=self.process_group)
              tensor = tensor.narrow(0, 0, param_numel).reshape(param.size())
              nonsharded_tensors.append(tensor)
  
          # Create a new flat_param from the loaded, non-sharded tensors.
          flat_param = self.module.flat_param
          loaded_flat_param = FlatParameter(nonsharded_tensors, requires_grad=False)
  
          # Get the chunk from the loaded flat_param for the local rank.
          loaded_flat_param, num_to_pad = self._get_shard(loaded_flat_param)
          assert flat_param.numel() == loaded_flat_param.numel(), (
              f"The loaded local chunk has different numel({flat_param.numel()}) "
              f"from the local chunk {flat_param.numel()}."
          )
          assert flat_param.num_padded == num_to_pad, (
              f"The loaded local chunk has different padding({num_to_pad}) "
              f"from the local chunk {flat_param.num_padded}."
          )
          state_dict[f"{prefix}_fsdp_wrapped_module.flat_param"] = loaded_flat_param
  
      @staticmethod
      def _pre_load_state_dict_hook(
          module: nn.Module,
          state_dict: Dict[str, Any],
          prefix: str,
          *args: Any,
      ) -> None:
          """
          ``_pre_state_dict_hook` is called before ``self._load_from_state_dict()``
          is called. ``self._state_dict_type`` is used to decide what preprocessing
          will be done.
          """
          # Code that is common for all state_dict impls
          self = cast(FullyShardedDataParallel, module)
          if torch.cuda.is_available():
              torch.cuda.synchronize()
          # Dispatch into state_dict specific implementation of pre-hook.
          self._pre_load_state_dict_hook_fn[self._state_dict_type](state_dict, prefix)
  
      @staticmethod
      def _post_load_state_dict_hook(module: nn.Module, *args: Any) -> None:
          # Code that is common for all state_dict impls
          self = cast(FullyShardedDataParallel, module)
          # Dispatch into state_dict type specific implementation of post-hook for
          # loading state_dict.
          self._post_load_state_dict_hook_fn[self._state_dict_type]()
  
      def load_state_dict(
          self,
          state_dict: Mapping[str, Any],
          *args,
      ) -> NamedTuple:
          """
          The entry point of all three FSDP ``load_state_dict`` APIs. By default,
          calling ``load_state_dict`` on an FSDP module will result in FSDP
          attempting to load a "full" state_dict, i.e. a state_dict consisting of
          full, unsharded, unflattened original module parameters. This requires
          FSDP to load the full parameter context on each rank which could result
          in GPU OOM. As a result, :func:`state_dict_type` API is available to
          configure between ``load_state_dict`` implementations. User can thus use
          ``with self.state_dict_type(self, StateDictType.LOCAL_STATE_DICT)`` context
          manager to load a local state dict checkpoint that will restore only
          local shards of the module. Currently, the only supported
          implementations are ``StateDictType.LOCAL_STATE_DICT`` and
          ``StateDictType.FULL_STATE_DICT`` (default). Please see :func:`state_dict`
          for documentation around creating an FSDP checkpoint.
  
          Example::
  
          >>> import torch
          >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
          >>> from torch.distributed.fsdp import StateDictType
          >>> torch.cuda.set_device(device_id)
          >>> my_module = nn.Linear(...)
          >>> sharded_module = FSDP(my_module)
          >>> checkpoint = torch.load(PATH)
          >>> full_state_dict = checkpoint['full_state_dict']
          >>> with FSDP.state_dict_type(sharded_module, StateDictType.FULL_STATE_DICT):
          >>>     sharded_module.load_state_dict(full_state_dict)
          >>> full_dict.keys()
          >>> odict_keys(['weight', 'bias'])
          >>> # using local state dict
          >>> local_state_dict = checkpoint['local_state_dict']
          >>> with FSDP.state_dict_type(sharded_module, StateDictType.LOCAL_STATE_DICT):
          >>>     sharded_module.load_state_dict(local_state_dict)
          >>> local_dict.keys()
          >>> odict_keys(['flat_param', 'inner.flat_param'])
  
          .. warning:: This needs to be called on all ranks, since synchronization
              primitives may be used.
          """
          return super().load_state_dict(state_dict, *args)
  
      def _load_local_state_dict(
          self,
          state_dict: Mapping[str, Any],
          *args,
      ) -> NamedTuple:
          """
          Load states from a flattened, sharded state dictionary.
          """
          with self.state_dict_type(self, StateDictType.LOCAL_STATE_DICT):
              return self.load_state_dict(state_dict, *args)
  
      def _load_sharded_state_dict(
          self,
          state_dict: Union[Dict[str, torch.Tensor], "OrderedDict[str, torch.Tensor]"],
          strict: bool = True,
      ) -> NamedTuple:
          """
          Load states from a unflattened, sharded state dictionary.
          """
          with self.set_state_dict_type(StateDictType.SHARDED_STATE_DICT):
              return self.load_state_dict(state_dict, strict)
  
      def forward(self, *args: Any, **kwargs: Any) -> Any:
          with torch.autograd.profiler.record_function("FullyShardedDataParallel.forward"):
              self._lazy_init()
  
              # Start of a forward pass.
              self.training_state = TrainingState_.FORWARD
              if self._is_root:
                  # TODO: disabling side stream for tensor copies for now, investigate
                  # perf with it on / off.
                  # Place inputs on compute_device. This is a noop if inputs are already
                  # on compute_device. Note that when device_id is specified,
                  # device_id == self.compute_device is guaranteed.
                  # TODO: for mixed precision, move inputs to right device + cast might
                  # be done in one go for performance.
                  args, kwargs = _to_kwargs(args, kwargs, self.compute_device.index, False)
                  args = args[0]
                  kwargs = kwargs[0]
  
              # Cast inputs to their mixed precision type.
              if (
                  self._is_root
                  and self._mixed_precision_enabled_for_params()
              ):
                  input_dtype = self.mixed_precision.param_dtype
                  args, kwargs = self._cast_fp_inputs_to_precision(
                      input_dtype, *args, **kwargs
                  )
  
              # All-gather full parameters, moving them to compute_device if
              # necessary.
              self._rebuild_full_params()
              # Wait for all_gather full parameters to finish before computation
              torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
  
              # Register backward hooks to reshard params and reduce-scatter grads.
              # These need to be re-registered every forward pass in some cases where grad_fn
              # is mutated.
              self._register_post_backward_hooks()
              outputs = self.module(*args, **kwargs)
  
              if self not in self._fsdp_graph_order:
                  self._my_fsdp_idx_in_graph = len(self._fsdp_graph_order)
                  self._fsdp_graph_order.append(self)
  
              if self.reshard_after_forward:
                  self._free_full_params()
                  if (
                      self._mixed_precision_enabled_for_params()
                  ):
                      self._free_mp_shard(self.params)
              # Switch to original local shards of params. We maintain this invariant throughout
              # the code, i.e., ``p.data == p._local_shard`` after each function. This
              # also ensures that after the first forward, the optimizer state will be
              # initialized with the correct dtype and (sharded) size, since optimizer
              # state is typically initialized lazily in ``optim.step()``. Note that
              # when CPU offload is enabled, _use_param_local_shard implicitly
              # offloads the local shard to CPU by making p.data point to
              # p._local_shard, which would reside on CPU.
              self._use_param_local_shard()
  
              # Register pre-backward hooks to all-gather the params for the backward
              # pass (if output's grad was needed). This won't register anything if
              # we are in eval mode.
              outputs = self._register_pre_backward_hooks(outputs)
  
              # Done with a forward pass.
              self.training_state = TrainingState_.IDLE
  
          return outputs
  
      @torch.no_grad()
      def _write_back_current_shard(self, full_params):
          """
          Writes back full_params into self.params.
          """
          for p, (full_param, _) in zip(self.params, full_params):
              if not p._is_sharded:  # type: ignore[attr-defined]
                  continue  # Already copied because no sharding.
  
              # TODO: Might be able to refactor to use _get_shard.
              chunks = full_param.chunk(self.world_size)  # type: ignore[attr-defined]
              assert len(chunks) > self.rank
              chunk = chunks[self.rank]
              p._local_shard.copy_(chunk)  # type: ignore[attr-defined]
  
      @contextlib.contextmanager
      def _summon_full_params(
          self,
          recurse: bool = True,
          writeback: bool = True,
          rank0_only: bool = False,
          offload_to_cpu: bool = False,
      ):
          if writeback and rank0_only:
              raise ValueError(
                  "writeback=True and rank0_only=True is not supported, as model "
                  "parameter shapes will be different across ranks, and writing "
                  "to them can lead to inconsistencies across ranks when the "
                  "context is exited."
              )
  
          if offload_to_cpu and not rank0_only:
              warnings.warn(
                  "offload_to_cpu and rank0_only=False will result in "
                  "full parameters being redundantly copied to CPU memory for "
                  "GPUs that reside on the same machine, which may incur the risk of "
                  "CPU OOM. It is recommended to use ``offload_to_cpu`` with "
                  "rank0_only=True."
              )
  
          def _free_full_params_and_use_local_shard(params_to_free):
              # We may not always be able to free the full param, for example in
              # the case where world_size == 1 and the shard actually points to
              # the full parameter.
              for (param, can_free) in params_to_free:
                  if can_free:
                      current_stream = torch.cuda.current_stream()
                      # Don't let PyTorch reuse this memory until all work in the
                      # current stream is complete
                      param.record_stream(current_stream)
                      _free_storage(param)
  
              # when CPU offload is enabled, _use_param_local_shard implicitly
              # offloads the local shard to CPU by making p.data point to
              # p._local_shard, which would reside on CPU.
              self._use_param_local_shard()
  
          if recurse:
              with contextlib.ExitStack() as stack:
                  # Summon all params for any nested FSDP instances.
                  for module in self.fsdp_modules(self):
                      stack.enter_context(
                          module._summon_full_params(
                              recurse=False,
                              writeback=writeback,
                              rank0_only=rank0_only,
                              offload_to_cpu=offload_to_cpu,
                          )
                      )
                  # Yield to the caller, with full params in all nested instances.
                  yield
              # Exiting from the ExitStack will re-shard params.
              return
          else:
              torch.cuda.synchronize()
              self._lazy_init()
              self._assert_state([TrainingState_.IDLE])
              # Set the state so that we assert when trying to go into
              # forward/backward.
              self.training_state = TrainingState_.SUMMON_FULL_PARAMS
  
              # Even if rank0_only = True, we need to materialize all params here
              # and free them right after as full param materialization requires
              # collective comm.
              currently_local_params = self._rebuild_full_params()
              # Wait for all_gather to finish before computation
              torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
              my_rank = dist.get_rank(self.process_group)
              if offload_to_cpu and (not rank0_only or my_rank == 0):
                  for p in self.params:
                      if p._is_sharded:
                          with torch.no_grad():
                              # Note that we avoid using p._full_param_padded
                              # directly here as we may not be using that param
                              # as the full_param from _rebuild_full_params (i.e.)
                              # in mixed precision.
                              for p, (full_param, _) in zip(
                                  self.params, currently_local_params
                              ):
                                  full_param = full_param.to(torch.device("cpu"))
                                  self._update_p_data(p, output_tensor=full_param)
  
              if rank0_only and my_rank != 0:
                  _free_full_params_and_use_local_shard(currently_local_params)
                  try:
                      yield
                  finally:
                      self.training_state = TrainingState_.IDLE
              else:
                  # FSDP now has the full flattened parameter. Unflatten it to get the
                  # full parameters.
                  with contextlib.ExitStack() as stack:
                      # Invariant: rank == 0 or !rank0_only
                      stack.enter_context(self.module.unflatten_params())
                      try:
                          yield
                      finally:
                          if offload_to_cpu:
                              for p in self.params:
                                  if p._is_sharded:
                                      with torch.no_grad():
                                          # Note that we avoid using
                                          # p._full_param_padded directly here as
                                          # we may not be using that param
                                          # as the full_param from
                                          # _rebuild_full_params (i.e. in mixed
                                          # precision.
                                          for p, (full_param, _) in zip(
                                              self.params, currently_local_params
                                          ):
                                              full_param = full_param.to(self.compute_device)
                                              self._update_p_data(
                                                  p, output_tensor=full_param,
                                              )
  
                          if writeback:
                              self._write_back_current_shard(currently_local_params)
                          stack.close()
                          _free_full_params_and_use_local_shard(currently_local_params)
                          self.training_state = TrainingState_.IDLE
  
      @staticmethod
      @contextlib.contextmanager
      def summon_full_params(
          module,
          recurse: bool = True,
          writeback: bool = True,
          rank0_only: bool = False,
          offload_to_cpu: bool = False,
      ) -> Generator:
          r""" A context manager to expose full params for FSDP instances.
          Can be useful *after* forward/backward for a model to get
          the params for additional processing or checking. It can take a non-FSDP
          module and will summon full params for all contained FSDP modules as
          well as their children, depending on the ``recurse`` argument.
  
          .. note:: This can be used on inner FSDPs.
          .. note:: This can *not* be used within a forward or backward pass. Nor
              can forward and backward be started from within this context.
          .. note:: Parameters will revert to their local shards after the context
              manager exits, storage behavior is the same as forward.
          .. note:: The full parameters can be modified, but only the portion
              corresponding to the local param shard will persist after the
              context manager exits (unless ``writeback=False``, in which case
              changes will be discarded). In the case where FSDP does not shard
              the parameters, currently only when ``world_size == 1``, or ``NO_SHARD``
              config, the modification is persisted regardless of ``writeback``.
          .. note:: This method works on modules which are not FSDP themselves but
              may contain multiple independent FSDP units. In that case, the given
              arguments will apply to all contained FSDP units.
  
          .. warning:: Note that ``rank0_only=True`` in conjunction with
              ``writeback=True`` is not currently supported and will raise an
              error. This is because model parameter shapes would be different
              across ranks within the context, and writing to them can lead to
              inconsistency across ranks when the context is exited.
  
          .. warning:: Note that ``offload_to_cpu`` and ``rank0_only=False`` will
              result in full parameters being redundantly copied to CPU memory for
              GPUs that reside on the same machine, which may incur the risk of
              CPU OOM. It is recommended to use ``offload_to_cpu`` with
              ``rank0_only=True``.
  
          Args:
              recurse (bool, Optional): recursively summon all params for nested
                  FSDP instances (default: True).
              writeback (bool, Optional): if ``False``, modifications to params are
                  discarded after the context manager exists;
                  disabling this can be slightly more efficient (default: True)
              rank0_only (bool, Optional): if ``True``, full parameters are
                  materialized on only global rank 0. This means that within the
                  context, only rank 0 will have full parameters and the other
                  ranks will have sharded parameters. Note that setting
                  ``rank0_only=True`` with ``writeback=True`` is not supported,
                  as model parameter shapes will be different across ranks
                  within the context, and writing to them can lead to
                  inconsistency across ranks when the context is exited.
              offload_to_cpu (bool, Optional): If ``True``, full parameters are
                  offloaded to CPU. Note that this offloading currently only
                  occurs if the parameter is sharded (which is only not the case
                  for world_size = 1 or ``NO_SHARD`` config). It is recommended
                  to use ``offload_to_cpu`` with ``rank0_only=True`` to avoid
                  redundant copies of model parameters being offloaded to the same CPU memory.
          """
          # Note that we specify root_only as FSDP roots will handle summoning
          # child FSDP instances based on recurse argument.
          fsdp_modules = FullyShardedDataParallel.fsdp_modules(
              module, root_only=True
          )
          # Summon all params for all FSDP instances
          with contextlib.ExitStack() as stack:
              for module in fsdp_modules:
                  stack.enter_context(
                      module._summon_full_params(
                          recurse=recurse,
                          writeback=writeback,
                          rank0_only=rank0_only,
                          offload_to_cpu=offload_to_cpu,
                      )
                  )
              # Yield to the caller, with full params in all FSDP instances.
              yield
          # Exiting from the ExitStack will reshard all params.
          return
  
      def named_buffers(
          self,
          *args,
          **kwargs,
      ) -> Iterator[Tuple[str, torch.Tensor]]:
          """
          Overrides :meth:`named_buffers()` to intercept buffer names and
          remove all occurrences of the FSDP-specific flattened buffer prefix
          when inside the :meth:`summon_full_params` context manager.
          """
          in_summon_full_params = self.training_state == TrainingState_.SUMMON_FULL_PARAMS
          for buffer_name, buffer in super().named_buffers(*args, **kwargs):
              if in_summon_full_params:
                  # Remove any instances of the FSDP-specific prefix; there can
                  # be multiple in the case of nested FSDP modules
                  buffer_name = buffer_name.replace(FSDP_PREFIX, "")
              yield (buffer_name, buffer)
  
      def named_parameters(
          self,
          *args,
          **kwargs,
      ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
          """
          Overrides :meth:`named_parameters()` to intercept parameter names and
          remove all occurrences of the FSDP-specific flattened parameter prefix
          when inside the :meth:`summon_full_params` context manager.
          """
          # Determine which logic to use based on the context at call time
          in_summon_full_params = self.training_state == TrainingState_.SUMMON_FULL_PARAMS
          for param_name, param in super().named_parameters(*args, **kwargs):
              if in_summon_full_params:
                  # Remove any instances of the FSDP-specific prefix; there can
                  # be multiple in the case of nested FSDP modules
                  param_name = param_name.replace(FSDP_PREFIX, "")
              yield (param_name, param)
  
      def _register_pre_backward_hooks(self, outputs: Any) -> Any:
          """Register pre-backward hook to run before the wrapped module's
          backward. Hooks should be attached to all outputs from the forward.
          Returns:
              outputs: new outputs with hooks registered if they requires gradient.
          """
          # Reset before each backward pass
          self._need_rebuild_full_params = False
  
          if not torch.is_grad_enabled():
              return outputs  # don't register hooks if grad isn't enabled
  
          if self._is_root:
              # This actually means that only root instance has
              # _post_backward_callback_queued defined. Accidentally accessing this field
              # will assert on all other instances, giving us a nice bug checker.
              self._post_backward_callback_queued = False
  
          # Reset before each backward pass
          self._pre_backward_hook_has_run = False
  
          def _pre_backward_hook(*unused: Any) -> None:
              # Run ``_pre_backward_hook`` only once per backward pass
              if self._pre_backward_hook_has_run:
                  return
              # try to queue final backward callback only once for root, so
              # that final backward callback is attached to the outer most
              # backward graph task and called after all the backward
              # calls are completed.
              if self._is_root:
                  self._queue_wait_for_post_backward()
  
              if self._need_prefetch_pre_backward_hook():
                  # Always wait for all_gather before rebuilding full params, just
                  # in case full params have already been prefetched in previous layer's
                  # pre-backward hook.
                  torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
  
              # Start of a backward pass for the first time in an backward pass.
              self._assert_state([TrainingState_.IDLE])
              self.training_state = TrainingState_.BACKWARD_PRE
  
              # All-gather full parameters, moving them to compute device if
              # necessary.
              self._rebuild_full_params()
              # Wait for all_gather to finish before computation
              torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
  
              # Prefetch next layer's full params in backward pass,
              # since it is prefetching, no need to wait for all_gather stream.
              if self._need_prefetch_pre_backward_hook():
                  self._fsdp_graph_order[self._my_fsdp_idx_in_graph - 1]._rebuild_full_params()  # type: ignore[operator]
  
              self._pre_backward_hook_has_run = True
              # Prepare p.grad so that it is in the right shape, device, accumulated values, etc.
              self._prep_grads_for_backward()
  
          def _register_hook(t: torch.Tensor) -> torch.Tensor:
              if t.requires_grad:
                  t.register_hook(_pre_backward_hook)
                  self._need_rebuild_full_params = True
              return t
  
          # Attach hooks to Tensor outputs.
          outputs = _apply_to_tensors(_register_hook, outputs)
  
          return outputs
  
      def _register_post_backward_hooks(self) -> None:
          """
          Register backward hooks to reshard params and reduce-scatter grads.
          This is called during forward pass. The goal is to attach a hook
          on each of the parameter's gradient generating function (``grad_acc``
          below) so that the hook is called *after* all gradients for that
          param are computed.
          Goals:
          1. We want the hook to fire once and only once *after* all gradients
          are accumulated for a param.
          2. If it fires more than once, we end up incorrectly shard the grad
          multiple times. (could lead to dimension too small)
          3. If it fires once but too early or doesn't fire, we leave gradients
          unsharded. (could lead to dimension too large)
          Due to multiple-pass forward, this function can be called on
          the same parameter multiple times in a single forward pass. If we register
          the hook multiple time, we end up getting called multiple times. We
          could try to get a new hook every time and delete the previous one
          registered. However, due to *unknown reason* (I have debugged it for
          a long time!), in mixed precision mode, we get two different ``grad_acc``
          objects below during different calls of this function (in the same
          forward pass). If we keep the last one, the hook end up firing too
          early. In full precision mode, we luckily get the *same* ``grad_acc``
          object, so deleting and re-registering still ensured the hook fire
          once after all gradients are generated.
          Empirically, keep the first hook register per forward pass seems to
          work the best. We do need to remove the hook at the end of the
          backward pass. Otherwise, the next forward pass will not register
          a new hook, which is needed for a new forward pass.
          """
          if not torch.is_grad_enabled():
              return  # don't register grad hooks if grad isn't enabled
          for p in self.params:
              if p.requires_grad:
                  if hasattr(p, "_shard_bwd_hook"):
                      continue
                  # Register a hook on the first call, empirically, autograd
                  # fires it at the end for this param, which makes sense.
                  p_tmp = p.expand_as(p)  # Get a grad_fn on p_tmp.
                  assert (
                      p_tmp.grad_fn is not None
                  ), "p_tmp grad_fn should not be None, it is used to access \
                      p's AccumulateGrad object and register post hook on it."
                  grad_acc = p_tmp.grad_fn.next_functions[0][
                      0
                  ]  # Gets its AccumulateGrad object.
                  handle = grad_acc.register_hook(
                      functools.partial(self._post_backward_hook, p)
                  )
                  p._shard_bwd_hook = (grad_acc, handle)  # type: ignore[attr-defined]
  
      @torch.no_grad()
      def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
          """
          At the start of :func:`_post_backward_hook`, ``param.grad`` contains the
          full gradient for the local batch. The reduce-scatter op will replace
          ``param.grad`` with a single shard of the summed gradient across all
          GPUs. This shard will align with the current GPU rank. For example::
              before reduce_scatter:
                  param.grad (GPU #0): [1, 2, 3, 4]
                  param.grad (GPU #1): [5, 6, 7, 8]
              after reduce_scatter:
                  param.grad (GPU #0): [6, 8]    # 1+5, 2+6
                  param.grad (GPU #1): [10, 12]  # 3+7, 4+8
          The local GPU's ``optim.step`` is responsible for updating a single
          shard of params, also corresponding to the current GPU's rank. This
          alignment is created by :func:`_shard_parameters`, which ensures that
          the local optimizer only sees the relevant parameter shard.
          """
          # First hook callback will see PRE state. If we have multiple params,
          # then subsequent hook callbacks will see POST state.
          self._assert_state([TrainingState_.BACKWARD_PRE, TrainingState_.BACKWARD_POST])
          self.training_state = TrainingState_.BACKWARD_POST
          if param.grad is None:
              return
  
          if param.grad.requires_grad:
              raise RuntimeError(
                  "FSDP only works with gradients that don't require gradients"
              )
  
          if self._require_backward_grad_sync or \
                  self.sharding_strategy == ShardingStrategy.FULL_SHARD:
              # We free full parameters unless we are in `no_sync()` (i.e. when
              # `_require_backward_grad_sync=False`) and not using the
              # `FULL_SHARD` strategy. If we are not using the `FULL_SHARD`
              # strategy (e.g. instead using `SHARD_GRAD_OP`), then we keep the
              # full parameters in memory and save network overhead.
              self._free_full_params(cast(List[FlatParameter], [param]))
  
          if self._mixed_precision_enabled_for_params():
              # Noop if reshard_after_forward=True because we'd free the param
              # shard when rebuilding the full params in the pre_beckward_hook.
              self._free_mp_shard(cast(List[FlatParameter], [param]))
  
          # Switch to local shard after backward. Note that
          # when CPU offload is enabled, _use_param_local_shard implicitly
          # offloads the local shard to CPU by making p.data point to
          # p._local_shard, which would reside on CPU.
          self._use_param_local_shard(cast(List[FlatParameter], [param]))
  
          # Prefetch previous layer's full params in backward pass post backward hook,
          # If next layer's backward computation is done and full params are freed,
          # no need to prefetch the full params again.
          # Only prefetch full params if any of the next layer's outputs requires grad
          if self._need_prefetch_post_backward_hook():
              self._fsdp_graph_order[self._my_fsdp_idx_in_graph - 1]._rebuild_full_params()  # type: ignore[operator]
              # Next layer's computation will start right after this all_gather,
              # Wait for all_gather to finish before computation.
              torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
  
          if not self._require_backward_grad_sync:
              return
  
          # Wait for all work in the current stream to finish, then start the
          # reductions in post_backward stream.
          self._streams["post_backward"].wait_stream(torch.cuda.current_stream())
  
          with torch.cuda.stream(self._streams["post_backward"]):
              orig_grad_data = param.grad.data
              if (
                  self._mixed_precision_enabled_for_reduce()
              ):
                  # Cast gradient to precision in which it should be communicated.
                  # TODO: Make this a communication hook when communication hooks
                  # are implemented for FSDP. Note that this is a noop if the
                  # reduce_dtype matches the param dtype.
                  param.grad.data = param.grad.data.to(self.mixed_precision.reduce_dtype)
  
              if self.gradient_predivide_factor > 1:
                  # Average grad by world_size for consistency with PyTorch DDP.
                  param.grad.div_(self.gradient_predivide_factor)
  
              grad = param.grad.data
              if param._is_sharded:  # type: ignore[attr-defined]
                  # We clear `param.grad` to permit repeated gradient
                  # computations when this FSDP module is called multiple times.
                  # This is to avoid a race among multiple re-entrant backward
                  # passes. For example, the second backward pass computation
                  # precedes ahead of the first backward pass reduction, which is
                  # possible since the reduction is in a different stream and is
                  # async. Then, the first backward pass may be incorrectly
                  # reducing the second backward pass's `param.grad`.
                  # The reduced gradients are accumulated in
                  # `param._saved_grad_shard`, and the gradient reductions can
                  # happen in arbitrary order, though we tolerate this due to the
                  # (approximate) commutativity of floating-point addition.
                  param.grad = None
                  grad_flatten = torch.flatten(grad)
                  chunks = list(grad_flatten.chunk(self.world_size))
                  num_pad = self.world_size * chunks[0].numel() - grad.numel()
                  input_flattened = F.pad(grad_flatten, [0, num_pad])
                  output = torch.zeros_like(chunks[0])
                  dist._reduce_scatter_base(
                      output, input_flattened, group=self.process_group
                  )
                  if self.gradient_postdivide_factor > 1:
                      # Average grad by world_size for consistency with PyTorch DDP.
                      output.div_(self.gradient_postdivide_factor)
  
                  # Note that we need to cast grads back to the full precision if
                  # 1) parameters were in reduced precision during fwd, as grads
                  # would thus be in this reduced precision, or
                  # 2) parameters did not have precision reduced, but grads
                  # had reduced precision for communication.
                  if (
                      self._mixed_precision_enabled_for_params() or self._mixed_precision_enabled_for_reduce()
                  ):
                      # Cast gradients back to the full parameter precision so that
                      # optimizer.step() happens in full precision.
                      orig_param_grad_data = output
                      output.data = output.data.to(dtype=param.data.dtype)
                      # Don't let this memory get reused until after the transfer.
                      orig_param_grad_data.record_stream(torch.cuda.current_stream())
  
                  # To support gradient accumulation outside `no_sync()`, we save
                  # the gradient data to `param._saved_grad_shard` before the
                  # backward pass, accumulate gradients into it here, and set
                  # `param.grad` with the accumulated value at the end of the
                  # backward pass in preparation for the optimizer step.
                  accumulate_grad = hasattr(param, "_saved_grad_shard")
                  if accumulate_grad:
                      p_assert(
                          param._saved_grad_shard.shape == output.shape,  # type: ignore[attr-defined]
                          "Shape mismatch when accumulating gradients: "  # type: ignore[attr-defined]
                          f"existing grad shape={param._saved_grad_shard.shape} "
                          f"new grad shape={output.shape}"  # type: ignore[attr-defined]
                      )
                      p_assert(
                          param._saved_grad_shard.device == output.device,  # type: ignore[attr-defined]
                          "Device mismatch when accumulating gradients: "  # type: ignore[attr-defined]
                          f"existing grad device={param._saved_grad_shard.device} "
                          f"new grad device={output.device}"  # type: ignore[attr-defined]
                      )
                      param._saved_grad_shard += output  # type: ignore[attr-defined]
                  else:
                      param._saved_grad_shard = output  # type: ignore[attr-defined]
                  grad = param._saved_grad_shard  # type: ignore[attr-defined]
              else:
                  # Currently the way for _is_sharded to be False is if
                  # world_size == 1 or sharding_strategy is NO_SHARD.
                  assert (
                      self.world_size == 1 or self.sharding_strategy == ShardingStrategy.NO_SHARD
                  ), "Currently the way for _is_sharded to be False is \
                      world_size == 1 or sharding_stratagy is set to be NO_SHARD"
                  if self.sharding_strategy == ShardingStrategy.NO_SHARD:
                      dist.all_reduce(param.grad, group=self.process_group)
                      if self.gradient_postdivide_factor > 1:
                          # Average grad by world_size for consistency with PyTorch DDP.
                          param.grad.div_(self.gradient_postdivide_factor)
                  # Note that we need to cast grads back to the full precision if
                  # 1) parameters were in reduced precision during fwd, as grads
                  # would thus be in this reduced precision, or
                  # 2) parameters did not have precision reduced, but grads
                  # had reduced precision for communication.
                  if (
                      self._mixed_precision_enabled_for_params() or self._mixed_precision_enabled_for_reduce()
                  ):
                      # Cast gradients back to the full parameter precision so that
                      # optimizer.step() happens in full precision.
                      orig_param_grad_data = param.grad.data
                      param.grad.data = param.grad.data.to(dtype=param.data.dtype)
                      # Don't let this memory get reused until after the transfer.
                      orig_param_grad_data.record_stream(torch.cuda.current_stream())
  
              # Regardless of sharding or not, offload the grad to CPU if we are
              # offloading params. This is so param and grad reside on same device
              # which is needed for the optimizer step.
              if self.cpu_offload.offload_params:
                  # We specify non_blocking=True
                  # and ensure the appropriate synchronization is done by waiting
                  # streams in _wait_for_post_backward.
                  param._cpu_grad.copy_(  # type: ignore[attr-defined]
                      grad.detach(), non_blocking=True
                  )
                  # Don't let this memory get reused until after the transfer.
                  grad.data.record_stream(torch.cuda.current_stream())
  
              # After _post_backward_hook returns, orig_grad_data will eventually
              # go out of scope, at which point it could otherwise be freed for
              # further reuse by the main stream while the div/reduce_scatter/copy
              # are underway in the post_backward stream. See:
              # github.com/NVIDIA/apex/blob/master/apex/parallel/distributed.py
              orig_grad_data.record_stream(self._streams["post_backward"])
  
      def _queue_wait_for_post_backward(self) -> None:
          """Try to queue a `wait_for_post_backward` callback.
          Only called on root and only queue one callback at the beginning of
          outer most backward.
          """
          assert (
              self._is_root
          ), "_queue_wait_for_post_backward can only be called on root."
          if not self._post_backward_callback_queued:
              self._assert_state([TrainingState_.IDLE])
              self._post_backward_callback_queued = True
              Variable._execution_engine.queue_callback(self._wait_for_post_backward)
  
      @torch.no_grad()
      def _wait_for_post_backward(self) -> None:
          """Wait for post-backward to finish. Only called on root instance."""
          assert self._is_root, "_wait_for_post_backward can only be called on root."
          # Check if the root module has params and if any of them has
          # the `requires_grad` field set. If `requires_grad=False` for
          # all the params, the post_backward hook will not fire and the
          # state will remain in `TrainingState_.BACKWARD_PRE`.
          if any([p.requires_grad for p in self.params]):
              self._assert_state(TrainingState_.BACKWARD_POST)
          else:
              self._assert_state(TrainingState_.BACKWARD_PRE)
  
          if self._require_backward_grad_sync:
              torch.cuda.current_stream().wait_stream(self._streams["post_backward"])
              if self.cpu_offload.offload_params:
                  # We need to wait for the non-blocking GPU ->
                  # CPU grad transfers to finish. We need to do this for GPU -> CPU
                  # copies because when grad is on CPU, it won't wait for any CUDA
                  # stream to finish GPU -> CPU copies unless we explicitly block the
                  # host-side with synchronize().
                  torch.cuda.current_stream().synchronize()
  
          # A backward pass is done, clean up below.
          self._exec_order_data.reset()
  
          def _finalize_params(fsdp_module: FullyShardedDataParallel) -> None:
              """Helper used below on all fsdp modules."""
              for p in fsdp_module.params:
                  if p.requires_grad:
                      if hasattr(p, "_shard_bwd_hook"):
                          assert len(p._shard_bwd_hook) == 2 and len(  # type: ignore[attr-defined]
                              p._shard_bwd_hook  # type: ignore[attr-defined]
                          ), (  # type: ignore[attr-defined]
                              "p._shard_bwd_hook fields are not valid."
                          )
                          p._shard_bwd_hook[1].remove()  # type: ignore[attr-defined]
                          delattr(p, "_shard_bwd_hook")
                      # Preserve the gradient accumulation state if not
                      # synchronizing: `p.grad` remains the unsharded gradient
                      # accumulated from prior `no_sync()` iterations, and
                      # `p._saved_grad_shard` remains the sharded gradient from
                      # the last synchronized iteration
                      if not self._require_backward_grad_sync:
                          continue
                      # Set `p.grad` as needed to ensure optimizer correctness
                      # since optimizers operate on the `grad` attribute
                      if hasattr(p, "_cpu_grad"):
                          p_assert(
                              p.device == torch.device("cpu"),
                              f"Device mismatch: p={p.device} "  # type: ignore[attr-defined]
                              f"p._cpu_grad={p._cpu_grad}"
                          )
                          p.grad = p._cpu_grad  # type: ignore[attr-defined]
                      elif hasattr(p, "_saved_grad_shard"):
                          p_assert(
                              p.device == p._saved_grad_shard.device,  # type: ignore[attr-defined]
                              f"Device mismatch: p={p.device} "  # type: ignore[attr-defined]
                              f"p._saved_grad_shard={p._saved_grad_shard.device}"
                          )
                          p.grad = p._saved_grad_shard  # type: ignore[attr-defined]
                      else:
                          p_assert(
                              not p._is_sharded, "All sharded parameters should "
                              "use `_saved_grad_shard`"
                          )
                      if hasattr(p, "_saved_grad_shard"):
                          delattr(p, "_saved_grad_shard")
  
          # Update root and nested FSDP's hooks and flags.
          for m in self.modules():  # includes self
              if isinstance(m, FullyShardedDataParallel):
                  _finalize_params(m)
                  m._pre_backward_hook_has_run = False
                  if any(p.requires_grad for p in m.parameters()):
                      # Check if the module has params and if any of them has
                      # the `requires_grad` field set. If `requires_grad=False` for
                      # all the params, the post_backward hook will not fire and the
                      # state will remain in `TrainingState_.BACKWARD_PRE`.
                      if any([p.requires_grad for p in m.params]):
                          m._assert_state(TrainingState_.BACKWARD_POST)
                      else:
                          m._assert_state(TrainingState_.BACKWARD_PRE)
                  else:
                      # When `m` and its children have no non-ignored params or
                      # have non-ignored params but none with `requires_grad==True`,
                      # there are two cases:
                      # 1. output tensors are `requires_grad==True`. In this case,
                      # pre-backward hook is still registered, so it is in BACKWARD_PRE state.
                      # 2. output tensors are `requires_grad==False`. In this case,
                      # pre-backward hook is not registered, so it is in IDLE state.
                      m._assert_state([TrainingState_.BACKWARD_PRE, TrainingState_.IDLE])
                  m.training_state = TrainingState_.IDLE
  
                  if m._is_root:
                      # reset this flag for cases like "one forward pass + multiple backward passes"
                      self._post_backward_callback_queued = False
  
      def _update_p_data(self, p, output_tensor: torch.Tensor) -> None:
          """
          Helper function to update p.data pointer.
          Args:
              output_tensor (torch.Tensor): this tensor contains the data we just gathered.
          """
          p.data = output_tensor
          # Trim any padding and reshape to match original size.
          p.data = p.data[: p._orig_size.numel()].view(p._orig_size)  # type: ignore[attr-defined]
  
      @torch.no_grad()
      def _rebuild_full_params(self) -> List[Tuple[torch.Tensor, bool]]:
          """
          Gather all shards of params.
          """
          # _summon_full_params must do a full precision rebuild even under mixed
          # precision, because it is used for e.g. checkpoint where we'd like to
          # checkpoint in full precision.
          force_full_precision = (self.training_state == TrainingState_.SUMMON_FULL_PARAMS)
          # full param output tensors and a flag indicating whether
          # _summon_full_params can free them or not. It is possible that we can't
          # free the full param, which currently occurs when the returned
          # parameter points to the unsharded param when world_size == 1, or when
          # we're returning the full parameter and reshard_after_forward=False
          # (because we need to ensure p._full_param_padded stays intact)
          output_tensors: List[Tuple[torch.Tensor, bool]] = []
          with torch.cuda.stream(self._streams["all_gather"]):
              for p in self.params:
                  mixed_precision_cast_ran = (
                      self._mixed_precision_enabled_for_params()
                      and not force_full_precision
                  )
                  if mixed_precision_cast_ran:
                      self._cast_param_shards_to_dtype()
                      # TODO: remove below
                      for p in self.params:
                          assert p.dtype == self.mixed_precision.param_dtype
                  # We can skip moving params to GPU if mixed precision, as p.data
                  # would then be pointing to p._mp_shard which is already on
                  # self.compute_device.
                  if self.cpu_offload.offload_params and not mixed_precision_cast_ran:
                      # Move params to GPU if needed. Note that we don't use
                      # self._full_param_padded.device here because the attr is
                      # not set always, i.e. when world_size=1 and
                      # p._is_sharded = False. However when it is set, the
                      # device is always self.compute_device.
                      p.data = p.data.to(self.compute_device, non_blocking=True)
                  # Check the validity of this `_rebuild_full_params()` call in
                  # terms of execution order (regardless of if FSDP actually
                  # needs to all-gather or not)
                  self._check_rebuild_full_params(p)
                  # e.g., when world_size == 1
                  if not p._is_sharded:  # type: ignore[attr-defined]
                      if mixed_precision_cast_ran:
                          # p.data should be the same type as p._mp_shard, and it
                          # is safe to free.
                          assert p.data.dtype == p._mp_shard.dtype
                          # Safe to free because p.data points to the mp shard.
                          output_tensors.append((p.data, True))
                      else:
                          # p.data points to the unsharded parameter, so not safe to
                          # free.
                          output_tensors.append((p.data, False))
                      continue
                  # If full param has been rebuilt or has not been freed, no need to call all gather
                  elif (
                      p._full_param_padded.storage().size()  # type: ignore[attr-defined]
                      == p._full_param_padded.size().numel()  # type: ignore[attr-defined]
                  ):
                      # Check that the full param is in the expected precision, if
                      # training with mixed precision
                      if mixed_precision_cast_ran:
                          if p._full_param_padded.dtype != self.mixed_precision.param_dtype:
                              raise ValueError(
                                  "_rebuild_full_params: Expected full param to be "
                                  f"of type {self.mixed_precision.param_dtype}, "
                                  f"but got {p._full_param_padded.dtype}!"
                              )
                      # output is full_param_padded which can be freed depending
                      # on reshard_after_forward (this path is exercised by tests
                      # in test_fsdp_summon_full_params).
                      output_tensors.append((p._full_param_padded, self.reshard_after_forward))
  
                      self._update_p_data(p, output_tensor=p._full_param_padded)  # type: ignore[attr-defined]
                      continue
                  else:
                      # If full param has not been rebuilt or has been freed, call all gather
                      p_data = p.data  # type: ignore[attr-defined]
                      p_full_size = p._full_param_padded.size()  # type: ignore[attr-defined]
                      assert (
                          p_full_size.numel() == p_data.numel() * self.world_size
                      ), "Param full size should be equal to its shard size multiply world_size."
                      assert (
                          p._full_param_padded.storage().size() == 0  # type: ignore[attr-defined]
                      ), "Full param's storage should have been freed before if all gather is needed."  # type: ignore[attr-defined]
                      if (
                          self._mixed_precision_enabled_for_params()
                          and force_full_precision
                      ):
                          # p._full_param_padded has the reduced precision type,
                          # but we need full precision rebuild as we're in
                          # _summon_full_params. Note that this is why
                          # _summon_full_params collects locally used params from
                          # _rebuild_full_params instead of relying on
                          # p._full_param_padded, as it may not always be
                          # allocated such as during mixed precision.
                          output_tensor = p_data.new_zeros(p_full_size)
                      else:
                          # Allocate based on full size from all shards.
                          _alloc_storage(p._full_param_padded, size=p_full_size)  # type: ignore[attr-defined]
                          output_tensor = p._full_param_padded  # type: ignore[attr-defined]
                      # Fill output_tensor with (p.data for each shard in self.world_size)
                      dist._all_gather_base(
                          output_tensor, p_data, group=self.process_group
                      )
  
                      # The full parameter, which can be freed. Note that we
                      # append here before update_p_data so as to not saved the
                      # tensor with padding trimmed, which causes issues with
                      # writeback in _summon_full_params.
                      output_tensors.append((output_tensor, True))
                      # Set p.data = output_tensor (with padding trimmed)
                      self._update_p_data(p, output_tensor=output_tensor)
                      # We can free the reduced precision shard as we have the
                      # full precision parameter.
                      if (
                          self._mixed_precision_enabled_for_params()
                      ):
                          self._free_mp_shard(cast(List[FlatParameter], [p]))
          return output_tensors
  
      def _check_rebuild_full_params(self, param: FlatParameter):
          """
          Checks the validity of a call to :meth:`_rebuild_full_params` in terms
          of the execution order. If on the first iteration, this uses an
          all-gather to check that all ranks are running ``forward()`` with the
          same parameter, erroring if not, and on subsequent iterations, if the
          forward order differs from that of the first iteration (meaning that we
          can no longer guarantee correct execution since all-gathers may be
          mismatched), then we issue a warning to the user. This only issues
          warnings on the first deviating iteration and stops checking
          thereafter.
  
          Only the :meth:`_rebuild_full_params` calls in the forward pass are
          checked since a correct forward order should imply a correct
          pre-backward order for typical cases.
  
          Executing in ``no_sync()`` does not affect this check for
          ``FULL_SHARD`` and ``SHARD_GRAD_OP``: (1) Being in ``no_sync()`` in the
          first iteration does not yield a different forward
          :meth:`_rebuild_full_params()` sequence, and (2) being in ``no_sync()``
          in a later iteration does not give false positive warnings since the
          forward :meth:`_rebuild_full_params()` sequence still matches the first
          iteration sequence (for ``FULL_SHARD``) or the first iteration
          sequence's prefix (for ``SHARD_GRAD_OP``).
          """
          # Only check when rebuilding the full parameters in the forward pass,
          # and skip the check (1) when in eval mode since then there is not a
          # safe point at which to reset the execution order data and (2) if
          # world size is 1 since then there is no chance of desynchronization
          if self.training_state != TrainingState_.FORWARD or \
                  not self.training or self.world_size == 1:
              return
          eod = self._exec_order_data
          param_index = eod.get_param_index(param)
          if not eod.is_first_iter:
              # Only issue warnings on the first deviating iteration and stop
              # checking thereafter to avoid flooding the console
              if eod.warn_status == _ExecOrderWarnStatus.WARNED:
                  return
              # However, we may issue multiple warnings on the first deviating
              # iteration to help debugging, where either:
              # 1. This iteration sees an extra `_rebuild_full_params()` in
              # `forward()` compared to the first iteration
              msg_prefix = curr_param_order = None  # non-`None` means we warn
              if eod.index >= len(eod.param_order):
                  msg_prefix = "Expected to not rebuild any more parameters " \
                      "in `forward()` for this module but trying to rebuild " \
                      "parameters for "
                  curr_param_order = eod.param_order + [param_index]
              else:
                  expected_param_index = eod.param_order[eod.index]
                  # 2. This iteration sees the same number of
                  # `_rebuild_full_params()` (so far) but the current parameter
                  # differs
                  if param_index != expected_param_index:
                      expected_param_names = eod.get_unflat_param_names(expected_param_index)
                      assert len(expected_param_names) > 0, \
                          "Expected parameter should always be valid"
                      msg_prefix = "Expected to rebuild parameters in " \
                          f"`forward()` for {expected_param_names} but " \
                          "instead trying to rebuild parameters for "
                      curr_param_order = eod.param_order[:eod.index - 1] + [param_index]
              to_issue_warning = msg_prefix is not None
              if to_issue_warning:
                  assert curr_param_order is not None
                  param_names = eod.get_unflat_param_names(param_index)
                  is_added_param = len(param_names) == 0
                  if is_added_param:
                      msg_suffix = "a newly-added parameter since construction time"
                  else:
                      msg_suffix = f"{param_names}"
                  sub_msg = msg_prefix + msg_suffix
                  first_iter_param_names = [
                      eod.get_unflat_param_names(index) for index in eod.param_order
                  ]
                  curr_iter_param_names = [
                      eod.get_unflat_param_names(index) for index in curr_param_order
                  ]
                  warnings.warn(
                      "Forward order differs from that of the first iteration "
                      f"on rank {self.rank} -- collectives are unchecked and may "
                      "give incorrect results or hang\n" + sub_msg + "\n" +
                      f"First iteration's forward order: {first_iter_param_names}"
                      "\nThis iteration's forward order (so far): "
                      f"{curr_iter_param_names}"
                  )
                  eod.warn_status = _ExecOrderWarnStatus.WARNING
              eod.index += 1
          else:
              # Use `compute_device` instead of the parameter's device in case it
              # is offloaded on CPU and we are using NCCL backend, which requires
              # communicated tensors be on GPU
              device = self.compute_device
              indices = torch.zeros(self.world_size, dtype=torch.int32, device=device)
              index = torch.tensor([param_index], dtype=torch.int32, device=device)
              dist._all_gather_base(indices, index, group=self.process_group)
              # Check that all ranks plan to all-gather the same parameter index
              for (r1, i1), (r2, i2) in itertools.combinations(
                  ((rank, indices[rank]) for rank in range(self.world_size)), 2,
              ):
                  if not torch.equal(i1, i2):
                      r1_param_names = eod.get_unflat_param_names(i1)
                      r2_param_names = eod.get_unflat_param_names(i2)
                      raise RuntimeError(
                          f"Forward order differs across ranks: rank {r1} is "
                          "rebuilding full parameters in `forward()` for "
                          f"{r1_param_names} while rank {r2} is rebuilding full "
                          f"parameters in `forward()` for {r2_param_names}"
                      )
              eod.param_order.append(param_index)
  
      @torch.no_grad()
      def _prep_grads_for_backward(self) -> None:
          """Make sure p.grad has the correct size/device, otherwise set it to None."""
          for p in self.params:
              if p.grad is not None and (
                  p.grad.size() != p._orig_size  # type: ignore[attr-defined]
                  or p.grad.device != p.device
              ):
                  offloaded: bool = p.grad.device != p.device
                  if offloaded:
                      assert self.cpu_offload.offload_params, \
                          "`p.grad.device` and `p.device` should be the same " \
                          "if not offloading parameters to CPU"
                  prev_iter_outside_no_sync: bool = \
                      p.grad.size() == p._local_shard.shape  # type: ignore[attr-defined]
                  # As long as the previous iteration was outside `no_sync()`,
                  # then we must save the gradient in `_saved_grad_shard`, even
                  # if the current iteration is inside `no_sync()`. This is to
                  # prepare for the next iteration outside `no_sync()`, which may
                  # try to accumulate gradients. FSDP accumulates gradients in
                  # the separate variable `p._saved_grad_shard` to leave `p.grad`
                  # for the per-iteration gradient.
                  if prev_iter_outside_no_sync:
                      # FSDP currently does not support gradient accumulation
                      # outside `no_sync()` when using CPU offloading (see the
                      # warning in the class's docstring).
                      if not offloaded:
                          p._saved_grad_shard = p.grad.data  # type: ignore[attr-defined]
                  p.grad = None
  
      @torch.no_grad()
      def _free_full_params(self, params: Optional[List[FlatParameter]] = None) -> None:
          """
          Free up storage for full parameters.
          """
          if params is None:
              params = self.params
          current_stream = torch.cuda.current_stream()
          for p in params:
              # e.g., world_size == 1 or self.sharding_strategy = NO_SHARD
              if not p._is_sharded:  # type: ignore[attr-defined]
                  if (
                      self._mixed_precision_enabled_for_params()
                  ):
                      self._free_mp_shard(cast(List[FlatParameter], [p]))
                  continue
              # Don't let PyTorch reuse this memory until all work in the current
              # stream is complete.
              p._full_param_padded.record_stream(current_stream)  # type: ignore[attr-defined]
              # There may be external references to the Tensor Storage that we
              # can't modify, such as references that are created by
              # ctx.save_for_backward in the forward pass. Thus when we
              # unshard parameters, we should reuse the original Tensor
              # Storage object and unshard it in-place. For now, just resize
              # the Storage to 0 to save memory.
              _free_storage(p._full_param_padded)  # type: ignore[attr-defined]
  
      @torch.no_grad()
      def _use_param_local_shard(
          self, params: Optional[List[FlatParameter]] = None
      ) -> None:
          """Use local shard for a list of params. Also implicitly offloads
          parameters back to CPU if we are CPU offloading."""
          if params is None:
              params = self.params
          for p in params:
              if self.cpu_offload.offload_params:
                  # Ensure local_shard resides in CPU if we are offloading params.
                  assert p._local_shard.device == torch.device(  # type: ignore[attr-defined]
                      "cpu"
                  ), "Expected p._local_shard to be on CPU"
              p.data = p._local_shard  # type: ignore[attr-defined]
  
      def _assert_state(self, state: Union[TrainingState_, List[TrainingState_]]) -> None:
          """Assert we are in the given state."""
          # Since assert can be turned off and this error checking
          # is really important, we use explicit error checking
          # and raise a ValueError if needed.
          if isinstance(state, TrainingState_):
              state = [state]
          if self.training_state not in state:
              msg = (
                  f"expected to be in states {state} but current state "
                  f"is {self.training_state}"
              )
              # In case we are failing in the context of autograd hook, asserting
              # may not generate useful msg. So, let's print it to be sure.
              if self.rank == 0:
                  print(f"Asserting FSDP instance is: {self}")
                  print(f"ERROR: {msg}")
                  traceback.print_stack()
              raise ValueError(msg)
  
      @contextmanager
      def no_sync(self) -> Generator:
          """
          A context manager to disable gradient synchronizations across FSDP
          instances. Within this context, gradients will be accumulated in module
          variables, which will later be synchronized in the first
          forward-backward pass after exiting the context. This should only be
          used on the root FSDP instance and will recursively apply to all
          children FSDP instances.
  
          .. note:: This likely results in higher memory usage because FSDP will
              accumulate the full model gradients (instead of gradient shards)
              until the eventual sync.
  
          .. note:: When used with CPU offloading, the gradients will not be
              offloaded to CPU when inside the context manager. Instead, they
              will only be offloaded right after the eventual sync.
          """
          self._lazy_init()
          assert self._is_root, "`no_sync()` on inner FSDP instances is not supported"
          self._assert_state(TrainingState_.IDLE)
          old_flags = []
          for m in self.modules():
              if isinstance(m, FullyShardedDataParallel):
                  old_flags.append((m, m._require_backward_grad_sync))
                  m._require_backward_grad_sync = False
          try:
              yield
          finally:
              for m, old_flag in old_flags:
                  assert not m._require_backward_grad_sync, (
                      "`_require_backward_grad_sync` was incorrectly set to "
                      "`True` while in the `no_sync()` context manager"
                  )
                  m._require_backward_grad_sync = old_flag
  
      @property
      def params_with_grad(self) -> List[Parameter]:
          """
          Recursively returns a list of all module parameters that have a gradient.
          """
          return [p for p in self.parameters() if p.grad is not None]
  
      @torch.no_grad()
      def clip_grad_norm_(
          self, max_norm: Union[float, int], norm_type: Union[float, int] = 2.0
      ) -> None:
          """
          Clip all gradients at this point in time. The norm is computed over all
          gradients together, as if they were concatenated into a single vector.
          Gradients are modified in-place.
  
          Args:
              max_norm (float or int): max norm of the gradients
              norm_type (float or int): type of the used p-norm. Can be ``'inf'``
                  for infinity norm.
  
          Returns:
              Total norm of the parameters (viewed as a single vector).
  
          .. note:: This is analogous to ``torch.nn.utils.clip_grad_norm_`` but
              handles the partitioning and multiple devices per rank under the
              hood. The default torch util is not applicable here, because each
              rank only has a partial view of all the grads in the model, so
              calling it for FSDP models would lead to different scaling being
              applied per subset of model parameters.
  
          .. warning:: This needs to be called on all ranks, since synchronization
              primitives will be used.
          """
          # Call `_lazy_init` to ensure the stream synchronization is done appropriately.
          self._lazy_init()
          assert self._is_root, "clip_grad_norm should only be called on the root (parent) instance"
          self._assert_state(TrainingState_.IDLE)
  
          max_norm = float(max_norm)
          norm_type = float(norm_type)
          # Computes the max norm for this shard's gradients and sync's across workers
          local_norm = _calc_grad_norm(self.params_with_grad, norm_type).cuda()  # type: ignore[arg-type]
          if norm_type == math.inf:
              total_norm = local_norm
              dist.all_reduce(total_norm, op=torch.distributed.ReduceOp.MAX, group=self.process_group)
          else:
              total_norm = local_norm ** norm_type
              dist.all_reduce(total_norm, group=self.process_group)
              total_norm = total_norm ** (1.0 / norm_type)
  
          if self.cpu_offload:
              total_norm = total_norm.cpu()
  
          clip_coef = torch.tensor(max_norm, dtype=total_norm.dtype, device=total_norm.device) / (total_norm + 1e-6)
          if clip_coef < 1:
              # multiply by clip_coef, aka, (max_norm/total_norm).
              for p in self.params_with_grad:
                  assert p.grad is not None
                  p.grad.detach().mul_(clip_coef.to(p.grad.device))
  
      @staticmethod
      def full_optim_state_dict(
          model: torch.nn.Module,
          optim: torch.optim.Optimizer,
          optim_input: Optional[Union[
              List[Dict[str, Any]], Iterable[torch.nn.Parameter],
          ]] = None,
          rank0_only: bool = True,
      ) -> Dict[str, Any]:
          """
          Consolidates the full optimizer state on rank 0 and returns it
          as a :class:`dict` following the convention of
          :meth:`torch.optim.Optimizer.state_dict`, i.e. with keys ``"state"``
          and ``"param_groups"``. The flattened parameters in ``FSDP`` modules
          contained in ``model`` are mapped back to their unflattened parameters.
  
          .. warning:: This needs to be called on all ranks since synchronization
              primitives are used. However, if ``rank0_only=True``, then the
              state dict is only populated on rank 0, and all other ranks return
              an empty :class:`dict`.
  
          .. warning:: Unlike ``torch.optim.Optimizer.state_dict()``, this method
              uses full parameter names as keys instead of parameter IDs.
  
          .. warning:: If you do not pass ``model.parameters()`` as the first
              argument to the optimizer, then you should pass that same value to
              this method as ``optim_input``.
  
          .. note:: Like in :meth:`torch.optim.Optimizer.state_dict`, the tensors
              contained in the optimizer state dict are not cloned, so there may
              be aliasing surprises. For best practices, consider saving the
              returned optimizer state dict immediately, e.g. using
              ``torch.save()``.
  
          Args:
              model (torch.nn.Module): Root module (which may or may not be a
                  :class:`FullyShardedDataParallel` instance) whose parameters
                  were passed into the optimizer ``optim``.
              optim (torch.optim.Optimizer): Optimizer for ``model`` 's
                  parameters.
              optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]):
                  Input passed into the optimizer ``optim`` representing either a
                  :class:`list` of parameter groups or an iterable of parameters;
                  if ``None``, then this method assumes the input was
                  ``model.parameters()``. (Default: ``None``)
              rank0_only (bool): If ``True``, saves the populated :class:`dict`
                  only on rank 0; if ``False``, saves it on all ranks. (Default:
                  ``True``)
  
          Returns:
              Dict[str, Any]: A :class:`dict` containing the optimizer state for
              ``model`` 's original unflattened parameters and including keys
              "state" and "param_groups" following the convention of
              :meth:`torch.optim.Optimizer.state_dict`. If ``rank0_only=False``,
              then nonzero ranks return an empty :class:`dict`.
          """
          osd = optim.state_dict()
          osd_state, osd_param_groups = osd["state"], osd["param_groups"]  # alias
  
          group = model.process_group if hasattr(model, "process_group") \
              else None  # not all `torch.nn.Module`s have `process_group`
          rank = dist.get_rank(group)
          to_save = not rank0_only or rank == 0
          full_osd: Dict = {"state": {}, "param_groups": []} if to_save else {}
          full_osd_state = full_osd["state"] if to_save else None  # alias
  
          # Handle the "state" part of the optimizer state dict
          param_to_unflat_param_names = _get_param_to_unflat_param_names(model)
          flat_param_id_to_param = _get_param_id_to_param(model, optim_input)
          flat_param_to_fsdp_module = _get_flat_param_to_fsdp_module(model)
          for flat_param_id, param in enumerate(flat_param_id_to_param):  # type: ignore[assignment]
              # Do not include parameters without state to avoid empty mappings
              if flat_param_id not in osd_state:
                  continue
              assert param in param_to_unflat_param_names, \
                  "Check the `param_to_unflat_params` construction\n" \
                  f"param: {param}"
              unflat_param_names = param_to_unflat_param_names[param]
              # For FSDP parameters, we need to unflatten
              if isinstance(param, FlatParameter):
                  assert param in flat_param_to_fsdp_module, \
                      "Check the `flat_param_to_fsdp_module` construction\n" \
                      f"param: {param}"
                  unflat_state = _unflatten_optim_state(
                      flat_param_to_fsdp_module[param], param,
                      osd_state[flat_param_id], to_save,
                  )
                  if to_save:
                      assert len(unflat_state) == len(unflat_param_names) and \
                          len(unflat_state) == param._num_unflattened_params, \
                          f"{len(unflat_state)} {len(unflat_param_names)} " \
                          f"{param._num_unflattened_params}"
                      for unflat_param_name, unflat_param_state in zip(
                          unflat_param_names, unflat_state,
                      ):
                          full_osd_state[unflat_param_name] = unflat_param_state
              # For parameters from non-FSDP modules, we do not need to unflatten
              elif to_save:
                  assert len(unflat_param_names) == 1
                  unflat_param_name = unflat_param_names[0]
                  # Do not `deepcopy()` to avoid unnecessarily duplicating
                  # tensor storage
                  full_osd_state[unflat_param_name] = \
                      copy.copy(osd_state[flat_param_id])
                  # Move all tensor state to CPU
                  param_state = full_osd_state[unflat_param_name]
                  for state_name, value in param_state.items():
                      if torch.is_tensor(value):
                          param_state[state_name] = value.cpu()
  
          # Non-target ranks may return since there is no more communication
          if not to_save:
              return full_osd
  
          # Handle the "param_groups" part of the optimizer state dict
          full_osd_param_groups = full_osd["param_groups"]  # alias
          for flat_param_group in osd_param_groups:
              unflat_param_group = copy.deepcopy(flat_param_group)
              param_group_params = [
                  flat_param_id_to_param[flat_param_id]
                  for flat_param_id in flat_param_group["params"]
              ]
              nested_unflat_param_names = [
                  param_to_unflat_param_names[param]
                  for param in param_group_params
              ]
              unflat_param_group["params"] = [
                  unflat_param_name
                  for unflat_param_names in nested_unflat_param_names
                  for unflat_param_name in unflat_param_names
              ]  # flatten the list of lists
              full_osd_param_groups.append(unflat_param_group)
          return full_osd
  
      @staticmethod
      def shard_full_optim_state_dict(
          full_optim_state_dict: Dict[str, Any],
          model: torch.nn.Module,
          optim_input: Optional[Union[
              List[Dict[str, Any]], Iterable[torch.nn.Parameter],
          ]] = None,
      ) -> Dict[str, Any]:
          """
          Shards the full optimizer state dict ``full_optim_state_dict`` by
          remapping the state to flattened parameters instead of unflattened
          parameters and restricting to only this rank's part of the optimizer
          state. The first argument should be the return value of
          :meth:`full_optim_state_dict`.
  
          Example::
  
              >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
              >>> model, optim = ...
              >>> full_osd = FSDP.full_optim_state_dict(model, optim)
              >>> torch.save(full_osd, PATH)
              >>> # Define new model with possibly different world size
              >>> new_model, new_optim = ...
              >>> full_osd = torch.load(PATH)
              >>> sharded_osd = FSDP.shard_full_optim_state_dict(full_osd, new_model)
              >>> new_optim.load_state_dict(sharded_osd)
  
          .. warning:: If you do not pass ``model.parameters()`` as the first
              argument to the optimizer, then you should pass that same value to
              this method as ``optim_input``.
  
          .. note:: Both :meth:`shard_full_optim_state_dict` and
              :meth:`scatter_full_optim_state_dict` may be used to get the
              sharded optimizer state dict to load. Assuming that the full
              optimizer state dict resides in CPU memory, the former requires
              each rank to have the full dict in CPU memory, where each rank
              individually shards the dict without any communication, while the
              latter requires only rank 0 to have the full dict in CPU memory,
              where rank 0 moves each shard to GPU memory (for NCCL) and
              communicates it to ranks appropriately. Hence, the former has
              higher aggregate CPU memory cost, while the latter has higher
              communication cost.
  
          Args:
              full_optim_state_dict (Dict[str, Any]): Optimizer state dict
                  corresponding to the unflattened parameters and holding the
                  full non-sharded optimizer state.
              model (torch.nn.Module): Root module (which may or may not be a
                  :class:`FullyShardedDataParallel` instance) whose parameters
                  correspond to the optimizer state in ``full_optim_state_dict``.
              optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]):
                  Input passed into the optimizer representing either a
                  :class:`list` of parameter groups or an iterable of parameters;
                  if ``None``, then this method assumes the input was
                  ``model.parameters()``. (Default: ``None``)
  
          Returns:
              Dict[str, Any]: The full optimizer state dict now remapped to
              flattened parameters instead of unflattened parameters and
              restricted to only include this rank's part of the optimizer state.
          """
          return _flatten_full_optim_state_dict(
              full_optim_state_dict, model, True, optim_input,
          )[0]
  
      @staticmethod
      def scatter_full_optim_state_dict(
          full_optim_state_dict: Optional[Dict[str, Any]],
          model: torch.nn.Module,
          optim_input: Optional[Union[
              List[Dict[str, Any]], Iterable[torch.nn.Parameter],
          ]] = None,
          group: Optional[Any] = None,
      ) -> Dict[str, Any]:
          """
          Scatters the full optimizer state dict from rank 0 to all other ranks,
          returning the sharded optimizer state dict on each rank. The return
          value is the same as :meth:`shard_full_optim_state_dict`, and on rank
          0, the first argument should be the return value of
          :meth:`full_optim_state_dict`.
  
          Example::
  
              >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
              >>> model, optim = ...
              >>> full_osd = FSDP.full_optim_state_dict(model, optim)  # only non-empty on rank 0
              >>> # Define new model with possibly different world size
              >>> new_model, new_optim, new_group = ...
              >>> sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, new_model, group=new_group)
              >>> new_optim.load_state_dict(sharded_osd)
  
          .. note:: Both :meth:`shard_full_optim_state_dict` and
              :meth:`scatter_full_optim_state_dict` may be used to get the
              sharded optimizer state dict to load. Assuming that the full
              optimizer state dict resides in CPU memory, the former requires
              each rank to have the full dict in CPU memory, where each rank
              individually shards the dict without any communication, while the
              latter requires only rank 0 to have the full dict in CPU memory,
              where rank 0 moves each shard to GPU memory (for NCCL) and
              communicates it to ranks appropriately. Hence, the former has
              higher aggregate CPU memory cost, while the latter has higher
              communication cost.
  
          Args:
              full_optim_state_dict (Optional[Dict[str, Any]]): Optimizer state
                  dict corresponding to the unflattened parameters and holding
                  the full non-sharded optimizer state if on rank 0; the argument
                  is ignored on nonzero ranks.
              model (torch.nn.Module): Root module (which may or may not be a
                  :class:`FullyShardedDataParallel` instance) whose parameters
                  correspond to the optimizer state in ``full_optim_state_dict``.
              optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]):
                  Input passed into the optimizer representing either a
                  :class:`list` of parameter groups or an iterable of parameters;
                  if ``None``, then this method assumes the input was
                  ``model.parameters()``; the argument is ignored on nonzero
                  ranks. (Default: ``None``)
              group (Optional[Any]): Model's process group or ``None`` if using
                  the default process group. (Default: ``None``)
  
          Returns:
              Dict[str, Any]: The full optimizer state dict now remapped to
              flattened parameters instead of unflattened parameters and
              restricted to only include this rank's part of the optimizer state.
          """
          # Try to use the passed-in process group, the model's process group,
          # or the default process group (i.e. ``None``) in that priority order
          if group is None and hasattr(model, "process_group"):
              group = model.process_group
          rank = dist.get_rank(group)
          world_size = dist.get_world_size(group)
          # Check for a valid broadcast device, preferring GPU when available
          using_nccl = dist.distributed_c10d._check_for_nccl_backend(group)
          broadcast_device = torch.device("cuda") if torch.cuda.is_available() \
              else torch.device("cpu")
          if using_nccl and not torch.cuda.is_available():
              raise RuntimeError("NCCL requires a GPU for collectives")
          # Flatten the optimizer state dict and construct a copy with the
          # positive-dimension tensors' shapes in place of the tensors themselves
          # since those tensors will be broadcast separately to avoid copying
          if rank == 0:
              if full_optim_state_dict is None:
                  raise ValueError("Rank 0 must pass in the full optimizer state dict")
              flat_osd, fsdp_flat_param_ids = _flatten_full_optim_state_dict(
                  full_optim_state_dict, model, False, optim_input,
              )
              processed_osd = _process_pos_dim_tensor_state(
                  flat_osd, fsdp_flat_param_ids, world_size,
              )
          # Broadcast the optim state dict without positive-dimension tensor
          # state and the FSDP parameter IDs from rank 0 to all ranks
          processed_osd, fsdp_flat_param_ids = \
              _broadcast_processed_optim_state_dict(
                  processed_osd if rank == 0 else None,
                  fsdp_flat_param_ids if rank == 0 else None, rank, group,
                  broadcast_device,
              )
          # Broadcast positive-dimension tensor state (both sharded tensors for
          # FSDP parameters and unsharded tensors for non-FSDP parameters)
          sharded_osd = _broadcast_pos_dim_tensor_states(
              processed_osd, fsdp_flat_param_ids,
              flat_osd if rank == 0 else None, rank, world_size, group,
              broadcast_device,
          )
          return sharded_osd
  
      @staticmethod
      def rekey_optim_state_dict(
          optim_state_dict: Dict[str, Any],
          optim_state_key_type: OptimStateKeyType,
          model: torch.nn.Module,
          optim_input: Optional[Union[
              List[Dict[str, Any]], Iterable[torch.nn.Parameter],
          ]] = None,
      ) -> Dict[str, Any]:
          """
          Re-keys the optimizer state dict ``optim_state_dict`` to use the key
          type ``optim_state_key_type``. This can be used to achieve
          compatibility between optimizer state dicts from models with FSDP
          instances and ones without.
  
          To re-key an FSDP full optimizer state dict (i.e. from
          :meth:`full_optim_state_dict`) to use parameter IDs and be loadable to
          a non-wrapped model::
  
              >>> wrapped_model, wrapped_optim = ...
              >>> full_osd = FSDP.full_optim_state_dict(wrapped_model, wrapped_optim)
              >>> nonwrapped_model, nonwrapped_optim = ...
              >>> rekeyed_osd = FSDP.rekey_optim_state_dict(full_osd, OptimStateKeyType.PARAM_ID, nonwrapped_model)
              >>> nonwrapped_optim.load_state_dict(rekeyed_osd)
  
          To re-key a normal optimizer state dict from a non-wrapped model to be
          loadable to a wrapped model::
  
              >>> nonwrapped_model, nonwrapped_optim = ...
              >>> osd = nonwrapped_optim.state_dict()
              >>> rekeyed_osd = FSDP.rekey_optim_state_dict(osd, OptimStateKeyType.PARAM_NAME, nonwrapped_model)
              >>> wrapped_model, wrapped_optim = ...
              >>> sharded_osd = FSDP.shard_full_optim_state_dict(rekeyed_osd, wrapped_model)
              >>> wrapped_optim.load_state_dict(sharded_osd)
  
          Returns:
              Dict[str, Any]: The optimizer state dict re-keyed using the
              parameter keys specified by ``optim_state_key_type``.
          """
          assert optim_state_key_type in \
              (OptimStateKeyType.PARAM_NAME, OptimStateKeyType.PARAM_ID)
          osd = optim_state_dict  # alias
          # Validate that the existing parameter keys are uniformly typed
          uses_param_name_mask = [
              type(param_key) is str for param_key in osd["state"]
          ]
          uses_param_id_mask = [
              type(param_key) is int for param_key in osd["state"]
          ]
          if (any(uses_param_name_mask) and not all(uses_param_name_mask)) or \
                  (any(uses_param_id_mask) and not all(uses_param_id_mask)):
              error_msg = f"Invalid parameter keys: {osd['state'].keys()}"
              raise ValueError(error_msg)
          # Return directly if the existing key type matches the target key type
          if (optim_state_key_type == OptimStateKeyType.PARAM_NAME and
              all(uses_param_name_mask)) or \
              (optim_state_key_type == OptimStateKeyType.PARAM_ID and
                  all(uses_param_id_mask)):
              return osd
          # Otherwise, actually perform the re-keying
          new_osd = {}
          if optim_state_key_type == OptimStateKeyType.PARAM_NAME:  # ID -> name
              param_id_to_param = _get_param_id_to_param(model, optim_input)
              param_to_param_name = _get_param_to_param_name(model)
              param_id_to_param_name: List[str] = [
                  param_to_param_name[param] for param in param_id_to_param
              ]
              new_osd["state"] = {
                  param_id_to_param_name[param_id]: param_state
                  for param_id, param_state in osd["state"].items()
              }
              new_osd["param_groups"] = copy.deepcopy(osd["param_groups"])
              for param_group in new_osd["param_groups"]:
                  param_group["params"] = sorted([
                      param_id_to_param_name[param_id]
                      for param_id in param_group["params"]
                  ])
              return new_osd
          elif optim_state_key_type == OptimStateKeyType.PARAM_ID:  # name -> ID
              param_name_to_param = _get_param_name_to_param(model)
              param_to_param_id = _get_param_to_param_id(model, optim_input)
              # Because not all model parameters may be passed as the optimizer
              # input, we may need to drop some parameters from this mapping
              param_name_to_param_id = {
                  param_name: param_to_param_id[param]
                  for param_name, param in param_name_to_param.items()
                  if param in param_to_param_id
              }
              new_osd["state"] = {
                  param_name_to_param_id[param_name]: param_state
                  for param_name, param_state in osd["state"].items()
              }
              new_osd["param_groups"] = copy.deepcopy(osd["param_groups"])
              for param_group in new_osd["param_groups"]:
                  param_group["params"] = sorted([
                      param_name_to_param_id[param_name]
                      for param_name in param_group["params"]
                  ])
              return new_osd
          return new_osd  # should never reach here
  
  
  def _get_default_cuda_device(module: nn.Module) -> torch.device:
      """Try to infer CUDA device from module parameters."""
      try:
          compute_device = next(module.parameters()).device
          if compute_device.type == "cuda":
              return compute_device
      # e.g., if module does not have parameters, it will throw StopIteration,
      # in this case, instead of raising exception, return cuda device.
      except StopIteration:
          pass
      # Fall back to current CUDA device
      return torch.device("cuda", torch.cuda.current_device())
  
  
  def _free_storage(data: torch.Tensor) -> None:
      """Free underlying storage of a Tensor."""
      if data.storage().size() > 0:
          # Since we're modifying the Tensor's Storage directly, make sure the Tensor
          # is the sole occupant of the Storage.
          assert (
              data.storage_offset() == 0
          ), "The tensor is not the sole occupant of the storage."
          data.storage().resize_(0)  # type: ignore[attr-defined]
  
  
  @torch.no_grad()
  def _alloc_storage(data: torch.Tensor, size: torch.Size) -> None:
      """Allocate storage for a tensor."""
      if data.storage().size() == size.numel():  # no need to reallocate
          return
      assert (
          data.storage().size() == 0
      ), "Then tensor storage should have been resized to be 0."
      data.storage().resize_(size.numel())  # type: ignore[attr-defined]
  
  def p_assert(cond: Any, s: Any) -> None:
      """This is used as an alternate to ``assert`` when in the backward context
      to print the error message ``s`` since otherwise, it is swallowed."""
      if not cond:
          print(s)
          raise AssertionError
  
  def _calc_grad_norm(parameters: List[torch.nn.Parameter], p: float) -> torch.Tensor:
      r"""Calculate gradient norm of an iterable of parameters.
      Returns:
          Total norm of the parameters (viewed as a single vector).
      """
      parameters = [p for p in parameters if p.grad is not None]
  
      if len(parameters) == 0:
          return torch.tensor(0.0)
      if p == math.inf:
          local_norm = torch.tensor(max(par.grad.detach().abs().max() for par in parameters))
      else:
          # Compute the norm in full precision no matter what
          local_norm = torch.linalg.vector_norm(
              torch.stack(
                  [
                      torch.linalg.vector_norm(par.grad.detach(), p, dtype=torch.float32)
                      for par in parameters
                  ]
              ),
              p,
          )
      local_norm.to(dtype=parameters[0].dtype)
      return local_norm
  
  
  def _get_param_to_unflat_param_names(
      model: torch.nn.Module,
      dedup_shared_params: bool = True,
  ) -> Dict[torch.nn.Parameter, List[str]]:
      """
      Constructs a mapping from flattened parameter (including non-FSDP-module
      parameters) to its unflattened parameter names. For non-FSDP-module
      parameters, these mapped-to lists always contain a single element. The
      unflattened parameter names should match the keys of the model state dict.
  
      For shared parameters, only the first parameter name is included (following
      the ``torch.nn.Module.parameters()`` order).
  
      Args:
          model (torch.nn.Module): Root module (which may or may not be a
              :class:`FullyShardedDataParallel` instance).
          dedup_shared_params (bool): If ``True``, only includes the first
              list of unflattened parameter names corresponding to a parameter
              in the module walk order; if ``False``, then includes all of the
              unflattened parameter names.
      """
      def _clean_param_name(prefix, param_info):
          """This replicates the parameter name cleaning logic in model state
          dict but avoids gathering any parameters."""
          name = clean_tensor_name(
              prefix + param_info.module_name + "." + param_info.param_name
          )
          return name
  
      def module_fn(module, prefix, param_to_unflat_param_names):
          # For FSDP modules, only add the entry when considering the contained
          # `FlattenParamsWrapper` to avoid duplication
          if not isinstance(module, FullyShardedDataParallel):
              for param_name, param in module.named_parameters(recurse=False):
                  prefixed_param_names = [
                      _clean_param_name(prefix, param_info)
                      for param_info in param._param_infos
                  ] if isinstance(param, FlatParameter) else [prefix + param_name]
                  # If this parameter has already been visited, then it is a
                  # shared parameter; then, only take the first parameter name
                  is_shared_param = param in param_to_unflat_param_names
                  if not is_shared_param:
                      param_to_unflat_param_names[param] = prefixed_param_names
                  elif not dedup_shared_params:
                      param_to_unflat_param_names[param].extend(prefixed_param_names)
  
      def return_fn(param_to_unflat_param_names):
          return param_to_unflat_param_names
  
      param_to_unflat_param_names: Dict[torch.nn.Parameter, List[str]] = {}
      return _apply_to_modules(
          model, module_fn, return_fn, param_to_unflat_param_names,
      )
  
  
  def _get_param_to_param_name(
      model: torch.nn.Module,
  ) -> Dict[torch.nn.Parameter, str]:
      """
      Constructs a mapping from parameters to their parameter names. ``model``
      should not contain any :class:`FullyShardedDataParallel` instances, which
      means that none of the parameters should be ``FlatParameter`` s. As a
      result, compared to :meth:`_get_param_to_unflat_param_names`, the mapped
      values may be flattened from singleton :class:`list` s to the contained
      names themselves.
  
      Args:
          model (torch.nn.Module): Root module, which should not contain any
              :class:`FullyShardedDataParallel` instances.
      """
      param_to_param_names = _get_param_to_unflat_param_names(model)
      for param_names in param_to_param_names.values():
          assert len(param_names) > 0, "`_get_param_to_unflat_param_names()` " \
              "should not construct empty lists"
          if len(param_names) > 1:
              raise RuntimeError(
                  "Each parameter should only map to one parameter name but got "
                  f"{len(param_names)}: {param_names}"
              )
      param_to_param_name = {
          param: param_names[0]
          for param, param_names in param_to_param_names.items()
      }
      return param_to_param_name
  
  
  def _get_param_name_to_param(
      model: torch.nn.Module,
  ) -> Dict[str, torch.nn.Parameter]:
      """Constructs the inverse mapping of :meth:`_get_param_to_param_name`."""
      param_to_param_name = _get_param_to_param_name(model)
      return dict(zip(param_to_param_name.values(), param_to_param_name.keys()))
  
  
  def clean_tensor_name(tensor_name: str) -> str:
      """Cleans the parameter or buffer name by removing any FSDP-related
      prefixes."""
      # FSDP full tensor names may not have both (i.e. `FSDP_PREFIX`), so we
      # call `replace()` twice separately
      tensor_name = tensor_name.replace(FSDP_WRAPPED_MODULE + ".", "")
      tensor_name = tensor_name.replace(FPW_MODULE + ".", "")
      return tensor_name
  