Vehicle-cpp/hkpc/software/torch/distributed/fsdp/_runtime_utils.py

import functools
from typing import (
    Any,
    Callable,
    Dict,
    Iterable,
    List,
    no_type_check,
    Optional,
    Set,
    Tuple,
)

import torch
import torch.distributed.fsdp._traversal_utils as traversal_utils
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.distributed.algorithms._comm_hooks import default_hooks, LOW_PRECISION_HOOKS
from torch.distributed.fsdp._common_utils import (
    _assert_in_training_states,
    _FSDPState,
    _get_module_fsdp_state,
    _get_sharding_strategy,
    _is_composable,
    TrainingState,
)
from torch.distributed.fsdp._init_utils import HYBRID_SHARDING_STRATEGIES
from torch.distributed.fsdp._utils import (
    _apply_to_tensors,
    _no_dispatch_record_stream,
    p_assert,
)
from torch.distributed.fsdp.api import BackwardPrefetch
from torch.distributed.fsdp.flat_param import (
    _HandlesKey,
    FlatParameter,
    FlatParamHandle,
    HandleShardingStrategy,
    HandleTrainingState,
)
from torch.distributed.utils import _to_kwargs

RESHARD_AFTER_FORWARD_STRATEGIES = {
    HandleShardingStrategy.FULL_SHARD,
    HandleShardingStrategy.HYBRID_SHARD,
}

# Do not include "process_group" to enable hybrid shard and MoE cases
HOMOGENEOUS_ATTR_NAMES = (
    "_use_orig_params",
    "limit_all_gathers",
)


def _get_fsdp_root_states_with_modules(
    module: nn.Module,
) -> Tuple[List[_FSDPState], List[nn.Module]]:
    """
    Returns a tuple containing:
    1. A list of the root ``_FSDPState`` instances in the module tree rooted at
    ``module`` without any duplicates and following the ``module.modules()``
    traversal order (which is assumed to be depth-first).
    2. A corresponding list of the root modules owning the states in the first
    list.

    This is similar to :func:`_get_fsdp_states_with_modules` except that we
    must call :func:`_is_fsdp_root` to force a lazy initialization to determine
    the FSDP root in case lazy initialization has not yet happened.
    """
    fsdp_root_states: List[_FSDPState] = []
    fsdp_root_modules: List[nn.Module] = []
    visited_fsdp_states: Set[_FSDPState] = set()
    # NOTE: This function assumes that `module.modules()` proceeds top-down.
    for submodule in module.modules():
        optional_state = _get_module_fsdp_state(submodule)
        if (
            optional_state is not None
            and optional_state not in visited_fsdp_states
            and _is_fsdp_root(optional_state, submodule)
        ):
            visited_fsdp_states.add(optional_state)
            fsdp_root_states.append(optional_state)
            fsdp_root_modules.append(submodule)
    return fsdp_root_states, fsdp_root_modules


def _get_fsdp_root_states(module: nn.Module) -> List[_FSDPState]:
    """See :func:`_get_fsdp_root_states_with_modules`."""
    fsdp_root_states, _ = _get_fsdp_root_states_with_modules(module)
    return fsdp_root_states


def _is_fsdp_root(state: _FSDPState, module: nn.Module) -> bool:
    """
    Returns if ``state`` corresponds to that of an FSDP root.

    For the wrapper code path, ``state`` and ``module`` should be the same. For
    the non-wrapper code path, ``state`` should be ``module`` 's state.
    """
    # Force a lazy initialization to determine the FSDP root
    _lazy_init(state, module)
    assert state._is_root is not None  # mypy
    return state._is_root


@no_type_check
def _validate_and_get_hybrid_shard_state(
    root_module: nn.Module,
) -> default_hooks.DefaultState:
    """
    Precondition: ``root_module`` is a ``FullyShardedDataParallel`` instance.

    This checks that all instances using a hybrid sharding strategy have the
    same intra- and inter-node process groups.

    Returns:
        DefaultState: One of the instances' inter-node state (does not
        matter which since they will share the same one).
    """
    intra_node_pgs = set()
    inter_node_pgs = set()
    inter_node_states = set()
    for fsdp_module in traversal_utils._get_fsdp_states(root_module):
        # TODO: Change this to handle's sharding strategy if we deprecate
        # `ShardingStrategy` internally.
        # https://github.com/pytorch/pytorch/issues/90857
        if fsdp_module.sharding_strategy in HYBRID_SHARDING_STRATEGIES:
            intra_node_pgs.add(fsdp_module.process_group)
            inter_node_pgs.add(fsdp_module._inter_node_pg)
            inter_node_states.add(fsdp_module._inter_node_state)
    if len(intra_node_pgs) == 0 and len(inter_node_pgs) == 0:
        # No instances use a hybrid sharding strategy
        return None
    error_prefix = "At least one instance uses a hybrid sharding strategy but has no "
    if len(intra_node_pgs) > 0 and len(inter_node_pgs) == 0:
        raise AssertionError(error_prefix + "inter-node proces group set")
    if len(intra_node_pgs) == 0 and len(inter_node_pgs) > 0:
        raise AssertionError(error_prefix + "intra-node process group set")
    error_prefix = "Some instances use a hybrid sharding strategy, but "
    if len(intra_node_pgs) != 1:
        raise ValueError(error_prefix + "intra-node process groups do not match")
    if len(inter_node_pgs) != 1:
        raise ValueError(error_prefix + "inter-node process groups do not match")
    return next(iter(inter_node_states))


