# mypy: allow-untyped-decorators
# mypy: allow-untyped-defs

from __future__ import annotations

import functools
from contextlib import contextmanager
from typing import Any, cast, Literal, NoReturn, overload, TYPE_CHECKING
from typing_extensions import deprecated

import torch
import torch.nn as nn
from torch.distributed._composable import contract

from ._fsdp_api import (
    AllGather,
    DataParallelMeshDims,
    MixedPrecisionPolicy,
    OffloadPolicy,
    ReduceScatter,
)
from ._fsdp_common import FSDPMeshInfo, ShardPlacementFnResult
from ._fsdp_init import (
    _apply_to_module,
    _get_device_from_mesh,
    _get_mesh_info,
    _get_modules_and_states,
    _get_post_forward_mesh_info,
    _init_default_mesh,
    _init_param_group,
    _validate_mesh,
    _validate_module,
)
from ._fsdp_state import _get_module_fsdp_state, FSDPState


if TYPE_CHECKING:
    from collections.abc import Callable, Iterable, Iterator

    from torch.distributed.tensor import DeviceMesh

    from ._fsdp_param_group import FSDPParamGroup

__all__ = [
    "fully_shard",
    "FSDPModule",
    "UnshardHandle",
    "register_fsdp_forward_method",
    "get_cls_to_fsdp_cls",
    "disable_fsdp_module_new_init",
    "share_comm_ctx",
]


cls_to_fsdp_cls: dict[type, type] = {}


def get_cls_to_fsdp_cls() -> dict[type, type]:
    return cls_to_fsdp_cls


@overload
# pyrefly: ignore [inconsistent-overload]
def fully_shard(
    module: nn.Module,
    *,
    mesh: DeviceMesh | None = ...,
    reshard_after_forward: bool | int | None = ...,
    shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None = ...,
    mp_policy: MixedPrecisionPolicy = ...,
    offload_policy: OffloadPolicy = ...,
    ignored_params: set[nn.Parameter] | None = ...,
    dp_mesh_dims: DataParallelMeshDims | None = ...,
) -> FSDPModule: ...


@overload
# pyrefly: ignore [inconsistent-overload]
def fully_shard(
    module: list[nn.Module],
    *,
    mesh: DeviceMesh | None = ...,
    reshard_after_forward: bool | int | None = ...,
    shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None = ...,
    mp_policy: MixedPrecisionPolicy = ...,
    offload_policy: OffloadPolicy = ...,
    ignored_params: set[nn.Parameter] | None = ...,
    dp_mesh_dims: DataParallelMeshDims | None = ...,
) -> list[FSDPModule]: ...


