Diff Coverage

        """
        self._validate_flags()
        self._validate_dp()
        self._validate_pp()
        self._validate_ac()

    def _validate_ac(self) -> None:
        """Validate the activation checkpointing mode and layer selectors."""
        if not isinstance(self.ac_mode, str) or self.ac_mode not in AC_MODES:
            raise ValueError(f"ac_mode must be one of {AC_MODES}, got {self.ac_mode!r}")
        if not isinstance(self.ac_save_forward_all_gather, bool):
            raise ValueError("ac_save_forward_all_gather must be a bool")
        self._validate_ac_layer_count()
        self._validate_ac_layer_indices()
        if self.ac_layer_count is not None and self.ac_layer_indices is not None:
            raise ValueError(
                "ac_layer_count and ac_layer_indices are mutually exclusive; "
                "set exactly one"
            )

            )

    def _validate_ac_layer_count(self) -> None:
        """Validate the optional prefix layer selector."""
        if self.ac_layer_count is not None and (
            isinstance(self.ac_layer_count, bool)
            or not isinstance(self.ac_layer_count, int)
            or self.ac_layer_count < 0
        ):
            raise ValueError(
                "ac_layer_count must be None or a non-negative integer, "
                f"got {self.ac_layer_count!r}"
            )

            )

    def _validate_ac_layer_indices(self) -> None:
        """Validate the optional explicit layer index selector."""
        indices = self.ac_layer_indices
        if indices is None:
            return
        invalid = (
            not isinstance(indices, list)
            or not indices
            or any(
                isinstance(index, bool) or not isinstance(index, int) or index < 0

                isinstance(index, bool) or not isinstance(index, int) or index < 0
                for index in indices
            )
        )
        if invalid:
            raise ValueError(
                "ac_layer_indices must be None or a non-empty list of "
                f"non-negative integers, got {indices!r}"
            )
        if len(set(indices)) != len(indices):
            raise ValueError(
                f"ac_layer_indices must not contain duplicates, got {indices!r}"
            )

    def _validate_flags(self) -> None:

            )

    def _validate_flags(self) -> None:
        """Validate boolean flags and the overlap memory budget."""
        for name in (
            "fsdp_bucketing",
            "fsdp_allgather_prefetch",
            "fsdp_overlap_grad_reduce",
        ):

        ):
            value = getattr(self, name)
            if not isinstance(value, bool):
                raise ValueError(f"{name} must be a bool")
        if self.fsdp_bucket_modules is not None and not isinstance(
            self.fsdp_bucket_modules, (list, tuple)
        ):
            raise ValueError(
                "fsdp_bucket_modules must be None or a list/tuple of module groups"
            )
        value = self.max_in_flight_gb
        if (

            or config.fsdp_config.edp_shard_size > 1
        )
    fsdp_degree = config.fsdp_config.dp_shard_size if fsdp_enabled else None

    ac_mode = "off"
    ac_layer_count = None
    ac_layer_indices = None
    ac_config = getattr(config, "activation_checkpoint", None)
    if ac_config is not None and getattr(ac_config, "mode", None) not in (None, "off"):
        ac_mode = ac_config.mode
        selection = getattr(ac_config, "selection", None)
        if getattr(selection, "source", "default") == "model_adapter_safe_regions":
            ac_layer_count = getattr(selection, "layer_count", None)
            ac_layer_indices = getattr(selection, "layer_indices", None)
        if getattr(ac_config, "swap_inputs", False):
            _LOG.warning(
                "activation_checkpoint.swap_inputs is not supported in graph "
                "mode; input swapping will be disabled."
            )
        if getattr(ac_config, "allow_cache_entry_mutation", False):
            _LOG.info(
                "activation_checkpoint.allow_cache_entry_mutation has no "
                "effect in graph mode; ignoring it."
            )

import torch
from torch import fx
from torch._functorch import compile_utils
from torch._functorch.partitioners import must_recompute
from torch._ops import OpOverload
from torch.utils.checkpoint import CheckpointPolicy

# Same-package reuse of the eager AC machinery: operator save/recompute sets,
# HF transformer-block discovery and model-adapter recompute specs. They are
# module-level privates because the eager module never exported them, but the

# Same-package reuse of the eager AC machinery: operator save/recompute sets,
# HF transformer-block discovery and model-adapter recompute specs. They are
# module-level privates because the eager module never exported them, but the
# coupling is the point: one policy vocabulary for both modes.
from hyper_parallel.distributed.activation_checkpoint import (
    _SELECTIVE_AC_COMM_OP_NAMES,
    _SELECTIVE_AC_FORCE_RECOMPUTE_OPS,
    _SELECTIVE_AC_MUST_SAVE_OPS,
    _adapter_recompute_spec,

    _select_layer_region_roots,
    _selected_checkpoint_layer_fqns,
)

from ..pass_config import PassConfig
from .base import GraphPass

_LOG = logging.getLogger(__name__)

# FSDP all-gathers are inserted after this pass, but forward (TP/SP/CP)
# all-gathers are already in the traced forward; saving them avoids a second
# collective during backward recompute. ``PassConfig.ac_save_forward_all_gather``

# FSDP all-gathers are inserted after this pass, but forward (TP/SP/CP)
# all-gathers are already in the traced forward; saving them avoids a second
# collective during backward recompute. ``PassConfig.ac_save_forward_all_gather``
# selects that superset.
_GRAPH_AC_COMM_OPS = _existing_ops(
    *(_resolve_torch_op(name) for name in _SELECTIVE_AC_COMM_OP_NAMES),
    _resolve_op_attr(torch.ops, "deepep.dispatch.default"),
    _resolve_op_attr(torch.ops, "deepep.combine.default"),
    _resolve_op_attr(torch.ops, "hybridep.dispatch.default"),

    _resolve_op_attr(torch.ops, "deepep.combine.default"),
    _resolve_op_attr(torch.ops, "hybridep.dispatch.default"),
    _resolve_op_attr(torch.ops, "hybridep.combine.default"),
)
_GRAPH_FORWARD_ALL_GATHER_OPS = _existing_ops(
    _resolve_torch_op("aten.all_gather_into_tensor"),
    _resolve_torch_op("_c10d_functional.all_gather_into_tensor"),
)

    _resolve_torch_op("_c10d_functional.all_gather_into_tensor"),
)


def _resolve_wait_tensor() -> Optional[Any]:
    """Return the c10d ``wait_tensor`` overload when this torch ships it."""
    functional = getattr(torch.ops, "_c10d_functional", None)
    wait = getattr(functional, "wait_tensor", None)
    if wait is None:
        return None
    return getattr(wait, "default", None)


# getitem / wait_tensor extract or await their producer's value: their
# save/recompute decision is inherited, and remat walks through them.
_TRANSPARENT_OPS: Tuple[Any, ...] = tuple(
    target
    for target in (operator.getitem, _resolve_wait_tensor())
    if target is not None
)

    for target in (operator.getitem, _resolve_wait_tensor())
    if target is not None
)

_RECOMPUTE_POLICIES = (
    CheckpointPolicy.PREFER_RECOMPUTE,
    CheckpointPolicy.MUST_RECOMPUTE,
)

    CheckpointPolicy.MUST_RECOMPUTE,
)


def _node_stack_fqns(node: fx.Node) -> List[str]:
    """Module FQNs on ``node``'s ``nn_module_stack``, innermost first."""
    stack = node.meta.get("nn_module_stack") or {}
    return [fqn for fqn, _ in reversed(list(stack.values()))]