@no_type_check
def _lazy_init(
    state: _FSDPState,
    root_module: nn.Module,
) -> _FSDPState:
    """
    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.

    For the non-composable code path, ``state`` and ``root_module`` should be
    the same, namely the FSDP instance itself.
    """
    if state._is_root is not None:
        return  # no-op: already lazily initialized
    if not torch.cuda.is_available():
        # Allow the FSDP constructor to run even without CUDA but check this
        # once we start real execution
        raise RuntimeError("FSDP does not support CPU only execution")
    # The following logic is only run on the root FSDP instance since it will
    # set `_is_root=False` for the non-root instances
    state._is_root = True
    _assert_in_training_states(state, [TrainingState.IDLE])
    _check_flat_params_on_expected_device(state, root_module)
    _init_streams(state)
    buffers, buffer_dtypes = _get_buffers_and_dtypes_for_computation(state, root_module)
    _cast_buffers_to_dtype_and_device(buffers, buffer_dtypes, state.compute_device)
    state._exec_order_data.init(state, root_module, state.process_group)
    _share_state_and_init_handle_attrs(state, root_module)
    return state


def _check_flat_params_on_expected_device(state: _FSDPState, module: nn.Module):
    """
    Checks that all ``FlatParameter``s in ``module`` 's tree managed by
    ``state`` are on the expected device for *lazy initialization*.
    """
    cpu_device = torch.device("cpu")
    for handle in traversal_utils._get_fsdp_handles(module):
        if (
            not handle._offload_params
            and handle.flat_param.device != state.compute_device
        ):
            raise RuntimeError(
                "An FSDP-managed module unexpectedly has parameters on "
                f"{handle.flat_param.device}. Make sure to move the module to "
                f"{state.compute_device} before training."
            )
        elif handle._offload_params and handle.flat_param.device != cpu_device:
            raise RuntimeError(
                "An FSDP-managed module with parameter CPU offloading enabled "
                f"has parameters on {handle.flat_param.device}. Make sure to "
                f"not move the module from CPU when offloading parameters."
            )


@no_type_check
def _share_state_and_init_handle_attrs(
    root_state: _FSDPState,
    root_module: nn.Module,
) -> None:
    """
    Shares data structure state from the ``root_state`` to all FSDP states in
    ``root_module`` 's module tree, and initializes handle attributes. These
    are done together to require a single loop over the states.
    """
    for handle in root_state._handles:
        handle.init_flat_param_attributes()
    inter_node_state = _validate_and_get_hybrid_shard_state(root_module)
    attr_name_to_values: Dict[str, Set[Any]] = {}
    for attr_name in HOMOGENEOUS_ATTR_NAMES:
        attr_name_to_values[attr_name] = set()
    for fsdp_state in traversal_utils._get_fsdp_states(root_module):
        for attr_name in HOMOGENEOUS_ATTR_NAMES:
            p_assert(
                hasattr(fsdp_state, attr_name),
                f"FSDP state missing attribute {attr_name}",
            )
            attr_name_to_values[attr_name].add(getattr(fsdp_state, attr_name))
        if fsdp_state is root_state:
            continue
        handle_sharding_strategy = _get_sharding_strategy(fsdp_state._handles)
        if handle_sharding_strategy in (
            HandleShardingStrategy.HYBRID_SHARD,
            HandleShardingStrategy._HYBRID_SHARD_ZERO2,
        ):
            # Share the all-reduce state across FSDP units. This is not strictly necessary
            # as each one already uses the same process group, but can slightly save memory
            # since other FSDP units allreduce state can be garbage collected.
            assert inter_node_state is not None, (
                "`_validate_and_get_hybrid_shard_state()` should have returned "
                "a valid inter-node state if there exists an FSDP instance "
                "using a hybrid sharding strategy"
            )
            fsdp_state._inter_node_state = inter_node_state

        # 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)
        p_assert(
            fsdp_state._is_root is None or not fsdp_state._is_root,
            "Non-root FSDP instance's `_is_root` should not have been "
            "set yet or should have been set to `False`",
        )
        fsdp_state._is_root = False
        fsdp_state._streams = root_state._streams
        fsdp_state._stream_to_name = root_state._stream_to_name
        fsdp_state._exec_order_data = root_state._exec_order_data
        fsdp_state._free_event_queue = root_state._free_event_queue
        fsdp_state._handles_prefetched = root_state._handles_prefetched
        fsdp_state._needs_pre_backward_unshard = root_state._needs_pre_backward_unshard
        for handle in fsdp_state._handles:
            handle.init_flat_param_attributes()
    for attr_name, attr_values in attr_name_to_values.items():
        if len(attr_values) != 1:
            raise ValueError(
                f"Expects one homogeneous value for {attr_name} but got {attr_values}"
            )


@no_type_check
def _init_streams(
    state: _FSDPState,
) -> _FSDPState:
    """
    Initializes CUDA streams for overlapping communication, computation, and
    data transfers. The streams should be shared across FSDP instances.
    """
    assert state._is_root
    assert torch.cuda.is_available()
    # Stream for unshard logic, including allocating the all-gather destination
    # tensors and the all-gathers themselves.
    state._streams["unshard"] = torch.cuda.Stream()
    # Stream for overlapping gradient reduction with the backward pass gradient
    # computation.
    state._streams["post_backward"] = torch.cuda.Stream()
    # Stream for pre-unshard logic, namely allocations and writes for CPU
    # offloading (H2D copy) and mixed precision (low precision cast).
    state._streams["pre_unshard"] = torch.cuda.Stream()
    # Default stream for computation
    state._streams["default"] = torch.cuda.current_stream()
    state._stream_to_name = {
        torch.cuda.current_stream(): "default",
        state._streams["unshard"]: "unshard",
        state._streams["pre_unshard"]: "pre_unshard",
        state._streams["post_backward"]: "post_backward",
    }


