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

        """Acquire the pinned CPU buffer that receives one bucket's D2H copy."""
        numel = bucket["total_numel"]
        if bucket["cpu_pool"] is None:
            return _get_cpu_pinned_buf(bucket["dtype_key"], numel, bucket["dtype"])
        raw_buf = bucket["cpu_pool"].acquire(bucket["total_bytes"])
        try:
            return raw_buf.view(bucket["dtype"])
        except Exception as exc:
            bucket["cpu_pool"].release(raw_buf)
            raise RuntimeError(
                "Failed to create a typed CPU-pool view for packed activation bucket: "
                f"group={self.group_name!r}, bucket={bucket_key!r}, "
                f"requested_dtype={bucket['dtype']}, total_numel={numel}, "
                f"total_bytes={bucket['total_bytes']}, raw_buffer_shape={tuple(raw_buf.shape)}, "

            ) from exc

    def _release_offloaded_bucket_bufs(self, group_cpu_bufs, copy_stream):
        """Give back every buffer already filled before a failed D2H loop."""
        release_event = _backend.new_event()
        release_event.record(copy_stream)
        for bucket_key, cpu_buf in group_cpu_bufs.items():
            cpu_pool = self._packed_buckets[bucket_key]["cpu_pool"]
            if cpu_pool is not None:
                cpu_pool.release(cpu_buf, event=release_event)
            else:
                _return_cpu_pinned_buf(cpu_buf)

    def _offload_buckets_d2h(self, group_device_bufs, copy_stream):
        """One-shot D2H per packed bucket."""
        group_cpu_bufs = {}

                cpu_buf[:bucket["total_numel"]].copy_(
                    group_device_bufs[bucket_key], non_blocking=True
                )
        except Exception as exc:
            self._release_offloaded_bucket_bufs(group_cpu_bufs, copy_stream)
            failed_bucket = self._packed_buckets.get(active_bucket_key, {})
            raise RuntimeError(
                "Failed to offload packed activation bucket from device to CPU: "
                f"group={self.group_name!r}, bucket={active_bucket_key!r}, "

        Units without a gradient are skipped.  ``max_exp_avg_sqs`` stays empty
        unless the group is amsgrad, and ``state_steps`` holds ``None`` for new
        AdamW, which advances the step counter itself.
        """
        args = GroupArgs([], [], [], [], [], [])
        for unit in units:
            if unit.grad is None:
                continue
            state = self.optimizer.state[unit.param]
            args.params.append(unit.param)
            args.grads.append(unit.grad)
            args.exp_avgs.append(self._slot_tensor(unit, "exp_avg", state["exp_avg"]))
            args.exp_avg_sqs.append(self._slot_tensor(unit, "exp_avg_sq", state["exp_avg_sq"]))
            if group.get("amsgrad", False):
                args.max_exp_avg_sqs.append(
                    self._slot_tensor(unit, "max_exp_avg_sq", state["max_exp_avg_sq"])
                )
            if self.is_new_adamw:
                args.state_steps.append(None)
            else:
                args.state_steps.append(state["step"])
        return args

    def step_batch(self, batch: List[UpdateUnit], step_context: Dict[str, Any]) -> None:
        """Run Torch functional Adam/AdamW for one batch."""
        del step_context

            self._step_group(self.optimizer.param_groups[group_index], units)

    def _step_group(self, group: Dict[str, Any], units: List[UpdateUnit]) -> None:
        """Run one group's parameters through the matching functional Adam/AdamW."""
        args = self._collect_group_args(units, group)
        params = args.params

        if not params:
            return

        if not params:
            return

        if self.is_new_adamw:
            self._step_new_adamw(
                group,
                args.params,
                args.grads,
                args.exp_avgs,

        }
        if self.functional_name == "adam":
            if "decoupled_weight_decay" in inspect.signature(func).parameters:
                kwargs["decoupled_weight_decay"] = self._decoupled_weight_decay(group)
        func(
            args.params,
            args.grads,
            args.exp_avgs,
            args.exp_avg_sqs,

    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,