def _at_or_under(fqn: str, root: str) -> bool:
    """Whether ``fqn`` equals ``root`` or lives below it."""
    return fqn == root or fqn.startswith(f"{root}.")


def _first_container(fqns: List[str], roots: frozenset) -> Optional[str]:
    """Return the first root that contains one of ``fqns`` (innermost wins).

    ``roots`` is sorted so a (currently impossible) nested-root match is still
    deterministic. ``fqns`` is innermost-first, so the first hit is the

    ``roots`` is sorted so a (currently impossible) nested-root match is still
    deterministic. ``fqns`` is innermost-first, so the first hit is the
    innermost containing root.
    """
    for fqn in fqns:
        for root in sorted(roots):
            if _at_or_under(fqn, root):
                return root
    return None


class _CheckpointScope:
    """Resolved recompute scope for one compilation.

    ``_resolve_scope`` always fills every field; the ``None`` defaults exist
    for direct construction in tests and mean "no region" / "no boundary

        unit_fqns: Boundary units for the cross-unit save rule. Empty (or
            ``None``) disables the rule.
    """

    def __init__(
        self,
        region_fqns: Optional[frozenset] = None,
        exclude_fqns: frozenset = frozenset(),
        unit_fqns: Optional[frozenset] = None,

        exclude_fqns: frozenset = frozenset(),
        unit_fqns: Optional[frozenset] = None,
    ) -> None:
        """Store the resolved scope membership sets."""
        self.region_fqns = region_fqns
        self.exclude_fqns = exclude_fqns
        self.unit_fqns = unit_fqns