@no_type_check
def _unshard(
    state: _FSDPState,
    handles: List[FlatParamHandle],
    unshard_stream: torch.cuda.Stream,
    pre_unshard_stream: torch.cuda.Stream,
) -> 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 handles:
        return
    any_ran_pre_unshard = False
    with torch.cuda.stream(pre_unshard_stream):
        for handle in handles:
            ran_pre_unshard = handle.pre_unshard()
            any_ran_pre_unshard = any_ran_pre_unshard or ran_pre_unshard
    if any_ran_pre_unshard:
        unshard_stream.wait_stream(pre_unshard_stream)
    if state.limit_all_gathers:
        event = state._free_event_queue.dequeue_if_needed()
        if event:
            event.synchronize()
    with torch.cuda.stream(unshard_stream):
        for handle in handles:
            handle.unshard()
            handle.post_unshard()


@no_type_check
def _reshard(
    state: _FSDPState,
    handles: List[FlatParamHandle],
    free_unsharded_flat_params: List[bool],
):
    """
    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.
    """
    if not handles:
        return
    p_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 state.limit_all_gathers and free_unsharded_flat_param:
            free_event = torch.cuda.Event()
            free_event.record()
            state._free_event_queue.enqueue(free_event)
        handle.post_reshard()
    # Since we prefetch entire handles keys at a time, conservatively mark
    # the entire key as no longer prefetched once we free at least one
    handles_key = tuple(handles)
    if any(free_unsharded_flat_params):
        state._handles_prefetched.pop(handles_key, None)


def _unshard_grads(
    handles: List[FlatParamHandle],
) -> None:
    for handle in handles:
        handle.unshard_grad()


def _reshard_grads(
    handles: List[FlatParamHandle],
) -> None:
    for handle in handles:
        handle.reshard_grad()


