From f595cf659ce1f13bab1033c01d748535c0a25437 Mon Sep 17 00:00:00 2001
From: fandawei <fandawei2@huawei.com>
Date: Fri, 23 Jan 2026 16:20:08 +0800
Subject: [PATCH] checkpoint support swap_input && component decoupling

---
 .../activation_checkpoint/__init__.py}        | 22 +-----
 .../activation_checkpoint.py                  | 66 +++++++++++++++++
 .../activation_checkpoint/swap.py             | 72 ++++++++++---------
 hyper_parallel/platform/mindspore/platform.py | 69 ++++++++++++++++++
 hyper_parallel/platform/platform.py           | 49 +++++++++++++
 .../torch/activation_checkpoint/__init__.py   |  3 -
 .../activation_checkpoint/activation_swap.py  |  2 +-
 .../torch/activation_checkpoint/sac.py        | 34 +--------
 hyper_parallel/platform/torch/platform.py     | 45 ++++++++++++
 tests/torch/activation_checkpoint/test_ac.py  | 10 +--
 .../activation_checkpoint/test_act_swap.py    |  4 +-
 11 files changed, 281 insertions(+), 95 deletions(-)
 rename hyper_parallel/{platform/torch/activation_checkpoint/activation_checkpoint.py => core/activation_checkpoint/__init__.py} (49%)
 create mode 100644 hyper_parallel/core/activation_checkpoint/activation_checkpoint.py
 rename hyper_parallel/{platform/torch => core}/activation_checkpoint/swap.py (88%)

diff --git a/hyper_parallel/platform/torch/activation_checkpoint/activation_checkpoint.py b/hyper_parallel/core/activation_checkpoint/__init__.py
similarity index 49%
rename from hyper_parallel/platform/torch/activation_checkpoint/activation_checkpoint.py
rename to hyper_parallel/core/activation_checkpoint/__init__.py
index bcb4263..a826689 100644
--- a/hyper_parallel/platform/torch/activation_checkpoint/activation_checkpoint.py
+++ b/hyper_parallel/core/activation_checkpoint/__init__.py
@@ -12,22 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""Activation Checkpointing Wrapper for Hyper Parallel."""
-import functools
-
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import \
-    checkpoint_wrapper as ckpt_wrapper
-from torch.utils.checkpoint import checkpoint, noop_context_fn
-
-from .sac import create_selective_checkpoint_contexts
-
-
-def _checkpoint(function, *args, policy_fn=None):
-    if policy_fn:
-        context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn)
-    else:
-        context_fn = noop_context_fn
-    return checkpoint(function, *args, context_fn=context_fn, use_reentrant=False)
-
-
-checkpoint_wrapper = functools.partial(ckpt_wrapper, checkpoint_fn=_checkpoint)
+"""Activation checkpointing related interfaces"""
+from .activation_checkpoint import CheckpointPolicy, checkpoint, checkpoint_wrapper
+from .swap import SwapManager
diff --git a/hyper_parallel/core/activation_checkpoint/activation_checkpoint.py b/hyper_parallel/core/activation_checkpoint/activation_checkpoint.py
new file mode 100644
index 0000000..4dab890
--- /dev/null
+++ b/hyper_parallel/core/activation_checkpoint/activation_checkpoint.py
@@ -0,0 +1,66 @@
+# Copyright 2026 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Activation checkpointing related interfaces"""
+import contextlib
+import enum
+from functools import partial
+from hyper_parallel.platform import get_platform
+plat = get_platform()
+
+
+class CheckpointPolicy(enum.Enum):
+    """
+    Enum for specifying the policy for checkpointing during backpropagation.
+
+    This enum extends PyTorch's selective activation checkpointing policies
+    by introducing a SWAP-based strategy, which allows activation tensors
+    to be offloaded during the forward pass and loaded back before backward
+    computation.
+
+    For PyTorch native policies (SAVE / RECOMPUTE semantics and MUST vs PREFER),
+    see: https://docs.pytorch.org/docs/2.6/checkpoint.html#torch.utils.checkpoint.CheckpointPolicy
+
+    Additional policy:
+
+    - ``MUST_SWAP``: The operation's output is offloaded to host memory during the
+      forward pass and loaded back asynchronously before backward computation. The backward
+      pass reuses the loaded activations without recomputation.
+
+      This policy must be used together with :class:`SwapManager` to coordinate
+      asynchronous offload/load and stream synchronization.
+
+    .. note::
+        ``MUST_SWAP`` is typically applied to operations that are either
+        computationally expensive or have large memory footprints. Note that
+        swapping very small outputs may introduce additional overhead and
+        reduce the effectiveness of asynchronous copy.
+    """
+    MUST_SAVE = 0
+    PREFER_SAVE = 1
+    MUST_RECOMPUTE = 2
+    PREFER_RECOMPUTE = 3
+
+    # Offload during forward, reload before backward. Requires SwapManager.
+    MUST_SWAP = 4
+
+
+def checkpoint(function, *args, swap_inputs=False, policy_fn=None, **kwargs):
+    context_fn = partial(plat.create_selective_checkpoint_contexts, policy_fn) if policy_fn else plat.noop_context_fn
+    context = plat.async_save_on_cpu if swap_inputs else contextlib.nullcontext
+    with context():
+        return plat.checkpoint(function, *args, context_fn=context_fn, use_reentrant=False, **kwargs)
+
+
+checkpoint_wrapper = partial(plat.ckpt_wrapper, checkpoint_fn=checkpoint)
diff --git a/hyper_parallel/platform/torch/activation_checkpoint/swap.py b/hyper_parallel/core/activation_checkpoint/swap.py
similarity index 88%
rename from hyper_parallel/platform/torch/activation_checkpoint/swap.py
rename to hyper_parallel/core/activation_checkpoint/swap.py
index f047827..c386130 100644
--- a/hyper_parallel/platform/torch/activation_checkpoint/swap.py
+++ b/hyper_parallel/core/activation_checkpoint/swap.py
@@ -22,8 +22,9 @@ import weakref
 from collections import defaultdict
 from typing import Any, Dict, List, Optional
 
