Diff Coverage

                else:
                    source_tensor = self._storage_tensor(source)
                    host_view.copy_(source_tensor.detach().reshape(-1).view(slot.shape), non_blocking=False)
                    if source_tensor.device.type != "cpu" and source is slot.tensor:
                        self.release_device_storage(slot)
                slot.host_offset = host_offset
                slot.cpu_tensor = host_view
                slot.bind_tensor(host_view)
                slot.state = "host"

        if not isinstance(storage_tensor, torch.Tensor):
            raise ValueError(f"Expected torch.Tensor for CPU mirror, got {type(tensor)!r}.")
        try:
            cpu_tensor = torch.empty_like(storage_tensor, device="cpu", pin_memory=True)
        except RuntimeError:
            cpu_tensor = torch.empty_like(storage_tensor, device="cpu")
        cpu_tensor.zero_()
        return cpu_tensor

    def make_device_tensor_like(self, param: Any, saved_tensor: Any) -> Any:

        self._packed_offload_events = {}

    def enqueue_packed_prefetch(self, batch_index: int, staging_index: int) -> None:
        """Enqueue one packed H2D and record its ready event."""
        copy_stream = self._get_copy_stream()
        with self.stream_context(copy_stream):
            self._copy_packed_to_device(batch_index, staging_index)
            ready_event = self.record_event(copy_stream)
        self._packed_ready_events[batch_index] = ready_event
        self._packed_tail_event = ready_event

    def wait_packed_prefetch(self, batch_index: int, staging_index: int) -> None:
        """Order the compute stream after the batch's packed transfer chain."""
        del staging_index

        del staging_index
        ready_event = self._packed_ready_events.get(batch_index)
        if ready_event is None:
            raise RuntimeError(f"Packed optimizer batch {batch_index} has no ready event.")
        self.wait_event(ready_event, self.current_stream())

    def activate_packed_batch(self, batch_index: int, staging_index: int) -> None:
        """Bind each swap slot to its slice of one staging arena."""
        batch_plan = self._packed_batch_plans[batch_index]

            next_index: Optional[int],
            staging_index: int,
    ) -> None:
        """Serialize current D2H and next same-parity H2D on the copy stream."""
        copy_stream = self._get_copy_stream()
        compute_event = self._record_current_stream_event()
        with self.stream_context(copy_stream):
            self.wait_event(compute_event, copy_stream)
            self._copy_packed_to_host(batch_index, staging_index)
            if next_index is not None:
                self._copy_packed_to_device(next_index, staging_index)
            chain_event = self.record_event(copy_stream)
        self._packed_offload_events[batch_index] = chain_event
        self._packed_tail_event = chain_event
        if next_index is not None:
            self._packed_ready_events[next_index] = chain_event

    def wait_packed_offload(self, batch_index: int) -> None:
        """Order the compute stream after a packed transfer chain."""
        offload_event = self._packed_offload_events.get(batch_index)

        """Order the compute stream after a packed transfer chain."""
        offload_event = self._packed_offload_events.get(batch_index)
        if offload_event is None:
            raise RuntimeError(f"Packed optimizer batch {batch_index} has no offload event.")
        self.wait_event(offload_event, self.current_stream())

    def finish_packed_offload(self, batch_index: int) -> None:
        """Make persistent pinned views authoritative after D2H completion."""
        batch_plan = self._packed_batch_plans[batch_index]

        if self._packed_tail_event is not None:
            # All D2H copies are serialized on the copy stream.  Waiting only
            # for its tail is sufficient before detaching views and releasing
            # the step-local device allocation.
            self.wait_event(self._packed_tail_event, None)
        active_slots = {
            id(slot): slot
            for batch_plan in self._packed_batch_plans
            for region in batch_plan.regions.values()

        if active_slots:
            compute_stream = self.current_stream()
            synchronize = getattr(compute_stream, "synchronize", None)
            if synchronize is not None:
                synchronize()
            for slot in active_slots.values():
                slot.cpu_tensor.copy_(self._storage_tensor(slot.tensor).detach(), non_blocking=False)
                slot.bind_tensor(slot.cpu_tensor)
                slot.state = "host"

                slot.state = "host"
                slot.event = None
        for arena in self._staging_arenas:
            if arena is None:
                continue
            # Drop views before shrinking the raw storage; otherwise a view
            # can keep the device allocation alive after the step ends.
            arena.dtype_views = {}
            arena.layout_signature = None

    @classmethod
    def matches(cls, optimizer: Any) -> bool:
        """Return whether this adapter supports ``optimizer``."""
        return isinstance(optimizer, cls.supported_cls)

    def validate(self) -> None:
        """Validate unsupported optimizer flags."""
        for group in self.optimizer.param_groups:

        units = []
        for group_index, group in enumerate(self.optimizer.param_groups):
            if self.is_new_adamw:
                group["step"] = (group.get("step") or 0) + 1
            for param in group["params"]:
                grad = getattr(param, "grad", None)
                if grad is None:
                    continue

                grad = getattr(param, "grad", None)
                state = self.optimizer.state.get(param)
                if grad is not None:
                    if getattr(grad, "is_sparse", False):
                        raise ValueError("Swap optimizer only supports dense Adam/AdamW gradients.")
                    state = self.optimizer.state[param]
                    self._init_param_state(param, grad, group)
                if state:
                    self._register_present_slots(param, state)

                if state:
                    self._register_present_slots(param, state)
                has_slots = any((id(param), key) in self._slots for key in self._configured_state_keys())
                if grad is None and not has_slots:
                    continue
                records.append((group_index, param, grad))

        self.runtime.prepare_packed_host(self._ordered_slots())
        self.publish_packed_state()

        for group_index, param, grad in records:
            state = self.optimizer.state[param]
            slots = self._build_slots(param, state)
            if grad is None and not any(slot.swappable and slot.packed for slot in slots):
                continue
            units.append(UpdateUnit(
                adapter_index=group_index,
                param=param,
                grad=grad,

        return {"units": units}

    def iter_update_units(self, step_context: Dict[str, Any]) -> List[UpdateUnit]:
        """Return units collected in ``prepare_step``."""
        return step_context["units"]

    def initial_slots(self) -> Iterable[SwapSlot]:
        """Discover optimizer states materialized before the swap wrapper was created."""
        slots = []

        return tuple(slots)

    def step_batch(self, batch: List[UpdateUnit], step_context: Dict[str, Any]) -> None:
        """Run Torch functional Adam/AdamW for one batch."""
        del step_context
        by_group: Dict[int, List[UpdateUnit]] = defaultdict(list)
        for unit in batch:
            by_group[unit.adapter_index].append(unit)
        for group_index, units in by_group.items():
            group = self.optimizer.param_groups[group_index]
            state_steps = []
            params = []
            grads = []
            exp_avgs = []
            exp_avg_sqs = []
            max_exp_avg_sqs = []
            for unit in units:
                if unit.grad is None:
                    continue
                state = self.optimizer.state[unit.param]
                params.append(unit.param)
                grads.append(unit.grad)
                exp_avgs.append(self._slot_tensor(unit, "exp_avg", state["exp_avg"]))
                exp_avg_sqs.append(self._slot_tensor(unit, "exp_avg_sq", state["exp_avg_sq"]))
                if group.get("amsgrad", False):
                    max_exp_avg_sqs.append(
                        self._slot_tensor(unit, "max_exp_avg_sq", state["max_exp_avg_sq"])
                    )
                if self.is_new_adamw:
                    state_steps.append(None)
                else:
                    state_steps.append(state["step"])

            if not params:
                continue

            if self.is_new_adamw:
                if params and params[0].device.type == "cpu":
                    # torch.optim._functional.adamw increments tensor state_steps
                    # internally. New AdamW already advanced group["step"] in
                    # prepare_step(), so feed step - 1 to preserve outer-step
                    # semantics for CPU-only tests.
                    step_tensor = torch.tensor(float(group["step"] - 1), dtype=torch.float32)
                    torch.optim._functional.adamw(
                        params,
                        grads,
                        exp_avgs,
                        exp_avg_sqs,

                        found_inf=None,
                        has_complex=False,
                    )
                else:
                    _new_adamw_func()(
                        params,
                        grads,
                        exp_avgs,
                        exp_avg_sqs,

                        weight_decay=group["weight_decay"],
                        eps=group["eps"],
                        maximize=group["maximize"],
                    )
                continue

            func = getattr(torch.optim._functional, self.functional_name)
            kwargs = {
                "amsgrad": group["amsgrad"],
                "beta1": group["betas"][0],
                "beta2": group["betas"][1],
                "lr": group["lr"],

                "grad_scale": getattr(self.optimizer, "grad_scale", None),
                "found_inf": getattr(self.optimizer, "found_inf", None),
                "has_complex": False,
            }
            if self.functional_name == "adam":
                if "decoupled_weight_decay" in inspect.signature(func).parameters:
                    kwargs["decoupled_weight_decay"] = self.decoupled_weight_decay or group.get(
                        "decoupled_weight_decay", False
                    )
            func(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, **kwargs)

    def finish_step(self, step_context: Any) -> Any:
        """Finish one outer optimizer step."""
        del step_context

    def all_slots(self) -> Iterable[SwapSlot]:
        """Iterate known swap slots."""
        return tuple(self._ordered_slots())

            ids_in_order.extend(group["params"])
        for param, param_id in zip(params_in_order, ids_in_order):
            saved_state = state_dict.get("state", {}).get(param_id)
            if not saved_state:
                continue
            exported_state = {}
            for key in self._state_keys_for_param(param):
                slot = self._slots.get((id(param), key))
                if slot is not None and slot.state == "host" and slot.cpu_tensor is not None and key in saved_state:
                    exported_state[key] = slot.cpu_tensor.detach().clone()
            for key, value in saved_state.items():
                if key not in exported_state:
                    exported_state[key] = copy.deepcopy(value)
            exported["state"][param_id] = exported_state

        removed: Dict[int, Dict[str, Any]] = {}
        swappable_keys = self._configured_state_keys()
        for param_id, saved_state in list(stripped.get("state", {}).items()):
            if not isinstance(saved_state, dict):
                continue
            for key in swappable_keys:
                if key in saved_state:
                    removed.setdefault(param_id, {})[key] = saved_state.pop(key)
        return stripped, removed

            current_params.extend(current_group["params"])
        for saved_id, param in zip(saved_ids, current_params):
            key_to_tensor = removed.get(saved_id, {})
            if not key_to_tensor:
                continue
            state = self.optimizer.state[param]
            for key, saved_tensor in key_to_tensor.items():
                cpu_tensor = self._cast_swappable_tensor_to_cpu(param, saved_tensor)
                if self.runtime.packed_enabled and self.runtime.is_packable_template(param, self.config.min_numel):
                    if self.runtime.is_distributed_tensor(param):
                        logical_tensor = torch.zeros_like(
                            param,
                            memory_format=torch.preserve_format,
                        )
                        slot = self._make_slot(key, logical_tensor)
                        state[key] = logical_tensor
                        self.runtime.release_device_storage(slot)
                    else:
                        slot = self._make_slot(key, None, template=param)
                        state[key] = cpu_tensor
                        slot.tensor = cpu_tensor
                    slot.cpu_tensor = cpu_tensor
                    slot.state = "host"
                    self._slots[(id(param), key)] = slot
                    continue
                device_tensor = self.runtime.make_empty_device_tensor_like(param, cpu_tensor)
                slot = self._make_slot(key, device_tensor)
                if slot.swappable:
                    state[key] = device_tensor
                    slot.cpu_tensor = self.runtime.make_cpu_tensor(cpu_tensor)
                    slot.state = "host"
                    self._slots[(id(param), key)] = slot
                    self.runtime.release_device_storage(slot)
                else:
                    device_tensor = self._cast_state_tensor_like_torch(
                        param,
                        saved_tensor,

                self._slots[(id(param), key)] = slot
                self.runtime.release_device_storage(slot)
                continue
            if key in configured_keys and self.runtime.is_packable_template(param, self.config.min_numel):
                self._slots[(id(param), key)] = self._make_slot(key, None, template=param)
                continue
            state[key] = torch.zeros_like(param, memory_format=torch.preserve_format)

    def _register_present_slots(self, param: Any, state: Dict[str, Any]) -> None:
        """Register configured state tensors that already exist in an optimizer state mapping."""

        slots = []
        for key in self._state_keys_for_param(param):
            tensor = state.get(key)
            if tensor is None:
                continue
            slot = self._slots.get((id(param), key))
            if slot is None:
                slot = self._make_slot(key, tensor)
                self._slots[(id(param), key)] = slot

        for key in keys:
            if key == "master_param":
                if self.config.state_keys is not None:
                    raise ValueError(f"Requested state key '{key}' is not available for {type(self.optimizer)!r}.")
                continue
            result.append(key)
        return tuple(result)

    def _cast_state_tensor_like_torch(

            key: str,
    ) -> Any:
        """Cast a loaded state tensor using PyTorch optimizer load semantics."""
        if not isinstance(saved_tensor, torch.Tensor):
            raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
        process = getattr(torch.optim.Optimizer, "_process_value_according_to_param_policy", None)
        if process is not None:
            return process(param, saved_tensor, saved_id, saved_groups, key).detach().clone()
        if key == "step":
            return saved_tensor.detach().clone()
        if param.is_floating_point():
            return saved_tensor.detach().to(dtype=param.dtype, device=param.device).clone()
        return saved_tensor.detach().to(device=param.device).clone()

    def _cast_swappable_tensor_to_cpu(self, param: Any, saved_tensor: Any) -> Any:
        """Cast swappable state dtype like PyTorch while keeping values on CPU."""
        if not isinstance(saved_tensor, torch.Tensor):
            raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
        if param.is_floating_point():
            return saved_tensor.detach().to(dtype=param.dtype, device="cpu")
        return saved_tensor.detach().to(device="cpu")

    @staticmethod
    def _default_state_keys() -> Tuple[str, ...]:
        return ("exp_avg", "exp_avg_sq", "max_exp_avg_sq")

def _new_adamw_func() -> Callable[..., Any]:
    """Return hyper-parallel's functional AdamW without importing it at module load."""
    from hyper_parallel.core.optimizer.adamw import (  # pylint: disable=import-outside-toplevel
        adamw as new_adamw_func,
    )
    return new_adamw_func


class SwapOptimizer(torch.optim.Optimizer):
    """Torch optimizer wrapper for Adam/AdamW state swap."""

        self.adapter.publish_packed_state()

    def __getattr__(self, name: str) -> Any:
        """Delegate unknown attributes to the base optimizer."""
        return getattr(self.optimizer, name)

    @property
    def param_groups(self):
        """Proxy parameter groups."""

        return self.optimizer.param_groups

    @param_groups.setter
    def param_groups(self, value) -> None:
        self.optimizer.param_groups = value

    @property
    def state(self):
        """Proxy optimizer state."""

    @property
    def defaults(self):
        """Proxy optimizer defaults."""
        return self.optimizer.defaults

    def add_param_group(self, param_group: Dict[str, Any]) -> None:
        """Proxy param group addition."""
        self.optimizer.add_param_group(param_group)

    def zero_grad(self, set_to_none: bool = True) -> None:
        """Proxy gradient clearing."""
        self.optimizer.zero_grad(set_to_none=set_to_none)

    def step(self, closure: Optional[Any] = None) -> Any:
        """Run one optimizer step with pipeline state swap."""
        if closure is not None:
            raise ValueError("Swap optimizer does not support closure.")
        with self._no_grad_context():
            step_context = self.adapter.prepare_step()
            units = self.adapter.iter_update_units(step_context)
            batches = self.runtime.partition(units)
            self.runtime.run_pipeline(batches, step_context, self.adapter.step_batch)
            return self.adapter.finish_step(step_context)

    def state_dict(self) -> Dict[str, Any]:
        """Return optimizer state dict using CPU mirrors for swappable tensors."""
        return self.adapter.checkpoint_state_dict()

        )

    @contextlib.contextmanager
    def _no_grad_context(self) -> Iterable[None]:
        with torch.no_grad():
            yield