# mypy: allow-untyped-defs
from __future__ import annotations

import contextlib
import logging
from typing import Any, cast, Literal, NamedTuple, TYPE_CHECKING
from typing_extensions import TypeVarTuple, Unpack

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.device_mesh import _get_device_handle
from torch.distributed.fsdp._common_utils import (
    _named_parameters_with_duplicates,
    collect_grad_tensors,
    replace_grad_tensors,
)
from torch.profiler import record_function
from torch.utils.hooks import RemovableHandle

from ._fsdp_api import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy
from ._fsdp_collectives import (
    AllGather,
    AllGatherResult,
    DefaultAllGather,
    DefaultReduceScatter,
    foreach_all_gather,
    foreach_all_gather_copy_out,
    foreach_reduce,
    ProcessGroupAllocAllGather,
    ProcessGroupAllocReduceScatter,
    ReduceScatter,
    SymmMemAllGather,
    SymmMemReduceScatter,
)
from ._fsdp_common import (
    _disable_functorch_if_active,
    _dynamo_disable,
    DataParallelMeshInfo,
    DDPMeshInfo,
    FSDPMeshInfo,
    HSDPMeshInfo,
    is_bw,
    ShardPlacementFnResult,
    TrainingState,
)
from ._fsdp_param import alloc_storage, FSDPParam, ParamModuleInfo, ShardedState


if TYPE_CHECKING:
    from collections.abc import Callable


logger = logging.getLogger("torch.distributed.fsdp.fully_shard")

_ModuleToHandleDict = dict[nn.Module, RemovableHandle]  # for state dict
_GradInputs = TypeVarTuple("_GradInputs")


"""
[Note: Overlapping all-gather copy-in and all-gather]
For implicit forward prefetching, we want to overlap the next copy-in with the
current all-gather. We do so using a separate copy-in stream. However, since
we have the all-gather input as a view into the output, we must make sure to
copy into different memory from the current all-gather's output. Thus, we keep
a reference to the current all-gather's output and have the next FSDP parameter
group free it after its copy-in. Finally, we have the last FSDP state flush the
reference to avoid holding onto memory after forward.
"""


class FSDPCommContext:
    """This has the communication state shared across FSDP states/parameter groups."""

    def lazy_init(self, device: torch.device):
        self.device_handle = _get_device_handle(device.type)
        # Setting the all-gather/reduce-scatter streams to be higher priority
        # can help avoid some issues where their copies in/out are delayed and
        # block computation (this is different from high-pri NCCL streams)
        high_priority = -1
        # All-gather state and copy-in stream allow overlapping the next
        # copy-in with the current all-gather in forward; copy-in overlaps with
        # reduce-scatter in backward without the separate copy-in stream
        self.all_gather_copy_in_stream = self.device_handle.Stream(
            priority=high_priority
        )
        # All-gather stream allows overlapping next all-gather with current
        # forward compute
        self.all_gather_stream = self.device_handle.Stream(priority=high_priority)
        # Reduce-scatter stream gives separate execution "thread" for post-
        # backward logic like pre/post-gradient division and reduce-scatter
        self.reduce_scatter_stream = self.device_handle.Stream(priority=high_priority)
        # The most recent post-reduce event across all param groups, recorded
        # on reduce_scatter_stream (FSDP) or all_reduce_stream (HSDP). Since
        # later events subsume earlier ones on the same stream, a single wait
        # on the last event recorded on each stream used in the post backward
        # hook covers all groups.
        self._last_post_reduce_events: dict[torch.Stream, torch.Event] = dict()
        # Run the HSDP all-reduces concurrently with all-gather/reduce-scatter
        # since collectives use different network resources and can overlap
        # in the typical intra-node sharding / inter-node replication case
        self.all_reduce_stream = self.device_handle.Stream()
        # All-gather/reduce-scatter states keep references to collective
        # tensors produced in one stream and used in another and accompanying
        # CUDA events for synchronization
        self.all_gather_state: AllGatherState | None = None
        self.reduce_scatter_states: list[ReduceScatterState] = []
        # Effective cap on retained reduce_scatter_states, resolved from the
        # per-group reduce_scatter_max_input_buffers in
        # FSDPState._init_shared_state. It lives here, not per group, because
        # reduce_scatter_states is a single shared list governed by one cap.
        self.reduce_scatter_max_input_buffers: int = 1
        # Post-forward order for explicit backward prefetching
        self.post_forward_order: list[FSDPParamGroup] = []  # will cause ref cycles

    def get_all_gather_streams(
        self, async_op: bool, training_state: TrainingState
    ) -> tuple[torch.Stream, torch.Stream]:
        if not async_op and training_state in (
            TrainingState.FORWARD,
            TrainingState.PRE_BACKWARD,
        ):
            # Use separate streams for implicit prefetching
            return self.all_gather_copy_in_stream, self.all_gather_stream
        current_stream = self.device_handle.current_stream()
        return current_stream, current_stream


# See [Note: Overlapping all-gather copy-in and all-gather]
class AllGatherState(NamedTuple):
    all_gather_result: AllGatherResult
    event: torch.Event | None  # all-gather copy-out


class ReduceScatterState(NamedTuple):
    reduce_scatter_input: torch.Tensor
    event: torch.Event | None  # reduce-scatter event


class AllReduceState(NamedTuple):
    # Holding all_reduce_input (the reduce-dtype AR buffer) keeps the
    # caching allocator from reusing the block across layers. This is a
    # structural invariant, not bookkeeping: without it, the next layer's
    # RS can reuse the same physical block before this layer's AR finishes
    # under slow AR, causing gradient aliasing. See PR #140044, PR #180900.
    all_reduce_input: torch.Tensor
    event: torch.Event | None  # all-reduce event