-import torch
-from torch.utils._pytree import tree_map
+from hyper_parallel.platform import get_platform
+
+platform = get_platform()
 
 
 class SwapTensor:
@@ -36,11 +37,10 @@ class SwapTensor:
 
     def __init__(self, val: Any) -> None:
         self.val = val
-
-        if isinstance(val, torch.Tensor) and not val.is_cpu:
+        if isinstance(val, platform.Tensor) and str(val.device).lower() != 'cpu':
             self._state = self.STATE_DEVICE
             self.is_slice_tensor = val.storage().size() != val.numel()
-            self.val_cpu: Optional[torch.Tensor] = torch.empty_like(
+            self.val_cpu = platform.empty_like(
                 val, device="cpu", pin_memory=True
             )
             self.storage_size = val.storage().size()
@@ -151,9 +151,9 @@ class Storage:
                 x.async_load()
             return x
 
-        for stroage_list in self.swap_storage.values():
-            for item in stroage_list:
-                tree_map(_async_load, item)
+        for storage_list in self.swap_storage.values():
+            for item in storage_list:
+                platform.tree_map(_async_load, item)
 
     def wait_load(self):
         """wait load for all tensors in swap storage"""
@@ -162,9 +162,9 @@ class Storage:
                 x.wait_load()
             return x
 
-        for stroage_list in self.swap_storage.values():
-            for item in stroage_list:
-                tree_map(_wait_load, item)
+        for storage_list in self.swap_storage.values():
+            for item in storage_list:
+                platform.tree_map(_wait_load, item)
 
     def wait_offload(self):
         """wait offload for all tensors in swap storage"""
@@ -173,9 +173,9 @@ class Storage:
                 x.wait_offload()
             return x
 
-        for stroage_list in self.swap_storage.values():
-            for item in stroage_list:
-                tree_map(_wait_offload, item)
+        for storage_list in self.swap_storage.values():
+            for item in storage_list:
+                platform.tree_map(_wait_offload, item)
 
     def launch_offload(self):
         """launch async offload for all tensors in swap storage"""
@@ -184,9 +184,9 @@ class Storage:
                 x.async_offload()
             return x
 
-        for stroage_list in self.swap_storage.values():
-            for item in stroage_list:
-                tree_map(_async_offload, item)
+        for storage_list in self.swap_storage.values():
+            for item in storage_list:
+                platform.tree_map(_async_offload, item)
 
 
 class SwapGroup:
@@ -204,10 +204,11 @@ class SwapGroup:
 
     def launch_offload(self, copy_stream):
         """Launch async offload for all storages in the group."""
-        compute_event = torch.npu.Event()
-        compute_event.record(torch.npu.current_stream())
-        self._offload_event = torch.npu.Event()
-        with torch.no_grad(), torch.npu.stream(copy_stream):
+        compute_event = platform.new_event()
+        compute_event.record(platform.get_current_stream())
+        self._offload_event = platform.new_event()
+        stream_context = platform.get_stream_context()
+        with platform.no_grad(), stream_context(copy_stream):
             compute_event.wait(copy_stream)
             for storage in self._storages:
                 storage.launch_offload()
@@ -215,8 +216,9 @@ class SwapGroup:
 
     def wait_offload(self):
         """Wait for offload to complete for all storages in the group."""
-        compute_stream = torch.npu.current_stream()
-        with torch.no_grad(), torch.npu.stream(compute_stream):
+        compute_stream = platform.get_current_stream()
+        stream_context = platform.get_stream_context()
+        with platform.no_grad(), stream_context(compute_stream):
             self._offload_event.wait(compute_stream)
             self._offload_event = None
             for storage in self._storages:
@@ -224,10 +226,11 @@ class SwapGroup:
 
     def launch_load(self, copy_stream):
         """Launch async load for all storages in the group."""
-        compute_event = torch.npu.Event()
-        compute_event.record(torch.npu.current_stream())
-        self._load_event = torch.npu.Event()
-        with torch.no_grad(), torch.npu.stream(copy_stream):
+        compute_event = platform.new_event()
+        compute_event.record(platform.get_current_stream())
+        self._load_event = platform.new_event()
+        stream_context = platform.get_stream_context()
+        with platform.no_grad(), stream_context(copy_stream):
             compute_event.wait(copy_stream)
             for storage in self._storages:
                 storage.launch_load()
@@ -235,8 +238,9 @@ class SwapGroup:
 
     def wait_load(self):
         """Wait for load to complete for all storages in the group."""
-        compute_stream = torch.npu.current_stream()
-        with torch.no_grad(), torch.npu.stream(compute_stream):
+        compute_stream = platform.get_current_stream()
+        stream_context = platform.get_stream_context()
+        with platform.no_grad(), stream_context(compute_stream):
             self._load_event.wait(compute_stream)
             self._load_event = None
             for storage in self._storages:
@@ -430,7 +434,7 @@ class SwapManager:
         def _register_hooks_once(module, group_name):
             hooks = [
                 ("_swap_forward_pre_hook_handle",
-                 lambda h: module.register_forward_pre_hook(h, prepend=True),
+                 lambda h: platform.register_forward_pre_hook(module, h, prepend=True),
                  functools.partial(_forward_pre_hook, group_name)),
 
                 ("_swap_forward_hook_handle",
@@ -438,11 +442,11 @@ class SwapManager:
                  functools.partial(_forward_hook, group_name)),
 
                 ("_swap_backward_pre_hook_handle",
-                 lambda h: module.register_full_backward_pre_hook(h, prepend=True),
+                 lambda h: platform.register_full_backward_pre_hook(module, h, prepend=True),
                  functools.partial(_backward_pre_hook, group_name)),
 
                 ("_swap_backward_hook_handle",
-                 module.register_full_backward_hook,
+                 lambda h: platform.register_full_backward_hook(module, h),
                  functools.partial(_backward_hook, group_name)),
             ]
 
@@ -455,7 +459,7 @@ class SwapManager:
         _register_hooks_once(second_layer, second_name)
 
     def _get_copy_stream(self):
-        """Return a singleton NPU copy stream, created on first access."""
+        """Return a singleton copy stream, created on first access."""
         if self._copy_stream is None:
-            self._copy_stream = torch.npu.Stream()
+            self._copy_stream = platform.new_stream()
         return self._copy_stream
diff --git a/hyper_parallel/platform/mindspore/platform.py b/hyper_parallel/platform/mindspore/platform.py
index 2ef4e99..899873a 100644
--- a/hyper_parallel/platform/mindspore/platform.py
+++ b/hyper_parallel/platform/mindspore/platform.py
@@ -15,8 +15,10 @@
 """MindSpore platform api"""
 import mindspore as ms
 import mindspore.common.dtype as mstype
+
 from mindspore.nn import Cell
 from mindspore import mint
+from mindspore.common.api import _no_grad
 from mindspore.common.dtype import type_size_in_bytes
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
@@ -36,6 +38,8 @@ _tensor_transform = TensorTransform.get_instance()
 
 
 # pylint: disable=C0103
+
+
 class MindSporePlatform(Platform):
     """MindSpore platform api"""
     Tensor = Tensor
@@ -338,3 +342,68 @@ class MindSporePlatform(Platform):
                 in Ascend. Default: ``None``.
         """
         dist.all_gather_object(object_list, obj, group)
+
+    @staticmethod
+    def no_grad():
+        return _no_grad()
+
+    @staticmethod
+    def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
+        return mint.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
+
+    def get_current_stream(self):
+        return ms.runtime.current_stream()
+
+    def new_event(self):
+        return ms.runtime.Event()
+
+    def tree_map(self, fn, tree):
+        """
+        Apply fn to each leaf in a nested structure (list / tuple / dict),
+        preserving the original structure.
+        """
+        if isinstance(tree, dict):
+            return type(tree)(
+                (k, self.tree_map(fn, v)) for k, v in tree.items()
+            )
+
+        if isinstance(tree, tuple):
+            return tuple(self.tree_map(fn, v) for v in tree)
+
+        if isinstance(tree, list):
+            return [self.tree_map(fn, v) for v in tree]
+
+        # leaf
+        return fn(tree)
+
+    @staticmethod
+    def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
+        return module.register_forward_pre_hook(hook, with_kwargs)
+
+    @staticmethod
+    def register_full_backward_hook(module, hook, prepend=False):
+        return module.register_backward_hook(hook)
+
+    @staticmethod
+    def register_full_backward_pre_hook(module, hook, prepend=False):
+        return module.register_backward_pre_hook(hook)
+
+    @property
+    def checkpoint(self):
+        return ms.recompute
+
+    @staticmethod
+    def ckpt_wrapper(module, checkpoint_fn=None, **checkpoint_fn_kwargs):
+        raise NotImplementedError("ckpt_wrapper is not supported on MindSpore platform")
+
+    @property
+    def noop_context_fn(self):
+        raise NotImplementedError("noop_context_fn is not supported on MindSpore platform")
+
+    @staticmethod
+    def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False):
+        raise NotImplementedError("create_selective_checkpoint_contexts is not supported on MindSpore platform")
+
+    @staticmethod
+    def async_save_on_cpu(policy_fn=None):
+        raise NotImplementedError("async_save_on_cpu is not supported on MindSpore platform")
diff --git a/hyper_parallel/platform/platform.py b/hyper_parallel/platform/platform.py
index 89c8355..86fcb3f 100644
--- a/hyper_parallel/platform/platform.py
+++ b/hyper_parallel/platform/platform.py
@@ -310,3 +310,52 @@ class Platform:
         Aggregates all Python objects objs in a specified communication group into object_list.
         """
         raise NotImplementedError("Platform subclasses must implement all_gather_object")
+
+    @staticmethod
+    def no_grad():
+        raise NotImplementedError("Platform subclasses must implement no_grad")
+
+    @staticmethod
+    def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
+        raise NotImplementedError("Platform subclasses must implement empty_like")
+
+    def get_current_stream(self):
+        raise NotImplementedError("Platform subclasses must implement get_current_stream")
+
+    def new_event(self):
+        raise NotImplementedError("Platform subclasses must implement new_event")
+
+    def tree_map(self, fn, tree):
+        raise NotImplementedError("Platform subclasses must implement tree_map")
+
+    @staticmethod
+    def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
+        return module.register_forward_pre_hook(hook, prepend=prepend, with_kwargs=with_kwargs)
+
+    @staticmethod
+    def register_full_backward_hook(module, hook, prepend=False):
+        return module.register_full_backward_hook(hook, prepend)
+
+    @staticmethod
+    def register_full_backward_pre_hook(module, hook, prepend=False):
+        return module.register_full_backward_pre_hook(hook, prepend)
+
+    @property
+    def checkpoint(self):
+        raise NotImplementedError("Platform subclasses must implement checkpoint")
+
+    @staticmethod
+    def ckpt_wrapper(module, checkpoint_fn=None, **checkpoint_fn_kwargs):
+        raise NotImplementedError("Platform subclasses must implement ckpt_wrapper")
+
+    @property
+    def noop_context_fn(self):
+        raise NotImplementedError("Platform subclasses must implement ckpt_wrapper")
+
+    @staticmethod
+    def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False):
+        raise NotImplementedError("Platform subclasses must implement create_selective_checkpoint_contexts")
+
+    @staticmethod
+    def async_save_on_cpu(policy_fn=None):
+        raise NotImplementedError("Platform subclasses must implement async_save_on_cpu")
diff --git a/hyper_parallel/platform/torch/activation_checkpoint/__init__.py b/hyper_parallel/platform/torch/activation_checkpoint/__init__.py
index 27aa939..a086540 100644
--- a/hyper_parallel/platform/torch/activation_checkpoint/__init__.py
+++ b/hyper_parallel/platform/torch/activation_checkpoint/__init__.py
@@ -13,7 +13,4 @@
 # limitations under the License.
 # ============================================================================
 """Activation checkpointing related interfaces"""