@no_type_check
def _pre_forward(
    state: _FSDPState,
    handles: List[FlatParamHandle],
    unshard_fn: Callable,
    module: nn.Module,
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
    """
    Runs the pre-forward logic. This includes an opportunity to unshard
    currently sharded parameters such as those for the current forward and
    registering post-backward hooks for these current parameters. This function
    also converts forward ``args`` and ``kwargs`` to the given precision.

    Args:
        handles (List[FlatParamHandle]): Handles giving the parameters used in
            the current forward.
        unshard_fn (Optional[Callable]): A callable to unshard any currently
            sharded parameters or ``None`` to not do any unsharding.
        module (nn.Module): Module whose forward this method runs right before;
            expected by the hook signature.
        args (Tuple[Any, ...]): Module forward ``args``.
        kwargs (Dict[str, Any]): Module forward ``kwargs``.
    """
    state.training_state = TrainingState.FORWARD_BACKWARD
    state._exec_order_data.record_pre_forward(handles, module.training)
    for handle in handles:
        handle._training_state = HandleTrainingState.FORWARD
    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.
    _register_post_backward_hooks(state, handles)

    # Recursively convert args and kwargs to specified precision.
    input_dtype: Optional[torch.dtype] = state.mixed_precision.param_dtype
    if state.mixed_precision.cast_forward_inputs:
        args, kwargs = _cast_forward_inputs(input_dtype, *args, **kwargs)
    return args, kwargs


@no_type_check
def _pre_forward_unshard(
    state: _FSDPState,
    handles: List[FlatParamHandle],
) -> None:
    """Unshards parameters in the pre-forward."""
    if not handles:
        return
    _unshard(state, handles, state._streams["unshard"], state._streams["pre_unshard"])
    handles_key = tuple(handles)
    state._needs_pre_forward_unshard[handles_key] = False
    torch.cuda.current_stream().wait_stream(state._streams["unshard"])
    _prefetch_handles(state, handles_key)


@no_type_check
def _post_forward(
    state: _FSDPState,
    handles: List[FlatParamHandle],
    reshard_fn: 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
    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.
        module (nn.Module): Module whose forward just ran, which should be a
            fully sharded module (see [Note: Fully Sharded Module]); 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.
    """
    state._exec_order_data.record_post_forward(handles)
    if reshard_fn is not None:
        reshard_fn()
    # Register pre-backward hooks to unshard the flattened parameters
    # for the gradient computation (if needed)
    output = _register_pre_backward_hooks(state, module, output, handles)
    state.training_state = TrainingState.IDLE
    for handle in handles:
        handle._training_state = HandleTrainingState.IDLE
    return output


@no_type_check
def _post_forward_reshard(
    state: _FSDPState,
    handles: List[FlatParamHandle],
) -> None:
    """Reshards parameters in the post-forward."""
    if not handles:
        return
    # 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 state._is_root
        and handle._sharding_strategy in RESHARD_AFTER_FORWARD_STRATEGIES
        for handle in handles
    ]
    _reshard(state, handles, free_unsharded_flat_params)


@no_type_check
def _root_pre_forward(
    state: _FSDPState,
    module: nn.Module,
    args,
    kwargs,
) -> None:
    """
    Runs pre-forward logic specific to the root FSDP instance, which should run
    before any individual module's pre-forward. This starts with an attempt at
    lazy initialization (which only runs non-vacuously once). Otherwise, if
    this is called on a non-root FSDP instance, then it returns directly.

    Args:
        module (nn.Module): Module for which this logic tries to run. It may or
            may not be the root. If not, then this method does not do anything.
    """
    _lazy_init(state, module)
    p_assert(state._is_root is not None, "Expects a root FSDP to have been set")
    if not state._is_root:
        return args, kwargs
    if state.forward_prefetch:
        handles_keys = []
        if _is_composable(state):
            # TODO: This assumes singleton handles keys.
            handles_keys = [tuple(handle) for handle in state._handles]
        else:
            for fsdp_module in traversal_utils._get_fsdp_states(state):
                handles_key = tuple(fsdp_module._handles)
                handles_keys.append(handles_key)
        for handles_key in handles_keys:
            state._needs_pre_forward_unshard[handles_key] = True
    _wait_for_computation_stream(
        torch.cuda.current_stream(),
        state._streams["unshard"],
        state._streams["pre_unshard"],
    )
    _clear_grads_if_needed(traversal_utils._get_fsdp_handles(module))

    # Prepares the forward inputs by moving them to ``compute_device``
    # TODO: Do not use the side stream for tensor copies for now; investigate
    # the perf with/without it.
    args_tuple, kwargs_tuple = _to_kwargs(
        args, kwargs, state.compute_device.index, False
    )
    args = args_tuple[0]
    kwargs = kwargs_tuple[0]

    input_dtype: Optional[torch.dtype] = state.mixed_precision.param_dtype

    if state.mixed_precision.cast_root_forward_inputs:
        args, kwargs = _cast_forward_inputs(input_dtype, *args, **kwargs)
    return args, kwargs


def _cast_forward_inputs(
    input_dtype: Optional[torch.dtype],
    *args: Any,
    **kwargs: Any,
) -> Tuple[Any, Any]:
    """
    Prepares the forward inputs by casting them to ``input_dtype`` if it is not ``None``.
    """
    # TODO: For mixed precision, cast to reduced-precision in a single `to()` call.
    if input_dtype is not None:
        args, kwargs = _cast_fp_inputs_to_dtype(input_dtype, *args, **kwargs)
    return args, kwargs


def _cast_fp_inputs_to_dtype(
    dtype: torch.dtype,
    *args: Any,
    **kwargs: Any,
) -> Tuple[Any, Any]:
    """
    Casts floating point tensors in ``args`` and ``kwargs`` to ``input_dtype``.
    This respects the existing ``requires_grad`` on the tensors.
    """

    def cast_fn(x: torch.Tensor) -> torch.Tensor:
        if not torch.is_floating_point(x) or x.dtype == dtype:
            return x
        return x.to(dtype)

    return (_apply_to_tensors(cast_fn, args), _apply_to_tensors(cast_fn, kwargs))


@no_type_check
def _pre_backward_hook(
    state: _FSDPState,
    module: nn.Module,
    _handles: List[FlatParamHandle],
    *unused: Any,
) -> Any:
    """
    Prepares ``_handles`` 's ``FlatParameter`` s for gradient computation.

    Args:
        module (nn.Module): Fully sharded module (see [Note: Fully Sharded
            Module]).
    """
    _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 _handles_key and state._ran_pre_backward_hook.get(_handles_key, False):
        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 state._is_root and not state._post_backward_callback_queued:
            _register_post_backward_final_callback(state, module)
            _clear_grads_if_needed(traversal_utils._get_fsdp_handles(module))
        elif _handles_key:
            allowed_states = [TrainingState.IDLE]
            if _is_composable(state):
                allowed_states.append(TrainingState.FORWARD_BACKWARD)
            _assert_in_training_states(state, allowed_states)
        state.training_state = TrainingState.FORWARD_BACKWARD
        # Queueing the post-backward callback is the only logic that is not
        # per-handle in the pre-backward hook, so we can return early here if
        # there are no handles.
        if not _handles_key:
            return
        for handle in _handles:
            handle._training_state = HandleTrainingState.BACKWARD_PRE

        # If the handles have been prefetched, this `_unshard()` simply
        # switches to using the unsharded parameter
        _unshard(
            state, _handles, state._streams["unshard"], state._streams["pre_unshard"]
        )
        torch.cuda.current_stream().wait_stream(state._streams["unshard"])

        # Set this to `False` to ensure that a mistargeted prefetch does not
        # actually unshard these handles
        state._needs_pre_backward_unshard[_handles_key] = False
        _prefetch_handles(state, _handles_key)
        for handle in _handles:
            handle.prepare_gradient_for_backward()
        state._ran_pre_backward_hook[_handles_key] = True


@no_type_check
@torch.no_grad()
def _post_backward_hook(
    state: _FSDPState,
    handle: FlatParamHandle,
    *unused: Any,
):
    """
    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).
    """
    flat_param = handle.flat_param
    flat_param._post_backward_called = True
    with torch.autograd.profiler.record_function(
        "FullyShardedDataParallel._post_backward_hook"
    ):
        _assert_in_training_states(state, [TrainingState.FORWARD_BACKWARD])
        # For multiple applications of reentrant AC across submodules sharing
        # the same `FlatParameter`, the post-backward hook may run multiple
        # times in one backward, in which case we permit the state to already
        # be in `BACKWARD_POST`.
        p_assert(
            handle._training_state
            in (HandleTrainingState.BACKWARD_PRE, HandleTrainingState.BACKWARD_POST),
            f"Expects `BACKWARD_PRE` or `BACKWARD_POST` state but got {handle._training_state}",
        )
        handle._training_state = HandleTrainingState.BACKWARD_POST

        if flat_param.grad is None:
            return
        if flat_param.grad.requires_grad:
            raise RuntimeError("FSDP does not support gradients of gradients")

        free_unsharded_flat_param = _should_free_in_backward(state, handle)
        _reshard(state, [handle], [free_unsharded_flat_param])

        # TODO: Post-backward prefetching does not support the multiple handles
        # per module case since the post-backward hook runs per handle, not per
        # group of handles.
        handles_key = (handle,)
        _prefetch_handles(state, handles_key)

        if not state._sync_gradients:
            if handle._use_orig_params:
                handle._use_unsharded_grad_views()
            return

        # Wait for all ops in the current stream (e.g. gradient
        # computation) to finish before reduce-scattering the gradient
        state._streams["post_backward"].wait_stream(torch.cuda.current_stream())

        with torch.cuda.stream(state._streams["post_backward"]):
            autograd_computed_grad = flat_param.grad.data
            if state._exec_order_data.is_first_iter:  # only check once
                _check_comm_hook(
                    state._communication_hook, state._communication_hook_state
                )
            if (
                not _low_precision_hook_enabled(state)
                and flat_param.grad.dtype != handle._reduce_dtype
            ):
                flat_param.grad.data = flat_param.grad.to(handle._reduce_dtype)

            if handle.uses_sharded_strategy:
                # We clear `.grad` to permit multiple backwards. This avoids a
                # race where the second backward pass computation precedes
                # ahead of the first backward pass reduction, which is possible
                # since the reduction is issued in a separate stream and is
                # async and would result in reducing the wrong gradient.
                unsharded_grad = flat_param.grad.data
                flat_param.grad = None
                chunks = list(unsharded_grad.chunk(state.world_size))
                numel_to_pad = (
                    state.world_size * chunks[0].numel() - unsharded_grad.numel()
                )
                padded_unsharded_grad = (
                    F.pad(unsharded_grad, [0, numel_to_pad])
                    if numel_to_pad > 0
                    else unsharded_grad
                )
                new_sharded_grad = torch.empty_like(chunks[0])  # padded
                state._communication_hook(
                    state._communication_hook_state,
                    padded_unsharded_grad,
                    new_sharded_grad,
                )
                if handle._sharding_strategy in (
                    HandleShardingStrategy.HYBRID_SHARD,
                    HandleShardingStrategy._HYBRID_SHARD_ZERO2,
                ):
                    default_hooks.allreduce_hook(
                        state=state._inter_node_state,
                        grad=new_sharded_grad,
                    )
                _cast_grad_to_param_dtype(state, new_sharded_grad, flat_param)
                # Save the sharded gradient in `_saved_grad_shard` to support
                # gradient accumulation -- for multiple backwards, the gradient
                # reductions may happen in arbitrary order
                accumulate_grad = hasattr(flat_param, "_saved_grad_shard")
                if accumulate_grad:
                    _check_grad_to_accumulate(
                        new_sharded_grad, flat_param._saved_grad_shard
                    )
                    flat_param._saved_grad_shard += new_sharded_grad
                else:
                    flat_param._saved_grad_shard = new_sharded_grad
                grad_to_offload = flat_param._saved_grad_shard
            else:
                state._communication_hook(
                    state._communication_hook_state, flat_param.grad
                )
                # For `NO_SHARD`, we can keep the low precision gradients by
                # simply omitting the cast altogether
                if not handle._keep_low_precision_grads:
                    _cast_grad_to_param_dtype(state, flat_param.grad, flat_param)
                grad_to_offload = flat_param.grad.data

            if handle._offload_params:
                # Offload the gradient to CPU to ensure parameters and
                # gradients are on the same device as required by the optimizer
                # TODO: Investigate why `NO_SHARD` breaks correctness when
                # using `non_blocking=True` here.
                non_blocking = handle.uses_sharded_strategy
                flat_param._cpu_grad.copy_(  # type: ignore[attr-defined]
                    grad_to_offload.detach(), non_blocking=non_blocking
                )  # synchronized in the post-backward callback
                # Since the gradient being offloaded may have been produced in
                # the computation stream and is being consumed here in the
                # post-backward stream, inform the caching allocator
                _no_dispatch_record_stream(
                    grad_to_offload.data,
                    state._streams["post_backward"],
                )

            # Since the unsharded gradient is produced in the computation
            # stream and consumed in the post-backward stream, inform the
            # caching allocator (before it goes out of scope)
            _no_dispatch_record_stream(
                autograd_computed_grad, state._streams["post_backward"]
            )

            if handle._use_orig_params:
                # Since the handle's `FlatParameter` completed its gradient
                # computation, we should reset the gradient noneness mask
                handle._reset_is_grad_none()
                # Delay using sharded gradient views until after the
                # reduce-scatter instead of immediately after resharding
                handle._use_sharded_grad_views()


@no_type_check
def _should_free_in_backward(
    state: _FSDPState,
    handle: FlatParamHandle,
) -> bool:
    """
    Returns whether FSDP should free the unsharded flattened parameter in the
    post-backward or not.
    """
    # We always free if we are syncing gradients (i.e. not in no_sync) and parameters
    # are sharded.
    free_unsharded = state._sync_gradients and handle.uses_sharded_strategy
    # For NO_SHARD we don't need to free full parameters, for ZeRO-2 strategies, we skip
    # freeing in backward.
    return free_unsharded or (
        handle._sharding_strategy in RESHARD_AFTER_FORWARD_STRATEGIES
    )


@no_type_check
def _cast_grad_to_param_dtype(
    state: _FSDPState,
    sharded_grad: torch.Tensor,
    param: FlatParameter,
):
    """
    Casts ``sharded_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.
    """
    _assert_in_training_states(state, [TrainingState.FORWARD_BACKWARD])
    if not _low_precision_hook_enabled(state) and sharded_grad.dtype != param.dtype:
        low_prec_grad_data = sharded_grad.data
        sharded_grad.data = sharded_grad.data.to(dtype=param.dtype)
        # Since for `NO_SHARD`, the gradient is produced in the computation
        # stream and consumed here in the post-backward stream, inform the
        # caching allocator; for the sharded strategies, the gradient is
        # produced in the post-backward stream, so this `record_stream()`
        # should be a no-op
        _no_dispatch_record_stream(low_prec_grad_data, torch.cuda.current_stream())


def _check_comm_hook(
    comm_hook: Any,
    comm_hook_state: Any,
) -> None:
    p_assert(comm_hook is not None, "Communication hook should not be `None`")
    p_assert(
        comm_hook_state is not None, "Communication hook state should not be `None`"
    )


def _check_grad_to_accumulate(
    new_sharded_grad: torch.Tensor,
    accumulated_grad: torch.Tensor,
) -> None:
    p_assert(
        accumulated_grad.shape == new_sharded_grad.shape,
        "Shape mismatch when accumulating gradients: "
        f"existing gradient shape={accumulated_grad.shape} "
        f"new gradient shape={new_sharded_grad.shape}",
    )
    p_assert(
        accumulated_grad.device == new_sharded_grad.device,
        "Device mismatch when accumulating gradients: "
        f"existing gradient device={accumulated_grad.device} "
        f"new gradient device={new_sharded_grad.device}",
    )


@no_type_check
def _low_precision_hook_enabled(state: _FSDPState) -> bool:
    return state._communication_hook in LOW_PRECISION_HOOKS


@no_type_check
@torch.no_grad()
def _post_backward_final_callback(
    state: _FSDPState,
    module: nn.Module,
):
    """
    This waits for the post-backward to finish and performs some final cleanup.
    This runs at the end of the entire backward pass and should only be called
    on the root FSDP instance.
    """
    p_assert(
        state._is_root,
        "The post-backward callback should only be called on the root FSDP instance",
    )
    root_state = state

    if root_state._sync_gradients:
        torch.cuda.current_stream().wait_stream(root_state._streams["post_backward"])
        if root_state.cpu_offload.offload_params:
            # Wait for non-blocking GPU -> CPU sharded gradient copies from the
            # post-backward hooks to finish explicitly since CPU gradients do
            # not automatically synchronize with the GPU
            torch.cuda.current_stream().synchronize()
    root_state._exec_order_data.next_iter()

    for fsdp_state in traversal_utils._get_fsdp_states(module):
        _catch_all_reshard(fsdp_state)
        _finalize_params(fsdp_state)
        fsdp_state._ran_pre_backward_hook.clear()
        fsdp_state.training_state = TrainingState.IDLE
        for handle in fsdp_state._handles:
            handle._training_state = HandleTrainingState.IDLE
        fsdp_state._handles_prefetched.clear()
    # Reset for cases like one forward and multiple backwards
    root_state._post_backward_callback_queued = False


@no_type_check
def _catch_all_reshard(
    state: _FSDPState,
) -> None:
    """
    Reshards the parameters that may not have been resharded in the
    post-backward hook. This can happen when a module's output is used in the
    forward pass, meaning that its pre-backward hook runs (unsharding the
    parameter), but the post-backward hook does not run because the output was
    not jused in the loss computation corresponding to this backward pass.
    """
    # Wrap with a try-except to provide a more informative traceback if an
    # error is raised
    try:
        free_unsharded_flat_params: List[bool] = []
        handles_to_reshard: List[FlatParamHandle] = []
        for handle in state._handles:
            # TODO: This already-resharded check is brittle:
            # https://github.com/pytorch/pytorch/issues/83956
            already_resharded = (
                handle.flat_param.data_ptr()
                == handle.flat_param._local_shard.data_ptr()
            )
            if already_resharded:
                continue
            free_unsharded_flat_params.append(_should_free_in_backward(state, handle))
            handles_to_reshard.append(handle)
        if handles_to_reshard:
            _reshard(state, handles_to_reshard, free_unsharded_flat_params)
    except Exception as e:
        p_assert(
            False,
            f"Got exception in the catch-all reshard for {state}: {str(e)}",
            raise_assertion_error=False,
        )
        raise e


@no_type_check
def _finalize_params(
    state: _FSDPState,
) -> None:
    """Finalizes the parameters before the next iteration."""
    for handle in state._handles:
        flat_param = handle.flat_param
        if flat_param.requires_grad:
            if hasattr(flat_param, "_post_backward_hook_state"):
                p_assert(
                    len(flat_param._post_backward_hook_state) == 2,
                    f"Invalid: ``_post_backward_hook_state``: {flat_param._post_backward_hook_state}",
                )
                flat_param._post_backward_hook_state[1].remove()
                delattr(flat_param, "_post_backward_hook_state")
            if not state._sync_gradients:
                # Preserve the gradient accumulation state if not synchronizing
                # gradients: `.grad` remains the unsharded gradient  from prior
                # `no_sync()` iterations, and `_saved_grad_shard` remains the
                # sharded gradient from the last synchronized iteration
                continue
            handle.prepare_gradient_for_optim()
            p_assert(
                hasattr(flat_param, "_post_backward_called"),
                "Expects `_post_backward_called` to be set on the `FlatParameter`",
            )
            flat_param._post_backward_called = False


@no_type_check
def _prefetch_handles(
    state: _FSDPState,
    current_handles_key: _HandlesKey,
) -> None:
    """
    Prefetches the next handles if needed (without synchronization). An empty
    handles key cannot prefetch.
    """
    if not current_handles_key:
        return
    handles_to_prefetch = _get_handles_to_prefetch(state, current_handles_key)
    for handles_key in handles_to_prefetch:
        # Prefetch the next set of handles without synchronizing to allow
        # the sync to happen as late as possible to maximize overlap
        _unshard(
            state, handles_key, state._streams["unshard"], state._streams["pre_unshard"]
        )
        state._handles_prefetched[handles_key] = True


@no_type_check
def _get_handles_to_prefetch(
    state: _FSDPState,
    current_handles_key: _HandlesKey,
) -> List[_HandlesKey]:
    """
    Returns a :class:`list` of the handles keys to prefetch for the next
    module(s), where ``current_handles_key`` represents the current module.

    "Prefetching" refers to running the unshard logic early (without
    synchronization), and the "next" modules depend on the recorded execution
    order and the current training state.
    """
    training_state = _get_training_state(current_handles_key)
    valid_training_states = (
        HandleTrainingState.BACKWARD_PRE,
        HandleTrainingState.BACKWARD_POST,
        HandleTrainingState.FORWARD,
    )
    p_assert(
        training_state in valid_training_states,
        f"Prefetching is only supported in {valid_training_states} but "
        f"currently in {training_state}",
    )
    eod = state._exec_order_data
    target_handles_keys: List[_HandlesKey] = []
    if (
        training_state == HandleTrainingState.BACKWARD_PRE
        and state.backward_prefetch == BackwardPrefetch.BACKWARD_PRE
    ) or (
        training_state == HandleTrainingState.BACKWARD_POST
        and state.backward_prefetch == BackwardPrefetch.BACKWARD_POST
    ):
        target_handles_keys = [
            target_handles_key
            for target_handles_key in eod.get_handles_to_backward_prefetch(
                current_handles_key
            )
            if state._needs_pre_backward_unshard.get(target_handles_key, False)
            and not state._handles_prefetched.get(target_handles_key, False)
        ]
    elif training_state == HandleTrainingState.FORWARD and state.forward_prefetch:
        target_handles_keys = [
            target_handles_key
            for target_handles_key in eod.get_handles_to_forward_prefetch(
                current_handles_key
            )
            if state._needs_pre_forward_unshard.get(target_handles_key, False)
            and not state._handles_prefetched.get(target_handles_key, False)
        ]
    return target_handles_keys


def _get_training_state(
    handles_key: _HandlesKey,
) -> 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 = {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))


@no_type_check
def _register_pre_forward_hooks(
    state: _FSDPState,
    modules: Iterable[nn.Module],
) -> None:
    """
    Registers pre-forward hooks on all modules in ``modules``. The pre-forward
    hooks are partially applied based on the current ``FlatParamHandle``
    construction, meaning that they must be re-registered if the construction
    changes.
    """
    for forward_handle in state._pre_forward_handles:
        forward_handle.remove()
    state._pre_forward_handles.clear()
    for module in modules:
        module_param_handles = state._fully_sharded_module_to_handles.get(module, [])
        if module_param_handles:
            unshard_fn = functools.partial(
                _pre_forward_unshard,
                state,
                module_param_handles,
            )
            hook = functools.partial(
                _pre_forward, state, module_param_handles, unshard_fn
            )
            state._pre_forward_handles.append(
                module.register_forward_pre_hook(hook, prepend=True, with_kwargs=True)
            )


@no_type_check
def _register_post_forward_hooks(
    state: _FSDPState,
    modules: Iterable[nn.Module],
) -> None:
    """
    Registers post-forward hooks on all modules in ``modules``. The
    post-forward hooks are partially applied based on the current
    ``FlatParamHandle`` construction, meaning that they must be re-registered
    if the construction changes.
    """
    for forward_handle in state._post_forward_handles:
        forward_handle.remove()
    state._post_forward_handles.clear()
    for module in modules:
        module_param_handles = state._fully_sharded_module_to_handles.get(module, [])
        if module_param_handles:
            reshard_fn = functools.partial(
                _post_forward_reshard,
                state,
                module_param_handles,
            )
            hook = functools.partial(
                _post_forward,
                state,
                module_param_handles,
                reshard_fn,
            )
            state._post_forward_handles.append(module.register_forward_hook(hook))


@no_type_check
def _register_root_pre_forward_hook(
    state: _FSDPState,
    module: nn.Module,
):
    """
    Registers root pre-forward hook on ``module``, which should be the local
    FSDP root.

    NOTE: For the current composable FSDP design, we have each application of
    ``fully_shard()`` to a module to indicate that that module is the local
    FSDP root. We may remove this assumption in the future, in which case we
    will need to register this root pre-forward hook on any candidate module
    that may be the local FSDP root.
    """
    for forward_handle in state._root_pre_forward_handles:
        forward_handle.remove()
    state._root_pre_forward_handles.clear()
    hook = functools.partial(_root_pre_forward, state)
    state._root_pre_forward_handles.append(
        module.register_forward_pre_hook(hook, prepend=True, with_kwargs=True)
    )


@no_type_check
def _register_pre_backward_hooks(
    state: _FSDPState,
    module: nn.Module,
    outputs: Any,
    handles: List[FlatParamHandle],
) -> None:
    """
    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``.

    Args:
        module (nn.Module): Fully sharded module (see [Note: Fully Sharded
            Module]).

    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 state._is_root:
        state._post_backward_callback_queued = False  # only defined on the root

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

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

    return _apply_to_tensors(_register_hook, outputs)


def _register_post_backward_hooks(
    state: _FSDPState,
    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)
        p_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]
        assert acc_grad is not None
        hook_handle = acc_grad.register_hook(
            functools.partial(_post_backward_hook, state, handle)
        )
        flat_param._post_backward_hook_state = (acc_grad, hook_handle)  # type: ignore[attr-defined]


