fully_sharded_data_parallel_fix.py

import collections
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 (
    Any,
    Callable,
    Deque,
    Dict,
    Generator,
    Iterable,
    Iterator,
    List,
    Mapping,
    NamedTuple,
    Optional,
    Set,
    Tuple,
    Union,
    cast,
)

import torch
import torch.distributed as dist
import torch.distributed.algorithms._checkpoint.checkpoint_wrapper as checkpoint_wrapper
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.algorithms._checkpoint.checkpoint_wrapper import (
    _CHECKPOINT_PREFIX,
)
from torch.distributed.algorithms._comm_hooks import (
    LOW_PRECISION_HOOKS,
    default_hooks,
)
from torch.distributed.distributed_c10d import _get_default_group
from torch.distributed.utils import (
    _replace_by_prefix,
    _sync_params_and_buffers,
    _to_kwargs,
)
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,
    _OptimStateKey,
    _process_pos_dim_tensor_state,
    _rekey_sharded_optim_state_dict,
    _unflatten_optim_state,
)
from ._utils import (
    _apply_to_modules,
    _apply_to_tensors,
    _contains_batchnorm,
    _override_batchnorm_mixed_precision,
)
from .flat_param import (
    FlatParameter,
    FlatParamHandle,
    HandleConfig,
    HandleShardingStrategy,
    HandleTrainingState,
)
from .flatten_params_wrapper import (
    FLAT_PARAM,
    FPW_MODULE,
    FlattenParamsWrapper,
)
from .wrap import (
    ParamExecOrderWrapPolicy,
    _or_policy,
    _recursive_wrap,
    _wrap_batchnorm_individually,
)

_TORCHDISTX_AVAIL = True
try:
    from torchdistx import deferred_init, fake
except ImportError:
    _TORCHDISTX_AVAIL = False

_TORCH_FX_AVAIL = True
if not hasattr(torch, "fx"):
    _TORCH_FX_AVAIL = False
if _TORCH_FX_AVAIL:
    from ._symbolic_trace import (
        TracingConfig,
        _init_execution_info,
        _patch_tracer,
    )


__all__ = [
    "FullyShardedDataParallel",
    "ShardingStrategy",
    "MixedPrecision",
    "CPUOffload",
    "BackwardPrefetch",
    "StateDictType",
    "StateDictConfig",
    "FullStateDictConfig",
    "LocalStateDictConfig",
    "ShardedStateDictConfig",
    "OptimStateKeyType",
    "TrainingState_",
    "p_assert",
    "clean_tensor_name",
]


FSDP_WRAPPED_MODULE = "_fsdp_wrapped_module"
FSDP_PREFIX = FSDP_WRAPPED_MODULE + "." + FPW_MODULE + "."

_PARAM_BROADCAST_BUCKET_SIZE = int(250 * 1024 * 1024)