-from .activation_checkpoint import checkpoint_wrapper
-from .sac import CheckpointPolicy
-from .swap import SwapManager
 from .activation_swap import swap_wrapper, ActivationPolicy
diff --git a/hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py b/hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py
index c44a10c..16608e2 100644
--- a/hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py
+++ b/hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py
@@ -25,7 +25,7 @@ from typing import Optional, Callable, Any
 import torch
 from torch import nn
 from torch.distributed.utils import _replace_by_prefix
-from .swap import Storage, SwapManager, SwapTensor
+from hyper_parallel.core.activation_checkpoint.swap import SwapManager, SwapTensor, Storage
 
 
 _SWAP_WRAPPED_MODULE = "_swap_wrapped_module"
diff --git a/hyper_parallel/platform/torch/activation_checkpoint/sac.py b/hyper_parallel/platform/torch/activation_checkpoint/sac.py
index f345f51..270a20e 100644
--- a/hyper_parallel/platform/torch/activation_checkpoint/sac.py
+++ b/hyper_parallel/platform/torch/activation_checkpoint/sac.py
@@ -17,7 +17,6 @@
 # ============================================================================
 """enhanced with selective checkpoint support swap"""
 # pylint: disable=W0212, W0613, C0115, C0116, C0103, R1705
-import enum
 from typing import Any, Optional, Union
 
 import torch