class FSDPParamGroup:
    """This class represents a parameter group to communicate together."""

    _orig_dtype: torch.dtype | None
    _reduce_dtype: torch.dtype | None

    def __init__(
        self,
        params: list[nn.Parameter],
        modules: tuple[nn.Module, ...],
        mesh_info: DataParallelMeshInfo,
        post_forward_mesh_info: FSDPMeshInfo | None,
        device: torch.device,
        shard_placement_fn: Callable[[nn.Parameter], ShardPlacementFnResult] | None,
        mp_policy: MixedPrecisionPolicy,
        offload_policy: OffloadPolicy,
    ):
        self.modules = modules  # permit ref cycle because 1:1 lifetime
        param_module_infos = _get_param_module_infos(params, modules)

        self.fsdp_params = [
            FSDPParam(
                param,
                module_info,
                mesh_info,
                post_forward_mesh_info,
                device,
                shard_placement_fn,
                mp_policy,
                offload_policy,
            )
            for param, module_info in zip(params, param_module_infos)
        ]
        self.mesh_info = mesh_info
        self.post_forward_mesh_info = post_forward_mesh_info
        self.device = device
        self.device_handle = _get_device_handle(device.type)
        self.mp_policy = mp_policy
        self.offload_policy = offload_policy
        self._training_state = TrainingState.IDLE
        # Group's sharded state always matches its parameters' sharded states
        self._sharded_state = ShardedState.SHARDED
        self._module_fqn: str | None = None  # prefixed from root module
        # Only consider resetting sharded parameters once in lazy init since it
        # can incur nontrivial overhead to reset them
        self._reset_sharded_params: bool = False

        # - Hook state
        self._module_to_pre_save_state_dict_hook_handle: _ModuleToHandleDict = {}
        self._module_to_pre_load_state_dict_hook_handle: _ModuleToHandleDict = {}
        self._all_reduce_hook: Callable[[torch.Tensor], None] | None = None
        self._all_gather_comm: AllGather = DefaultAllGather()
        self._all_gather_output = torch.empty(0, device=self.device)
        self._reduce_scatter_comm: ReduceScatter = DefaultReduceScatter()
        # Optional stream to run the user-defined all-reduce hook in
        # Saved here and not in the comm. context because we allow the user to
        # specify it, possibly at construction time before lazy init
        self._all_reduce_hook_stream: torch.cuda.Stream | None = None

        # - Communication and communication/computation overlap
        self.comm_ctx = FSDPCommContext()
        self._param_group_index: int = 0
        self._num_param_groups: int = 1
        # Group's indices in the shared post-forward order
        self._post_forward_indices: list[int] = []
        # Whether to reduce gradients at all (whether for FSDP or HSDP)
        self.reduce_grads: bool = True
        # Whether to all-reduce gradients for HSDP; only used if
        # `self.reduce_grads` is true, in which case setting this to false
        # means reduce-scatter but no all-reduce
        self.all_reduce_grads: bool = True
        # Whether to reshard parameters after backward (only useful for
        # gradient accumulation)
        self.reshard_after_backward: bool = True
        # Per-group input for the reduce-scatter copy-in (chunk_cat) buffer
        # cap-K, set via FSDPModule.set_reduce_scatter_max_input_buffers (see
        # its docstring). FSDP keeps 1 by default. These per-group values are
        # resolved into the single shared
        # comm_ctx.reduce_scatter_max_input_buffers that post_backward reads.
        self.reduce_scatter_max_input_buffers: int = 1
        # Optional custom factor for the gradient reduction op (e.g. to divide
        # by a factor other than the world size)
        self.gradient_divide_factor: float | None = None
        # Whether to include zero gradients for parameters that did not
        # receive a gradient in backward (e.g. due to conditional parameter
        # usage across ranks). This ensures all ranks participate in the same
        # reduce-scatter collectives, avoiding collective mismatch errors.
        self.reduce_scatter_unused_params: bool = False
        # Whether reduce-scatter and all-reduce should be issued using only
        # summations, potentially with separate pre-/post-scaling.
        self.force_sum_reduction_for_comms: bool = False
        # `async_op` arg used for pre-forward/pre-backward unshard; can be
        # overridden to only do explicit prefetching and avoid inter-stream
        # fragmentation from using separate unshard streams
        self.unshard_async_op: bool = False
        # Whether to unshard in backward: can be overridden by the user if the
        # parameters in this group are not needed for backward (e.g. embedding)
        self.unshard_in_backward: bool = True

        # - CUDA events for stream synchronization
        # Holds the all-gather output buffer, sync objects, and metadata
        self._all_gather_result: AllGatherResult | None = None
        # Holds the reduce-scatter/all-reduce view-out CUDA event that marks the end of
        # the group's post-backward (e.g. reduce-scatter, all-reduce and div), which
        # should be waited on at the end of backward
        self._post_reduce_event: torch.Event | None = None
        # Holds the reshard-after-forward CUDA event when resharding to a
        # different world size, which should be waited on in the next unshard
        self._reshard_after_forward_event: torch.Event | None = None

        # Only for HSDP, if accumulating gradients without all-reduce, save the
        # partial reduce output (only reduce-scattered but not all-reduced)
        self._partial_reduce_output: torch.Tensor | None = None
        # Holds the reduce-dtype AR buffer + completion event across
        # layers in HSDP+AR with reduce_dtype != orig_dtype (e.g., bf16
        # reduce + fp32 params). Structural invariant: the live Python
        # ref keeps the buffer off the caching allocator's free list,
        # preventing the next layer's RS from reusing the same physical
        # block while this layer's AR is still in flight. See
        # AllReduceState docstring and regression test PR #180900.
        self._all_reduce_state: AllReduceState | None = None

    # Initialization #
    def _init_mp_dtypes(self) -> None:
        for fsdp_param in self.fsdp_params:
            fsdp_param.init_dtype_attrs(self.mp_policy)
        trainable_params: list[FSDPParam] = [
            p for p in self.fsdp_params if p.sharded_param.requires_grad
        ]
        orig_dtypes = {p.orig_dtype for p in trainable_params}
        reduce_dtypes = {p.reduce_dtype for p in trainable_params}
        if len(trainable_params) > 0 and len(orig_dtypes) != 1:
            # Models may have no grad params
            raise AssertionError(
                f"FSDP expects uniform original parameter dtype but got {orig_dtypes}"
            )
        self._orig_dtype = next(iter(orig_dtypes)) if trainable_params else None
        if len(trainable_params) > 0 and len(reduce_dtypes) != 1:
            # This can be relaxed if we issue one reduce-scatter per reduce
            # dtype (but we would need a way for users to specify multiple
            # reduce dtypes)
            raise AssertionError(
                f"FSDP expects uniform reduce dtype but got {reduce_dtypes}"
            )
        self._reduce_dtype = next(iter(reduce_dtypes)) if trainable_params else None

    def lazy_init(self):
        # Lazy init should be idempotent
        # Users may change or register parameters after construction time.
        # For example, DoRA (https://arxiv.org/abs/2402.09353) initializes linear magnitudes based on
        # other parameters (e.g. loaded from the state dict).
        if not hasattr(self.comm_ctx, "device_handle"):
            self.comm_ctx.device_handle = _get_device_handle(self.device.type)
        if self.is_sharded and not self._reset_sharded_params:
            for fsdp_param in self.fsdp_params:
                fsdp_param.reset_sharded_param()
                fsdp_param._init_extensions()  # allow monkey patch after init
            self._reset_sharded_params = True
        self._validate_no_meta_params()
        self._validate_cpu_offload_params()
        self._validate_reduce_scatter_max_input_buffers()
        # Initialize mixed precision attributes lazily in case the user changes
        # the parameter dtypes after construction time but before forward
        self._init_mp_dtypes()
        self._register_state_dict_hooks()

    def set_symm_mem(self, backend: Literal["NCCL"] = "NCCL") -> None:
        if not isinstance(self._all_gather_comm, (DefaultAllGather | SymmMemAllGather)):
            raise AssertionError(
                "cannot call set_symm_mem() "
                f"when all gather comm is custom: {self._all_gather_comm.__class__.__name__}"
            )
        self._all_gather_comm = SymmMemAllGather(
            self._all_gather_process_group, backend
        )
        if not isinstance(
            self._reduce_scatter_comm, (DefaultReduceScatter | SymmMemReduceScatter)
        ):
            raise AssertionError(
                "cannot call set_symm_mem() "
                f"when reduce scatter comm is custom: {self._reduce_scatter_comm.__class__.__name__}"
            )
        if self.force_sum_reduction_for_comms:
            # As of NCCL 2.29.3, NCCL symmetric reduce-scatter only supports SUM reduction
            self._reduce_scatter_comm = SymmMemReduceScatter(
                self._reduce_scatter_process_group, backend
            )

    def set_allocate_memory_from_process_group(self, enable: bool) -> None:
        """
        Whether to (try to) use the ProcessGroup's allocate_tensor method for
        the staging buffers for collective comms.
        """
        if not isinstance(
            self._all_gather_comm, (DefaultAllGather | ProcessGroupAllocAllGather)
        ):
            raise AssertionError(
                "cannot call set_allocate_memory_from_process_group() "
                f"when all gather comm is custom: {self._all_gather_comm.__class__.__name__}"
            )
        self._all_gather_comm = (
            ProcessGroupAllocAllGather(self._all_gather_process_group)
            if enable
            else DefaultAllGather()
        )

        if not isinstance(
            self._reduce_scatter_comm,
            (DefaultReduceScatter | ProcessGroupAllocReduceScatter),
        ):
            raise AssertionError(
                "cannot call set_allocate_memory_from_process_group() "
                f"when reduce scatter comm is custom: {self._reduce_scatter_comm.__class__.__name__}"
            )
        self._reduce_scatter_comm = (
            ProcessGroupAllocReduceScatter(self._reduce_scatter_process_group)
            if enable
            else DefaultReduceScatter()
        )

    # Runtime #
    @_disable_functorch_if_active
    def unshard(self, async_op: bool = False):
        if self._all_gather_result is not None:  # already called, pending wait
            return
        if self.is_unsharded:
            return  # no-op
        if (
            not self.unshard_in_backward
            and self._training_state == TrainingState.PRE_BACKWARD
        ):
            return
        if self._reshard_after_forward_event is not None:
            # Resharded parameter data is allocated in the default stream and
            # used in the all-gather streams
            self._wait_all_gather_streams_on_event(self._reshard_after_forward_event)
            self._reshard_after_forward_event = None

        if isinstance(self.mesh_info, FSDPMeshInfo):
            world_size = self._all_gather_process_group.size()
        else:
            world_size = 1
        if world_size == 1:
            # can't skip due to early return in wait_for_unshard if
            # no self._all_gather_result
            self._all_gather_result = AllGatherResult(
                all_gather_output=self._all_gather_output,
                all_gather_event=self.device_handle.Event().record(),
                all_gather_work=None,
                param_all_gather_input_dtypes=[],
                param_all_gather_input_numels=[],
                all_gather_input_split_sizes=[],
            )

            return

        with record_function(self._with_fqn("FSDP::all_gather")):
            self._all_gather_result = foreach_all_gather(
                self.fsdp_params,
                self._all_gather_process_group,
                async_op,
                *self.comm_ctx.get_all_gather_streams(async_op, self._training_state),
                self.device,
                self._all_gather_comm,
            )

    @_disable_functorch_if_active
    def wait_for_unshard(self):
        """
        1. In forward with implicit prefetching, to overlap the current copy-out
        with the next all-gather, we save a reference to the current all-gather
        result to free after the next copy-out.
        2. Otherwise (explicit prefetching or in backward), we free the
        all-gather result immediately after the current copy-out since we can
        already overlap the current copy-out with the previous reduce-scatter.
        """
        if not self._all_gather_result:
            return  # no preceding unshard
        async_op = self._all_gather_result.all_gather_work is not None
        if self._training_state == TrainingState.FORWARD:  # implicit prefetch
            if prev_all_gather_state := self.comm_ctx.all_gather_state:
                self._wait_all_gather_streams_on_event(prev_all_gather_state.event)
                self.comm_ctx.all_gather_state = None  # free the all-gather result
        if isinstance(self.mesh_info, FSDPMeshInfo):
            world_size = self._all_gather_process_group.size()
        else:
            world_size = 1
        if world_size == 1:
            # directly initialize unsharded parameters from sharded parameters

            for fsdp_param in self.fsdp_params:
                # Use all_gather_inputs which already handles conversion to param_dtype
                # This is consistent with the world_size > 1 path
                all_gather_input = fsdp_param.all_gather_inputs[0]

                # Make sure the all_gather_outputs has proper storage size before using it
                # First ensure we have at least one tensor in all_gather_outputs
                fsdp_param.init_all_gather_outputs(
                    [all_gather_input.numel()],
                    [all_gather_input.dtype],
                    world_size,
                    self.device,
                )

                tensor = fsdp_param.all_gather_outputs[0]
                alloc_storage(tensor)

                # find alternative way to check if tensor.is_inference
                with torch.autograd._unsafe_preserve_version_counter(tensor):
                    tensor.copy_(all_gather_input)

        else:
            with record_function(self._with_fqn("FSDP::all_gather_copy_out")):
                foreach_all_gather_copy_out(
                    self._all_gather_result,
                    self.fsdp_params,
                    self._all_gather_process_group,
                )

        for fsdp_param in self.fsdp_params:
            fsdp_param.init_unsharded_param()

        self._to_unsharded()
        all_gather_copy_out_event = self.device_handle.Event()
        all_gather_copy_out_event.record()

        if (
            not async_op
            and self._training_state == TrainingState.FORWARD
            and world_size > 1
        ):
            # Defer free to allow for overlap of this copy-out with next
            # all-gather collective
            self.comm_ctx.all_gather_state = AllGatherState(
                self._all_gather_result, all_gather_copy_out_event
            )
        else:
            self._wait_all_gather_streams_on_event(all_gather_copy_out_event)

        self._all_gather_result = None  # free unless saved in `all_gather_state`

    def _wait_all_gather_streams_on_event(self, event: torch.Event | None):
        # Calling `unshard` before lazy init means streams are not initialized
        if hasattr(self.comm_ctx, "all_gather_copy_in_stream") and event is not None:
            self.comm_ctx.all_gather_copy_in_stream.wait_event(event)
        if hasattr(self.comm_ctx, "all_gather_stream") and event is not None:
            self.comm_ctx.all_gather_stream.wait_event(event)

    @_disable_functorch_if_active
    def reshard(self):
        if self._training_state == TrainingState.FORWARD:
            if not self._reshard_after_forward:
                return
            if self._use_post_forward_mesh:
                self._to_sharded_post_forward()
                self._reshard_after_forward_event = self.device_handle.Event()
                if self._reshard_after_forward_event is not None:
                    self._reshard_after_forward_event.record()
                return
        self._to_sharded()

    def _reset_iter_state(self) -> None:
        # See FSDPState._reset_iter_state for semantics. Waits on any
        # in-flight collective events owned by this group, discards
        # accumulated grad-reduction state, and restores sharded params.
        current_stream = self.device_handle.current_stream()
        if self._all_gather_result is not None:
            if (event := self._all_gather_result.all_gather_event) is not None:
                current_stream.wait_event(event)
            work = self._all_gather_result.all_gather_work
            if isinstance(work, dist.distributed_c10d.Work):
                work.wait()
            self._all_gather_result = None
        if self._post_reduce_event is not None:
            current_stream.wait_event(self._post_reduce_event)
            self._post_reduce_event = None
        if (
            self._all_reduce_state is not None
            and self._all_reduce_state.event is not None
        ):
            current_stream.wait_event(self._all_reduce_state.event)
        self._all_reduce_state = None
        if self._reshard_after_forward_event is not None:
            self._wait_all_gather_streams_on_event(self._reshard_after_forward_event)
            self._reshard_after_forward_event = None
        self._partial_reduce_output = None
        self._post_forward_indices.clear()
        self._training_state = TrainingState.IDLE
        self._to_sharded()

    def pre_forward(
        self, module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]
    ) -> tuple[tuple[Any, ...], dict[str, Any]]:
        logger.debug("%s", self._with_fqn("FSDP::pre_forward"))
        with record_function(self._with_fqn("FSDP::pre_forward")):
            # FORWARD at entry means another module in the grouped
            # ``fully_shard([a, b])`` already registered post_backward this
            # pass; skip to avoid duplicate ``RegisterPostBackwardFunction``
            # autograd nodes.
            entering_forward_pass = self._training_state != TrainingState.FORWARD
            self._training_state = TrainingState.FORWARD
            self.unshard(self.unshard_async_op)
            self.wait_for_unshard()
            if entering_forward_pass:
                args, kwargs = self._register_post_backward_hook(args, kwargs)
            return args, kwargs

    def post_forward(self, module: nn.Module, input: Any, output: Any):
        logger.debug("%s", self._with_fqn("FSDP::post_forward"))
        with record_function(self._with_fqn("FSDP::post_forward")):
            # for AC(fully_shard(model)), AC runs fsdp's _pre_forward
            # it shouldn't change post_forward_order
            if not is_bw():
                self.reshard()
                self._record_post_forward()
            self._training_state = TrainingState.IDLE
            return output

    def _record_post_forward(self) -> None:
        # Since a group has one pre-backward unshard for each forward call
        # before the backward, we record each usage (with multiplicity)
        post_forward_index = len(self.comm_ctx.post_forward_order)
        self.comm_ctx.post_forward_order.append(self)
        self._post_forward_indices.append(post_forward_index)

    @_dynamo_disable
    def pre_backward(self, default_prefetch: bool, *unused: Any):
        if self._training_state == TrainingState.PRE_BACKWARD:
            return
        logger.debug("%s", self._with_fqn("FSDP::pre_backward"))
        with record_function(self._with_fqn("FSDP::pre_backward")):
            self._training_state = TrainingState.PRE_BACKWARD
            self.unshard(self.unshard_async_op)  # no-op if prefetched
            self.wait_for_unshard()
            if default_prefetch:
                self._backward_prefetch()

    @_dynamo_disable
    def post_backward(self, *unused: Any):
        # This method should be idempotent and safe to call even when this
        # FSDP parameter group was not used in backward (should be a no-op)
        logger.debug("%s", self._with_fqn("FSDP::post_backward"))
        # Partial-group-forward detection: grouped ``fully_shard([a, b, ...])``
        # where the forward ran only a subset of the group's modules (chunked
        # loss, 1F1B, etc.). The wrapped post-hook took its partial path, so
        # ``state._post_forward`` never ran and no ``_pre_backward`` hook was
        # registered — ``pre_backward`` therefore never fired and state stays
        # at FORWARD instead of transitioning IDLE → PRE_BACKWARD.
        is_partial_group_backward = (
            len(self.modules) > 1  # grouped (structural)
            and self._training_state == TrainingState.FORWARD  # partial path taken
        )
        self._training_state = TrainingState.POST_BACKWARD
        with record_function(self._with_fqn("FSDP::post_backward_accumulate")):
            for fsdp_param in self.fsdp_params:
                fsdp_param.accumulate_unsharded_grad_if_needed()
        with record_function(self._with_fqn("FSDP::post_backward_reshard")):
            if not self.reduce_grads:
                if self.reshard_after_backward:
                    self.reshard()
                for fsdp_param in self.fsdp_params:
                    fsdp_param.to_accumulated_grad_if_needed()
                return
            # Save the autograd-computed gradients before resharding to only
            # access the unsharded parameters when their data is present
            fsdp_params_with_grad: list[FSDPParam] = []
            unsharded_grads: list[torch.Tensor] = []

            for fsdp_param in self.fsdp_params:
                if not hasattr(fsdp_param, "_unsharded_param"):
                    continue
                # May have an accumulated gradient of the reduce dtype if the
                # previous backward did not reduce-scatter
                if fsdp_param.unsharded_accumulated_grad is not None:
                    fsdp_params_with_grad.append(fsdp_param)
                    unsharded_grads.append(fsdp_param.unsharded_accumulated_grad_data)
                    fsdp_param.unsharded_accumulated_grad = None
                elif fsdp_param.unsharded_param.grad is not None:
                    fsdp_params_with_grad.append(fsdp_param)
                    unsharded_grads.append(fsdp_param.unsharded_grad_data)
                    fsdp_param.unsharded_param.grad = None
                elif (
                    self.reduce_scatter_unused_params
                    and fsdp_param.unsharded_param.requires_grad
                ):
                    fsdp_params_with_grad.append(fsdp_param)
                    unsharded_grads.append(fsdp_param.unsharded_zero_grad_data)
            if self.reshard_after_backward:
                self.reshard()
        # Recycle prior modules' reduce-scatter input buffers, keeping at most
        # `max_input_buffers` in flight: reclaim the oldest (wait on its
        # reduce-scatter, then drop the keepalive ref that was deferring the
        # allocator's reuse of its memory) until fewer than that remain, before
        # this module's reduce appends one more. See
        # set_reduce_scatter_max_input_buffers for the memory/overlap tradeoff.
        # (Assumes backward fires groups N-1 first; if not, overlap degrades but
        # correctness is preserved.)
        max_input_buffers = self.comm_ctx.reduce_scatter_max_input_buffers
        states = self.comm_ctx.reduce_scatter_states
        if (
            self._param_group_index == self._num_param_groups - 1
            and len(states) >= max_input_buffers
        ):
            with record_function(f"FSDP::post_backward_rs_wait ({self._module_fqn})"):
                while len(states) >= max_input_buffers:
                    oldest = states.pop(0)
                    if oldest.event is not None:
                        self.device_handle.current_stream().wait_event(oldest.event)
                    del oldest
        if len(fsdp_params_with_grad) == 0:
            return
        with record_function(self._with_fqn("FSDP::post_backward_reduce")):
            all_reduce_pg = (
                self._all_reduce_process_group
                if isinstance(self.mesh_info, DDPMeshInfo)
                else None
            )
            all_reduce_stream: torch.cuda.Stream
            if all_reduce_pg is None and self._all_reduce_hook_stream is not None:
                # this means the native HSDP is not enabled,
                # but user may want to have a custom HSDP setup
                if self._all_reduce_hook is None:
                    raise AssertionError(
                        "all reduce hook stream is specified but hook itself is missing."
                    )
                all_reduce_stream = self._all_reduce_hook_stream
            else:
                all_reduce_stream = self.comm_ctx.all_reduce_stream

            self._wait_for_post_backward()
            (
                reduce_scatter_input,
                reduce_scatter_event,
                post_reduce_stream,
                self._post_reduce_event,
                all_reduce_input,
                all_reduce_event,
                self._partial_reduce_output,
            ) = foreach_reduce(
                fsdp_params_with_grad,
                unsharded_grads,
                (
                    # pyrefly: ignore [bad-argument-type]
                    self._reduce_scatter_process_group
                    if isinstance(self.mesh_info, FSDPMeshInfo)
                    else None  # pyre-fixme[6]
                ),
                self.comm_ctx.reduce_scatter_stream,
                self._reduce_scatter_comm,
                self._orig_dtype,
                self._reduce_dtype,
                self.device,
                self.gradient_divide_factor,
                (
                    self._all_reduce_process_group
                    if isinstance(self.mesh_info, DDPMeshInfo)
                    else None
                ),
                all_reduce_stream,
                self.all_reduce_grads,
                self._partial_reduce_output,
                self._all_reduce_hook,
                self.force_sum_reduction_for_comms,
            )
            self.comm_ctx._last_post_reduce_events[post_reduce_stream] = (
                self._post_reduce_event
            )
            self.comm_ctx.reduce_scatter_states.append(
                ReduceScatterState(reduce_scatter_input, reduce_scatter_event)
            )
            if is_partial_group_backward:
                # Serialize the default stream on this invocation's
                # post-accumulate event before returning to autograd.
                # Otherwise the next partial-group invocation's autograd
                # kernels queue on the default stream concurrently with
                # this one's accumulate on the RS stream; the caching
                # allocator then hands an autograd kernel a block that is
                # a live input/output of this accumulate, corrupting grads
                # on non-zero ranks. The post-accumulate event is load-
                # bearing (pre-accumulate reduce_scatter_event is NOT, per
                # MI350X); waiting at the next post_backward's entry is
                # too late because autograd runs between post_backwards.
                # No CPU sync.
                #
                # TODO(#181218): open questions on scope.
                #   1. Conditional vs unconditional. The same cross-stream
                #      hazard exists structurally in regular FSDP but has
                #      not been observed to fire, plausibly because
                #      ``reduce_scatter_states.append(...)`` ref-holds
                #      ``reduce_scatter_input`` and cross-layer allocator
                #      pressure is looser than within a group. Dropping
                #      the gate needs a real-model overlap-loss measurement.
                #   2. ROCm-specific vs cross-platform. Only observed on
                #      ROCm/RCCL/MI350X; CUDA FSDP passes without this fix.
                #      The standalone repro shows the stream-ordering
                #      hazard (vector 1) is cross-platform, but FSDP's
                #      ref-hold closes it on both. The residual FSDP-side
                #      race on ROCm (vector 2) is on a non-Python-reachable
                #      buffer — RCCL workspace or allocator fragment.
                #      Whether vector 2 exists on CUDA FSDP but is timing-
                #      masked is unresolved. See
                #      ``fsdp2_chunked_loss_rocm_race.md``.
                self.device_handle.current_stream().wait_event(self._post_reduce_event)
            if all_reduce_input is not None:
                if self.device.type != "cpu":
                    if all_reduce_event is None:
                        raise AssertionError(
                            "Expected all_reduce_event to be set for non-CPU device"
                        )
                self._all_reduce_state = AllReduceState(
                    all_reduce_input, all_reduce_event
                )

    def finalize_backward(self):
        for event in self.comm_ctx._last_post_reduce_events.values():
            self.device_handle.current_stream().wait_event(event)
        self.comm_ctx._last_post_reduce_events = dict()
        self._post_reduce_event = None
        self._all_reduce_state = None
        for fsdp_param in self.fsdp_params:
            if fsdp_param.grad_offload_event is not None:
                fsdp_param.grad_offload_event.synchronize()
                fsdp_param.grad_offload_event = None
        if self._all_gather_result is not None:
            # If there was a mistargeted unshard without a corresponding wait,
            # then we wait here and clear the unshard
            if (event := self._all_gather_result.all_gather_event) is not None:
                torch.accelerator.current_stream().wait_event(event)
            work = self._all_gather_result.all_gather_work
            if isinstance(work, dist.distributed_c10d.Work):
                work.wait()
            self._all_gather_result = None
        self._post_forward_indices.clear()

    def _wait_for_post_backward(self):
        if self._post_reduce_event is not None:
            self.device_handle.current_stream().wait_event(self._post_reduce_event)
            self._post_reduce_event = None
        if (
            self._all_reduce_state is not None
            and self._all_reduce_state.event is not None
        ):
            self.device_handle.current_stream().wait_event(self._all_reduce_state.event)
        self._all_reduce_state = None

    def _backward_prefetch(self) -> None:
        if self._training_state == TrainingState.PRE_BACKWARD:
            if not self._post_forward_indices:
                # Can be cleared if running multiple `backward`s
                return
            curr_index = self._post_forward_indices.pop()
            if self._num_param_groups > 1:
                # Backward fires groups in reverse forward order:
                # N-1, N-2, ..., 1, 0.  Index 1 is always the
                # penultimate group regardless of N.  Prefetching here
                # lets the next module's AG overlap with group 0's RS
                # without holding unsharded params too long (as would
                # happen if we prefetched from N-1).
                if self._param_group_index != 1:
                    return
                # E.g. fully_shard(block, shard_placement_fn=...) creates two
                # param groups per block (dense + moe), giving
                # post_forward_order = [block0, block0.moe, block1, block1.moe].
                # block1.moe walks back past block1 to prefetch block0.moe then block0.
                curr_modules = self.modules
                target_modules: tuple[nn.Module, ...] | None = None
                for step in range(1, curr_index + 1):
                    target = self.comm_ctx.post_forward_order[curr_index - step]
                    if target.modules is curr_modules:
                        continue
                    if target_modules is None:
                        target_modules = target.modules
                    elif target.modules is not target_modules:
                        break
                    # Prefetch all groups of the target module in
                    # reverse forward order (highest index first),
                    # matching the explicit path in _pre_backward.
                    self._prefetch_unshard(target, "backward")
            elif curr_index > 0:
                target = self.comm_ctx.post_forward_order[curr_index - 1]
                self._prefetch_unshard(target, "backward")

    @staticmethod
    def _prefetch_unshard(
        target_fsdp_param_group: FSDPParamGroup, pass_type: str
    ) -> None:
        if pass_type == "backward":
            training_state = TrainingState.PRE_BACKWARD
        elif pass_type == "forward":
            training_state = TrainingState.FORWARD
        else:
            raise ValueError(f"Unknown pass type: {pass_type}")
        target_fqn = target_fsdp_param_group._module_fqn
        with (
            record_function(f"FSDP::{pass_type}_prefetch for {target_fqn}"),
            target_fsdp_param_group.use_training_state(training_state),
        ):
            async_op = target_fsdp_param_group.unshard_async_op
            target_fsdp_param_group.unshard(async_op)

    # Utilities #
    def _to_sharded(self):
        if not self.is_sharded:
            for fsdp_param in self.fsdp_params:
                fsdp_param.to_sharded()
            self._sharded_state = ShardedState.SHARDED

    def _to_sharded_post_forward(self):
        if not self.is_sharded_post_forward:
            for fsdp_param in self.fsdp_params:
                fsdp_param.to_sharded_post_forward()
            self._sharded_state = ShardedState.SHARDED_POST_FORWARD

    def _to_unsharded(self):
        if not self.is_unsharded:
            for fsdp_param in self.fsdp_params:
                fsdp_param.to_unsharded()
            self._sharded_state = ShardedState.UNSHARDED

    @property
    def is_sharded(self) -> bool:
        return self._sharded_state == ShardedState.SHARDED

    @property
    def is_sharded_post_forward(self) -> bool:
        return self._sharded_state == ShardedState.SHARDED_POST_FORWARD

    @property
    def is_unsharded(self) -> bool:
        return self._sharded_state == ShardedState.UNSHARDED

    @contextlib.contextmanager
    def use_training_state(self, training_state: TrainingState):
        old_training_state = self._training_state
        self._training_state = training_state
        try:
            yield
        finally:
            self._training_state = old_training_state

    # Hook Registration #
    def _register_post_backward_hook(
        self, args: tuple[Any, ...], kwargs: dict[str, Any]
    ) -> tuple[tuple[Any, ...], dict[str, Any]]:
        if not torch.is_grad_enabled():
            return args, kwargs

        # Collect all tensors that require gradients (including from dataclasses)
        inp_tensors = collect_grad_tensors((args, kwargs))
        if not inp_tensors:
            return args, kwargs

        # Apply RegisterPostBackwardFunction to all tensors at once
        out_tensors = RegisterPostBackwardFunction.apply(self, *inp_tensors)

        # Replace tensors in the structure (iterator order matches collect order)
        new_args, new_kwargs = replace_grad_tensors((args, kwargs), iter(out_tensors))
        return new_args, new_kwargs

    def _register_state_dict_hooks(self) -> None:
        num_pre_save_hooks = len(self._module_to_pre_save_state_dict_hook_handle)
        num_pre_load_hooks = len(self._module_to_pre_load_state_dict_hook_handle)
        if num_pre_save_hooks != num_pre_load_hooks:
            raise AssertionError(
                f"Pre-save: {num_pre_save_hooks} pre-load: {num_pre_load_hooks}"
            )
        if num_pre_save_hooks > 0:
            return  # already registered
        modules_with_fsdp_params: set[nn.Module] = {
            fsdp_param._module_info.module for fsdp_param in self.fsdp_params
        }

        def to_sharded_hook(*args: Any, **kwargs: Any) -> None:
            self._to_sharded()

        for module in modules_with_fsdp_params:
            self._module_to_pre_save_state_dict_hook_handle[module] = (
                module.register_state_dict_pre_hook(to_sharded_hook)
            )
            self._module_to_pre_load_state_dict_hook_handle[module] = (
                module._register_load_state_dict_pre_hook(to_sharded_hook)
            )

    # Properties #
    @property
    def _reshard_after_forward(self) -> bool:
        return self.post_forward_mesh_info is not None

    @property
    def _use_post_forward_mesh(self) -> bool:
        return (
            self._reshard_after_forward
            and self.mesh_info != self.post_forward_mesh_info
        )

    @property
    def _is_hsdp(self) -> bool:
        return isinstance(self.mesh_info, HSDPMeshInfo)

    @property
    def _all_gather_process_group(self) -> dist.ProcessGroup:
        mesh_info = (
            cast(FSDPMeshInfo, self.post_forward_mesh_info)
            if self.is_sharded_post_forward
            else self.mesh_info
        )
        if not isinstance(mesh_info, FSDPMeshInfo):
            raise AssertionError(
                f"Expected mesh_info to be FSDPMeshInfo, got {type(mesh_info)}"
            )
        return mesh_info.shard_process_group

    @property
    def _reduce_scatter_process_group(self) -> dist.ProcessGroup:
        if not isinstance(self.mesh_info, FSDPMeshInfo):
            raise AssertionError(
                f"Expected mesh_info to be FSDPMeshInfo, got {type(self.mesh_info)}"
            )
        # Falls back to the shard PG (shared with all-gather) unless
        # set_separate_reduce_scatter_group opted in to a dedicated PG.
        return (
            self.mesh_info.reduce_scatter_process_group
            or self.mesh_info.shard_process_group
        )

    def _set_separate_reduce_scatter_group(
        self, enable: bool, new_groups: dict[tuple[int, ...], dist.ProcessGroup]
    ) -> None:
        # Give reduce-scatter its own process group (separate from the shared
        # all-gather/shard PG) so the two can overlap in backward, or reset to
        # the shared group. ``new_groups`` caches groups by shard ranks within
        # one set_separate_reduce_scatter_group call so meshes sharding over the
        # same ranks share one communicator (typically one total) rather than one
        # per mesh. HSDP ranks may create only their local shard subgroup, so
        # use local synchronization to avoid waiting for nonmember ranks in
        # new_group's post-init barrier.
        mesh_info = self.mesh_info
        if not isinstance(mesh_info, FSDPMeshInfo):
            raise AssertionError(
                f"Expected mesh_info to be FSDPMeshInfo, got {type(mesh_info)}"
            )
        if not enable:
            mesh_info.reduce_scatter_process_group = None
            return
        ranks = tuple(dist.get_process_group_ranks(mesh_info.shard_process_group))
        if ranks not in new_groups:
            # Reuse the mesh's existing group if already enabled (idempotent),
            # else create one for this rank set.
            existing = mesh_info.reduce_scatter_process_group
            if existing is not None:
                new_groups[ranks] = existing
            else:
                new_groups[ranks] = dist.new_group(
                    list(ranks),
                    use_local_synchronization=True,
                    group_desc="fsdp_reduce_scatter",
                )
        mesh_info.reduce_scatter_process_group = new_groups[ranks]

    @property
    def _all_reduce_process_group(self) -> dist.ProcessGroup:
        if not isinstance(self.mesh_info, DDPMeshInfo):
            raise AssertionError(
                f"Expected mesh_info to be DDPMeshInfo or HSDPMeshInfo, got {type(self.mesh_info)}"
            )
        return self.mesh_info.replicate_process_group

    def _with_fqn(self, label: str) -> str:
        if self._module_fqn:
            label = f"{label} ({self._module_fqn})"
        if self._num_param_groups > 1 and isinstance(self.mesh_info, FSDPMeshInfo):
            label = f"{label} [pg={self.mesh_info.shard_mesh_size}]"
        return label

    def __repr__(self):
        return f"FSDPParamGroup(fqn={self._module_fqn})"

    def _validate_no_meta_params(self):
        param_names_on_meta = [
            fsdp_param._param_fqn
            for fsdp_param in self.fsdp_params
            if fsdp_param.sharded_param.device.type == "meta"
        ]
        if param_names_on_meta:
            raise RuntimeError(
                "FSDP parameters should be materialized from meta device before training, "
                f"but the following were still on meta device: {param_names_on_meta}\n"
                "For example, call module.to_empty(device) to materialize to device and "
                "call module.reset_parameters() on each module to initialize values."
            )

    def _validate_cpu_offload_params(self):
        if not isinstance(self.offload_policy, CPUOffloadPolicy):
            return
        fsdp_params_not_on_cpu = [
            fsdp_param
            for fsdp_param in self.fsdp_params
            if fsdp_param.sharded_param.device.type != "cpu"
        ]
        if fsdp_params_not_on_cpu:
            raise RuntimeError(
                "FSDP parameters should be materialized on CPU when enabling CPU offloading. "
                'For example, load a CPU state dict or call module.to_empty(device="cpu"). '
                "Found following parameters on non-CPU device: "
                f"{[(fsdp_param._param_fqn, fsdp_param.sharded_param.device) for fsdp_param in fsdp_params_not_on_cpu]}\n"
            )

    def _validate_reduce_scatter_max_input_buffers(self):
        # Reject max_input_buffers > 1 with the symmetric-memory reduce-scatter
        # comm: retaining >1 input buffers relies on the allocator handing out a
        # fresh buffer while prior ones are ref-held, but the symmetric-memory
        # pool reuses a single buffer, so retaining K would force K symmetric
        # segments and can exhaust symmetric memory. Read the per-group cap (set
        # in __init__) so this is safe to call before the comm context is lazily
        # initialized (e.g. unshard before the first forward).
        if self.reduce_scatter_max_input_buffers > 1 and isinstance(
            self._reduce_scatter_comm, SymmMemReduceScatter
        ):
            raise ValueError(
                "set_reduce_scatter_max_input_buffers(>1) is not supported with "
                "the symmetric-memory reduce-scatter comm, but got "
                f"{self.reduce_scatter_max_input_buffers} for module "
                f"'{self._module_fqn}' using "
                f"{self._reduce_scatter_comm.__class__.__name__}. Keep "
                "max_input_buffers=1 with symmetric memory, or use the default "
                "reduce-scatter comm."
            )