class ShardingStrategy(Enum):
    """
    This specifies the sharding strategy to be used for distributed training by
    :class:`FullyShardedDataParallel`.
    FULL_SHARD: Parameters, gradients, and optimizer states are sharded. For
                the parameters, this algorithm all-gathers before the forward,
                reshards after the forward, all-gathers before the backward
                computation, and reshards after the backward computation. The
                gradients are synchronized and sharded via reduce-scatter after
                the backward computation. The sharded optimizer states are
                updated locally.
    SHARD_GRAD_OP: Gradients and optimizer states are sharded during
                   computation, and additionally parameters are sharded outside
                   computation. For the parameters, this algorithm all-gathers
                   before the forward, does not reshard after the forward, and
                   only reshards after the backward computation. The gradients
                   are synchronized and sharded via reduce-scatter after the
                   backward computation. The sharded optimizer states are
                   updated locally. Inside ``no_sync()``, the parameters are
                   not resharded after the backward computation.
    NO_SHARD: Parameters, gradients, and optimizer states are not sharded but
              instead replicated across ranks, similar to PyTorch's
              ``DistributedDataParallel`` API. The gradients are synchronized
              via all-reduce after the backward computation. The unsharded
              optimizer states are updated locally.
    HYBRID_SHARD(future support): Apply ``FULL_SHARD`` intra-node and
                                  ``NO_SHARD`` inter-node.

    """

    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:
                   >>> # xdoctest: +SKIP("undefined variables")
                   >>> 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:
        >>> # xdoctest: +SKIP("undefined variables")
        >>> 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 structures to track the execution order. We track
    the pre-forward order on the *first* iteration for forward prefetching
    (which thus assumes static graph) and the post-forward order on *every*
    iteration for backward prefetching (which thus does not assume static
    graph but may be provide an incorrect order).

    Additionally, if the distributed debug level is set to at least INFO, then
    this tracks additional data structures for checking the execution order
    across ranks on the first iteration and per rank across iterations.
    """

    def __init__(self, debug_level: dist.DebugLevel) -> None:
        self.handles_to_pre_forward_order_index: Dict[
            Tuple[FlatParamHandle, ...], int
        ] = {}
        self.handles_pre_forward_order: List[int] = []
        self.handles_to_post_forward_order_index: Dict[
            Tuple[FlatParamHandle, ...], int
        ] = {}
        self.handles_post_forward_order: List[int] = []
        self.is_first_iter: Optional[bool] = None
        self._checking_order: bool = debug_level in [
            dist.DebugLevel.INFO,
            dist.DebugLevel.DETAIL,
        ]

        # The following are only used if distributed debug level >= INFO
        self.process_group: Optional[dist.ProcessGroup] = None
        self.world_size: Optional[int] = None
        self.all_handles: List[FlatParamHandle] = []
        self.handle_to_handle_index: Dict[FlatParamHandle, int] = {}
        self.flat_param_to_prefixed_param_names: Dict[FlatParameter, List[str]] = {}
        self.current_order_index = 0
        self.warn_status = _ExecOrderWarnStatus.NONE

    def init(
        self,
        fsdp_root: "FullyShardedDataParallel",
        process_group: dist.ProcessGroup,
    ) -> None:
        """
        Initializes the data structures needed for checking the forward order.
        This is a no-op if the distributed debug level is less than INFO.
        Otherwise, this should be called after a root FSDP instance has been
        set during lazy initialization.
        """
        if not self._checking_order:
            return
        self.process_group = process_group
        self.rank = process_group.rank()
        self.world_size = process_group.size()
        # Fix an order over the handles, which should be the same across ranks
        for fsdp_module in fsdp_root.fsdp_modules(fsdp_root):
            for handle in fsdp_module._handles:
                index = len(self.all_handles)
                self.all_handles.append(handle)
                self.handle_to_handle_index[handle] = index
        self.flat_param_to_prefixed_param_names = cast(
            Dict[FlatParameter, List[str]],
            _get_param_to_unflat_param_names(fsdp_root),
        )
        # TODO (awgu): We can broadcast the metadata of rank 0's `all_handles`
        # to check that all ranks have the same handles in the same order.
        # https://github.com/pytorch/pytorch/issues/79620

    def get_handles_to_backward_prefetch(
        self,
        current_handles_key: Tuple[FlatParamHandle, ...],
    ) -> Optional[Tuple[FlatParamHandle, ...]]:
        """
        Returns the handles key of the handles to backward prefetch given the
        current handles key or ``None`` if there is no valid handles key to
        prefetch.
        """
        current_index = self.handles_to_post_forward_order_index.get(
            current_handles_key, None
        )
        if current_index is None:
            return None
        target_index = current_index - 1
        if target_index < 0:
            return None
        target_handles_key = self.handles_post_forward_order[target_index]
        return target_handles_key

    def record_post_forward(self, handles: List[FlatParamHandle]) -> None:
        """
        Records ``handles`` in the post-forward order, where ``handles`` should
        be a group of handles used in the same module's forward.

        Unlike :meth:`record_pre_forward`, this records the order *every*
        iteration with the expectation that the recorded order is reset in
        :meth:`next_iter`, and the recorded order includes empty handles keys.
        """
        handles_key = tuple(handles)
        if handles_key and handles_key in self.handles_to_post_forward_order_index:
            return
        index = len(self.handles_post_forward_order)
        if handles_key:
            self.handles_to_post_forward_order_index[handles_key] = index
        self.handles_post_forward_order.append(handles_key)

    def record_pre_forward(self, handles: List[FlatParamHandle]) -> None:
        """
        Records ``handles`` in the pre-forward order on the first iteration,
        where ``handles`` should be a group of handles used in the same
        module's forward. If ``handles`` is empty, then it is omitted.

        If the distributed debug level is at least INFO, then this additionally
        checks the execution order across ranks. See :meth:`_check_order` for
        details.
        """
        if not handles:
            return
        handles_key = tuple(handles)
        if self._checking_order:
            self._check_order(handles_key)
        # Fix the order after the first iteration
        # TODO (awgu): For now, only record the first usage of a module, which
        # is consistent with the existing implementation.
        if (
            not self.is_first_iter
            or handles_key in self.handles_to_pre_forward_order_index
        ):
            return
        index = len(self.handles_pre_forward_order)
        self.handles_to_pre_forward_order_index[handles_key] = index
        self.handles_pre_forward_order.append(handles_key)

    def _check_order(self, handles_key: Tuple[FlatParamHandle, ...]) -> None:
        """
        Checks the forward execution order. This should only be called if the
        distributed debug level is at least INFO.

        On the first iteration, this uses all-gathers to check that all ranks
        are all-gathering the same handles and hence ``FlatParameter`` s,
        raising an error if not.
        On subsequent iterations, this checks that each rank is locally
        consistent with its own forward order from the first iteration, issuing
        a warning if not. This issues a warning on the first deviating
        iteration and stops warning thereafter.
        """
        if self.is_first_iter:
            msg_prefix = "Forward order differs across ranks:"
            local_indices: Optional[Tuple[int, ...]] = self._get_handle_indices(
                handles_key
            )
            device = handles_key[0].device  # guaranteed to be non-CPU
            num_valid_indices = sum((index is not None) for index in local_indices)
            tensor_kwargs = {"dtype": torch.int32, "device": device}
            world_num_valid_indices = torch.zeros(self.world_size, **tensor_kwargs)
            local_num_valid_indices = torch.tensor([num_valid_indices], **tensor_kwargs)
            dist._all_gather_base(
                world_num_valid_indices,
                local_num_valid_indices,
                group=self.process_group,
            )
            # Check that all ranks plan to all-gather the same number of
            # parameters
            # TODO (awgu): Since every module has at most one handle in the
            # current implementation, this should never raise the error.
            for (r1, n1), (r2, n2) in itertools.combinations(
                (
                    (rank, world_num_valid_indices[rank])
                    for rank in range(self.world_size)
                ),
                2,
            ):
                if n1 != n2:
                    raise RuntimeError(
                        f"{msg_prefix} rank {r1} is all-gathering {n1} parameters "
                        f"while rank {r2} is all-gathering {n2} parameters"
                    )
            world_indices = torch.zeros(
                self.world_size * num_valid_indices, **tensor_kwargs
            )
            local_indices = torch.tensor(local_indices, **tensor_kwargs)
            dist._all_gather_base(
                world_indices, local_indices, group=self.process_group
            )
            # Check that all ranks plan to all-gather the same index parameters
            for (r1, i1), (r2, i2) in itertools.combinations(
                (
                    (
                        rank,
                        world_indices[
                            rank * num_valid_indices : (rank + 1) * num_valid_indices
                        ],
                    )
                    for rank in range(self.world_size)
                ),
                2,
            ):
                if i1 != i2:
                    r1_param_names = self._get_names_from_handle_indices(i1)
                    r2_param_names = self._get_names_from_handle_indices(i2)
                    raise RuntimeError(
                        f"{msg_prefix} rank {r1} is all-gathering parameters "
                        f"for {r1_param_names} while rank {r2} is all-gathering "
                        f"parameters for {r2_param_names}"
                    )
        else:
            # Only issue warnings on the first deviating iteration and stop
            # checking thereafter to avoid flooding the console
            if self.warn_status == _ExecOrderWarnStatus.WARNED:
                return
            msg_prefix = None  # non-`None` means we should warn
            if self.current_order_index >= len(self.handles_pre_forward_order):
                # This iteration sees extra all-gather(s) compared to the first
                msg_prefix = (
                    "Expected to not all-gather any more parameters in the "
                    "forward but trying to all-gather parameters for "
                )
            else:
                expected_handles_key = self.handles_pre_forward_order[
                    self.current_order_index
                ]
                if expected_handles_key != handles_key:
                    expected_param_names = self._get_names_from_handles(
                        expected_handles_key
                    )
                    msg_prefix = (
                        f"Expected to all-gather for {expected_param_names} "
                        "but trying to all-gather parameters for "
                    )
            if msg_prefix is not None:
                param_names = self._get_names_from_handles(handles_key)
                msg_suffix = (
                    f"{param_names}"
                    if param_names
                    else "a newly-added parameter since construction time"
                )
                warnings.warn(
                    "Forward order differs from that of the first iteration "
                    f"on rank {self.rank}. Collectives are unchecked and may "
                    f"give incorrect results or hang.\n{msg_prefix}{msg_suffix}"
                )
                self.warn_status = _ExecOrderWarnStatus.WARNING
            self.current_order_index += 1

    def _get_handle_indices(
        self,
        handles_key: Tuple[FlatParamHandle, ...],
    ) -> Tuple[Optional[int], ...]:
        """
        Returns the handle indices (i.e. indices into ``self.all_handles``)
        corresponding to the handles in ``handles_key``. An entry in the
        returned tuple is ``None`` if the handle is invalid.
        """
        indices: List[int] = []
        for handle in handles_key:
            if handle not in self.handle_to_handle_index:
                indices.append(None)
            else:
                indices.append(self.handle_to_handle_index[handle])
        return tuple(indices)

    def _get_names_from_handle_indices(
        self,
        handle_indices: Tuple[int, ...],
    ) -> List[List[str]]:
        """
        Returns a list of prefixed parameter names for each handle in
        ``handle_indices``. If a handle index is invalid, then its prefixed
        parameter names are omitted from the returned list.
        """
        prefixed_param_names: List[List[str]] = []
        for index in handle_indices:
            if index is None or index < 0 or index >= len(self.all_handles):
                continue
            handle = self.all_handles[index]
            flat_param = handle.flat_param
            prefixed_param_names.append(
                self.flat_param_to_prefixed_param_names[flat_param]
            )
        return prefixed_param_names

    def _get_names_from_handles(
        self,
        handles_key: Tuple[FlatParamHandle, ...],
    ) -> List[List[str]]:
        """
        Returns a list of prefixed parameter names for each handle in
        ``handles_key``. If a handle is invalid, then its prefixed parameter
        names are omitted from the returned list.
        """
        prefixed_param_names: List[List[str]] = []
        for handle in handles_key:
            flat_param = handle.flat_param
            if flat_param not in self.flat_param_to_prefixed_param_names:
                continue
            prefixed_param_names.append(
                self.flat_param_to_prefixed_param_names[flat_param]
            )
        return prefixed_param_names

    def next_iter(self):
        """
        Advances the internal data structures per iteration. This should be
        called in the root's pre-forward rather than in the post-backward
        callback since the backward may not run (e.g. inference).
        """
        if self.is_first_iter is None:
            self.is_first_iter = True
        else:
            self.is_first_iter = False
        self.handles_to_post_forward_order_index.clear()
        self.handles_post_forward_order.clear()
        if self._checking_order:
            self.current_order_index = 0
            if self.warn_status == _ExecOrderWarnStatus.WARNING:
                self.warn_status = _ExecOrderWarnStatus.WARNED


# TODO (awgu): Refactor this later
sharding_strategy_map = {
    ShardingStrategy.NO_SHARD: HandleShardingStrategy.NO_SHARD,
    ShardingStrategy.FULL_SHARD: HandleShardingStrategy.FULL_SHARD,
    ShardingStrategy.SHARD_GRAD_OP: HandleShardingStrategy.SHARD_GRAD_OP,
}


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::

        >>> # xdoctest: +SKIP("undefined variables")
        >>> 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[[nn.Module, bool, int], bool]]):
            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 transformer-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``, and return
            a ``bool`` specifying whether the passed in ``module``` should be wrapped
            (if ``recurse=False``) or whether we should recurse down the subgraph of ``module``
            children (if ``recurse=True``). 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
                >>> # Configure a custom min_num_params
                >>> my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=1e5)

        backward_prefetch (Optional[BackwardPrefetch]):
            This is an experimental feature that is subject to change in the
            the near future. It allows users to enable two different backward_prefetch
            algorithms to help backward communication and computation overlapping.
            Pros and cons of each algorithm is explained in the class ``BackwardPrefetch``.
        mixed_precision (Optional[MixedPrecision]): A ``MixedPrecision`` instance
            describing the mixed precision training config to be used. ``MixedPrecision``
            supports configuring parameter, buffer, and gradient communication dtype. Note
            that only floating point data is cast to the reduced precision. This allows
            users potential memory saving and training speedup while trading off
            accuracy during model training. If ``None``, no mixed precision is applied.
            Note that if ``mixed_precision`` is enabled for FSDP model that
            contains ``BatchNorm`` with ``auto_wrap_policy``, FSDP will take
            care to disable mixed precision for ``BatchNorm`` units by wrapping
            them separately in their own FSDP unit with ``mixed_precision=None``.
            This is done because several ``BatchNorm`` kernels do not implement
            reduced type support at the moment. If individually wrapping the model,
            users must take care to set ``mixed_precision=None`` for
            ``BatchNorm`` units.
            (Default: ``None``)
        ignored_modules (Optional[Iterable[torch.nn.Module]]): Modules whose
            own parameters and child modules' parameters and buffers are
            ignored by this instance. None of the modules directly in
            ``ignored_modules`` should be :class:`FullyShardedDataParallel`
            instances, and any child modules that are already-constructed
            :class:`FullyShardedDataParallel` instances will not be ignored if
            they are nested under this instance. This argument may be used to
            avoid sharding specific parameters at module granularity when using an
            ``auto_wrap_policy`` or if parameters' sharding is not managed by
            FSDP. (Default: ``None``)
        param_init_fn (Optional[Callable[[nn.Module], None]]):
            A ``Callable[torch.nn.Module] -> None`` that
            specifies how modules that are currently on the meta device should be initialized
            onto an actual device. Note that as of v1.12, we detect modules on the meta
            device via ``is_meta`` check and apply a default initialization that calls
            ``reset_parameters`` method on the passed in ``nn.Module`` if ``param_init_fn``
            is not specified, otherwise we run ``param_init_fn`` to initialize the passed
            in ``nn.Module``. In particular, this means that if ``is_meta=True`` for any
            module parameters for modules that will be wrapped with FSDP and ``param_init_fn``
            is not specified, we assume your module properly implements a ``reset_paramters()``
            and will throw errors if not. Note that additionally, we offer support for modules
            initialized with torchdistX's (https://github.com/pytorch/torchdistX)
            ``deferred_init`` API. In this case, deferred modules would be initialized
            by a default initialization function that calls torchdistX's
            ``materialize_module``, or the passed in ``param_init_fn``, if it is not
            ``None``. The same ``Callable`` is applied to initialize all meta modules.
            Note that this initialization function is applied before doing any FSDP sharding
            logic.

            Example::

                >>> # xdoctest: +SKIP("undefined variables")
                >>> 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``)
        limit_all_gathers (bool): If ``False``, then FSDP allows the CPU
            thread to schedule all-gathers without any extra synchronization.
            If ``True``, then FSDP explicitly synchronizes the CPU thread to
            prevent too many in-flight all-gathers. This ``bool`` only affects
            the sharded strategies that schedule all-gathers.
            TODO (awgu): Explain the implications on GPU memory.
    """

    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,
        limit_all_gathers: bool = True,
    ):
        if isinstance(auto_wrap_policy, ParamExecOrderWrapPolicy):
            self._init_param_exec_order_wrap_policy(
                module=module,
                process_group=process_group,
                sharding_strategy=sharding_strategy,
                cpu_offload=cpu_offload,
                auto_wrap_policy=auto_wrap_policy,
                backward_prefetch=backward_prefetch,
                mixed_precision=mixed_precision,
                ignored_modules=ignored_modules,
                param_init_fn=param_init_fn,
                device_id=device_id,
                sync_module_states=sync_module_states,
                limit_all_gathers=limit_all_gathers,
            )
            return

        torch._C._log_api_usage_once("torch.distributed.fsdp")
        super().__init__()
        
        # --- Just announce we are live
        print(f"--> ** running with high res rate limiter update! ** ")

        self._ignored_modules = self._get_ignored_modules(module, ignored_modules)
        ignored_params, self._ignored_param_names = self._get_ignored_params(
            module, self._ignored_modules
        )
        self._buffer_names = self._get_buffer_names(module)
        if auto_wrap_policy is not None:
            auto_wrap_kwargs = {
                "module": module,
                "auto_wrap_policy": auto_wrap_policy,
                "wrapper_cls": FullyShardedDataParallel,
                "ignored_modules": self._ignored_modules,
                "ignored_params": ignored_params,
                "only_wrap_children": True,  # avoid double wrapping the root
            }
            fsdp_kwargs = {
                "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,
                "limit_all_gathers": limit_all_gathers,
            }
            self._auto_wrap(auto_wrap_kwargs, fsdp_kwargs)

        self.process_group = process_group or _get_default_group()
        self.rank = self.process_group.rank()
        self.world_size = self.process_group.size()
        self.training_state = TrainingState_.IDLE
        self.cpu_offload = cpu_offload or CPUOffload()
        self.backward_prefetch = backward_prefetch
        self.limit_all_gathers = limit_all_gathers
        self._max_num_inflight_all_gathers = 2  # empirically chosen
        if self.world_size == 1:
            # World size of 1 is functionally equivalent to `NO_SHARD`
            sharding_strategy = ShardingStrategy.NO_SHARD
        self.sharding_strategy = sharding_strategy or ShardingStrategy.FULL_SHARD
        self.mixed_precision = mixed_precision or MixedPrecision()
        # Save a mapping from fully prefixed buffer name to its original dtype
        # since for mixed precision, buffers are restored to their original
        # dtype for model checkpointing
        self._buffer_name_to_orig_dtype: Dict[str, torch.dtype] = {}

        if not torch.cuda.is_available():
            raise RuntimeError("FSDP does not support CPU only execution")
        self._check_single_device_module(module, ignored_params)
        device_from_device_id: Optional[torch.device] = self._get_device_from_device_id(
            device_id
        )
        self._materialize_module(module, param_init_fn, device_from_device_id)
        self._move_module_to_device(module, ignored_params, device_from_device_id)
        self.compute_device = self._get_compute_device(
            module, ignored_params, device_from_device_id
        )
        params_to_flatten = list(self._get_orig_params(module, ignored_params))
        if sync_module_states:
            self._sync_module_states(module, params_to_flatten)

        # This FSDP instance's handles should inherit the same process group,
        # compute device, CPU offload, and mixed precision settings. However,
        # different sharding strategies are allowed.
        config = HandleConfig(
            sharding_strategy_map[self.sharding_strategy],
            self.cpu_offload.offload_params,
            self.mixed_precision.param_dtype,
            self.mixed_precision.reduce_dtype,
        )
        self._fsdp_wrapped_module = FlattenParamsWrapper(
            module,
            params_to_flatten,
            self.compute_device,
            config,
        )
        self._check_orig_params_flattened(ignored_params)
        # Invariant: `self.params` contains exactly the `FlatParameter`s of the
        # handles in `self._handles`
        self._handles: List[FlatParamHandle] = []
        self.params: List[FlatParameter] = []
        if self._fsdp_wrapped_module.has_params:
            handle = self._fsdp_wrapped_module.handle
            self.params.append(handle.flat_param)
            self._register_param_handle(handle)
            handle.shard(self.process_group)
            if self.cpu_offload.offload_params:
                with torch.no_grad():
                    handle._flat_param_to(torch.device("cpu"))

        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._sync_gradients = True
        self.communication_hook = self._get_default_comm_hook()
        self.communication_hook_state = self._get_default_comm_hook_state()
        self._hook_registered = False

        # Used to prevent running the pre-backward hook multiple times
        self._pre_backward_hook_has_run: Dict[Tuple[FlatParamHandle, ...], bool] = {}
        self._is_root: Optional[bool] = None  # `None` indicates not yet set
        # The following attributes are owned by the root FSDP instance and
        # shared with non-root FSDP instances
        self._streams: Dict[str, torch.cuda.Stream] = {}
        self._free_events: Deque[torch.cuda.Event] = collections.deque()
        self._debug_level = dist.get_debug_level()
        self._exec_order_data = _ExecOrderData(self._debug_level)
        self._handles_prefetched: Dict[Tuple[FlatParamHandle, ...], bool] = {}
        # Used for `BACKWARD_POST` prefetching
        self._need_pre_backward_unshard: Dict[Tuple[FlatParamHandle, ...], bool] = {}
        # The data structures use tuples of handles to generalize over the case
        # where a module's forward involves multiple handles. The two forward
        # order structures are populated and finalized in the first iteration.

        # `_state_dict_type` controls the `state_dict()` behavior, which is
        # implemented using post-save and pre-load hooks
        self._state_dict_type = StateDictType.FULL_STATE_DICT
        self._state_dict_config = FullStateDictConfig()
        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,
        }

    def _get_ignored_modules(
        self,
        root_module: nn.Module,
        _ignored_modules: Optional[Iterable[torch.nn.Module]],
    ) -> Set[nn.Module]:
        """
        Checks that ``_ignored_modules`` is an iterable of ``nn.Module`` s
        without any FSDP instances, and returns the modules contained in their
        module subtrees as a :class:`set`. Nested FSDP instances are excluded,
        but their already-computed ignored modules are included.
        """
        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 themselves
        ignored_modules = set(
            child
            for module in ignored_root_modules
            for child in module.modules()
            if not isinstance(child, (FullyShardedDataParallel, FlattenParamsWrapper))
        )
        if root_module in ignored_modules:
            warnings.warn(
                "Trying to ignore the top-level module passed into the FSDP "
                "constructor itself will result in all parameters being "
                f"ignored and is not well-supported: {module}"
            )
        # Include nested FSDP modules' ignored modules
        for submodule in root_module.modules():
            if isinstance(submodule, FullyShardedDataParallel):
                assert hasattr(submodule, "_ignored_modules")
                ignored_modules.update(submodule._ignored_modules)
        return ignored_modules

    def _get_ignored_params(
        self,
        root_module: torch.nn.Module,
        ignored_modules: Set[torch.nn.Module],
    ) -> Tuple[Set[torch.nn.Parameter], Set[str]]:
        """
        Returns the parameters of the modules in ``ignored_modules``,
        excluding any :class:`FlatParameter` s, and their fully prefixed names,
        both as :class:`set` s.
        """
        ignored_params = set(
            p
            for m in ignored_modules
            for p in m.parameters()
            if not isinstance(p, FlatParameter)
        )
        # Conservatively include all shared parameters' names
        param_to_unflat_param_names = _get_param_to_unflat_param_names(
            root_module,
            dedup_shared_params=False,
        )
        ignored_param_names = set()
        for param in ignored_params:
            unflat_param_names = param_to_unflat_param_names[param]
            clean_names = []
            for k in unflat_param_names:
                # Clean any module wrapper prefixes in case of nested wrapping
                clean_names.append(clean_tensor_name(k))
            ignored_param_names.update(clean_names)
        return ignored_params, ignored_param_names

    def _get_buffer_names(self, root_module: nn.Module) -> Set[str]:
        """
        Returns the fully prefixed names of all buffers in the module hierarchy
        rooted at ``root_module`` as a class:`set`.
        """

        def module_fn(module: nn.Module, prefix: str, buffer_names: Set[str]):
            # 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):
                    # Clean module wrapper prefixes in case of nested wrapping
                    prefixed_buffer_name = clean_tensor_name(prefix + buffer_name)
                    buffer_names.add(prefixed_buffer_name)

        def return_fn(buffer_names: Set[str], *args):
            return buffer_names

        buffer_names: Set[str] = set()
        return _apply_to_modules(
            root_module,
            module_fn,
            return_fn,
            buffer_names,
        )

    def _auto_wrap(
        self,
        auto_wrap_kwargs: Dict[str, Any],
        fsdp_kwargs: Dict[str, Any],
    ) -> None:
        """
        Recursively auto wraps the root module given by the key "module" in
        ``auto_wrap_kwargs`` with the arguments in ``auto_wrap_kwargs`` and
        ``fsdp_kwargs``.

        Precondition: ``auto_wrap_policy`` contains the arguments expected by
            ``_recursive_wrap()``, where ``auto_wrap_policy`` is not ``None``.
            ``fsdp_kwargs`` contains all FSDP arguments except ``module``.
        """
        auto_wrap_policy = auto_wrap_kwargs["auto_wrap_policy"]
        root_module = auto_wrap_kwargs["module"]
        assert auto_wrap_policy is not None
        # For auto wrapping, submodules should not already be wrapped with FSDP
        # since double wrapping is not supported
        for module_name, module in root_module.named_modules():
            if isinstance(module, FullyShardedDataParallel):
                raise ValueError(
                    f"Expected {module_name} to NOT be FullyShardedDataParallel "
                    "if using an `auto_wrap_policy`"
                )
        mixed_precision = fsdp_kwargs["mixed_precision"]
        if mixed_precision is not None and _contains_batchnorm(root_module):
            _override_batchnorm_mixed_precision(root_module)
            auto_wrap_policy = functools.partial(
                _or_policy, policies=[_wrap_batchnorm_individually, auto_wrap_policy]
            )
            warnings.warn(
                "Both mixed precision and an `auto_wrap_policy` were specified "
                "for FSDP, where the wrapped module has batch norm submodules. "
                "The batch norm submodules will be wrapped as separate FSDP "
                "instances with mixed precision disabled since some batch norm "
                "kernels do not support low precision."
            )
            auto_wrap_kwargs["auto_wrap_policy"] = auto_wrap_policy
        _recursive_wrap(**auto_wrap_kwargs, **fsdp_kwargs)

    def _check_single_device_module(
        self,
        module: nn.Module,
        ignored_params: Set[nn.Parameter],
    ) -> None:
        """
        Raises an error if ``module`` has original parameters on multiple
        devices, ignoring the parameters in ``ignored_params``. Thus, after
        this method, the module must be either fully on the CPU or fully on a
        non-CPU device.
        """
        devices = set(
            param.device for param in self._get_orig_params(module, ignored_params)
        )
        if len(devices) > 1:
            raise RuntimeError(
                f"FSDP only supports single device modules but got params on {devices}"
            )

    def _get_device_from_device_id(
        self,
        device_id: Optional[Union[int, torch.device]],
    ) -> Optional[torch.device]:
        """ """
        if device_id is None:
            return None
        device = (
            device_id
            if isinstance(device_id, torch.device)
            else torch.device(device_id)
        )
        if device == torch.device("cuda"):
            warnings.warn(
                f"FSDP got the argument `device_id` {device_id} on rank "
                f"{self.rank}, which does not have an explicit index. "
                f"FSDP will use the current device {torch.cuda.current_device()}. "
                "If this is incorrect, please explicitly call `torch.cuda.set_device()` "
                "before FSDP initialization or pass in the explicit device "
                "index as the `device_id` argument."
            )
            device = torch.device("cuda", torch.cuda.current_device())
        return device

    def _materialize_module(
        self,
        module: nn.Module,
        param_init_fn: Optional[Callable[[nn.Module], None]],
        device_from_device_id: Optional[torch.device],
    ) -> None:
        """
        Materializes the wrapped module ``module`` in place if needed: either
        if the module has parameters that use meta device or are torchdistX
        fake tensors.

        This method uses ``param_init_fn`` to materialize the module if the
        function is not ``None`` and falls back to default behavior otherwise.
        For meta device, this moves the module to ``device_from_device_id`` if
        it is not ``None`` or the current device otherwise and calls
        ``reset_parameters()``, and for torchdistX fake tensors, this calls
        ``deferred_init.materialize_module()``.
        """
        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())
        )
        if (
            is_meta_module or is_torchdistX_deferred_init
        ) and param_init_fn is not None:
            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)
        elif is_meta_module:
            # Run default meta device initialization
            materialization_device = (
                device_from_device_id or torch.cuda.current_device()
            )
            module.to_empty(device=materialization_device)
            try:
                with torch.no_grad():
                    module.reset_parameters()
            except BaseException as e:
                warnings.warn(
                    "Unable to call `reset_parameters()` for module on meta "
                    f"device with error {str(e)}. Please ensure your "
                    "module implements a `reset_parameters()` method."
                )
                raise e
        elif is_torchdistX_deferred_init:
            # Run default torchdistX initialization
            deferred_init.materialize_module(
                module,
                check_fn=lambda k: not isinstance(k, FullyShardedDataParallel),
            )

    def _move_module_to_device(
        self,
        module: nn.Module,
        ignored_params: Set[nn.Parameter],
        device_from_device_id: Optional[torch.device],
    ):
        """
        Moves ``module`` depending on ``device_from_device_id`` and its current
        device.

        - If ``device_from_device_id`` is not ``None``, then this moves
        ``module`` to the device.
        - If ``device_from_device_id`` is ``None``, then this does not move
        ``module`` but warns the user if it is on CPU.

        Precondition: ``_check_single_device_module()``.
        """
        cpu_device = torch.device("cpu")
        param = next(self._get_orig_params(module, ignored_params), None)
        if param is None:
            return  # no original parameters to manage
        if device_from_device_id is not None:
            if param.device == cpu_device:
                # NOTE: This includes moving ignored parameters.
                module = module.to(device_from_device_id)
                # TODO (awgu): This is a temporary fix to move already-
                # constructed `FlatParameter`s back to CPU if needed. This is
                # needed to make CPU offload work with `device_id`.
                for submodule in module.modules():
                    if (
                        isinstance(submodule, FullyShardedDataParallel)
                        and submodule.cpu_offload.offload_params
                    ):
                        with torch.no_grad():
                            for handle in submodule._handles:
                                handle._flat_param_to(torch.device("cpu"))
        elif param.device == cpu_device:
            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."
            )

    def _get_compute_device(
        self,
        module: nn.Module,
        ignored_params: Set[nn.Parameter],
        device_from_device_id: Optional[torch.device],
    ) -> torch.device:
        """
        Determines and returns this FSDP instance's compute device. If the
        module is already on a non-CPU device, then the compute device is that
        non-CPU device. If the module is on CPU, then the compute device is the
        current device.

        Since this method should be called after materializing the module, any
        non-CPU device should not be meta device. For now, the compute device
        is always a CUDA GPU device with its explicit index.

        Precondition: ``_check_single_device_module()`` and
        ``_move_module_to_device()``.
        """
        # If the module is on GPU already, then that GPU device has priority
        # over the current device
        param = next(self._get_orig_params(module, ignored_params), None)
        if param is not None and param.device.type == "cuda":
            compute_device = param.device
        else:
            compute_device = torch.device("cuda", torch.cuda.current_device())
        if (
            device_from_device_id is not None
            and compute_device != device_from_device_id
        ):
            raise ValueError(
                "Inconsistent compute device and `device_id` on rank "
                f"{self.rank}: {compute_device} vs {device_from_device_id}"
            )
        return compute_device

    def _sync_module_states(
        self, module: nn.Module, params: List[nn.Parameter]
    ) -> None:
        """
        Synchronizes module states (i.e. parameters ``params`` and all
        not-yet-synced buffers) by broadcasting from rank 0 to all ranks.

        Precondition: ``sync_module_states == True`` and ``self.process_group``
        has been set.
        """
        if params and any(param.device == torch.device("cpu") for param in params):
            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."
            )
        module_states: List[torch.Tensor] = []
        # TODO (awgu): When exposing the original parameters, we need to also
        # use this attribute to prevent re-synchronizing parameters.
        for buffer in module.buffers():
            # Avoid re-synchronizing buffers in case of nested wrapping
            if not getattr(buffer, "_fsdp_synced", False):
                buffer._fsdp_synced = True
                module_states.append(buffer.detach())
        module_states.extend(param.detach() for param in params)
        _sync_params_and_buffers(
            self.process_group,
            module_states,
            _PARAM_BROADCAST_BUCKET_SIZE,
            src=0,
        )

    def _get_orig_params(
        self,
        module: nn.Module,
        ignored_params: Set[nn.Parameter],
    ) -> Iterator[nn.Parameter]:
        """
        Returns an iterator over the original parameters in ``module``,
        ignoring the parameters in ``ignored_params`` and any ``FlatParameter``
        s (which may be present due to nested FSDP wrapping).
        """
        param_gen = module.parameters()
        try:
            while True:
                param = next(param_gen)
                if param not in ignored_params and not isinstance(param, FlatParameter):
                    yield param
        except StopIteration:
            pass

    def _check_orig_params_flattened(self, ignored_params: Set[nn.Parameter]) -> None:
        """
        Checks that all original parameters have been flattened and hence made
        invisible to ``named_parameters()``. This should be called as a sanity
        check after flattening the wrapped module's parameters.
        """
        for param_name, param in self.named_parameters():
            if param not in ignored_params and not isinstance(param, FlatParameter):
                raise RuntimeError(
                    f"Found an unflattened parameter: {param_name}; "
                    f"{param.size()} {param.__class__}"
                )

    def _register_param_handle(self, handle: FlatParamHandle) -> None:
        """Registers the parameter handle to this FSDP instance."""
        if handle not in self._handles:
            self._handles.append(handle)

    @torch.no_grad()
    def _unshard(
        self,
        handles: List[FlatParamHandle],
    ) -> None:
        """
        Unshards the handles in ``handles``. If the handles are in
        :meth:`summon_full_params` and are using mixed precision, then they are
        forced to full precision.

        Postcondition: Each handle's ``FlatParameter`` 's data is the padded
        unsharded flattened parameter on the compute device.
        """
        if (
            not self.limit_all_gathers
            and len(self._free_events) >= self._max_num_inflight_all_gathers
        ):
            # Synchronize the current stream to block the CPU thread and
            # prevent it from over-all-gathering
            free_event = self._free_events.popleft()
            free_event.synchronize()
        with torch.cuda.stream(self._streams["all_gather"]):
            for handle in handles:
                handle.pre_unshard()
                handle.unshard()
                handle.post_unshard()

    def _reshard(
        self,  # unused
        handles: List[FlatParamHandle],
        free_unsharded_flat_params: List[bool],
    ) -> None:
        """
        Reshards the handles in ``handles``. ``free_unsharded_flat_params``
        should have the same length as ``handles``, and each element should
        give whether the corresponding handle should free its padded unsharded
        flattened parameter.
        """
        assert len(handles) == len(free_unsharded_flat_params), (
            "Expects both lists to have equal length but got "
            f"{len(handles)} and {len(free_unsharded_flat_params)}"
        )
        for handle, free_unsharded_flat_param in zip(
            handles,
            free_unsharded_flat_params,
        ):
            handle.reshard(free_unsharded_flat_param)
            if not self.limit_all_gathers and free_unsharded_flat_param:
                free_event = torch.cuda.Event()
                free_event.record()
                self._free_events.append(free_event)
            handle.post_reshard()

    @property
    def module(self) -> nn.Module:
        """
        Returns the wrapped module (like :class:`DistributedDataParallel`).
        """
        assert isinstance(self._fsdp_wrapped_module, FlattenParamsWrapper)
        return self._fsdp_wrapped_module.module

    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._fsdp_wrapped_module, name)

    def __getitem__(self, key: int) -> Any:
        """Forward indexing calls in case the module is a nn.Sequential."""
        return self._fsdp_wrapped_module.__getitem__(key)  # type: ignore[operator]

    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

    def _get_gradient_predivide_factor(self, world_size: int) -> float:
        """
        Returns a pre-divide factor to prevent underflow/overflow during
        gradient synchronization. The post-divide factor should be the
        reciprocal of the pre-divide factor.
        """
        factor: int = 1
        while world_size % factor == 0 and world_size / factor > factor:
            factor *= 2
        return float(factor)

    def _mixed_precision_enabled_for_params(self) -> bool:
        """
        Whether user explicitly enabled mixed precision for
        parameters or not.
        """
        return 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.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.reduce_dtype is not None

    def _low_precision_hook_enabled(self) -> bool:
        """
        Wether a low precision hook is registered or not.
        """
        return (
            self.communication_hook is not None
            and self.communication_hook in LOW_PRECISION_HOOKS
        )

    def _cast_fp_inputs_to_dtype(
        self, dtype: torch.dtype, *args: Any, **kwargs: Any
    ) -> Tuple[Any, Any]:
        """
        Casts floating point tensors in ``args`` and ``kwargs`` to the
        precision given by ``dtype``, while respecting the existing
        ``requires_grad`` on the tensors.
        """

        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` since this runs inside
            # `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),
            )

    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._buffer_name_to_orig_dtype:
                        self._buffer_name_to_orig_dtype[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)

    def _reset_lazy_init(self) -> None:
        """
        Reset instance so :func:`_lazy_init` will run on the next forward.
        """
        self._is_root: Optional[bool] = None
        for p in self.params:
            if hasattr(p, "_local_shard"):
                # TODO (awgu): We only need to `del` `_local_shard` because
                # `_init_param_attributes()` gates the logic based on its
                # existence (and not any of the other attributes).
                del p._local_shard  # type: ignore[attr-defined]

    def _lazy_init(self) -> None:
        """
        Performs initialization lazily, typically right before the first
        forward pass. The laziness is needed to ensure that the parameter
        device/dtype and the FSDP hierarchy have finalized.

        This method's actual logic only runs on the root FSDP instance, which
        performs initialization for all non-root FSDP instances to avoid
        partial initialization.
        """
        if self._is_root is not None:
            return  # no-op: already initialized
        # The following logic is only run on the root FSDP instance since it
        # will set `_is_root=False` for the non-root instances
        self._is_root = True
        self._assert_state(TrainingState_.IDLE)
        self._init_streams()
        self._cast_buffers(recurse=True)
        for handle in self._handles:
            self._init_param_attributes(handle)
        self._exec_order_data.init(self, self.process_group)
        # Initialize non-root FSDP instances and share attributes from the root
        # to non-root instances
        inconsistent_limit_all_gather = False
        for fsdp_module in self.fsdp_modules(self):
            if fsdp_module is not self:
                # Relax the assert for non-root FSDP instances in case the
                # nested initialized module is wrapped again in FSDP later (e.g.
                # after training to run inference)
                assert fsdp_module._is_root is None or not fsdp_module._is_root, (
                    "Non-root FSDP instance's `_is_root` should not have been "
                    "set yet or should have been set to `False`"
                )
                fsdp_module._is_root = False
                fsdp_module._streams = self._streams
                fsdp_module._exec_order_data = self._exec_order_data
                if fsdp_module.limit_all_gathers != self.limit_all_gathers:
                    # Prefer the root's value
                    inconsistent_limit_all_gather = True
                    fsdp_module.limit_all_gathers = self.limit_all_gathers
                fsdp_module._free_events = self._free_events
                fsdp_module._handles_prefetched = self._handles_prefetched
                fsdp_module._need_pre_backward_unshard = self._need_pre_backward_unshard
                for handle in fsdp_module._handles:
                    fsdp_module._init_param_attributes(handle)
        if inconsistent_limit_all_gather:
            warnings.warn(
                "Found inconsistent `limit_all_gathers` values across FSDP "
                f"instances on rank {self.rank}. Using the root FSDP's value "
                f"of {self.limit_all_gathers} for all instances."
            )

    # TODO (awgu): Move this to the `FlatParamHandle` class later
    @torch.no_grad()
    def _init_param_attributes(self, handle: FlatParamHandle) -> None:
        """
        We manage several attributes on each Parameter instance. The first is
        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.
        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`:
            ``_post_backward_hook_state``: it holds the parameter's AccumulateGrad object
                and the registered post hook handle.
        """
        p = handle.flat_param
        # 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 handle.uses_sharded_strategy:  # 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]

            if self._mixed_precision_enabled_for_params():
                p._full_prec_full_param_padded = torch.zeros(  # type: ignore[attr-defined]
                    p.numel() * self.world_size,
                    device=self.compute_device,
                    dtype=p.dtype,  # full precision
                )
                _free_storage(p._full_prec_full_param_padded)

        # Track whether the `FlatParameter`'s post-backward hook has been
        # called for validation in `_wait_for_post_backward()`
        p._post_backward_called = False

    def _init_streams(self) -> None:
        """Initializes CUDA streams for overlapping data transfer and
        computation. This should only be called on the root FSDP instance."""
        assert self._is_root
        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()

    def _wait_for_previous_optim_step(self) -> None:
        """
        The root :class:`FullyShardedDataParallel` instance needs to
        synchronize with the default stream to ensure that the previous
        optimizer step is done.
        """
        if not torch.cuda.is_available() or not self._is_root:
            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 _prefetch_handles(
        self,
        current_handles_key: Tuple[FlatParamHandle, ...],
    ) -> None:
        """
        Prefetches the next handles if needed (without synchronization). An
        empty handles key cannot prefetch.
        """
        if not current_handles_key:
            return
        out = self._get_handles_to_prefetch(current_handles_key)
        if out is not None:
            handles_to_prefetch, current_training_state = out
            # Prefetch the next set of parameters without synchronizing to
            # allow the sync to happen as late as possible to maximize overlap
            self._unshard(handles_to_prefetch)
            self._handles_prefetched[handles_to_prefetch] = True

    def _get_handles_to_prefetch(
        self,
        current_handles_key: Tuple[FlatParamHandle, ...],
    ) -> Optional[Tuple[Tuple[FlatParamHandle, ...], HandleTrainingState]]:
        """
        Returns a tuple of the handles to prefetch and the *current* handles'
        training state if this FSDP instance should prefetch the next FSDP
        instance's parameters and ``None`` otherwise.

        "Prefetching" refers to running the unshard logic early (without
        synchronization), and the "next" FSDP instance depends on the recorded
        execution order and the current training state.
        """
        training_state = self._get_training_state(current_handles_key)
        valid_training_states = (
            HandleTrainingState.PRE_BACKWARD,
            HandleTrainingState.POST_BACKWARD,
        )
        p_assert(
            training_state in valid_training_states,
            f"Prefetching is only supported in {valid_training_states} but "
            f"currently in {training_state}",
        )
        eod = self._exec_order_data
        if (
            training_state == HandleTrainingState.PRE_BACKWARD
            and self.backward_prefetch == BackwardPrefetch.BACKWARD_PRE
        ):
            target_handles_key = eod.get_handles_to_backward_prefetch(
                current_handles_key
            )
            if target_handles_key is not None:
                target_training_state = self._get_training_state(target_handles_key)
                if target_training_state != HandleTrainingState.POST_BACKWARD:
                    return target_handles_key, training_state
        elif (
            training_state == HandleTrainingState.POST_BACKWARD
            and self.backward_prefetch == BackwardPrefetch.BACKWARD_POST
        ):
            target_handles_key = eod.get_handles_to_backward_prefetch(
                current_handles_key
            )
            if target_handles_key is not None:
                target_training_state = self._get_training_state(target_handles_key)
                # TODO (awgu): I think we should check that the target training
                # state is not even pre-backward yet too or else it will have
                # already unsharded, which means the unshard call from the prefetch
                # is a no-op.
                if (
                    target_training_state != HandleTrainingState.POST_BACKWARD
                    and self._need_pre_backward_unshard.get(target_handles_key, False)
                ):
                    return target_handles_key, training_state
        return None

    def _get_training_state(
        self,
        handles_key: Tuple[FlatParamHandle, ...],
    ) -> HandleTrainingState:
        """Returns the training state of the handles in ``handles_key``."""
        p_assert(len(handles_key) > 0, "Expects a non-empty handles key")
        training_states = set(handle._training_state for handle in handles_key)
        p_assert(
            len(training_states) == 1,
            f"Expects uniform training state but got {training_states}",
        )
        return next(iter(training_states))

    @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::

            >>> # xdoctest: +SKIP("undefined variables")
            >>> 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

    @property
    def _param_fqns(self) -> Iterator[Tuple[str, str, str]]:
        for (
            param_name,
            module_name,
        ) in self._fsdp_wrapped_module.handle.parameter_module_names():
            module_name = module_name.replace(f"{FPW_MODULE}.", "")
            module_name = module_name.replace(f"{FPW_MODULE}", "")
            if module_name:
                module_name = f"{module_name}."
            # Activation checkpoint adds a prefix that has to be
            # removed as well.
            module_name = module_name.replace(
                f"{checkpoint_wrapper._CHECKPOINT_PREFIX}.", ""
            )
            fqn = f"{module_name}{param_name}"
            yield fqn, param_name, module_name

    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.
        """
        _replace_by_prefix(state_dict, prefix + f"{FSDP_WRAPPED_MODULE}.", prefix)
        self._assert_state([TrainingState_.SUMMON_FULL_PARAMS])
        # Return early for trivial cases
        if not state_dict or not self._fsdp_wrapped_module.has_params:
            return state_dict

        # If the `FlatParameter` is registered, then this rank only needed to
        # participate in the all-gather but does not actually save the state
        # dict (e.g. when `rank0_only=True` and `self.rank != 0`)
        if hasattr(self._fsdp_wrapped_module, "flat_param"):
            return state_dict

        offload_to_cpu = self._state_dict_config.offload_to_cpu
        cpu_device = torch.device("cpu")

        # Loop only the parameters saved in self._fsdp_wrapped_module to avoid
        # processing buffers.
        for fqn, param_name, module_name in self._param_fqns:
            fqn = f"{prefix}{fqn}"
            clean_key = fqn
            clean_prefix = clean_tensor_name(prefix)
            # Strip prefix out of key if needed as buffer names and param names
            # do not have prefix considered as they are not computed in `state_dict`
            # call.
            if clean_key.startswith(clean_prefix):
                clean_key = clean_key[len(clean_prefix) :]

            # Clone non-ignored parameters before exiting the
            # `_summon_full_params()` context
            assert fqn in state_dict, (
                f"FSDP assumes {fqn} is in the state_dict but the state_dict "
                f"only has {state_dict.keys()}. prefix={prefix}, "
                f"module_name={module_name} param_name={param_name} rank={self.rank}."
            )
            if clean_key not in self._ignored_param_names and not getattr(
                state_dict[fqn], "_has_been_cloned", False
            ):
                try:
                    state_dict[fqn] = state_dict[fqn].clone().detach()
                    state_dict[fqn]._has_been_cloned = True  # type: ignore[attr-defined]
                except BaseException as e:
                    warnings.warn(
                        f"Failed to clone() tensor with name {fqn}. 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 {fqn}. Error: {str(e)}"
                    )

        # 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:
            for clean_key in self._buffer_names:
                # This is a hack to support activation checkpoint.
                clean_key = clean_key.replace(
                    f"{checkpoint_wrapper._CHECKPOINT_PREFIX}.", ""
                )
                fqn = f"{prefix}{clean_key}"
                if state_dict[fqn].device != cpu_device:
                    state_dict[fqn] = state_dict[fqn].to(cpu_device)
        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 not self._fsdp_wrapped_module.has_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._fsdp_wrapped_module, FLAT_PARAM, None)
        # Construct a ShardedTensor from the flat_param.
        full_numel = flat_param._unsharded_size.numel()
        shard_offset = flat_param.numel() * self.rank
        valid_data_size = flat_param.numel() - flat_param._shard_numel_padded
        if valid_data_size > 0 and flat_param._shard_numel_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 not self._fsdp_wrapped_module.has_params:
            return state_dict

        assert self.training_state != TrainingState_.SUMMON_FULL_PARAMS, (
            "Inside _sharded_post_load_state_dict_hook, the training_state must "
            "not be SUMMON_FULL_PARAMS."
        )
        with self._summon_full_params(recurse=False, writeback=False):
            for fqn, _, _ in self._param_fqns:
                # Create a ShardedTensor for the unflattened, non-sharded parameter.
                param = functools.reduce(getattr, fqn.split("."), self.module)
                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)
                ]
                fqn = f"{prefix}{fqn}"
                state_dict[fqn] = init_from_local_shards(
                    local_shards, param.size(), process_group=self.process_group
                )  # type: ignore[assignment]
        state_dict.pop(f"{prefix}{FLAT_PARAM}")
        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::

            >>> # xdoctest: +SKIP("undefined variables")
            >>> 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._buffer_name_to_orig_dtype, 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
            or self._state_dict_type == StateDictType.SHARDED_STATE_DICT
        ):
            if (
                self._fsdp_wrapped_module.flat_param is not None
                and not self._fsdp_wrapped_module.handle.uses_sharded_strategy
            ):
                raise RuntimeError(
                    "sharded_state_dict/local_state_dict can only be called "
                    "when parameters are flatten and sharded."
                )
            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._fsdp_wrapped_module, FLAT_PARAM, None) is None
            ), "No flat parameter in state_dict but self._fsdp_wrapped_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._fsdp_wrapped_module.flat_param
        assert flat_param is not None
        if flat_param._shard_numel_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._shard_numel_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 not self._fsdp_wrapped_module.has_params:
            return

        if not self._fsdp_wrapped_module.handle.uses_sharded_strategy:
            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 (
            param_name,
            _,
            module_name,
        ) in self._fsdp_wrapped_module.handle.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._fsdp_wrapped_module.flat_param
        loaded_flat_param = FlatParamHandle.flatten_params(
            nonsharded_tensors, requires_grad=False
        )

        # Get the chunk from the loaded flat_param for the local rank.
        loaded_flat_param, num_to_pad = FlatParamHandle._get_shard(
            loaded_flat_param,
            self.rank,
            self.world_size,
        )
        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._shard_numel_padded == num_to_pad, (
            f"The loaded local chunk has different padding({num_to_pad}) "
            f"from the local chunk {flat_param._shard_numel_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::

            >>> # xdoctest: +SKIP("undefined variables")
            >>> 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:
        """
        Runs the forward pass for the wrapped module, inserting FSDP-specific
        pre- and post-forward sharding logic.
        """
        with torch.autograd.profiler.record_function(
            "FullyShardedDataParallel.forward"
        ):
            self._lazy_init()
            args, kwargs = self._fsdp_root_pre_forward(*args, **kwargs)
            unused = None
            unshard_fn = functools.partial(
                self._pre_forward_unshard,
                handles=self._handles,
            )
            # Do not free the root's parameters in the post-forward for
            # `FULL_SHARD` with the intention that they are immediately used
            # for backward computation (though this may not be true)
            free_unsharded_flat_params = [
                not self._is_root
                and handle._config.sharding_strategy
                == HandleShardingStrategy.FULL_SHARD
                for handle in self._handles
            ]
            reshard_fn = functools.partial(
                self._reshard,
                self._handles,
                free_unsharded_flat_params,
            )
            self._pre_forward(self._handles, unused, unshard_fn, unused, unused)
            for handle in self._handles:
                assert handle.flat_param.device == self.compute_device
            output = self._fsdp_wrapped_module(*args, **kwargs)
            return self._post_forward(
                self._handles, reshard_fn, unused, unused, unused, output
            )

    def _pre_forward(
        self,
        handles: List[FlatParamHandle],
        reshard_fn: Optional[Callable],
        unshard_fn: Optional[Callable],
        module: nn.Module,
        input: Any,
    ):
        """
        Runs the pre-forward logic. This includes an opportunity to reshard
        currently unsharded parameters such as those from previous forwards; an
        opportunity to unshard currently sharded parameters such as those for
        the current forward; and registering post-backward hooks for these
        current parameters.

        Args:
            handles (List[FlatParamHandle]): Handles giving the parameters
                used in the current forward.
            reshard_fn (Optional[Callable]): A callable to reshard any
                currently unsharded parameters (e.g. from previous forwards) or
                ``None`` to not do any resharding.
            unshard_fn (Optional[Callable]): A callable to unshard any
                currently sharded parameters or ``None`` to not do any
                unsharding.
            module (nn.Module): Unused; expected by the hook signature.
            input (Any): Unused; expected by the hook signature.
        """
        with torch.autograd.profiler.record_function(
            "FullyShardedDataParallel.forward"
        ):
            self.training_state = TrainingState_.FORWARD
            self._exec_order_data.record_pre_forward(handles)
            for handle in handles:
                handle._training_state = HandleTrainingState.FORWARD
            if reshard_fn is not None:
                reshard_fn()
            if unshard_fn is not None:
                unshard_fn()
            # Register post-backward hooks to reshard the parameters and
            # reduce-scatter their gradients. They must be re-registered every
            # forward pass in case the `grad_fn` is mutated.
            self._register_post_backward_hooks(handles)

    def _pre_forward_unshard(
        self,
        handles: List[FlatParamHandle],
    ) -> None:
        """Unshards parameters in the pre-forward including logic for forward
        prefetching."""
        handles_key = tuple(handles)
        self._unshard(handles)
        torch.cuda.current_stream().wait_stream(self._streams["all_gather"])

    def _post_forward(
        self,
        handles: List[FlatParamHandle],
        reshard_fn: Optional[Callable],
        unshard_fn: Optional[Callable],
        module: nn.Module,
        input: Any,
        output: Any,
    ) -> Any:
        """
        Runs the post-forward logic. This includes an opportunity to reshard
        currently unsharded parameters such as those used in the current
        forward; an opportunity to unshard currently sharded parameters such as
        those for the next forward; and registering pre-backward hooks on the
        forward outputs.

        Args:
            handles (List[FlatParamHandle]): Handles giving the parameters
                used in the current forward.
            reshard_fn (Optional[Callable]): A callable to reshard any
                currently unsharded parameters (e.g. from the current forward)
                or ``None`` to not do any resharding.
            unshard_fn (Optional[Callable]): A callable to unshard any
                currently sharded parameters or ``None`` to not do any
                unsharding.
            module (nn.Module): Unused; expected by the hook signature.
            input (Any): Unused; exepcted by the hook signature.
            output (Any): Forward pass output; pre-backward hooks are
                registered on the tensors that require gradients in this
                output.

        Postcondition: Each ``FlatParameter`` 's data points to the sharded
        flattened parameter.
        """
        with torch.autograd.profiler.record_function(
            "FullyShardedDataParallel.forward"
        ):
            self._exec_order_data.record_post_forward(handles)
            if reshard_fn is not None:
                reshard_fn()
            if unshard_fn is not None:
                unshard_fn()
            # Register pre-backward hooks to unshard the flattened parameters
            # for the gradient computation (if needed)
            output = self._register_pre_backward_hooks(output, handles)
            self.training_state = TrainingState_.IDLE
            for handle in handles:
                handle._training_state = HandleTrainingState.IDLE
            return output

    def _cast_forward_inputs(self, *args, **kwargs):
        """Moves the forward inputs to the compute device and casts them to the
        appropriate dtype if needed."""
        # TODO: Do not use the side stream for tensor copies for now;
        # investigate the perf with/without it
        # TODO: For mixed precision, move the inputs to the compute device and
        # cast to reduced-precision in a single `to()` call
        args, kwargs = _to_kwargs(args, kwargs, self.compute_device.index, False)
        args = args[0]
        kwargs = kwargs[0]
        if self._mixed_precision_enabled_for_params():
            input_dtype = self.mixed_precision.param_dtype
            args, kwargs = self._cast_fp_inputs_to_dtype(
                input_dtype,
                *args,
                **kwargs,
            )
        return args, kwargs

    def _fsdp_root_pre_forward(self, *args, **kwargs):
        """
        Runs pre-forward logic specific to the root FSDP instance, which should
        run before any individual module's pre-forward. This includes
        synchronizing with the previous iteration and casting the forward
        inputs appropriately. If this is called on a non-root FSDP instance,
        then the forward inputs are returned directly.
        """
        assert self._is_root is not None, "Expects a root FSDP to have been set"
        if not self._is_root:
            return args, kwargs
        self._exec_order_data.next_iter()
        self._wait_for_previous_optim_step()
        args, kwargs = self._cast_forward_inputs(*args, **kwargs)
        return args, kwargs

    @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 exits;
                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.
        root_fsdp_modules = FullyShardedDataParallel.fsdp_modules(
            module, root_only=True
        )
        # Summon all params for all FSDP instances
        with contextlib.ExitStack() as stack:
            for module in root_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

    @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."
            )

        if recurse:
            with contextlib.ExitStack() as stack:
                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
            return

        torch.cuda.synchronize()
        self._lazy_init()
        self._assert_state([TrainingState_.IDLE])
        for handle in self._handles:
            assert handle._training_state == HandleTrainingState.IDLE
        self.training_state = TrainingState_.SUMMON_FULL_PARAMS
        for handle in self._handles:
            handle._training_state = HandleTrainingState.SUMMON_FULL_PARAMS

        free_unsharded_flat_params = [
            handle.needs_unshard() for handle in self._handles
        ]
        self._unshard(self._handles)
        torch.cuda.current_stream().wait_stream(self._streams["all_gather"])

        if rank0_only and self.rank != 0:
            # Free the unsharded flattened parameter early
            self._reshard(self._handles, free_unsharded_flat_params)
            try:
                yield
            finally:
                self.training_state = TrainingState_.IDLE
                for handle in self._handles:
                    handle._training_state = HandleTrainingState.IDLE
        else:
            # Unflatten the unsharded flattened parameters
            with contextlib.ExitStack() as stack:
                # Invariant: rank == 0 or !rank0_only
                for handle in self._handles:
                    if offload_to_cpu and handle.uses_sharded_strategy:
                        stack.enter_context(handle.to_cpu())
                # TODO (awgu): This FPW call assumes 1 `FlatParameter`
                stack.enter_context(self._fsdp_wrapped_module.unflatten_as_params())
                try:
                    yield
                finally:
                    stack.close()
                    if writeback:
                        self._write_back_to_local_shard(self._handles)
                    self._reshard(self._handles, free_unsharded_flat_params)
                    self.training_state = TrainingState_.IDLE
                    for handle in self._handles:
                        handle._training_state = HandleTrainingState.IDLE

    @torch.no_grad()
    def _write_back_to_local_shard(self, handles: List[FlatParamHandle]):
        """
        For each handle, writes back the this rank's shard of the unsharded
        flattened parameter to the sharded flattened parameter.

        Precondition: Each handle's ``FlatParameter`` 's data points to the
        padded unsharded flattened parameter.
        """
        for handle in handles:
            # For `NO_SHARD`, `_local_shard` is the unsharded flattened
            # parameter as well
            if not handle.uses_sharded_strategy:
                continue
            assert (
                handle.flat_param.ndim == 1
            ), f"Expects `flat_param` to be flattened but got {handle.flat_param.shape}"
            # Get the unpadded shard instead of the padded shard to persist
            # user changes to the padding (though FSDP does not explicitly
            # support this)
            shard, _ = FlatParamHandle._get_unpadded_shard(
                handle.flat_param, handle.rank, handle.world_size
            )
            handle.flat_param._local_shard[: shard.numel()].copy_(shard)

    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,
        handles: List[FlatParamHandle],
    ) -> Any:
        """
        Registers pre-backward hooks on the tensors that require gradients in
        the forward pass outputs ``outputs``, which were computed using the
        ``FlatParameter`` s of ``handles``.

        Returns:
            Forward pass outputs with pre-backward hooks registered to tensors
            that require gradients.
        """
        # If there is no gradient computation, then there is no need for
        # pre-backward logic
        if not torch.is_grad_enabled():
            return outputs

        if self._is_root:
            self._post_backward_callback_queued = False  # only defined on the root

        # Since these handles' `FlatParameter`s participated in a forward,
        # we conservatively assume that they will be used in the backward
        handles_key = tuple(handles)
        self._need_pre_backward_unshard[handles_key] = False
        self._pre_backward_hook_has_run[handles_key] = False

        def _pre_backward_hook(_handles: List[FlatParamHandle], *unused: Any) -> None:
            """Prepares ``_handles`` 's ``FlatParameter`` s for gradient
            computation."""
            _handles_key = tuple(_handles)  # avoid shadowing `handles_key`
            # Only run the pre-backward hook once per group of handles involved
            # in the same module forward computation
            if self._pre_backward_hook_has_run[_handles_key]:
                return

            with torch.autograd.profiler.record_function(
                "FullyShardedDataParallel._pre_backward_hook"
            ):
                # Queue the post-backward callback once for the root FSDP
                # instance to attach it to the outermost backward graph task so
                # that it is called after all backward calls complete
                if self._is_root:
                    self._queue_wait_for_post_backward()

                self._assert_state([TrainingState_.IDLE])
                self.training_state = TrainingState_.BACKWARD_PRE
                for handle in _handles:
                    handle._training_state = HandleTrainingState.PRE_BACKWARD

                params_prefetched = self._handles_prefetched.get(_handles_key, False)
                if not params_prefetched:
                    self._unshard(_handles)
                torch.cuda.current_stream().wait_stream(self._streams["all_gather"])

                # TODO (awgu): This is not necessarily the best place to reset
                # the flag, but it works because the pre-backward hook logic is
                # guaranteed to only run once per iteration.
                self._handles_prefetched.pop(_handles_key, None)
                self._prefetch_handles(_handles_key)
                for handle in self._handles:
                    handle.prepare_gradient()
                self._pre_backward_hook_has_run[_handles_key] = True

        def _register_hook(t: torch.Tensor) -> torch.Tensor:
            if t.requires_grad:
                t.register_hook(functools.partial(_pre_backward_hook, handles))
                self._need_pre_backward_unshard[handles_key] = True
            return t

        return _apply_to_tensors(_register_hook, outputs)

    def _register_post_backward_hooks(
        self,
        handles: List[FlatParamHandle],
    ) -> None:
        """
        Registers post-backward hooks on the ``FlatParameter`` s'
        ``AccumulateGrad`` objects to reshard and to reduce-scatter gradients.

        The ``AccumulateGrad`` object represents the last function that
        finalizes the ``FlatParameter`` 's gradient, so it only runs after its
        entire gradient computation has finished.

        We register the post-backward hook only once in the *first* forward
        that a ``FlatParameter`` participates in. This relies on the
        ``AccumulateGrad`` object being preserved through multiple forwards.
        """
        # If there is no gradient computation, then there is no need for
        # post-backward logic
        if not torch.is_grad_enabled():
            return
        for handle in handles:
            flat_param = handle.flat_param
            already_registered = hasattr(flat_param, "_post_backward_hook_state")
            if already_registered or not flat_param.requires_grad:
                continue
            # Get the `AccumulateGrad` object
            temp_flat_param = flat_param.expand_as(flat_param)
            assert (
                temp_flat_param.grad_fn is not None
            ), "The `grad_fn` is needed to access the `AccumulateGrad` and register the post-backward hook"
            acc_grad = temp_flat_param.grad_fn.next_functions[0][0]
            hook_handle = acc_grad.register_hook(
                functools.partial(self._post_backward_hook, handle)
            )
            # TODO (awgu): We should probably store this state on the handle
            # and not on the `FlatParameter` itself.
            flat_param._post_backward_hook_state = (acc_grad, hook_handle)  # type: ignore[attr-defined]

    @torch.no_grad()
    def _post_backward_hook(
        self,
        handle: FlatParamHandle,
        *unused: Any,
    ) -> None:
        """
        Reduce-scatters the gradient of ``handle`` 's ``FlatParameter``.

        Precondition: The ``FlatParameter`` 's ``.grad`` attribute contains the
        unsharded gradient for the local batch.

        Postcondition:
        - If using ``NO_SHARD``, then the ``.grad`` attribute is the reduced
        unsharded gradient.
        - Otherwise, the ``_saved_grad_shard`` attribute is the reduced sharded
        gradient (accumulating with any existing gradient).
        """
        param = handle.flat_param
        param._post_backward_called = True
        with torch.autograd.profiler.record_function(
            "FullyShardedDataParallel._post_backward_hook"
        ):
            # 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
            handle._training_state = HandleTrainingState.POST_BACKWARD

            if (
                self._use_param_exec_order_policy()
                and self._param_exec_order_prep_stage
            ):
                # In self._fsdp_params_exec_order, the parameters are ordered based on
                # the execution order in the backward pass in the first iteration.
                self._fsdp_params_exec_order.append(param)

            if param.grad is None:
                return
            if param.grad.requires_grad:
                raise RuntimeError(
                    "FSDP only works with gradients that don't require gradients"
                )

            free_unsharded_flat_param = (
                (self._sync_gradients and handle.uses_sharded_strategy)
                or handle._config.sharding_strategy == HandleShardingStrategy.FULL_SHARD
            )
            self._reshard([handle], [free_unsharded_flat_param])

            # TODO (awgu): Post-backward prefetching does not support the
            # multiple handles per module case (which was why we keyed by
            # *tuple*). The post-backward hook runs per handle, not per group
            # of handles. To generalize this, we may need a 2-level mapping,
            # where we map each individual handle to its groups of handles and
            # then from the groups of handles to their indices in the order.
            handles_key = (handle,)
            self._prefetch_handles(handles_key)

            if not self._sync_gradients:
                return

            # Wait for all ops in the current stream (e.g. gradient
            # computation) to finish before reduce-scattering the gradient
            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()
                    and not self._low_precision_hook_enabled()
                ):
                    # Cast gradient to precision in which it should be communicated.
                    # If a low precision hook is registered and reduce_dtype is specified
                    # in `MixedPrecision`, communication hook will take care of
                    # casting to lower precision and back.
                    # 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
                    and self.communication_hook is None
                ):
                    # Average grad by pre-division factor. Together pre- and post-division factors
                    # lead to an overall averaging by world_size, required for consistency with PyTorch DDP.
                    # This is a two-step process to avoid potential underflow and overflow.
                    param.grad.div_(self.gradient_predivide_factor)

                grad = param.grad.data
                if handle.uses_sharded_strategy:
                    # 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 pre-division factor. Together pre- and post-division factors
                        # lead to an overall averaging by world_size, required for consistency with PyTorch DDP.
                        # This is a two-step process to avoid potential underflow and overflow.
                        output.div_(self.gradient_postdivide_factor)

                    self._cast_grad_to_param_dtype(output, param)

                    # 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:
                    if self.sharding_strategy == ShardingStrategy.NO_SHARD:
                        # if a communication hook was not registered,
                        # then a default hook (`all_reduce`) will be used
                        self.communication_hook(
                            self.communication_hook_state, param.grad
                        )

                    self._cast_grad_to_param_dtype(param.grad, param)

                # 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 handle._config.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 _cast_grad_to_param_dtype(
        self,
        grad: torch.Tensor,
        param: FlatParameter,
    ):
        """
        Casts gradient ``grad`` back to the full parameter dtype so that the
        optimizer step runs with that dtype. This performs an actual cast if
        1. parameters were in reduced precision during the forward since then
        gradients would be in that reduced precision, or
        2. parameters were not in reduced precision but gradients were in
        reduced precision for communication.
        However, if a low precision communication hook is registered, then this
        dtype cast happens in the hook instead.
        """
        self._assert_state(TrainingState_.BACKWARD_POST)
        if not self._low_precision_hook_enabled() and (
            self._mixed_precision_enabled_for_params()
            or self._mixed_precision_enabled_for_reduce()
        ):
            low_prec_grad_data = grad.data
            grad.data = grad.data.to(dtype=param.dtype)
            # Do not let the low precision gradient memory get reused until
            # the cast to full parameter precision completes
            low_prec_grad_data.record_stream(torch.cuda.current_stream())

    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:
            # TODO (awgu): I think it makes more sense for the state to
            # transition to pre-backward the moment the execution enters the
            # pre-backward. Right now, we call this method and then transition,
            # which is why this assert is for `IDLE`.
            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._sync_gradients:
            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.
        def _finalize_params(fsdp_module: FullyShardedDataParallel) -> None:
            """Helper used below on all fsdp modules."""
            for handle in fsdp_module._handles:
                p = handle.flat_param
                if p.requires_grad:
                    if hasattr(p, "_post_backward_hook_state"):
                        assert len(p._post_backward_hook_state) == 2 and len(  # type: ignore[attr-defined]
                            p._post_backward_hook_state  # type: ignore[attr-defined]
                        ), (  # type: ignore[attr-defined]
                            "p._post_backward_hook_state fields are not valid."
                        )
                        p._post_backward_hook_state[1].remove()  # type: ignore[attr-defined]
                        delattr(p, "_post_backward_hook_state")
                    # 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._sync_gradients:
                        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 handle.uses_sharded_strategy
                            or not p._post_backward_called,
                            "All sharded parameters that received a gradient "
                            "should use `_saved_grad_shard`",
                        )
                    if hasattr(p, "_saved_grad_shard"):
                        delattr(p, "_saved_grad_shard")

                    p_assert(
                        hasattr(p, "_post_backward_called"),
                        "Expected flag _post_backward_called to be set on param.",
                    )
                    # Reset _post_backward_called in preparation for the next iteration.
                    p._post_backward_called = False

        # Update root and nested FSDP's hooks and flags.
        for m in self.modules():  # includes self
            if isinstance(m, FullyShardedDataParallel):
                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`.
                    managed_param_requires_grad = any(p.requires_grad for p in m.params)
                    if managed_param_requires_grad:
                        p_assert(
                            all(hasattr(p, "_post_backward_called") for p in m.params),
                            "Expected all params to have flag _post_backward_called set!",
                        )
                        post_backward_hook_called = any(
                            p._post_backward_called for p in m.params
                        )
                        if post_backward_hook_called:
                            m._assert_state(TrainingState_.BACKWARD_POST)
                        else:
                            # post backward hook was not called, meaning param
                            # did not have a gradient computed. It was either unused
                            # in forward, or unused in loss computation so it did
                            # not get gradient
                            m._assert_state(
                                [TrainingState_.BACKWARD_PRE, TrainingState_.IDLE]
                            )
                    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])

                _finalize_params(m)
                m._pre_backward_hook_has_run.clear()
                m.training_state = TrainingState_.IDLE
                for handle in m._handles:
                    handle._training_state = HandleTrainingState.IDLE

                # TODO (awgu): We could replace this with an assert, but it may
                # trigger for dynamic execution order across ranks. If the
                # backward prefetch mismatches the pre-backward order, we may
                # incorrectly have some parameters marked as prefetched (from
                # the pre-backward) when they have been resharded in the
                # post-backward already.
                m._handles_prefetched.clear()

                if m._is_root:
                    # reset this flag for cases like "one forward pass + multiple backward passes"
                    self._post_backward_callback_queued = False

        if self._use_param_exec_order_policy() and self._param_exec_order_prep_stage:
            self._param_exec_order_policy_second_iter_init()

    def _param_exec_order_policy_second_iter_init(self) -> None:
        self._param_exec_order_prep_stage = False
        # Let the parameters in self._fsdp_params_exec_order ordered based on
        # the execution order in the forward pass.
        self._fsdp_params_exec_order.reverse()
        for m in self.modules():
            if m is not self and isinstance(m, FullyShardedDataParallel):
                assert hasattr(
                    m, "_param_exec_order_policy"
                ), "Non-root FSDP modules should also have _param_exec_order_policy attribute"
                assert hasattr(
                    m, "_param_exec_order_prep_stage"
                ), "Non-root FSDP modules should also have _param_exec_order_prep_stage attribute"
                m._param_exec_order_prep_stage = False
        # TODO (linjianma): Construct a fsdp_wrap_map whose keys are all children modules with a FSDP wrap,
        # and values are its FSDP wraps. These children FSDP wraps will be detached from the root FSDP module
        # and will be used to schedule the parameters (rebuild_full_params and reshard).
        # TODO (linjianma): Remove all internal FSDP wraps from the root FSDP module.
        # TODO (linjianma): Based on self._fsdp_params_exec_order, get the information
        # needed to patch the forward() function of each key in the fsdp_wrap_map. The rules are as follows:
        # 1: Before each forward(), rebuild_full_params of all parameters that are currently sharded and
        # will be used in the forward, and reshard all parameters that are currently full and will not be
        # used in the next forward()
        # 2: After each forward(), reshard all parameters just used in the forward, and rebuild_full_params of
        # all parameters that will be used next.
        # TODO (linjianma): Patch the forward of each model in the keys
        # of fsdp_wrap_map based on the information above.

    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._sync_gradients))
                m._sync_gradients = False
        try:
            yield
        finally:
            for m, old_flag in old_flags:
                assert not m._sync_gradients, (
                    "`_sync_gradients` was incorrectly set to "
                    "`True` while in the `no_sync()` context manager"
                )
                m._sync_gradients = 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.
        """
        self._lazy_init()
        self._wait_for_previous_optim_step()
        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,
        group=None,
    ) -> 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``)
            group (dist.ProcessGroup): Model's process group or ``None`` if using
                the default process group. (Default: ``None``)

        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"]
        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

        # Construct the local mapping between unflattened parameter names
        # (`_OptimStateKey`s) and parameter IDs and broadcast rank 0's mapping
        param_to_unflat_param_names: Dict[
            torch.nn.Parameter, List[str]
        ] = _get_param_to_unflat_param_names(model)
        flat_param_id_to_param: List[torch.nn.Parameter] = _get_param_id_to_param(
            model, optim_input
        )
        optim_state_key_to_flat_param_id: Dict[_OptimStateKey, int] = {}  # local
        r0_flat_param_id_to_optim_state_key: Dict[
            int, _OptimStateKey
        ] = collections.OrderedDict()  # rank 0
        for flat_param_id, param in enumerate(flat_param_id_to_param):
            # Do not include parameters without state to avoid empty mappings
            # just like in normal `torch.optim.Optimizer.state_dict()`
            if flat_param_id not in osd_state:
                continue
            optim_state_key = _OptimStateKey(
                unflat_param_names=tuple(param_to_unflat_param_names[param]),
                is_flat_param=isinstance(param, FlatParameter),
            )
            if rank == 0:
                r0_flat_param_id_to_optim_state_key[flat_param_id] = optim_state_key
            optim_state_key_to_flat_param_id[optim_state_key] = flat_param_id
        obj_list = [r0_flat_param_id_to_optim_state_key] if rank == 0 else [None]
        dist.broadcast_object_list(obj_list, src=0, group=group)
        r0_flat_param_id_to_optim_state_key = obj_list[0]

        # Ensure that all ranks have at least the optimizer states needed by
        # rank 0's optimizer
        missing_keys: List[_OptimStateKey] = []
        for r0_optim_state_key in r0_flat_param_id_to_optim_state_key.values():
            if r0_optim_state_key not in optim_state_key_to_flat_param_id:
                # A parameter from rank 0's optimizer does not exist for this
                # rank's optimizer
                missing_keys.append(r0_optim_state_key)
                continue
            flat_param_id = optim_state_key_to_flat_param_id[r0_optim_state_key]
            assert flat_param_id >= 0 and flat_param_id < len(
                flat_param_id_to_param
            ), "Check the `flat_param_id_to_param` construction"
        device = torch.device("cuda", torch.cuda.current_device())
        num_missing = torch.tensor(
            [len(missing_keys)], dtype=torch.int32, device=device
        )
        dist.all_reduce(num_missing, group=group)
        if num_missing.item() > 0:
            obj_list = [None for _ in range(dist.get_world_size(group))]
            dist.all_gather_object(obj_list, missing_keys, group=group)
            error_msg = (
                "FSDP currently requires each rank to have at least the "
                "optimizer states needed by rank 0's optimizer but some ranks "
                "are missing some of those states"
            )
            for rank, keys in enumerate(obj_list):
                if len(keys) > 0:
                    error_msg += (
                        f"\nRank {rank} is missing states for the parameters: "
                        f"{[key.unflat_param_names for key in keys]}"
                    )
            raise RuntimeError(error_msg)

        # Iterate in rank 0's flattened parameter ID order to ensure aligned
        # all-gathers across ranks
        flat_param_to_fsdp_module = _get_flat_param_to_fsdp_module(model)
        for r0_optim_state_key in r0_flat_param_id_to_optim_state_key.values():
            flat_param_id = optim_state_key_to_flat_param_id[r0_optim_state_key]
            param = flat_param_id_to_param[flat_param_id]
            if r0_optim_state_key.is_flat_param:
                fsdp_module = flat_param_to_fsdp_module[param]
                unflat_state = _unflatten_optim_state(
                    param,
                    osd_state[flat_param_id],
                    fsdp_module,
                    to_save,
                )
                if to_save:
                    assert len(unflat_state) == len(
                        r0_optim_state_key.unflat_param_names
                    )
                    for unflat_param_name, unflat_param_state in zip(
                        r0_optim_state_key.unflat_param_names,
                        unflat_state,
                    ):
                        full_osd_state[unflat_param_name] = unflat_param_state
            elif to_save:
                assert len(r0_optim_state_key.unflat_param_names) == 1
                unflat_param_name = r0_optim_state_key.unflat_param_names[0]
                full_osd_state[unflat_param_name] = copy.copy(osd_state[flat_param_id])
                for state_name, value in full_osd_state[unflat_param_name].items():
                    if torch.is_tensor(value):
                        full_osd_state[unflat_param_name][state_name] = value.cpu()

        if not to_save:
            return {}

        # 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::

            >>> # xdoctest: +SKIP("undefined variables")
            >>> 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.
        """
        sharded_osd = _flatten_full_optim_state_dict(
            full_optim_state_dict,
            model,
            True,
        )
        return _rekey_sharded_optim_state_dict(sharded_osd, model, optim_input)

    @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::

            >>> # xdoctest: +SKIP("undefined variables")
            >>> 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()``. (Default: ``None``)
            group (dist.ProcessGroup): 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 = _flatten_full_optim_state_dict(
                full_optim_state_dict, model, False
            )
            processed_osd = _process_pos_dim_tensor_state(flat_osd, 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 = _broadcast_processed_optim_state_dict(
            processed_osd 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,
            flat_osd if rank == 0 else None,
            rank,
            world_size,
            group,
            broadcast_device,
        )
        # Rekey the optimizer state dict to use parameter IDs according to this
        # rank's `optim_input`
        sharded_osd = _rekey_sharded_optim_state_dict(sharded_osd, model, optim_input)
        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::

            >>> # xdoctest: +SKIP("undefined variables")
            >>> 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::

            >>> # xdoctest: +SKIP("undefined variables")
            >>> 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_comm_hook(self) -> Any:
        r"""
        Returns a default communication hook based on a sharding strategy.
        """
        if self.sharding_strategy != ShardingStrategy.NO_SHARD:
            return None
        else:
            return default_hooks.allreduce_hook

    def _get_default_comm_hook_state(self) -> Any:
        r"""
        Returns a default communication hook state based on a sharding strategy.
        """
        if self.sharding_strategy != ShardingStrategy.NO_SHARD:
            return None
        else:
            return default_hooks.DefaultState(process_group=self.process_group)

    def register_comm_hook(self, state: object, hook: callable):
        """
        Registers a communication hook which is an enhancement that provides a
        flexible hook to users where they can specify how FSDP aggregates gradients
        across multiple workers.
        This hook can be used to implement several algorithms like
        `GossipGrad <https://arxiv.org/abs/1803.05880>`_ and gradient compression
        which involve different communication strategies for
        parameter syncs while training with :class:`FullyShardedDataParallel`.

        .. warning::
            FSDP only support communication hooks for a ``NO_SHARD`` strategy at this time.
            If other strategies are used, an error will be raised.

        .. warning ::
            FSDP communication hook should be registered before running an initial forward pass
            and only once.

        Args:
            state (object): Passed to the hook to maintain any state information during the training process.
                            Examples include error feedback in gradient compression,
                            peers to communicate with next in `GossipGrad <https://arxiv.org/abs/1803.05880>`_, etc.
                            It is locally stored by each worker
                            and shared by all the gradient tensors on the worker.
            hook (Callable): Callable with the following signature:
                            ``hook: Callable[torch.Tensor] -> None``:
                            This function takes in a Python tensor, which represents
                            the full, flattened, unsharded gradient with respect to all variables
                            corresponding to the model this FSDP unit is wrapping
                            (that are not wrapped by other FSDP sub-units).
                            It then performs all necessary processing and returns ``None``.

        """
        if not self.check_is_root():
            raise AssertionError(
                "register_comm_hook can only be called on a root instance."
            )
        if self.sharding_strategy != ShardingStrategy.NO_SHARD:
            raise NotImplementedError(
                "Communication hooks are currently only available for a NO_SHARD strategy."
            )
        else:
            # register same hook for root and all submodules
            for submodule in self.fsdp_modules(self):
                assert (
                    not submodule._hook_registered
                ), "communication hook can be only registered once"
                submodule._hook_registered = True
                assert (
                    submodule.communication_hook == self._get_default_comm_hook()
                ), f"communication hook should be default, but it is {submodule.communication_hook.__name__} instead"
                submodule.communication_hook_state = state
                submodule.communication_hook = hook

    def _init_param_exec_order_wrap_policy(self, *args, **kwargs) -> None:
        auto_wrap_policy = kwargs["auto_wrap_policy"]
        module = kwargs["module"]
        assert hasattr(auto_wrap_policy, "tracing_config")
        if not _TORCH_FX_AVAIL:
            assert (
                auto_wrap_policy.tracing_config is None
            ), "tracing_config should be None when torch.fx is not enabled"
        elif isinstance(auto_wrap_policy.tracing_config, TracingConfig):
            tracer = auto_wrap_policy.tracing_config.tracer
            execution_info = _init_execution_info(module)

            for m in module.modules():
                assert not isinstance(
                    m, FullyShardedDataParallel
                ), "The input module of _patch_tracer should not contain FSDP modules"

            with _patch_tracer(
                tracer=tracer,
                root_module=module,
                execution_info=execution_info,
            ):
                try:
                    tracer.trace(module, auto_wrap_policy.tracing_config.concrete_args)
                except BaseException as e:
                    raise RuntimeError(
                        "tracer.trace failed inside _init_param_exec_order_wrap_policy"
                        f" with the error: {e}."
                    )
        else:
            assert (
                auto_wrap_policy.tracing_config is None
            ), "tracing_config should either be an instance of TracingConfig or be None"
        # The initial FSDP wrapping is done with auto_wrap_policy.init_policy
        kwargs["auto_wrap_policy"] = auto_wrap_policy.init_policy
        self.__init__(*args, **kwargs)
        self._param_exec_order_policy: bool = True
        # self._param_exec_order_prep_stage is set to True before we get the execution order
        self._param_exec_order_prep_stage: bool = True
        # A list that stores the flatten parameters and its name based on the parameter execution order
        self._fsdp_params_exec_order: List[FlatParameter] = []
        if _TORCH_FX_AVAIL and isinstance(
            auto_wrap_policy.tracing_config, TracingConfig
        ):
            # Initialize a dict that maps each module to its parent FSDP wrap
            module_to_fsdp: Dict[nn.Module, FullyShardedDataParallel] = dict()
            for wrap in self.fsdp_modules(self):
                module_to_fsdp[wrap.module] = wrap
            # Set self._fsdp_params_exec_order based on execution_info.module_forward_order.
            # TODO (linjianma): self._fsdp_params_exec_order will be set based on
            # the parameter execution order rather than module_forward_order,
            # once the non-recursive wrapping policy is fully implemented.
            for m in execution_info.module_forward_order:
                if m in module_to_fsdp:
                    for flat_param in module_to_fsdp[m].params:
                        self._fsdp_params_exec_order.append(flat_param)
            self._param_exec_order_prep_stage = False

        for m in self.modules():
            if m is not self and isinstance(m, FullyShardedDataParallel):
                # Assignment by reference, so each children FSDP wrap has access to
                # the _fsdp_params_exec_order of the root module
                m._fsdp_params_exec_order = self._fsdp_params_exec_order
                m._param_exec_order_policy = self._param_exec_order_policy
                m._param_exec_order_prep_stage = self._param_exec_order_prep_stage

    def _use_param_exec_order_policy(self) -> bool:
        return (
            hasattr(self, "_param_exec_order_policy") and self._param_exec_order_policy
        )

    def _is_param_exec_order_prep_stage(self) -> bool:
        is_prep_stage = (
            hasattr(self, "_param_exec_order_prep_stage")
            and self._param_exec_order_prep_stage
        )
        if not is_prep_stage:
            for p in self.parameters():
                assert not hasattr(
                    p, "_params_exec_order_hook_handle"
                ), "When not in execution order prep stage, all _params_exec_order_hook_handle should be removed."
        return is_prep_stage


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 _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 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):
                module_prefixed_param_names = (
                    param._prefixed_param_names
                    if isinstance(param, FlatParameter)
                    else [param_name]
                )  # prefixed from `module`
                fully_prefixed_param_names = [
                    clean_tensor_name(prefix + name)
                    for name in module_prefixed_param_names
                ]  # fully prefixed from the top level including `prefix`
                # 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] = fully_prefixed_param_names
                elif not dedup_shared_params:
                    param_to_unflat_param_names[param].extend(
                        fully_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 module wrapper
    prefixes."""
    # Call `replace()` twice separately since the name may not have both
    tensor_name = tensor_name.replace(FSDP_WRAPPED_MODULE + ".", "")
    tensor_name = tensor_name.replace(FPW_MODULE + ".", "")
    # TODO: Explicitly replacing checkpoint_wrapper prefix is not ideal,
    # as it increases coupling between CheckpointWrapper and FSDP. This is also not
    # scalable for additional wrapped modules, we should come up with a general solution
    # for this issue.
    tensor_name = tensor_name.replace(_CHECKPOINT_PREFIX + ".", "")
    return tensor_name