@@ -25,7 +24,8 @@ import torch.fx.traceback as fx_traceback
 from torch._functorch._aot_autograd.functional_utils import is_fun
 from torch.utils._pytree import tree_map
 from torch.utils._python_dispatch import TorchDispatchMode
-from .swap import SwapManager, SwapTensor, Storage  # patch code
+from hyper_parallel.core.activation_checkpoint import CheckpointPolicy  # patch code
+from hyper_parallel.core.activation_checkpoint.swap import SwapManager, SwapTensor, Storage  # patch code
 
 def _is_compiling(func, args, kwargs):
     # Check if we are under AOTAutograd tracing
@@ -98,36 +98,6 @@ class SelectiveCheckpointContext:
         self.is_recompute = is_recompute
 
 
-class CheckpointPolicy(enum.Enum):
-    """
-    Enum for specifying the policy for checkpointing during backpropagation.
-
-    The following policies are supported:
-
-    - ``{MUST,PREFER}_SAVE``: The operation's output will be saved during the forward
-      pass and will not be recomputed during the backward pass
-    - ``{MUST,PREFER}_RECOMPUTE``: The operation's output will not be saved during the
-      forward pass and will be recomputed during the backward pass
-
-    Use ``MUST_*`` over ``PREFER_*`` to indicate that the policy should not be overridden
-    by other subsystems like `torch.compile`.
-
-    .. note::
-        A policy function that always returns ``PREFER_RECOMPUTE`` is
-        equivalent to vanilla checkpointing.
-
-        A policy function that returns ``PREFER_SAVE`` every op is
-        NOT equivalent to not using checkpointing. Using such a policy would
-        save additional tensors not limited to ones that are actually needed for
-        gradient computation.
-    """
-    MUST_SAVE = 0
-    PREFER_SAVE = 1
-    MUST_RECOMPUTE = 2
-    PREFER_RECOMPUTE = 3
-    MUST_SWAP = 4  # patch code
-
-
 def _policy_from_bool(b):
     # For backward compatibility
     return CheckpointPolicy.MUST_SAVE if b else CheckpointPolicy.PREFER_RECOMPUTE