# The decorator adds a state object to `module` that can be accessed via
# `fully_shard.state(module)`. The state object and module are 1:1.
# [1] Python runtime decorator does not play well with static type checking
# so suppressing some type checks to support type overloads
# such that caller can still get correct return types based on input type
@contract(state_cls=FSDPState)  # type: ignore[misc] # see [1]
def fully_shard(
    module,
    *,
    mesh: DeviceMesh | None = None,
    reshard_after_forward: bool | int | None = None,
    shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None = None,
    mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
    offload_policy: OffloadPolicy = OffloadPolicy(),
    ignored_params: set[nn.Parameter] | None = None,
    dp_mesh_dims: DataParallelMeshDims | None = None,
):
    """
    Apply fully sharded data parallelism (FSDP) to ``module``, where FSDP
    shards module parameters, gradients, and optimizer states across data
    parallel workers to save memory at the cost of communication.

    At initialization, FSDP shards the module's parameters across the data
    parallel workers given by ``mesh``. Before forward, FSDP all-gathers the
    sharded parameters across the data-parallel workers to get the unsharded
    parameters for forward computation. If ``reshard_after_forward`` is
    ``True``, then FSDP frees the unsharded parameters after forward and
    re-all-gathers them in backward before gradient computation. After gradient
    computation, FSDP frees the unsharded parameters and reduce-scatters the
    unsharded gradients across data-parallel workers.

    This implementation represents the sharded parameters as :class:`DTensor` s
    sharded on dim-0, while the unsharded parameters will be like the original
    parameters on ``module`` (e.g. :class:`torch.Tensor` if originally
    :class:`torch.Tensor`). A module
    `forward pre-hook <https://pytorch.org/docs/main/generated/torch.nn.Module.html#torch.nn.Module.register_forward_pre_hook>`_
    on ``module`` all-gathers the parameters, and a module
    `forward hook <https://pytorch.org/docs/main/generated/torch.nn.Module.html#torch.nn.Module.register_forward_hook>`_
    on ``module`` frees them (if needed). Similar backward hooks all-gather
    parameters and later free parameters and reduce-scatter gradients.

    Since grouping multiple tensors together for one collective is critical for
    communication efficiency, this implementation makes this grouping first
    class. Calling :meth:`fully_shard` on ``module`` constructs one group that
    includes the parameters in ``module.parameters()`` except those already
    assigned to a group from an earlier call on a submodule. This means that
    :meth:`fully_shard` should be called bottom-up on your model. Each group's
    parameters are all-gathered in one collective, and its gradients are
    reduce-scattered in one collective. Partitioning the model into multiple
    groups ("layer by layer") allows for peak memory savings and communication/computation
    overlap. Users generally should *not* call :meth:`fully_shard` only on the
    topmost root module.

    Args:
        module (Union[nn.Module, List[nn.Module]): The module or modules to
            shard with FSDP and group together for communication.
        mesh (Optional[DeviceMesh]): This data parallel mesh defines the
            sharding and device. If 1D, then parameters are fully sharded
            across the 1D mesh (FSDP) with ``(Shard(0),)`` placement. If 2D,
            then parameters are sharded across the 1st dim and replicated
            across the 0th dim (HSDP) with ``(Replicate(), Shard(0))``
            placement. The mesh's device type gives the device type used for
            communication; if a CUDA or CUDA-like device type, then we use the
            current device.
        reshard_after_forward (Optional[Union[bool, int]]): This controls the parameter
            behavior after forward and can trade off memory and communication:

            - If ``True``, then this reshards parameters after forward and
              re-all-gathers in backward.
            - If ``False``, then this keeps the unsharded parameters in memory
              after forward and avoids the all-gather in backward. For best performance,
              we usually set ``False`` for the root module, because the root module
              is typically required immediately when the backward pass begins.
            - If ``None``, it is set to ``True`` for non-root modules and ``False``
              for root modules.
            - If an ``int``, then this represents the world size to reshard to
              after forward. It should be a non-trivial divisor of the ``mesh``
              shard dim size (i.e. excluding 1 and the dim size itself). A
              choice may be the intra-node size (e.g. ``torch.cuda.device_count()``).
              This allows the all-gather in backward to be over a smaller world
              size at the cost of higher memory usage than setting to ``True``.
            - After forward, the parameters registered to the module depend on
              to this: The registered parameters are the sharded parameters if
              ``True``; unsharded parameters if ``False``; and the parameters
              resharded to the smaller mesh otherwise. To modify the parameters
              between forward and backward, the registered parameters must be
              the sharded parameters. For ``False`` or an ``int``, this can be
              done by manually resharding via :meth:`reshard`.
        shard_placement_fn (Optional[Callable[[nn.Parameter], Optional[Shard | ShardPlacementResult]]]):
            This callable can be used to override the sharding placement and/or
            mesh for a parameter. It can return:

            - ``None``: Use default sharding (Shard(0)) on the mesh passed to
              ``fully_shard``.
            - :class:`Shard`: Shard the parameter on the specified dimension
              using the mesh passed to ``fully_shard``.
            - :class:`ShardPlacementResult`: Specify both the shard placement
              and a custom :class:`FSDPMeshInfo`. This allows different
              parameters to be sharded across different process groups, enabling
              use cases like Mixture of Experts where expert params use a
              different mesh than regular params.

            If sharding on a nonzero dim, we currently require even sharding,
            i.e. the tensor dim size on that dim must be divisible by the FSDP
            shard mesh size.
        mp_policy (MixedPrecisionPolicy): This controls the mixed precision
            policy, which offers parameter/reduction mixed precision for this
            module. See :class:`MixedPrecisionPolicy` for details.
        offload_policy (OffloadPolicy): This controls the offloading policy,
            which offers parameter/gradient/optimizer state offloading. See
            :class:`OffloadPolicy` and its subclasses for details.
        ignored_params: Optional(Set[nn.Parameter]): The set of parameters to be
            ignored by FSDP. They will not be sharded, nor moved to the device
            during init, nor have their gradients reduced in backward.
        dp_mesh_dims (Optional[DataParallelMeshDims]): When provided,
            ``mesh`` is treated as the full SPMD mesh, and parameters should be
            DTensors on this mesh with ``Replicate()`` on all DP dimensions.
            The ``shard`` field names which dim(s) FSDP shards on (multiple
            dims are flattened). The ``replicate`` field names the HSDP
            replication dim(s) (multiple dims are flattened).

    Returns:
        FSDPModule: The module with FSDP applied (in-place).
    """
    torch._C._log_api_usage_once("torch.distributed.fsdp.fully_shard")
    _validate_module(module, "fully_shard")
    mesh = mesh or _init_default_mesh()
    _validate_mesh(mesh, dp_mesh_dims)
    mesh_info = _get_mesh_info(mesh, dp_mesh_dims)
    device = _get_device_from_mesh(mesh)
    auto_reshard_after_forward = reshard_after_forward is None
    # If the user does not provide ``reshard_after_forward``, we set it to True.
    # During lazy_init, we identify which module is the root and override its value to False
    if isinstance(mesh_info, FSDPMeshInfo):
        if (
            mesh_info.is_spmd_mesh
            and not isinstance(reshard_after_forward, bool)
            and isinstance(reshard_after_forward, int)
        ):
            raise NotImplementedError(
                "reshard_after_forward as int is not yet supported with "
                "SPMD mesh (dp_mesh_dims)"
            )
        post_forward_mesh_info = _get_post_forward_mesh_info(
            reshard_after_forward if not auto_reshard_after_forward else True,  # type: ignore[arg-type]
            mesh_info,
        )
    else:
        # DDPMeshInfo: no sharding, so no post-forward resharding needed
        post_forward_mesh_info = None
    arg_module, modules, managed_modules, params, buffers = _get_modules_and_states(
        module, device, ignored_params
    )
    state = fully_shard.state(modules[0])  # type: ignore[attr-defined]
    state.init(modules, device, mp_policy, auto_reshard_after_forward)

    _init_param_group(
        state,
        params,
        modules,
        mesh_info,
        post_forward_mesh_info,
        device,
        shard_placement_fn,
        mp_policy,
        offload_policy,
        reshard_after_forward=reshard_after_forward
        if not auto_reshard_after_forward
        else True,
    )

    # For Dynamo
    for managed_module in managed_modules:
        managed_module._is_fsdp_managed_module = True  # type: ignore[assignment]
        managed_module._fsdp_use_orig_params = True  # type: ignore[assignment]

    # Place FSDP leftmost for highest priority in the method resolution order
    _apply_to_module(
        modules, cls_to_fsdp_cls, FSDPModule, "FSDP", _unimplemented_deepcopy
    )
    return arg_module


