Diff Coverage

                    # their memory instead of aliasing the slot tensors.
                    self._localize_state_dict_tensors(sd)
                overlap = set(merged.keys()) & set(sd.keys())
                if overlap:
                    raise KeyError(
                        f"Key clash detected while merging state dict for optimizer '{name}': "
                        f"{', '.join(sorted(overlap))}"
                    )
                merged.update(sd)

                cleanup = getattr(opt, "cleanup_synced_state", None)
                if callable(cleanup):
                    cleanup()
                    continue
                finish = getattr(opt, "_finish_swap_checkpoint", None)
                if callable(finish):
                    finish()

    @staticmethod
    def _localize_state_dict_tensors(state_dict: Dict[str, Any]) -> None:
        """Replace exported tensors with CPU copies that own their storage.

                return to_local_if_dtensor(value).detach().cpu().clone()
            if isinstance(value, dict):
                return {key: localize(item) for key, item in value.items()}
            if isinstance(value, list):
                return [localize(item) for item in value]
            if isinstance(value, tuple):
                return tuple(localize(item) for item in value)
            return value

        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(AdamUpdateUnit(
                adapter_index=group_index,
                param=param,
                grad=grad,

        """
        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(_slot_tensor(unit, "exp_avg", state["exp_avg"]))

            args.grads.append(unit.grad)
            args.exp_avgs.append(_slot_tensor(unit, "exp_avg", state["exp_avg"]))
            args.exp_avg_sqs.append(_slot_tensor(unit, "exp_avg_sq", state["exp_avg_sq"]))
            if group.get("amsgrad", False):
                args.max_exp_avg_sqs.append(
                    _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

        args = self._collect_group_args(units, group)
        params = args.params

        if not params:
            return

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

                args.exp_avgs,
                args.exp_avg_sqs,
                args.max_exp_avg_sqs,
            )
            return

        func = getattr(torch.optim._functional, self.functional_name)
        kwargs = {
            "amsgrad": group["amsgrad"],

            "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(group)
        func(
            args.params,
            args.grads,
            args.exp_avgs,

        )

    def _decoupled_weight_decay(self, group: Dict[str, Any]) -> bool:
        """Resolve the decoupled weight decay flag for one parameter group."""
        return self.decoupled_weight_decay or group.get("decoupled_weight_decay", False)

    def _step_new_adamw(
            self,
            group: Dict[str, Any],

            exp_avg_sqs: Sequence[Any],
            max_exp_avg_sqs: Sequence[Any],
    ) -> None:
        """Run the new AdamW functional for one group."""
        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,

                grad_scale=None,
                found_inf=None,
                has_complex=False,
            )
            return
        _new_adamw_func()(
            params,
            grads,
            exp_avgs,
            exp_avg_sqs,

                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)

    @staticmethod
    def _is_active_device_slot(slot: SwapSlot, key: str) -> bool:

    @staticmethod
    def _is_active_device_slot(slot: SwapSlot, key: str) -> bool:
        """Return whether ``slot`` holds ``key`` live on the local device."""
        if slot.name != key or not slot.swappable or slot.state != "device":
            return False
        return slot.tensor is not None

    @staticmethod
    def default_state_keys() -> Tuple[str, ...]:
        """Return the swap state keys an Adam/AdamW optimizer owns."""

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


