Diff Coverage

                logical_bytes = self.val.numel() * _backend.get_element_size(self.val)
                self._cpu_pool_buffer = self.cpu_pool.acquire(logical_bytes)
                try:
                    self.val_cpu = self._cpu_pool_buffer.view(self.val.dtype).reshape(self.val.shape)
                except Exception as exc:
                    pool_buffer = self._cpu_pool_buffer
                    self.release_cpu_buffer()
                    raise RuntimeError(
                        "Failed to create a typed CPU-pool view for activation tensor "
                        f"from {self.funcname!r}: shape={tuple(self.val.shape)}, dtype={self.val.dtype}, "
                        f"logical_bytes={logical_bytes}, pool_buffer_shape={tuple(pool_buffer.shape)}, "
                        f"pool_buffer_dtype={pool_buffer.dtype}. Original error: {exc}"

            if self.cpu_pool is not None or self.is_slice_tensor:
                self.val_cpu.copy_(self.val, non_blocking=True)
            else:
                self.val_cpu.untyped_storage().copy_(self.val.untyped_storage(), non_blocking=True)
        except Exception as exc:
            if self.cpu_pool is not None and self._cpu_pool_buffer is not None:
                release_event = _backend.new_event()
                release_event.record(_backend.get_current_stream())
                self.release_cpu_buffer(release_event)

                release_event = _backend.new_event()
                release_event.record(_backend.get_current_stream())
                self.release_cpu_buffer(release_event)
            self.val_cpu = None
            copy_mode = "tensor" if self.cpu_pool is not None or self.is_slice_tensor else "storage"
            raise RuntimeError(
                "Failed to offload activation tensor from device to CPU: "
                f"source={self.funcname!r}, shape={tuple(self.val.shape)}, dtype={self.val.dtype}, "
                f"device={self.val.device}, copy_mode={copy_mode}, "
                f"cpu_pool={'enabled' if self.cpu_pool is not None else 'disabled'}. "

                cpu_pool = bucket["cpu_pool"]
                if cpu_pool is None:
                    cpu_buf = _get_cpu_pinned_buf(dtype_key, numel, bucket["dtype"])
                else:
                    raw_buf = cpu_pool.acquire(bucket["total_bytes"])
                    try:
                        cpu_buf = raw_buf.view(bucket["dtype"])
                    except Exception as exc:
                        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)}, "

                            f"raw_buffer_dtype={raw_buf.dtype}. Original error: {exc}"
                        ) from exc
                group_cpu_bufs[bucket_key] = cpu_buf
                cpu_buf[:numel].copy_(group_device_bufs[bucket_key], non_blocking=True)
        except Exception as exc:
            release_event = _backend.new_event()
            release_event.record(copy_stream)
            for completed_bucket_key, cpu_buf in group_cpu_bufs.items():
                bucket = self._packed_buckets[completed_bucket_key]
                if bucket["cpu_pool"] is not None:
                    bucket["cpu_pool"].release(cpu_buf, event=release_event)
                else:
                    _return_cpu_pinned_buf(cpu_buf)
            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}, "
                f"dtype={failed_bucket.get('dtype', 'unknown')}, "
                f"total_numel={failed_bucket.get('total_numel', 'unknown')}, "

            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"

        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.
        """
        params: List[Any] = []
        grads: List[Any] = []
        exp_avgs: List[Any] = []
        exp_avg_sqs: List[Any] = []
        max_exp_avg_sqs: List[Any] = []
        state_steps: List[Any] = []
        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"])
        return params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps

    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():
            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."""
        params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps = (
            self._collect_group_args(units, group)
        )

        if not params:
            return

        if self.is_new_adamw:
            self._step_new_adamw(group, params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs)
            return

        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(group)
        func(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, **kwargs)

    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,

        """Restore one removed checkpoint buffer as a host or device-resident slot."""
        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 = 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
            return
        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,

        container_by_compute_index[index]
        for index in range(fwd_leaf.compute_index + 1, bwd_leaf.compute_index)
    ]
    if not compute_between:
        return None

    stage_index = fwd_leaf.step.stage_index
    micro_index = fwd_leaf.step.micro_index
    # Always launch offload immediately after the FWD container so that the

        return self._boundary_p2p

    def __eq__(self, value):
        if not isinstance(value, MetaStep):
            return NotImplemented
        return (self.type == value.type
                and self.micro_index == value.micro_index
                and self.stage_index == value.stage_index
                and self.sub_steps == value.sub_steps)

                and self.sub_steps == value.sub_steps)

    def __ne__(self, value):
        if not isinstance(value, MetaStep):
            return NotImplemented
        return not self.__eq__(value)

    def __hash__(self):
        return hash((self.type, self.micro_index, self.stage_index))

    def __str__(self):
        if self.sub_steps:
            sub = ", ".join(str(s) for s in self.sub_steps)
            return (f"MetaStep(type={self.type}, micro_index={self.micro_index}, "
                    f"stage_index={self.stage_index}, sub_steps=[{sub}])")
        return f"MetaStep(type={self.type}, micro_index={self.micro_index}, stage_index={self.stage_index})"

    def __repr__(self):
        return self.__str__()

    @staticmethod
    def from_str(step_str):
        """Parse a MetaStep from its string representation."""