def _unimplemented_deepcopy(*args: Any, **kwargs: Any) -> NoReturn:
    raise AssertionError(
        "FSDP does not support deepcopy. Please use state dict for serialization."
    )


_enable_fsdp_module_new_init: bool = True


@contextmanager
def disable_fsdp_module_new_init() -> Iterator[None]:
    global _enable_fsdp_module_new_init
    prev, _enable_fsdp_module_new_init = _enable_fsdp_module_new_init, False
    try:
        yield
    finally:
        _enable_fsdp_module_new_init = prev


class FSDPModule:
    # Index in MRO where the original class is found.
    # For FSDP: [FSDP<Orig>, FSDPModule, Orig, ...] -> index 2
    # Subclasses like ReplicateModule override this.
    _orig_cls_mro_index: int = 2

    def __new__(cls, *args, **kwargs):
        """
        Override ``__new__`` to remove the FSDP class and directly construct
        the original class for cases like indexing into a container module.
        """
        orig_cls = cls.__mro__[cls._orig_cls_mro_index]
        self = orig_cls.__new__(orig_cls, *args, **kwargs)
        if _enable_fsdp_module_new_init:
            self.__init__(*args, **kwargs)
        return self

    def reshard(self) -> None:
        """
        Reshards the module's parameters, freeing the unsharded parameters if
        they are allocated and registering the sharded parameters to the
        module. This method is *not* recursive.
        """
        state = self._get_fsdp_state()
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group.reshard()

    def unshard(self, async_op: bool = False) -> UnshardHandle | None:
        """
        Unshards the module's parameters by allocating memory and all-gathering
        the parameters. This method is *not* recursive. The unshard follows the
        :class:`MixedPrecisionPolicy`, so it will all-gather following
        ``param_dtype`` if set.

        Args:
            async_op (bool): If ``True``, then returns a :class:`UnshardHandle`
                that has a :meth:`wait` method to wait on the unshard op. If
                ``False``, then returns ``None`` and waits on the handle inside
                this function.

        .. note:: If ``async_op=True``, then FSDP will wait on the pending
            unshard in the module's pre-forward for the user. The user only
            needs to call :meth:`wait` explicitly if the wait should happen
            before pre-forward.
        """
        state = self._get_fsdp_state()
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group.lazy_init()
            fsdp_param_group.unshard(async_op=async_op)
        handle = _UnshardHandleImpl(
            list(state._fsdp_param_groups) if state._fsdp_param_groups else None
        )
        if async_op:
            return handle
        handle.wait()
        return None

    def set_is_last_backward(self, is_last_backward: bool) -> None:
        """
        Sets whether the next backward is the last one. On the last backward,
        FSDP waits on pending gradient reduction and clears internal data
        data structures for backward prefetching. This can be useful for
        microbatching.
        """
        state = self._get_fsdp_state()
        state._state_ctx.is_last_backward = is_last_backward

    def set_requires_gradient_sync(
        self, requires_gradient_sync: bool, *, recurse: bool = True
    ) -> None:
        """
        Sets if the module should sync gradients. This can be used to implement
        gradient accumulation *without communication*. For HSDP, this controls
        both reduce-scatter and all-reduce together. This is the equivalence of
        `no_sync` in FSDP1.

        Args:
            requires_gradient_sync (bool): Whether to reduce gradients for the
                module's parameters.
            recurse (bool): Whether to set for all FSDP submodules or just the
                passed-in module.
        """
        self_module = cast(nn.Module, self)
        modules = list(self_module.modules()) if recurse else [self_module]
        for module in modules:
            if isinstance(module, FSDPModule):
                state = module._get_fsdp_state()
                for fsdp_param_group in state._fsdp_param_groups:
                    fsdp_param_group.reduce_grads = requires_gradient_sync
                    fsdp_param_group.all_reduce_grads = requires_gradient_sync

    def set_requires_all_reduce(
        self, requires_all_reduce: bool, *, recurse: bool = True
    ) -> None:
        """
        Sets if the module should all-reduce gradients. This can be used to
        implement gradient accumulation with only reduce-scatter but not
        all-reduce for HSDP.
        """
        self_module = cast(nn.Module, self)
        modules = list(self_module.modules()) if recurse else [self_module]
        for module in modules:
            if isinstance(module, FSDPModule):
                state = module._get_fsdp_state()
                for fsdp_param_group in state._fsdp_param_groups:
                    fsdp_param_group.all_reduce_grads = requires_all_reduce

    def set_reshard_after_forward(
        self, reshard_after_forward: bool, recurse: bool = True
    ) -> None:
        """
        Sets if the module should reshard parameters after forward. This can be
        used to change the ``reshard_after_forward`` FSDP arg at runtime. For
        example, this can be used to set the FSDP root module's value to
        ``True`` (since it is otherwise specially set to ``False``), or it can
        set an FSDP module's value to ``False`` for running evals and set back
        to ``True`` for training.

        Args:
            reshard_after_forward (bool): Whether to reshard parameters after
                forward.
            recurse (bool): Whether to set for all FSDP submodules or just the
                passed-in module.
        """
        if not isinstance(reshard_after_forward, bool):
            raise ValueError(
                f"reshard_after_forward should be a bool, got {type(reshard_after_forward)}"
            )
        self_module = cast(nn.Module, self)
        modules = list(self_module.modules()) if recurse else [self_module]
        for module in modules:
            if isinstance(module, FSDPModule):
                state = module._get_fsdp_state()
                state._auto_reshard_after_forward = False
                for fsdp_param_group in state._fsdp_param_groups:
                    if not isinstance(fsdp_param_group.mesh_info, FSDPMeshInfo):
                        raise AssertionError
                    fsdp_param_group.post_forward_mesh_info = (
                        _get_post_forward_mesh_info(
                            reshard_after_forward,
                            fsdp_param_group.mesh_info,
                        )
                    )

    def set_reshard_after_backward(
        self, reshard_after_backward: bool, *, recurse: bool = True
    ) -> None:
        """
        Sets if the module should reshard parameters after backward. This can
        be used during gradient accumulation to trade off higher memory for
        reduced communication since the unsharded parameters do not need to be
        re-all-gathered before the next forward.

        Args:
            reshard_after_backward (bool): Whether to reshard parameters after
                backward.
            recurse (bool): Whether to set for all FSDP submodules or just the
                passed-in module.
        """
        self_module = cast(nn.Module, self)
        modules = list(self_module.modules()) if recurse else [self_module]
        for module in modules:
            if isinstance(module, FSDPModule):
                state = module._get_fsdp_state()
                for fsdp_param_group in state._fsdp_param_groups:
                    fsdp_param_group.reshard_after_backward = reshard_after_backward

    def set_modules_to_forward_prefetch(self, modules: list[FSDPModule]) -> None:
        """
        Sets the FSDP modules for which this FSDP module should explicitly
        prefetch all-gathers in forward. The prefetching runs after this
        module's all-gather copy-out.

        Passing a singleton list containing the next FSDP module gives the same
        all-gather overlap behavior as the default overlap behavior, except the
        prefetched all-gather is issued earlier from the CPU. Passing a list
        with at least length two is required for more aggressive overlap and
        will use more reserved memory.

        Args:
            modules (List[FSDPModule]): FSDP modules to prefetch.
        """
        _assert_all_fsdp_modules(modules)
        self._get_fsdp_state()._states_to_forward_prefetch = [
            module._get_fsdp_state() for module in modules
        ]

    def set_modules_to_backward_prefetch(self, modules: list[FSDPModule]) -> None:
        """
        Sets the FSDP modules for which this FSDP module should explicitly
        prefetch all-gathers in backward. This overrides the default backward
        pretching implementation that prefetches the next FSDP module based on
        the reverse post-forward order.

        Passing a singleton list containing the previous FSDP module gives the
        same all-gather overlap behavior as the default overlap behavior.
        Passing a list with at least length two is required for more aggressive
        overlap and will use more reserved memory.

        Args:
            modules (List[FSDPModule]): FSDP modules to prefetch.
        """
        _assert_all_fsdp_modules(modules)
        self._get_fsdp_state()._states_to_backward_prefetch = [
            module._get_fsdp_state() for module in modules
        ]

    def set_custom_all_gather(self, comm: AllGather) -> None:
        """
        Overrides the default ``all_gather`` communication behavior,
        to have better control over the communication and memory usage.
        See `Comm` and `ReduceScatter` for details.

        Args:
            comm (AllGather): Custom all-gather communication.
        """
        state = self._get_fsdp_state()
        if len(state._fsdp_param_groups) > 1:
            raise ValueError(
                "set_custom_all_gather is not supported with multiple param "
                "groups (from per-param mesh via shard_placement_fn). "
                "The custom comm would be ambiguous across groups with different meshes."
            )
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group._all_gather_comm = comm

    def set_custom_reduce_scatter(self, comm: ReduceScatter) -> None:
        """
        Overrides the default ``reduce_scatter`` communication behavior,
        to have better control over the communication and memory usage.
        See `Comm` and `ReduceScatter` for details.

        Args:
            comm (ReduceScatter): Custom reduce_scatter communication.
        """
        state = self._get_fsdp_state()
        if len(state._fsdp_param_groups) > 1:
            raise ValueError(
                "set_custom_reduce_scatter is not supported with multiple param "
                "groups (from per-param mesh via shard_placement_fn). "
                "The custom comm would be ambiguous across groups with different meshes."
            )
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group._reduce_scatter_comm = comm

    def set_all_reduce_hook(
        self,
        hook: Callable[[torch.Tensor], None],
        *,
        stream: torch.cuda.Stream | None = None,
    ):
        """
        Args:
            hook (Callable[[torch.Tensor], None]): User-defined all-reduce hook
                with expected signature ``hook(reduce_output: torch.Tensor) -> None``
                where ``reduce_output`` is the reduce-scatter output if only
                using FSDP or the all-reduce output if using native HSDP.
            stream (Optional[torch.cuda.Stream]): Stream to run the all-reduce
                hook in. This should only be set if not using native HSDP. If
                using native HSDP, the hook will run in the internally defined
                all-reduce stream used by the native HSDP all-reduce.
        """
        state = self._get_fsdp_state()
        if len(state._fsdp_param_groups) > 1:
            raise ValueError(
                "set_all_reduce_hook is not supported with multiple param "
                "groups (from per-param mesh via shard_placement_fn). "
                "The hook would be ambiguous across groups with different meshes."
            )
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group._all_reduce_hook = hook
            if stream is not None:
                if fsdp_param_group._is_hsdp:
                    raise ValueError("stream cannot be set when using native HSDP")
                fsdp_param_group._all_reduce_hook_stream = stream

    def set_post_optim_event(self, event: torch.Event) -> None:
        """
        Sets a post-optimizer-step event for the root FSDP module to wait the
        all-gather streams on.

        By default, the root FSDP module waits the all-gather streams on the
        current stream to ensure that the optimizer step has finished before
        all-gathering. However, this may introduce false dependencies if
        there is unrelated computation after the optimizer step. This API
        allows the user to provide their own event to wait on. After the root
        waits on the event, the event is discarded, so this API should be
        called with a new event each iteration.

        Args:
            event (torch.Event): Event recorded after the optimizer step
                to wait all-gather streams on.
        """
        self._get_fsdp_state()._state_ctx.post_optim_event = event

    @deprecated("Use `set_gradient_divide_factor` instead")
    def set_reduce_scatter_divide_factor(self, factor: float) -> None:
        """Use :py:meth:`set_gradient_divide_factor` instead"""
        self.set_gradient_divide_factor(factor)

    def set_gradient_divide_factor(self, factor: float) -> None:
        """
        Sets a custom divide factor for the gradient reduction. This might use
        a custom reduce op using NCCL's PreMulSum, which allows multiplying by
        the factor before reduction.

        Args:
            factor (float): Custom divide factor.
        """
        state = self._get_fsdp_state()
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group.gradient_divide_factor = factor

    def set_force_sum_reduction_for_comms(self, enable: bool) -> None:
        """
        Sets whether to require the low-level collective communication
        primitives to exclusively use "sum"-type reductions, even if it comes
        at the cost of separate additional pre- or post-scaling operations.
        This is needed for example because NCCL currently supports zero-copy
        transfers only for this kind of collectives.

        NB: for MTIA devices, this is always implicitly enabled.

        NB: if `set_all_reduce_hook` is used under FSDP setup, the caller needs
        to ensure the custom all-reduce across FSDP units follow this strategy
        as well, as FSDP can no longer automatically handle that.

        Args:
            enable (bool): Whether to only ever use ReduceOp.SUM for comms.
        """
        state = self._get_fsdp_state()
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group.force_sum_reduction_for_comms = enable

    def set_unshard_in_backward(self, unshard_in_backward: bool) -> None:
        """
        Sets whether the FSDP module's parameters need to be unsharded in
        backward. This can be used in expert cases when the user knows that all
        parameters in this FSDP module's parameter group are not needed for
        backward computation (e.g. embedding).
        """
        state = self._get_fsdp_state()
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group.unshard_in_backward = unshard_in_backward

    def set_allocate_memory_from_process_group_for_comm(self, enable: bool) -> None:
        """
        Sets whether the temporary staging buffers used to send and receive data
        over collective communications should be allocated using the custom
        optimized allocator provided by the ProcessGroup itself (if any). This
        might allow the ProcessGroup to be more efficient. For example, when
        using NCCL, this enables it to leverage zero-copy transfers over SHARP
        (for NVLink and/or InfiniBand).

        This cannot be used together with :meth:`set_custom_all_gather` or
        :meth:`set_custom_reduce_scatter` as those APIs allow for
        finer-grained control over each communication, and this method cannot
        determine their staging buffer allocation strategy.

        Args:
            enable (bool): Whether to turn on ProcessGroup allocation.
        """
        state = self._get_fsdp_state()
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group.set_allocate_memory_from_process_group(enable)

    def set_symm_mem_for_comm(self, backend: Literal["NCCL"] = "NCCL") -> None:
        """
        Sets the symmetric memory (``symm_mem``) backend for allocating the
        staging buffers used in all-gather collectives. This allows NCCL to use
        optimized all-gather implementations via symmetric memory. Such
        optimization may depend on the topology of the system.  For single node,
        Copy Engine All-Gather may be used. For multi-node, Symmetric Kernel
        All-Gather may be used.

        To enable Copy Engine All-Gather, you need to set the NCCL process group
        with the zero-CTA policy.
        ```python
        opts = dist.ProcessGroupNCCL.Options()
        opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO
        dist.init_process_group(backend="nccl", pg_options=opts, device_id=device)
        ```
        Alternatively, you can set the environment variable `NCCL_CTA_POLICY` to 2.
        ```bash
        export NCCL_CTA_POLICY=2
        ```
        For more details, see [Copy Engine
        Collectives](https://docs.pytorch.org/docs/2.11/symmetric_memory.html#copy-engine-collectives).

        This cannot be used together with :meth:`set_custom_all_gather` or
        :meth:`set_custom_reduce_scatter`.

        Args:
            backend (str): The symmetric memory backend to use. Defaults to
                ``"NCCL"``. Currently, only ``"NCCL"`` is supported.
        """
        state = self._get_fsdp_state()
        for fsdp_param_group in state._fsdp_param_groups:
            fsdp_param_group.set_symm_mem(backend)

    def _set_unshard_async_op(self, async_op: bool):
        """
        Sets whether to use ``async_op=True`` or ``False`` for the pre-forward
        and pre-backward unshard op. This defaults to ``False`` but can be set
        to ``True`` with this method.

        Setting this to ``True`` allows the all-gather allocations to happen in
        the default stream, avoiding inter-stream memory fragmentation.
        However, you must use explicit prefetching (e.g. via :meth:`unshard`)
        in forward to still get overlap, and the pre-all-gather ops like dtype
        casting and copy-in will not overlap with compute.
        """
        self_module = cast(nn.Module, self)
        for module in self_module.modules():
            if isinstance(module, FSDPModule):
                state = module._get_fsdp_state()
                for fsdp_param_group in state._fsdp_param_groups:
                    fsdp_param_group.unshard_async_op = async_op

    def _get_fsdp_state(self) -> FSDPState:
        if (state := _get_module_fsdp_state(cast(nn.Module, self))) is None:
            raise AssertionError(f"No FSDP state found on {self}")
        return state

    def _apply(self, *args: Any, **kwargs: Any) -> Any:
        # Reshard to ensure that sharded parameters are registered
        self.reshard()
        ret = super()._apply(*args, **kwargs)  # type: ignore[misc]
        state = self._get_fsdp_state()
        if not state._fsdp_param_groups:
            return ret
        # TODO: Remove this padding logic once DTensor pads the local tensor:
        # https://github.com/pytorch/pytorch/issues/113045
        with torch.no_grad():
            for fsdp_param_group in state._fsdp_param_groups:
                for fsdp_param in fsdp_param_group.fsdp_params:
                    fsdp_param.reset_sharded_param()
        return ret