ADAM_ADAPTERS = (
    TorchNewAdamWAdapter,

        if not state:
            return
        buffer = state.get("momentum_buffer")
        if buffer is None:
            return
        if getattr(param, "grad", None) is not None and getattr(param.grad, "is_sparse", False):
            raise ValueError("Muon swap optimizer only supports dense gradients.")
        local = self._runtime_local_tensor(buffer)
        if not isinstance(local, torch.Tensor):
            raise ValueError(
                f"Muon swap optimizer requires tensor momentum state, got {type(buffer)!r}."
            )
        if local.layout != torch.strided:
            raise ValueError("Muon swap optimizer requires dense momentum state.")
        if not local.is_contiguous():
            raise ValueError("Muon swap optimizer requires contiguous momentum state.")

    @staticmethod

        scheduled: set[int] = set()
        for group_index, group in enumerate(self.optimizer.param_groups):
            info = schedule.get(group_index)
            if info is None:
                scheduled.update(id(param) for param in group["params"])
                continue
            scheduled.update(id(param) for param in info.get("no_comm", []))
            for batch_group in info.get("batch_groups", []):
                for assignment in batch_group.get("sub_batches", []):
                    scheduled.update(id(param) for param in assignment.owned_params)

    def _params_with_local_state(self, params: Iterable[Any]) -> Iterable[Any]:
        """Keep checkpoint-restored momentum only for local schedule owners."""
        scheduled_ids = self._scheduled_param_ids()
        if scheduled_ids is None:
            return params
        return tuple(param for param in params if id(param) in scheduled_ids)

    def initial_slots(self) -> Iterable[SwapSlot]:
        """Discover only local-owner momentum state before the first step."""

        no_comm" rule the bare step applies.
        """
        info = getattr(self.optimizer, "_hsdp_assignment_batches", {}).get(group_index)
        if not info:
            return list(self.optimizer.unshard_params_by_group.get(group_index, []))
        return list(info.get("no_comm", []))

    def _hsdp_linear_batches(self, group_index: int) -> List[HSDPGroupAssignment]:
        """Flatten one group's nested HSDP schedule into a linear batch list."""

        allocates momentum when the key is missing, so a registered slot is
        written in place and never replaced.
        """
        if getattr(param, "grad", None) is None:
            return
        if self._runtime_local_tensor(param.grad).layout != torch.strided:
            raise ValueError("Muon swap optimizer only supports dense gradients.")
        state = self.optimizer.state[param]
        for key in self._configured_state_keys():
            if (id(param), key) in self._slots:
                continue

            if (id(param), key) in self._slots:
                continue
            device_tensor = self.runtime.make_empty_state_like(param)
            if not self._runtime_local_tensor(device_tensor).is_contiguous():
                raise ValueError("Muon swap optimizer requires contiguous momentum state.")
            state[key] = device_tensor
            slot = self._make_slot(key, device_tensor)
            if not slot.swappable:
                # Too small or otherwise ineligible: keep ordinary device state.

        for param in params:
            for key in keys:
                slot = self._slots.get((id(param), key))
                if slot is None or id(slot) in seen:
                    continue
                slots.append(slot)
                seen.add(id(slot))
        return slots

            self._apply_unshard_update(group, ns_inputs)
        else:
            hsdp_assign = unit.hsdp_assign
            if hsdp_assign is None:
                raise RuntimeError("Muon swap HSDP unit is missing its group assignment.")
            if hsdp_assign.is_shard:
                self.optimizer._process_shard_params(
                    group, ns_inputs, [hsdp_assign], unit.group_index,
                    buffer_cache={},

    @property
    def copy_stream(self) -> Optional[Any]:
        """Return the transfer stream after its lazy initialization."""
        return self._copy_stream

    def set_packed_slot_bind_hook(self, hook: Callable[[SwapSlot], None]) -> None:
        """Install the adapter hook that republishes a packed slot's live view."""
        self._packed_slot_bind_hook = hook

            self, packed_slots: Sequence[SwapSlot],
    ) -> Tuple[Tuple[Any, ...], ...]:
        """Validate candidates and return their persistent layout signature."""
        if self._packed_batch_plans or self._packed_tail_event is not None:
            raise RuntimeError("Cannot rebuild packed optimizer layout during an active transfer.")
        slot_ids = [id(slot) for slot in packed_slots]
        if len(slot_ids) != len(set(slot_ids)):
            raise RuntimeError("Packed optimizer layout contains a duplicate state slot.")
        self.validate_packed_devices(packed_slots)

            self._copy_stream.synchronize()
        for event in _iter_unique_events(packed_slots):
            self.wait_event(event, None)
        if any(slot.state in ("h2d", "d2h") for slot in packed_slots):
            self.synchronize_cpu_mirrors(packed_slots)

    def invalidate_packed_layout(self) -> None:
        """Discard an idle packed layout after state slots have been rebound.

        swappable_slots = [slot for slot in slots if slot.swappable]
        if not swappable_slots:
            return False
        if len({id(slot) for slot in swappable_slots}) != len(swappable_slots):
            raise RuntimeError("Packed optimizer batches contain a duplicate state slot.")
        devices = {slot.device for slot in swappable_slots}
        eligible = len(devices) == 1 and None not in devices and all(
            slot.packed
            and slot.cpu_tensor is not None

        Raises:
            NotImplementedError: Always, unless a subclass overrides it.
        """
        raise NotImplementedError

    def __init__(self, optimizer: Any, config: Any, runtime: Any) -> None:
        """Bind a leaf optimizer to its state slots and transfer runtime."""
        self.optimizer = optimizer

    def on_packed_slot_bound(self, slot: SwapSlot) -> None:
        """Publish a packed host or staging view into optimizer state."""
        binding = self._packed_state_bindings.get(id(slot))
        if binding is None:
            raise RuntimeError(f"Packed optimizer slot {slot.name!r} has no live state binding.")
        param, key = binding
        state = self.optimizer.state.get(param)
        if state is None:
            raise RuntimeError(f"Packed optimizer slot {slot.name!r} has no optimizer state entry.")
        state[key] = slot.tensor

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

        else:
            # Identity verifies a missing map; it never infers a replacement.
            if any(tensor is tracked.tensor or tensor is tracked.cpu_tensor for tracked in self._slots.values()):
                raise ValueError("Replacing a swap parameter requires an explicit new_to_old_param mapping.")
            source = tensor
        source = self.runtime._storage_tensor(source)
        if source.numel() != local_param.numel():
            raise ValueError(
                f"Cannot rebind {key!r} with local shape {tuple(source.shape)} "

        cpu_tensor = self._cast_swappable_tensor_to_cpu(param, source).reshape(local_param.shape)
        device_tensor = self.runtime.make_empty_state_like(param)
        rebound = self._make_slot(key, device_tensor)
        if rebound.swappable:
            rebound.cpu_tensor = self.runtime.make_cpu_tensor(cpu_tensor)
            rebound.state = "host"
            self.runtime.release_device_storage(rebound)
            if rebound.packed:
                rebound.bind_tensor(rebound.cpu_tensor)
        else:
            self.runtime._storage_tensor(device_tensor).copy_(cpu_tensor)
        return rebound

        if isinstance(value, torch.Tensor):
            to_local = getattr(value, "to_local", None)
            local = to_local() if callable(to_local) else value
            if local is not value:
                return local.detach().cpu().clone()
            return copy.deepcopy(value)
        if isinstance(value, dict):
            return {
                key: StateSwapAdapter._copy_checkpoint_value(item)

        state = self.optimizer.state[param]
        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 = self.runtime.make_empty_state_like(param)
                slot = self._make_slot(key, logical_tensor)
                state[key] = logical_tensor
                self.runtime.release_device_storage(slot)
            else:

            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()

    @staticmethod
    def _cast_swappable_tensor_to_cpu(param: Any, saved_tensor: Any) -> Any:
        """Cast swappable state dtype like PyTorch while keeping values on CPU."""

    @staticmethod
    def _cast_swappable_tensor_to_cpu(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")

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

            self._export_active = previous_active

    def _finish_swap_checkpoint(self) -> None:
        """Return checkpoint-broadcast state to its normal swapped form."""
        self.adapter.finish_checkpoint()

    def _broadcast_state_fused_for_ckpt(self) -> None:
        """Broadcast wrapped optimizer state with swap storage materialized."""
        broadcast = getattr(self.optimizer, "_broadcast_state_fused_for_ckpt", None)