Diff Coverage

        """Resolve contexts and capture RNG state before the forward pass runs."""
        device_type = _infer_device_type(*args)
        device_module = _get_device_module(device_type)
        forward_context, recompute_context = _resolve_contexts(context_fn)
        device_autocast_kwargs, cpu_autocast_kwargs = _get_autocast_kwargs(device_type)

        # Device RNG is captured only when the device is already initialized;
        # otherwise the device indexes its state lazily at first use.
        device_initialized = preserve_rng_state and getattr(device_module, "_initialized", False)
        device_ids, device_states = _get_device_states(device_type, *args) if device_initialized else ([], [])

        return cls(
            device_type=device_type,
            device_module=device_module,
            forward_context=forward_context,
            recompute_context=recompute_context,

        )

    def check_device_state_preserved(self, preserve_rng_state: bool) -> None:
        """Fail when the device was first initialized inside the checkpointed forward."""
        device_initialized = getattr(self.device_module, "_initialized", False)
        if preserve_rng_state and not self.device_was_initialized and device_initialized:
            raise RuntimeError(
                "The device state was initialized inside a Hyper checkpoint forward, so its initial RNG state "
                "could not be preserved. Initialize the device before entering checkpoint."
            )

            )

    def recompute_fn(self, function: Callable, *inputs: Any) -> None:
        """Restore execution state and rerun the checkpointed ``function``."""
        function_kwargs, *function_args = inputs
        with torch.random.fork_rng(
            devices=self.device_ids,
            enabled=self.cpu_state is not None,
            device_type=self.device_type,
        ):
            if self.cpu_state is not None:
                torch.set_rng_state(self.cpu_state)
                if self.device_was_initialized:
                    _set_device_states(self.device_type, self.device_ids, self.device_states)

            device_autocast_context = contextlib.nullcontext()
            if self.device_autocast_kwargs is not None:
                device_autocast_context = torch.amp.autocast(
                    device_type=self.device_type, **self.device_autocast_kwargs
                )
            with device_autocast_context, torch.amp.autocast("cpu", **self.cpu_autocast_kwargs), self.recompute_context:
                function(*function_args, **function_kwargs)


def _resolve_contexts(context_fn: Callable) -> Tuple[Any, Any]:
    """Validate and unpack the (forward, recompute) context tuple from ``context_fn``."""

    """Validate and unpack the (forward, recompute) context tuple from ``context_fn``."""
    contexts = context_fn()
    if not isinstance(contexts, tuple) or len(contexts) != 2:
        raise ValueError("context_fn must return a (forward_context, recompute_context) tuple.")
    return contexts[0], contexts[1]


def _checkpoint_without_reentrant_generator(
    function: Callable,

    """Set up eager checkpoint state around the caller's forward execution."""
    metadata_fn = _resolve_metadata_fn(determinism_check)
    state = _ExecutionState.capture(context_fn, preserve_rng_state, args)

    frame = _CheckpointFrame(partial(state.recompute_fn, function), early_stop, metadata_fn)
    dummy = torch.empty((0,), requires_grad=True)
    frame.input_saver = _NoopSaveInputs.apply(dummy, kwargs, *args)

    if frame.input_saver.grad_fn is None:

    if frame.input_saver.grad_fn is None:
        yield
        return

    if _RECOMPUTE_SESSION.get() is not None:
        raise CheckpointError("Nested checkpoint is not supported during scheduled recomputation.")

    collector = _RECOMPUTE_COLLECTOR.get()
    if collector is not None:

    collector = _RECOMPUTE_COLLECTOR.get()
    if collector is not None:
        collector.append(frame)
    try:
        with _create_checkpoint_hooks(frame), state.forward_context:
            yield
        frame.forward_completed = True
        state.check_device_state_preserved(preserve_rng_state)
    except BaseException:
        if collector is not None and frame in collector:
            collector.remove(frame)
        raise

    return CheckpointPolicy.MUST_SAVE if b else CheckpointPolicy.PREFER_RECOMPUTE


def _make_swap_entry(x, *, has_alias, funcname, group_swap, cpu_pool):
    return _SwapCacheEntry(
        _maybe_detach(x, has_alias), funcname, group_swap=group_swap, cpu_pool=cpu_pool
    )

    recomputation may need to unshard them inside that region. These allocation,
    copy and collective operators manage parameters rather than model
    activations, and therefore must not be matched against the forward replay.
    """
    ignore_sac_ops(_FSDP_SAC_IGNORED_OPS)


def compile_selective_checkpoint_policy(
    ctx: Any,

    Returns:
        The checkpoint policy for ``func``.
    """
    del ctx, args, kwargs
    if func in _SELECTIVE_AC_FORCE_RECOMPUTE_OPS:
        return CheckpointPolicy.MUST_RECOMPUTE
    if func in _SELECTIVE_AC_MATMUL_OPS:
        return CheckpointPolicy.MUST_SAVE
    if func in _SELECTIVE_AC_MUST_SAVE_OPS:
        return CheckpointPolicy.MUST_SAVE
    return CheckpointPolicy.MUST_RECOMPUTE


def _make_selective_checkpoint_policy_fn() -> Callable:
    """Create an isolated eager selective activation checkpointing policy."""

    running fully eager and memory-unbounded.
    """
    if wrapped_count != 0 or not containers:
        return
    layer_count = sum(len(container.blocks) for container in containers)
    logger.warning(
        "%s activation checkpointing wrapped no module on %d layer(s) in %s; the "
        "model is running without activation checkpointing. Expected submodule "
        "names may not match this architecture, or the layers are already wrapped.",
        activation_checkpoint.capitalize(),

            sub_config is not config and not hasattr(sub_config, "use_cache")
        ):
            continue
        if getattr(sub_config, "use_cache", None) is not False:
            _try_disable_use_cache(sub_config)
    return False


def _try_disable_use_cache(sub_config: Any) -> None:

def _try_disable_use_cache(sub_config: Any) -> None:
    """Best-effort disable of ``use_cache`` on one config object."""
    try:
        sub_config.use_cache = False
    except Exception:  # pylint: disable=broad-exception-caught
        # Configuration objects may reject assignment with custom errors.
        pass


def _validate_activation_checkpoint_config(
    activation_checkpoint: Optional[str],

    Returns:
        The number of layers, or submodules in the KV-shared fallback, wrapped.
    """
    if has_kv_sharing:
        logger.warning(
            "Selective activation checkpointing is not supported for KV-shared models; "
            "falling back to submodule activation checkpointing."
        )
        return apply_submodule_checkpointing(
            ac_layers,
            has_kv_sharing,
            enable_compile=enable_compile,
            swap_inputs=swap_inputs,

    ac_layers = _flatten_layer_container_infos(containers)
    has_kv_sharing = _detect_kv_sharing_and_maybe_disable_cache(model)

    if hasattr(model, "gradient_checkpointing_disable"):
        model.gradient_checkpointing_disable()

    if activation_checkpoint == "selective":
        wrapped_count = _apply_selective_checkpointing(
            containers,