@no_type_check
def _register_post_backward_final_callback(
    state: _FSDPState, module: nn.Module
) -> None:
    """
    Registers the post-backward final callback that runs at the end of the
    backward pass. This should be called from the root FSDP instance at the
    beginning of the pre-backward.
    """
    p_assert(
        state._is_root,
        "Only the root FSDP instance should register the post-backward callback",
    )
    if state._post_backward_callback_queued:
        return
    _assert_in_training_states(state, [TrainingState.IDLE])
    state._post_backward_callback_queued = True
    Variable._execution_engine.queue_callback(
        functools.partial(_post_backward_final_callback, state, module)
    )


def _wait_for_computation_stream(
    computation_stream: torch.cuda.Stream,
    unshard_stream: torch.cuda.Stream,
    pre_unshard_stream: torch.cuda.Stream,
):
    """
    Has the unshard and pre-unshard streams wait for the computation stream.
    For example, this should be called in the FSDP root's pre-forward to
    respect optimizer step computation.
    """
    unshard_stream.wait_stream(computation_stream)
    # Having the pre-all-gather stream wait for the current stream even if we
    # do not leverage the pre-all-gather stream is tolerable since this only
    # runs once per iteration
    pre_unshard_stream.wait_stream(computation_stream)


def _clear_grads_if_needed(
    handles: List[FlatParamHandle],
):
    """
    Clears the original parameters' gradients if needed. This method's CPU
    overhead is minimal, so we may call it throughout FSDP methods, which serve
    as callsites to free the gradient memory earlier.
    """
    for handle in handles:
        if handle._use_orig_params:
            handle._clear_grads_if_needed()