def _resolve_scope(
    model: Optional[torch.nn.Module], pass_config: PassConfig
) -> _CheckpointScope:
    """Resolve the checkpoint scope from the dynamic-config vocabulary.

        ValueError: When the live model is missing, no transformer blocks
            are discovered, or the adapter declares no recompute region in
            the selected layers.
    """
    if model is None:
        raise ValueError(
            "ActivationCheckpointPass requires the live model via kwargs "
            "(model=...) so it can discover the transformer blocks to "
            "checkpoint"
        )
    containers = _find_transformer_layer_container_infos(model)
    if not containers:
        raise ValueError(
            f"{type(model).__name__} has no module with a 'gradient_checkpointing' "
            "attribute and a non-empty repeated block container; activation "
            "checkpointing cannot resolve its layers"
        )
    block_fqns = frozenset(
        block.fqn for container in containers for block in container.blocks
    )
    has_count = pass_config.ac_layer_count is not None
    has_indices = pass_config.ac_layer_indices is not None
    if not has_count and not has_indices:
        return _CheckpointScope(region_fqns=block_fqns, unit_fqns=block_fqns)
    selection = SimpleNamespace(
        layer_count=pass_config.ac_layer_count,
        layer_indices=pass_config.ac_layer_indices,
    )
    selected_layers = frozenset(_selected_checkpoint_layer_fqns(containers, selection))
    if not selected_layers:
        # Zero selected layers is a no-op, not an error.
        _LOG.info("activation checkpoint layer selection matched zero layers")
        return _CheckpointScope(region_fqns=frozenset(), unit_fqns=block_fqns)
    recompute_spec = _adapter_recompute_spec(model, required=False)
    if recompute_spec is None or not recompute_spec.region_patterns:
        return _CheckpointScope(region_fqns=selected_layers, unit_fqns=block_fqns)
    module_by_fqn = dict(model.named_modules())
    roots = _select_layer_region_roots(module_by_fqn, selected_layers, recompute_spec)
    if not roots:
        raise ValueError(
            "model adapter recompute patterns matched no region in the selected layers"
        )
    region_fqns = frozenset(fqn for fqn, _ in roots)
    excluded = _select_excluded_targets(module_by_fqn, roots, recompute_spec)
    exclude_fqns = frozenset(fqn for fqn, _ in excluded)
    return _CheckpointScope(
        region_fqns=region_fqns,
        exclude_fqns=exclude_fqns,
        unit_fqns=block_fqns,
    )

        unit_fqns=block_fqns,
    )


def _is_eligible(node: fx.Node, scope: _CheckpointScope) -> bool:
    """Whether ``node`` may carry a recompute decision at all.

    Nodes outside the checkpoint regions stay saved: activation
    checkpointing targets transformer-block forwards, mirroring the eager

    Nodes outside the checkpoint regions stay saved: activation
    checkpointing targets transformer-block forwards, mirroring the eager
    path that wraps the discovered blocks only.
    """
    if not scope.region_fqns:
        return False
    fqns = _node_stack_fqns(node)
    if not fqns:
        return False
    if _first_container(fqns, scope.exclude_fqns) is not None:
        return False
    return _first_container(fqns, scope.region_fqns) is not None