class UnshardHandle:
    """
    A handle to wait on a :meth:`FSDPModule.unshard` op.
    """

    def wait(self) -> None:
        """
        Waits on the unshard op. This ensures that the current stream can use
        the unsharded parameters, which are now registered to the module.
        """
        return


class _UnshardHandleImpl(UnshardHandle):
    def __init__(self, fsdp_param_groups: list[FSDPParamGroup] | None):
        self._fsdp_param_groups = fsdp_param_groups

    def wait(self):
        if self._fsdp_param_groups is not None:
            for fsdp_param_group in self._fsdp_param_groups:
                fsdp_param_group.wait_for_unshard()
            # Avoid keeping a reference
            self._fsdp_param_groups = None


def register_fsdp_forward_method(module: nn.Module, method_name: str) -> None:
    """
    Registers a method on ``module`` to be considered a forward method for
    FSDP.

    FSDP all-gathers parameters pre-forward and optionally frees parameters
    post-forward (depending on ``reshard_after_forward``). FSDP only knows to
    do this for :meth:`nn.Module.forward` by default. This function patches a
    user-specified method to run the pre/post-forward hooks before/after the
    method, respectively. If ``module`` is not an :class:`FSDPModule`, then
    this is a no-op.

    Args:
        module (nn.Module): Module to register the forward method on.
        method_name (str): Name of the forward method.
    """
    if not isinstance(module, FSDPModule):
        # Make no-op to allow including both when using/not using FSDP
        return
    if not hasattr(module, method_name):
        raise ValueError(f"{type(module)} does not have a method {method_name}")
    orig_method = getattr(module, method_name)

    @functools.wraps(orig_method)
    def wrapped_method(self, *args, **kwargs):
        fsdp_state = self._get_fsdp_state()
        args, kwargs = fsdp_state._pre_forward(self, args, kwargs)
        out = orig_method(*args, **kwargs)
        return fsdp_state._post_forward(self, args, out)

    # Use `__get__` to make `wrapped_method` an instance method
    setattr(
        module,
        method_name,
        wrapped_method.__get__(module, type(module)),  # type:ignore[attr-defined]
    )