diff --git a/hyper_parallel/platform/torch/platform.py b/hyper_parallel/platform/torch/platform.py
index d34a5cc..b72a434 100644
--- a/hyper_parallel/platform/torch/platform.py
+++ b/hyper_parallel/platform/torch/platform.py
@@ -18,6 +18,8 @@ from torch import nn
 from torch import Tensor
 from torch.nn import Parameter, Module
 from torch._ops import OpOverload, OpOverloadPacket
+from torch.utils.checkpoint import noop_context_fn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper
 import torch.distributed.nn.functional as dist_func
 import torch.distributed as dist
 from hyper_parallel.platform.torch.dtensor import DTensorBase
@@ -320,3 +322,46 @@ class TorchPlatform(Platform):
             be unmodified.
         """
         dist.all_gather_object(object_list, obj, group)
+
+    @staticmethod
+    def no_grad():
+        return torch.no_grad()
+
+    @staticmethod
+    def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
+        return torch.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
+
+    def get_current_stream(self):
+        device = self.get_device_handle()
+        return device.current_stream()
+
+    def new_event(self):
+        device = self.get_device_handle()
+        return device.Event()
+
+    def tree_map(self, fn, tree):
+        return torch.utils._pytree.tree_map(fn, tree)  # pylint:disable=protected-access
+
+    @property
+    def checkpoint(self):
+        return torch.utils.checkpoint.checkpoint
+
+    @staticmethod
+    def ckpt_wrapper(module, checkpoint_fn=None, **checkpoint_fn_kwargs):
+        return checkpoint_wrapper(module, checkpoint_fn=checkpoint_fn, **checkpoint_fn_kwargs)
+
+    @property
+    def noop_context_fn(self):
+        return noop_context_fn
+
+    @staticmethod
+    def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False):
+        # pylint: disable=C0415
+        from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts
+        return create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation)
+
+    @staticmethod
+    def async_save_on_cpu(policy_fn=None):
+        # pylint: disable=C0415
+        from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import AsyncSaveOnCpu
+        return AsyncSaveOnCpu(policy_fn)
diff --git a/tests/torch/activation_checkpoint/test_ac.py b/tests/torch/activation_checkpoint/test_ac.py
index c77e5e9..1df8cbf 100644
--- a/tests/torch/activation_checkpoint/test_ac.py
+++ b/tests/torch/activation_checkpoint/test_ac.py
@@ -14,11 +14,13 @@
 # ============================================================================
 """Test Activation checkpoint memory comparison: None vs Recompute vs Save vs Swap"""
 import torch
+
+from hyper_parallel.core.activation_checkpoint import CheckpointPolicy, SwapManager, checkpoint_wrapper
 from tests.common.mark_utils import arg_mark
 from tests.torch.common_net import SimpleTransformer
-from hyper_parallel.platform.torch.activation_checkpoint import (
-    SwapManager, CheckpointPolicy, checkpoint_wrapper)
-from .utils import prepare_data, train_one_mode, seed_memory_time_context
+
+from .utils import prepare_data, seed_memory_time_context, train_one_mode
+
 
 def apply_recompute(model, mode):
     """Apply activation checkpointing based on the specified mode."""
@@ -40,7 +42,7 @@ def apply_recompute(model, mode):
             return CheckpointPolicy.MUST_RECOMPUTE
 
         for i, layer in enumerate(model.layers):
-            model.layers[i] = checkpoint_wrapper(layer, policy_fn=policy_fn)
+            model.layers[i] = checkpoint_wrapper(layer, policy_fn=policy_fn, swap_inputs=True)
 
         for i in range(len(model.layers) - 1):
             SwapManager().set_forward_prefetch_layer(model.layers[i], model.layers[i + 1])
diff --git a/tests/torch/activation_checkpoint/test_act_swap.py b/tests/torch/activation_checkpoint/test_act_swap.py
index 3ee6125..bae1633 100644
--- a/tests/torch/activation_checkpoint/test_act_swap.py
+++ b/tests/torch/activation_checkpoint/test_act_swap.py
@@ -15,8 +15,8 @@
 """Test Activation Swap memory comparison: None vs Swap"""
 from tests.common.mark_utils import arg_mark
 from tests.torch.common_net import SimpleTransformer
-from hyper_parallel.platform.torch.activation_checkpoint import (
-    SwapManager, ActivationPolicy, swap_wrapper)
+from hyper_parallel.core.activation_checkpoint import SwapManager
+from hyper_parallel.platform.torch.activation_checkpoint import ActivationPolicy, swap_wrapper
 from .utils import prepare_data, train_one_mode, seed_memory_time_context
 
 
-- 
Gitee