def _make_policy_fn(pass_config: PassConfig) -> Callable[[fx.Node], CheckpointPolicy]:
    """Build the per-node policy for ``ac_mode``.

    ``full`` recomputes every scoped node. ``selective`` saves the shared
    compute/communication operator set and forces recompute only on the

    compute/communication operator set and forces recompute only on the
    operators the eager policy must re-replay (in-place-unsafe outputs such
    as ``topk``); everything else prefers recompute.
    """
    if pass_config.ac_mode == "full":
        return lambda node: CheckpointPolicy.MUST_RECOMPUTE

    must_save = _SELECTIVE_AC_MUST_SAVE_OPS
    force_recompute = _SELECTIVE_AC_FORCE_RECOMPUTE_OPS

    def selective_policy(node: fx.Node) -> CheckpointPolicy:
        """Apply the shared selective-AC operator policy."""
        if node.target in force_recompute:
            return CheckpointPolicy.MUST_RECOMPUTE
        if node.target in must_save:
            return CheckpointPolicy.MUST_SAVE
        return CheckpointPolicy.PREFER_RECOMPUTE

    return selective_policy


def _unit_of(node: fx.Node, scope: _CheckpointScope) -> Optional[str]:
    """Return ``node``'s boundary unit (block or adapter region root), if any."""
    if not scope.unit_fqns:
        return None
    return _first_container(_node_stack_fqns(node), scope.unit_fqns)


def _tag_policy_pass(
    graph_module: fx.GraphModule,
    pass_config: PassConfig,
    scope: _CheckpointScope,
) -> None:

    ``full``): replaying a collective in backward changes communication
    semantics. All-gathers join the saved set when
    ``ac_save_forward_all_gather`` is enabled.
    """
    policy_fn = _make_policy_fn(pass_config)
    comm_ops = _GRAPH_AC_COMM_OPS
    if pass_config.ac_save_forward_all_gather:
        comm_ops = comm_ops | _GRAPH_FORWARD_ALL_GATHER_OPS
    for node in graph_module.graph.nodes:
        if node.op != "call_function" or node.meta.get("autograd_backward", False):
            continue
        if node.target in _TRANSPARENT_OPS:
            parent = node.args[0]
            if isinstance(parent, fx.Node) and "recompute" in parent.meta:
                node.meta["recompute"] = parent.meta["recompute"]
            continue
        if not _is_eligible(node, scope):
            continue
        if _is_rng_op(node) or node.target in comm_ops:
            node.meta["recompute"] = CheckpointPolicy.MUST_SAVE
            continue
        value = node.meta.get("val")
        if isinstance(value, torch.SymInt):
            node.meta["recompute"] = CheckpointPolicy.MUST_SAVE
            continue
        node.meta["recompute"] = policy_fn(node)


def _flatten_getitems(graph_module: fx.GraphModule) -> None:
    """Flatten tuple projections when this torch build provides the helper."""
    raise_getitems = getattr(compile_utils, "raise_getitems", None)
    if raise_getitems is not None:
        raise_getitems(graph_module)


def _is_mutable_op(target: Any) -> bool:
    """Return whether ``target`` is an in-place ``OpOverload``."""
    if not isinstance(target, OpOverload):
        return False
    return bool(getattr(getattr(target, "_schema", None), "is_mutable", False))


def _is_rng_op(node: fx.Node) -> bool:
    """Return True when a node's seeded nondeterminism cannot be ruled out."""
    return bool(
        hasattr(node.target, "tags")
        and torch.Tag.nondeterministic_seeded in node.target.tags
    )

        and torch.Tag.nondeterministic_seeded in node.target.tags
    )


def _tag_inplace_saves(graph_module: fx.GraphModule) -> None:
    """Force saves on in-place mutations and their forward consumers."""
    for node in graph_module.graph.nodes:
        if (
            node.op != "call_function"
            or node.meta.get("autograd_backward", False)
            or "recompute" not in node.meta
            or not _is_mutable_op(node.target)

            or node.meta.get("autograd_backward", False)
            or "recompute" not in node.meta
            or not _is_mutable_op(node.target)
        ):
            continue
        node.meta["recompute"] = CheckpointPolicy.MUST_SAVE
        for user in node.users:
            if (
                not user.meta.get("autograd_backward", False)
                and "recompute" in user.meta
            ):
                user.meta["recompute"] = CheckpointPolicy.MUST_SAVE