def share_comm_ctx(modules: list[FSDPModule]) -> None:
    """
    Share cuda streams for multiple FSDPModules

    Example usage:
        from torch.distributed.fsdp import share_comm_ctx
        share_comm_ctx([fsdp_model_1, fsdp_model_2, ...])

    For Pipeline Parallelism (PP), each model chunk is a FSDP root. We want
    to share cuda streams for all-gather, reduce-scatter, and all-reduce.
    This avoids allocating inter-stream memory framgmentation

    Args:
        modules (List[FSDPModule]): modules to share cuda streams
    """
    if len(modules) == 0:
        return
    for module in modules:
        if not isinstance(module, FSDPModule):
            raise ValueError(f"Expects list of FSDPModules but got {module}")
    fsdp_states = [module._get_fsdp_state() for module in modules]
    comm_ctx = fsdp_states[0]._comm_ctx
    for fsdp_state in fsdp_states[1:]:
        fsdp_state._comm_ctx = comm_ctx
        for fsdp_param_group in fsdp_state._fsdp_param_groups:
            fsdp_param_group.comm_ctx = comm_ctx


def _assert_all_fsdp_modules(modules: Iterable[Any]) -> None:
    for module in modules:
        if not isinstance(module, FSDPModule):
            raise ValueError(f"Expects FSDPModule but got {type(module)}: {module}")