@no_type_check
def _get_buffers_and_dtypes_for_computation(
    state: _FSDPState,
    root_module: nn.Module,
) -> Tuple[List[torch.Tensor], List[Optional[torch.dtype]]]:
    """
    Returns all buffers in the module tree rooted at ``root_module`` and a
    corresponding list of the buffer dtypes for computation. Each buffer dtype
    is either ``None`` if buffer mixed precision is not enabled or the buffer
    low precision dtype otherwise.
    """
    p_assert(state._is_root, "Expects the root to cast buffers")
    buffers: List[torch.Tensor] = []
    buffer_dtypes: List[Optional[torch.dtype]] = []
    if _is_composable(state):
        buffers = [
            buffer for module in root_module.modules() for buffer in module.buffers()
        ]
        buffer_dtypes = [
            state.mixed_precision.buffer_dtype for _ in range(len(buffers))
        ]
    else:
        visited_buffers = set()
        # Traverse the FSDP instances bottom-up so that we prefer the owning
        # FSDP instance's mixed precision setting for each buffer
        for fsdp_module in reversed(traversal_utils._get_fsdp_states(root_module)):
            for buffer in fsdp_module.buffers():
                if buffer in visited_buffers:
                    continue
                visited_buffers.add(buffer)
                buffers.append(buffer)
                buffer_dtypes.append(fsdp_module.mixed_precision.buffer_dtype)
    assert len(buffers) == len(buffer_dtypes), f"{len(buffers)} {len(buffer_dtypes)}"
    return buffers, buffer_dtypes


