Source code for torch.distributed.fsdp.fully_sharded_data_parallel

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 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)
        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:
            raise ValueError(
                "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}"
            )
        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]
            ignored_param_names.update(unflat_param_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
                m._ignored_param_names = self._ignored_param_names
                m._buffer_names = self._buffer_names

    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)
            # 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=True``,
            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 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