def _get_param_module_infos(
    params: list[nn.Parameter], modules: tuple[nn.Module, ...]
) -> list[ParamModuleInfo]:
    """
    Shared parameter: lin1.weight = lin2.weight
    Shared module: mlp.lin1 = mlp.lin2
    We do not remove duplicates when traversing both modules and parameters to
    find shared modules' parameters and shared parameters within a module.
    """
    params_set = set(params)
    param_to_module_info: dict[nn.Parameter, ParamModuleInfo] = {}
    for module in modules:
        for _, submodule in module.named_modules(remove_duplicate=False):
            for param_name, param in _named_parameters_with_duplicates(
                submodule, recurse=False
            ):
                if param in params_set:
                    if param not in param_to_module_info:
                        param_to_module_info[param] = ParamModuleInfo(
                            submodule, param_name
                        )
                    else:
                        param_to_module_info[param].shared_modules.append(submodule)
                        param_to_module_info[param].shared_param_names.append(
                            param_name
                        )
    if len(param_to_module_info) != len(params):
        raise AssertionError(f"Some parameters are not in the module tree of {modules}")
    return [param_to_module_info[param] for param in params]


class RegisterPostBackwardFunction(torch.autograd.Function):
    generate_vmap_rule = True

    @staticmethod
    # pyrefly: ignore [bad-override]
    def forward(param_group: FSDPParamGroup, *inputs: torch.Tensor):
        # All tensors in `inputs` should require gradient
        return inputs

    @staticmethod
    def setup_context(ctx, inputs: tuple[Any, ...], output: Any) -> None:
        param_group, *_ = inputs
        ctx.param_group = param_group

    @staticmethod
    def backward(ctx, *grads: torch.Tensor):
        ctx.param_group.post_backward()
        return (None,) + grads

    @staticmethod
    # pyrefly: ignore [bad-override]
    def jvp(
        ctx: Any,
        param_group_tangent: object,
        *grad_inputs: Unpack[_GradInputs],
    ) -> tuple[Unpack[_GradInputs]]:
        # Drop the non-tensor param_group tangent. The output pre-backward hook
        # queues final post-backward after all primal/tangent paths finish.
        return grad_inputs