def _tag_cross_unit_saves(
    graph_module: fx.GraphModule, scope: _CheckpointScope
) -> None:
    """Save recomputable nodes read forward across a unit boundary.

    Their original activation stays alive for the forward consumer anyway, so
    recomputing them in backward would be pure waste.
    """
    for node in graph_module.graph.nodes:
        if (
            node.op != "call_function"
            or node.meta.get("autograd_backward", False)
            or node.meta.get("recompute") not in _RECOMPUTE_POLICIES
        ):
            continue
        if _has_cross_unit_forward_consumer(node, scope):
            node.meta["recompute"] = CheckpointPolicy.MUST_SAVE


def _tag_sac_policy(
    graph_module: fx.GraphModule,
    pass_config: PassConfig,
    scope: _CheckpointScope,
) -> None:

    on in-place mutations and their forward consumers; pass 3 forces saves on
    recomputable nodes read forward across a unit boundary (their original
    stays alive anyway, so the backward recompute would be pure waste).
    """
    if not scope.region_fqns:
        return
    _tag_policy_pass(graph_module, pass_config, scope)
    _tag_inplace_saves(graph_module)
    if scope.unit_fqns:
        _tag_cross_unit_saves(graph_module, scope)


def _has_cross_unit_forward_consumer(node: fx.Node, scope: _CheckpointScope) -> bool:
    """Whether ``node`` is read forward outside its own boundary unit.

    Transparent (inherited) getitem/wait_tensor chains are walked through;
    any other forward consumer in a different unit (or outside every unit,

    Transparent (inherited) getitem/wait_tensor chains are walked through;
    any other forward consumer in a different unit (or outside every unit,
    e.g. the epilogue) means the original must stay alive anyway.
    """
    node_unit = _unit_of(node, scope)
    policy = node.meta["recompute"]
    seen = {node}
    stack = [node]
    while stack:
        current = stack.pop()
        for user in current.users:
            if user in seen:
                continue
            seen.add(user)
            if user.op == "output" or user.meta.get("autograd_backward", False):
                continue
            if (
                user.op == "call_function"
                and user.target in _TRANSPARENT_OPS
                and user.meta.get("recompute") == policy
            ):
                stack.append(user)
                continue
            if _unit_of(user, scope) != node_unit:
                return True
    return False


def _duplicate_subgraph_get_attrs(
    graph_module: fx.GraphModule, fwd_node: fx.Node, bwd_target: fx.Node
) -> Dict[fx.Node, fx.Node]:
    """Give a recompute duplicate private copies of HOP subgraph get_attrs.

    ``get_attr`` nodes. Sharing the original would let a later regional
    Inductor pass place the shared subgraph in only one partition; fresh
    get_attr nodes keep each region self-contained.
    """
    duplicates: Dict[fx.Node, fx.Node] = {}
    for arg in fwd_node.all_input_nodes:
        if arg.op != "get_attr":
            continue
        if not isinstance(getattr(graph_module, arg.target, None), fx.GraphModule):
            continue
        with graph_module.graph.inserting_before(bwd_target):
            dup_attr = graph_module.graph.get_attr(arg.target)
            dup_attr.meta.update(arg.meta)
        duplicates[arg] = dup_attr
    return duplicates


def _remat(graph_module: fx.GraphModule) -> int:
    """Duplicate tagged recompute nodes in front of their backward consumers.

    Returns the number of duplicated forward nodes. The graph is mutated in
    place; originals whose consumers were all redirected are erased, and

    Returns the number of duplicated forward nodes. The graph is mutated in
    place; originals whose consumers were all redirected are erased, and
    ``raise_getitems`` flattens output/tuple projections for better memory.
    """
    graph = graph_module.graph
    all_nodes = list(graph.nodes)
    order = {node: idx for idx, node in enumerate(all_nodes)}

    # Map each tagged forward node to its earliest backward consumer; the
    # transitive walk through tagged ancestors mirrors the partitioner's
    # recompute-closure gathering (saved ancestors stay live, so only tagged

    # Map each tagged forward node to its earliest backward consumer; the
    # transitive walk through tagged ancestors mirrors the partitioner's
    # recompute-closure gathering (saved ancestors stay live, so only tagged
    # ones need duplicating).
    remat_targets: Dict[fx.Node, fx.Node] = {}
    for bwd_node in all_nodes:
        if not bwd_node.meta.get("autograd_backward", False):
            continue
        seen: Set[fx.Node] = set()
        stack = list(bwd_node.all_input_nodes)
        while stack:
            current = stack.pop()
            if current in seen or not must_recompute(current):
                continue
            seen.add(current)
            remat_targets.setdefault(current, bwd_node)
            stack.extend(current.all_input_nodes)

    if not remat_targets:
        return 0

    recomputed: Dict[fx.Node, fx.Node] = {}

    def remat_arg(arg: Any, get_attr_dups: Dict[fx.Node, fx.Node]) -> Any:
        """Redirect duplicate args to their recompute copies."""
        if isinstance(arg, fx.Node):
            if arg in get_attr_dups:
                return get_attr_dups[arg]
            if arg in recomputed:
                return recomputed[arg]
        return arg

    for fwd_node in sorted(remat_targets, key=order.__getitem__):
        bwd_target = remat_targets[fwd_node]
        get_attr_dups = _duplicate_subgraph_get_attrs(
            graph_module, fwd_node, bwd_target
        )
        with graph.inserting_before(bwd_target):
            duplicate = graph.node_copy(
                fwd_node,
                functools.partial(remat_arg, get_attr_dups=get_attr_dups),
            )
        duplicate.name = f"{fwd_node.name}_recomputed"
        duplicate.meta["autograd_backward"] = True
        # node_copy shares meta, including the policy tag. Drop it so a second
        # pass run cannot mistake the duplicate for a recompute candidate.
        duplicate.meta.pop("recompute", None)
        recomputed[fwd_node] = duplicate

    def redirect(arg: Any) -> Any:
        """Arg transform for backward consumers of recomputed originals."""
        if isinstance(arg, fx.Node) and arg in recomputed:
            return recomputed[arg]
        return arg

    direct_backward_consumers = {
        user
        for fwd_node in recomputed
        for user in fwd_node.users
        if user.meta.get("autograd_backward", False)

        for fwd_node in recomputed
        for user in fwd_node.users
        if user.meta.get("autograd_backward", False)
    }
    for bwd_node in direct_backward_consumers:
        bwd_node.args = fx.map_arg(bwd_node.args, redirect)
        bwd_node.kwargs = fx.map_arg(bwd_node.kwargs, redirect)

    for original in reversed(list(recomputed)):
        if not original.users:
            graph.erase_node(original)

    _flatten_getitems(graph_module)
    graph.lint()
    return len(recomputed)


class ActivationCheckpointPass(GraphPass):
    """Static-graph activation checkpointing (tag + recompute-materialize).

    Configure through ``PassConfig.ac_mode`` (``"off"`` / ``"full"`` /
    ``"selective"``) and the optional ``ac_layer_count`` /

    ``activation_checkpoint`` section is projected onto these fields by
    ``build_pass_config_from_trainer_config``.
    """

    name = "activation_checkpoint"

    def run(
        self,
        graph_module: fx.GraphModule,
        pass_config: PassConfig,
        **kwargs: Any,

        Returns:
            The transformed graph module.
        """
        if pass_config.ac_mode == "off":
            return graph_module
        scope = _resolve_scope(kwargs.get("model"), pass_config)
        _tag_sac_policy(graph_module, pass_config, scope)
        duplicated = _remat(graph_module)
        _LOG.info(
            "Activation checkpointing (mode=%s) recomputed %d node(s)",
            pass_config.ac_mode,
            duplicated,
        )
        return graph_module

        # 2. Activation checkpointing: recompute duplicates must exist while
        # the graph is still plain (no collectives yet), so FSDPPass inserts
        # parameter all-gathers / backward re-gathers around them.
        if self.config.ac_mode != "off":
            self.passes.append(ActivationCheckpointPass())

        # 3. Execution layer: Parallel dimension partitioning
        if getattr(self.config, "fsdp_enabled", False):
            self.passes.append(FSDPPass(parallel_plan=self.parallel_plan))