@no_type_check
def _get_buffer_dtypes(
    state: _FSDPState,
    buffer_names: List[str],
) -> List[torch.dtype]:
    """
    Returns the original buffer types of the given buffer names.
    """
    buffer_dtypes: List[torch.dtype] = []
    for buffer_name in buffer_names:
        p_assert(
            buffer_name in state._buffer_name_to_orig_dtype,
            f"{buffer_name} is missing from pre-computed dict on rank "
            f"{state.rank}, which only has keys "
            f"{state._buffer_name_to_orig_dtype.keys()}",
        )
        buffer_dtypes.append(state._buffer_name_to_orig_dtype[buffer_name])
    return buffer_dtypes


def _cast_buffers_to_dtype_and_device(
    buffers: List[torch.Tensor],
    buffer_dtypes: List[Optional[torch.dtype]],
    device: torch.device,
) -> None:
    """
    Casts ``buffers`` to the dtypes given by ``buffer_dtypes`` and moves them
    to ``device``. If an element in ``buffer_dtypes`` is ``None``, then the
    corresponding buffer is only moved to ``device``.
    """
    p_assert(
        buffer_dtypes is None or len(buffers) == len(buffer_dtypes),
        f"Expects `buffers` and `buffer_dtypes` to have the same length if "
        f"`buffer_dtypes` is specified but got {len(buffers)} and "
        f"{len(buffer_dtypes)}",
    )
    for buffer, buffer_dtype in zip(buffers, buffer_dtypes):
        if not torch.is_floating_point(buffer) or buffer_dtype is None:
            buffer.data = buffer.to(device=device)
        else:
            buffer.data = buffer.to(device=device, dtype=buffer_dtype)