From dc71d9e8a45f504c9744eb273a7f5fab6f0a718b Mon Sep 17 00:00:00 2001 From: fandawei Date: Mon, 19 Jan 2026 16:19:07 +0800 Subject: [PATCH 1/2] select activation checkpoint support swap op output --- .../torch/activation_checkpoint/__init__.py | 17 + .../activation_checkpoint.py | 310 ++++++++++++++ .../torch/activation_checkpoint/swap.py | 396 ++++++++++++++++++ hyper_parallel/utils/__init__.py | 0 hyper_parallel/utils/logger.py | 105 +++++ tests/torch/activation_checkpoint/test_ac.py | 225 ++++++++++ 6 files changed, 1053 insertions(+) create mode 100644 hyper_parallel/platform/torch/activation_checkpoint/__init__.py create mode 100644 hyper_parallel/platform/torch/activation_checkpoint/activation_checkpoint.py create mode 100644 hyper_parallel/platform/torch/activation_checkpoint/swap.py create mode 100644 hyper_parallel/utils/__init__.py create mode 100644 hyper_parallel/utils/logger.py create mode 100644 tests/torch/activation_checkpoint/test_ac.py diff --git a/hyper_parallel/platform/torch/activation_checkpoint/__init__.py b/hyper_parallel/platform/torch/activation_checkpoint/__init__.py new file mode 100644 index 0000000..fb25bcd --- /dev/null +++ b/hyper_parallel/platform/torch/activation_checkpoint/__init__.py @@ -0,0 +1,17 @@ +# 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. +# ============================================================================ + +from .activation_checkpoint import CheckpointPolicy, checkpoint_wrapper +from .swap import SwapManager diff --git a/hyper_parallel/platform/torch/activation_checkpoint/activation_checkpoint.py b/hyper_parallel/platform/torch/activation_checkpoint/activation_checkpoint.py new file mode 100644 index 0000000..09b0cf0 --- /dev/null +++ b/hyper_parallel/platform/torch/activation_checkpoint/activation_checkpoint.py @@ -0,0 +1,310 @@ +# 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. +# ============================================================================ +from typing import * +import enum +import functools + +import torch +import hyper_parallel.utils.logger as logger +from torch.utils._pytree import tree_map +from torch.utils._python_dispatch import TorchDispatchMode +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper as ckpt_wrapper +from torch.utils.checkpoint import checkpoint, noop_context_fn +from .swap import Storage, SwapManager, SwapTensor + + +class _VersionWrapper: + # Check that cached tensors are not mutated. + def __init__(self, val): + self.version: Optional[int] = val._version if isinstance( + val, torch.Tensor) else None + self.val: Union[torch.Tensor, Any] = val + + def get_val(self, allow_cache_entry_mutation): + if self.version is not None and not allow_cache_entry_mutation: + if self.val._version != self.version: + # Can we give user a stack trace of where the mutation happened? + raise RuntimeError( + "Tensor cached during selective activation checkpoint has been mutated" + ) + return self.val + + +def _maybe_detach(x, any_ret_has_alias_info): + # We detach for two separate reasons: + # - For view ops, we need to ensure that when the tensor is returned from + # CachedDispatchMode, as_view sees that the AutogradMeta is nullptr + # - Avoid reference cycles + # For case 1, it is not enough to check whether x has differentiable dtype + # because non-differentiable dtype can have non-nullptr AutogradMeta, e.g. + # when the tensor is a view. + if isinstance(x, torch.Tensor) and (x.is_floating_point() or x.is_complex() or any_ret_has_alias_info): + with torch._C._SetExcludeDispatchKeyGuard(torch._C.DispatchKey.ADInplaceOrView, False): + # Ensure that view performed beneath autograd properly propagates + # version counter. TODO: Use reentrant_dispatch instead of + # manually manipulating dispatch keys. Using reentrant_dispatch + # would respect inference_mode, though that is not relevant for + # this case. + x = x.detach() + return x + + +class SelectiveCheckpointContext: + """ + Context passed to policy function during selective checkpointing. + + This class is used to pass relevant metadata to the policy function during + selective checkpointing. The metadata includes whether the current invocation + of the policy function is during recomputation or not. + + Example: + >>> # xdoctest: +SKIP(stub) + >>> + >>> def policy_fn(ctx, op, *args, **kwargs): + >>> print(ctx.is_recompute) + >>> + >>> context_fn = functools.partial(create_checkpoint_contexts, policy_fn) + >>> + >>> out = torch.utils.checkpoint.checkpoint( + >>> fn, x, y, + >>> use_reentrant=False, + >>> context_fn=context_fn, + >>> ) + """ + + def __init__(self, *, is_recompute): + 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 + + +SAC_IGNORED_OPS = { + # AC inserts different number of detach during forward and recompute. + torch.ops.aten.detach.default, + # AC's determinism check invokes additional metadata ops during forward. + # With subclasses involved, these metadata ops become dispatchable, this + # can result in incorrectness if these ops are selected cached. + torch.ops.prim.device.default, +} | set(torch._subclasses.functional_tensor.FunctionalTensor.metadata_fns) + + +class _CachingTorchDispatchMode(TorchDispatchMode): + # Used together with _CachedTorchDispatchMode to implement SAC. + def __init__(self, policy_fn, storage): + self.policy_fn = policy_fn + self.storage = storage + self.add_to_storage = False + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if func in SAC_IGNORED_OPS: + return func(*args, **kwargs) + + kwargs = {} if kwargs is None else kwargs + policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False), + func, *args, **kwargs) + + out = func(*args, **kwargs) + + any_ret_has_alias_info = any( + ret.alias_info is not None for ret in func._schema.returns) + + if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE): + self.storage.save_storage[func].append(tree_map( + lambda x: _VersionWrapper(_maybe_detach(x, any_ret_has_alias_info)), out)) + logger.debug("op is %s, policy is %s, save to device", func, policy) + elif policy == CheckpointPolicy.MUST_SWAP: + if not self.add_to_storage: + group_name = SwapManager().get_current_group_name() + SwapManager().add_storage(group_name, self.storage) + self.add_to_storage = True + self.storage.swap_storage[func].append( + tree_map(lambda x: SwapTensor(_maybe_detach(x, any_ret_has_alias_info)), out)) + logger.debug("op is %s, policy is %s, offload to host mem", func, policy) + else: + logger.debug("op is %s, policy is %s, recompute activation", func, policy) + return out + + +class _CachedTorchDispatchMode(TorchDispatchMode): + # Used together with _CachedTorchDispatchMode to implement SAC. + def __init__(self, policy_fn, storage, allow_cache_entry_mutation): + self.policy_fn = policy_fn + self.storage = storage + self.allow_cache_entry_mutation = allow_cache_entry_mutation + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if func in SAC_IGNORED_OPS: + return func(*args, **kwargs) + + kwargs = {} if kwargs is None else kwargs + policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=True), + func, *args, **kwargs) + + if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE): + storage = self.storage.save_storage.get(func) + if storage is None: + raise RuntimeError( + f"{func} encountered during backward, but not found in storage") + if len(storage) == 0: + raise RuntimeError( + "Trying to backward an extra time. You are only allowed to backward once " + "on any region computed under selective activation checkpoint." + ) + logger.debug("op is %s, policy is %s, direct return", func, policy) + out = tree_map(lambda x: x.get_val( + self.allow_cache_entry_mutation), storage.pop(0)) + elif policy == CheckpointPolicy.MUST_SWAP: + storage = self.storage.swap_storage.get(func) + if storage is None: + raise RuntimeError( + f"{func} encountered during backward, but not found in storage") + if len(storage) == 0: + raise RuntimeError( + "Trying to backward an extra time. You are only allowed to backward once " + "on any region computed under selective activation checkpoint." + ) + logger.debug("op is %s, policy is %s, load from host mem", func, policy) + out = tree_map(lambda x: x.get_val(), storage.pop(0)) + else: + logger.debug("op is %s, policy is %s, recompute activation", func, policy) + out = func(*args, **kwargs) + return out + + +def create_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False): + """ + Helper to avoid recomputing certain ops during activation checkpointing. + + Use this with `torch.utils.checkpoint.checkpoint` to control which + operations are recomputed during the backward pass. + + Args: + policy_fn_or_list (Callable or List): + - If a policy function is provided, it should accept a + :class:`SelectiveCheckpointContext`, the :class:`OpOverload`, args and + kwargs to the op, and return a :class:`CheckpointPolicy` enum value + indicating whether the execution of the op should be recomputed or not. + - If a list of operations is provided, it is equivalent to a policy + returning `CheckpointPolicy.MUST_SAVE` for the specified + operations and `CheckpointPolicy.PREFER_RECOMPUTE` for all other + operations. + allow_cache_entry_mutation (bool, optional): By default, an error is + raised if any tensors cached by selective activation checkpoint are + mutated in order to ensure correctness. If set to `True`, this check + is disabled. + Returns: + A tuple of two context managers. + + Example: + >>> # xdoctest: +REQUIRES(LINUX) + >>> import functools + >>> + >>> x = torch.rand(10, 10, requires_grad=True) + >>> y = torch.rand(10, 10, requires_grad=True) + >>> + >>> ops_to_save = [ + >>> torch.ops.aten.mm.default, + >>> ] + >>> + >>> def policy_fn(ctx, op, *args, **kwargs): + >>> if op in ops_to_save: + >>> return CheckpointPolicy.MUST_SAVE + >>> else: + >>> return CheckpointPolicy.PREFER_RECOMPUTE + >>> + >>> context_fn = functools.partial(create_checkpoint_contexts, policy_fn) + >>> + >>> # or equivalently + >>> context_fn = functools.partial(create_checkpoint_contexts, ops_to_save) + >>> + >>> def fn(x, y): + >>> return torch.sigmoid(torch.matmul(torch.matmul(x, y), y)) * y + >>> + >>> out = torch.utils.checkpoint.checkpoint( + >>> fn, x, y, + >>> use_reentrant=False, + >>> context_fn=context_fn, + >>> ) + """ + # NB: If grad_mode is disabled, checkpoint would not run forward under + # context_fn anyway, so proceed as usual. + if isinstance(policy_fn_or_list, list): + for op in policy_fn_or_list: + if not isinstance(op, torch._ops.OpOverload): + _extra_msg = ( + "Please update the OpOverloadPacket to a specific OpOverload." + "For example, if you have `torch.ops.aten.mm`, change it to `torch.ops.aten.mm.default`." + ) if isinstance(op, torch._ops.OpOverloadPacket) else "" + raise ValueError( + f"Expected op in `op_list` to be an OpOverload but got: {op} " + f"of type {type(op)}. {_extra_msg}" + ) + + def policy_fn(ctx, op, *args, **kwargs): + if op in policy_fn_or_list: + return CheckpointPolicy.MUST_SAVE + else: + return CheckpointPolicy.PREFER_RECOMPUTE + elif callable(policy_fn_or_list): + policy_fn = policy_fn_or_list + else: + raise TypeError( + "policy_fn_or_list must be either a function or a list of ops.") + + storage = Storage() + + return ( + _CachingTorchDispatchMode(policy_fn, storage), + _CachedTorchDispatchMode( + policy_fn, storage, allow_cache_entry_mutation), + ) + +def _checkpoint(function, *args, policy_fn=None): + if policy_fn: + context_fn = functools.partial(create_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) diff --git a/hyper_parallel/platform/torch/activation_checkpoint/swap.py b/hyper_parallel/platform/torch/activation_checkpoint/swap.py new file mode 100644 index 0000000..966617f --- /dev/null +++ b/hyper_parallel/platform/torch/activation_checkpoint/swap.py @@ -0,0 +1,396 @@ +# 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. +# ============================================================================ +from typing import * +import threading +import torch +import warnings +import weakref +import hyper_parallel.utils.logger as logger +import functools + +from collections import defaultdict +from torch.utils._pytree import tree_map + + +class Storage: + def __init__(self): + self.save_storage: Dict[Any, List[Any]] = defaultdict(list) + self.swap_storage: Dict[Any, List[Any]] = defaultdict(list) + + def launch_load(self): + def _async_load(x): + if isinstance(x, SwapTensor): + x.async_load() + return x + + for stroage_list in self.swap_storage.values(): + for item in stroage_list: + tree_map(_async_load, item) + + def wait_load(self): + def _wait_load(x): + if isinstance(x, SwapTensor): + x.wait_load() + return x + + for stroage_list in self.swap_storage.values(): + for item in stroage_list: + tree_map(_wait_load, item) + + def wait_offload(self): + def _wait_offload(x): + if isinstance(x, SwapTensor): + x.wait_offload() + return x + + for stroage_list in self.swap_storage.values(): + for item in stroage_list: + tree_map(_wait_offload, item) + + def launch_offload(self): + def _async_offload(x): + if isinstance(x, SwapTensor): + x.async_offload() + return x + + for stroage_list in self.swap_storage.values(): + for item in stroage_list: + tree_map(_async_offload, item) + + +class SwapManager: + _instance: Optional["SwapManager"] = None + _lock = threading.Lock() + + def __new__(cls): + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._groups: Dict[str, SwapGroup] = dict() + cls._instance._current_group_name: str = "" + cls._instance._counter_lock = threading.Lock() + cls._instance._layer_count = 0 + cls._copy_stream = None + return cls._instance + + def add_storage(self, group_name: str, storage: Storage) -> None: + logger.info("[SwapManager.add_storage] add storage in group '%s'", group_name) + if group_name not in self._groups: + self._groups[group_name] = SwapGroup(group_name) + self._groups[group_name].add(storage) + + def launch_offload(self, group_name: str, copy_stream=None): + logger.info("[SwapManager.launch_offload] '%s'", group_name) + group = self._groups.get(group_name) + if group is None: + raise RuntimeError(f"Group {group_name} does not exist.") + if copy_stream is None: + copy_stream = self._get_copy_stream() + group.launch_offload(copy_stream) + + def wait_offload(self, group_name: str): + logger.info("[SwapManager.wait_offload] '%s'", group_name) + group = self._groups.get(group_name) + if group is None: + raise RuntimeError(f"Group {group_name} does not exist.") + group.wait_offload() + + def launch_load(self, group_name: str, copy_stream=None): + logger.info("[SwapManager.launch_load] '%s'", group_name) + group = self._groups.get(group_name) + if group is None: + raise RuntimeError(f"Group {group_name} does not exist.") + if copy_stream is None: + copy_stream = self._get_copy_stream() + group.launch_load(copy_stream) + + def wait_load(self, group_name: str): + logger.info("[SwapManager.wait_load] '%s'", group_name) + group = self._groups.get(group_name) + if group is None: + raise RuntimeError(f"Group {group_name} does not exist.") + group.wait_load() + + def get_current_group_name(self): + return self._current_group_name + + def set_current_group_name(self, group_name): + self._current_group_name = group_name + + def set_forward_prefetch_layer(self, first_layer, second_layer): + """ + Configure prefetching and offloading order between two consecutive layers. + + Usage: + for i in range(len(model.layers) - 1): + set_forward_prefetch_layer(model.layers[i], model.layers[i + 1]) + + Ensures idempotency: safe to call multiple times on the same layer pair. + """ + + def _ensure_group_name(module): + """Assign a unique swap group name to the module if not already assigned.""" + if not hasattr(module, "_swap_group_name"): + name = f"swap_group_{self._layer_count}" + self._layer_count += 1 + module._swap_group_name = name + module._swap_group_order = {"prev": None, "next": None} + return module._swap_group_name + first_name = _ensure_group_name(first_layer) + second_name = _ensure_group_name(second_layer) + + if first_layer._swap_group_order["next"] is None: + first_layer._swap_group_order["next"] = second_name + if second_layer._swap_group_order["prev"] is None: + second_layer._swap_group_order["prev"] = first_name + + def _forward_pre_hook(group_name, module, _): + logger.info(f"enter {group_name} swap forward pre hook") + if getattr(module, "_swap_state", None) == "pre_backward": + logger.info(f"in pre backward, skip {group_name} forward pre hook") + return + SwapManager().set_current_group_name(group_name) + + def _forward_hook(group_name, module, args, output): + """ + compute stream: layer 0 | layer 1 | - | layer 2 + copy stream offlod: layer 0 | - | layer 1 + ├── launch offload layer 1 + └── wait offload layer 0 + after layer finish, if it has next layer offload current mem + after layer finish, if it has prev layer, waitoffload prev layer + """ + if getattr(module, "_swap_state", None) == "pre_backward": + logger.info(f"in pre backward, skip {group_name} swap forward hook") + return + next_name = module._swap_group_order.get('next', None) + if next_name: + SwapManager().launch_offload(group_name) + prev_name = module._swap_group_order.get('prev', None) + if prev_name: + SwapManager().wait_offload(prev_name) + logger.info(f"finished {group_name} forward hook") + + def _backward_pre_hook(group_name, module, grad_input): + module._swap_state = "pre_backward" + logger.info(f"enter {group_name} swap pre backward hook") + prev_name = module._swap_group_order.get('prev', None) + if prev_name: + SwapManager().launch_load(prev_name) + + next_name = module._swap_group_order.get('next', None) + if next_name: + SwapManager().wait_load(group_name) + + def _backward_hook(group_name, module, grad_input, grad_output): + module._swap_state = "backward" + logger.info(f"finished {group_name} swap backward hook") + + def _register_hooks_once(module, group_name): + hooks = [ + ("_swap_forward_pre_hook_handle", + lambda h: module.register_forward_pre_hook(h, prepend=True), + functools.partial(_forward_pre_hook, group_name)), + + ("_swap_forward_hook_handle", + module.register_forward_hook, + functools.partial(_forward_hook, group_name)), + + ("_swap_backward_pre_hook_handle", + lambda h: module.register_full_backward_pre_hook(h, prepend=True), + functools.partial(_backward_pre_hook, group_name)), + + ("_swap_backward_hook_handle", + module.register_full_backward_hook, + functools.partial(_backward_hook, group_name)), + ] + + for attr_name, register_func, hook in hooks: + if not hasattr(module, attr_name): + handle = register_func(hook) + setattr(module, attr_name, handle) + # Register for both layers + _register_hooks_once(first_layer, first_name) + _register_hooks_once(second_layer, second_name) + + def _get_copy_stream(self): + """Return a singleton NPU copy stream, created on first access.""" + if self._copy_stream is None: + self._copy_stream = torch.npu.Stream() + return self._copy_stream + +class SwapGroup: + def __init__(self, group_name: str): + self.group_name = group_name + self._storages = weakref.WeakSet() + self._load_event = None + self._offload_event = None + + def add(self, storage): + self._storages.add(storage) + + def launch_offload(self, copy_stream): + logger.info(f"[SwapGroup.launch_offload] Group '{self.group_name}'") + + 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.wait(copy_stream) + for storage in self._storages: + storage.launch_offload() + self._offload_event.record(copy_stream) + + def wait_offload(self): + compute_stream = torch.npu.current_stream() + logger.info(f"[SwapGroup.wait_offload] Group '{self.group_name}'") + + with torch.no_grad(), torch.npu.stream(compute_stream): + self._offload_event.wait(compute_stream) + self._offload_event = None + for storage in self._storages: + storage.wait_offload() + + def launch_load(self, copy_stream): + logger.info(f"[SwapGroup.launch_load] Group '{self.group_name}'") + + 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.wait(copy_stream) + for storage in self._storages: + storage.launch_load() + self._load_event.record(copy_stream) + + def wait_load(self): + compute_stream = torch.npu.current_stream() + logger.info(f"[SwapGroup.wait_load] Group '{self.group_name}'") + + with torch.no_grad(), torch.npu.stream(compute_stream): + self._load_event.wait(compute_stream) + self._load_event = None + for storage in self._storages: + storage.wait_load() + + +class SwapTensor: + STATE_DEVICE = "device" + STATE_HOST = "host" + STATE_D2H = "d2h" + STATE_H2D = "h2d" + STATE_NON_TENSOR = "non_tensor" + + def __init__(self, val: Any) -> None: + self.val = val + + if isinstance(val, torch.Tensor) and not val.is_cpu: + self._state = self.STATE_DEVICE + self.is_slice_tensor = val.untyped_storage().size() != val.numel() + self.val_cpu: Optional[torch.Tensor] = torch.empty_like( + val, device="cpu", pin_memory=True + ) + self.storage_size = val.untyped_storage().size() + else: + self._state = self.STATE_NON_TENSOR + self.val_cpu = None + + def get_val(self) -> Any: + if self._state == self.STATE_NON_TENSOR: + return self.val + if self._state != self.STATE_DEVICE: + raise RuntimeError( + f"Cannot call get_val(): tensor is in '{self._state}' state. " + f"Must be in 'device' state." + ) + return self.val + + def async_load(self): + if self._state == self.STATE_NON_TENSOR: + return + + if self._state != self.STATE_HOST: + warnings.warn( + f"[SwapTensor.async_load] Invalid state: current={self._state}, " + f"expected 'host'. Operation skipped." + ) + return + + assert self.val_cpu is not None + self.val.untyped_storage().resize_(self.storage_size) + logger.debug("loading, copy stream is %s", torch.npu.current_stream()) + if self.is_slice_tensor: + self.val.copy_(self.val_cpu, non_blocking=True) + else: + self.val.untyped_storage().copy_(self.val_cpu.untyped_storage(), non_blocking=True) + self._state = self.STATE_H2D + + def wait_load(self): + if self._state == self.STATE_NON_TENSOR: + return + + if self._state == self.STATE_DEVICE: + return # already loaded + if self._state != self.STATE_H2D: + warnings.warn( + f"[SwapTensor.wait_load] Called in invalid state: {self._state}. " + f"Expected 'h2d'. Skipped." + ) + return + self._state = self.STATE_DEVICE + + def async_offload(self): + if self._state == self.STATE_NON_TENSOR: + return + + if self._state != self.STATE_DEVICE: + warnings.warn( + f"[SwapTensor.async_offload] Invalid state: current={self._state}, " + f"expected 'device'. Operation skipped." + ) + return + + logger.debug("offloading, copy stream is %s", torch.npu.current_stream()) + if 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) + self._state = self.STATE_D2H + + def wait_offload(self): + if self._state == self.STATE_NON_TENSOR: + return + + if self._state == self.STATE_HOST: + return + if self._state != self.STATE_D2H: + warnings.warn( + f"[SwapTensor.wait_offload] Called in invalid state: {self._state}. " + f"Expected 'd2h'. Skipped." + ) + return + self.val.untyped_storage().resize_(0) + self._state = self.STATE_HOST + + @property + def state(self) -> str: + return self._state + + def __repr__(self): + if self._state == self.STATE_NON_TENSOR: + return f"" + else: + return f"" diff --git a/hyper_parallel/utils/__init__.py b/hyper_parallel/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hyper_parallel/utils/logger.py b/hyper_parallel/utils/logger.py new file mode 100644 index 0000000..24dd31a --- /dev/null +++ b/hyper_parallel/utils/logger.py @@ -0,0 +1,105 @@ +# 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. +# ============================================================================ +import os +import logging + + +def _get_rank_id(): + """Try to get rank from environment variables set by torchrun / torch.distributed.""" + # Priority: RANK (global rank) > LOCAL_RANK (node-local rank) + # If not found, assume single process with rank 0 + rank = os.getenv("RANK") + if rank is not None: + return int(rank) + local_rank = os.getenv("LOCAL_RANK") + if local_rank is not None: + return int(local_rank) + return 0 + + +def setup_logger(log_dir=None, log_level_env="0"): + """ + Configure the root logger. + + Args: + log_dir (str, optional): Directory to save per-rank log files. + If None, only console output is used. + log_level_env (str): Log level string ("0": WARNING, "1": INFO, "2": DEBUG) + """ + level_map = {"0": logging.WARNING, "1": logging.INFO, "2": logging.DEBUG} + log_level = level_map.get(log_level_env, logging.WARNING) + + # Get rank ID + rank_id = _get_rank_id() + + # Create formatter + formatter = logging.Formatter( + fmt=f"[PID %(process)d | Rank {rank_id}] %(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s", + datefmt="%H:%M:%S" + ) + + # Setup handlers + handlers = [] + + # Console handler + console_handler = logging.StreamHandler() + console_handler.setFormatter(formatter) + handlers.append(console_handler) + + # File handler (if log_dir specified) + if log_dir is not None: + os.makedirs(log_dir, exist_ok=True) + log_file = os.path.join(log_dir, f"rank_{rank_id}.log") + file_handler = logging.FileHandler(log_file, mode='w') + file_handler.setFormatter(formatter) + handlers.append(file_handler) + + # Configure root logger + logging.basicConfig( + level=log_level, + handlers=handlers, + force=True + ) + +# ----------------------------- +# Default initialization +# ----------------------------- + + +# Read LOG_LEVEL from env (for backward compatibility) +log_level_env = os.getenv("LOG_LEVEL", "0") + +# Optional: read LOG_DIR from environment variable +log_dir = os.getenv("LOG_DIR") # e.g., export LOG_DIR=./logs + +# Setup logger +setup_logger(log_dir=log_dir, log_level_env=log_level_env) + +# Get logger instance +logger = logging.getLogger(__name__) + +# Convenience functions with stacklevel=2 + + +def debug(msg, *args, **kwargs): + logger.debug(msg, *args, stacklevel=2, **kwargs) + + +def info(msg, *args, **kwargs): + logger.info(msg, *args, stacklevel=2, **kwargs) + + +def warning(msg, *args, **kwargs): + logger.warning(msg, *args, stacklevel=2, **kwargs) diff --git a/tests/torch/activation_checkpoint/test_ac.py b/tests/torch/activation_checkpoint/test_ac.py new file mode 100644 index 0000000..b50a011 --- /dev/null +++ b/tests/torch/activation_checkpoint/test_ac.py @@ -0,0 +1,225 @@ +# Copyright 2025 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. +# ============================================================================ +"""Test Activation checkpoint memory comparison: None vs Recompute vs Save vs Swap""" +from tests.common.mark_utils import arg_mark +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, TensorDataset +from hyper_parallel.platform.torch.activation_checkpoint import ( + checkpoint_wrapper, CheckpointPolicy, SwapManager +) +import time +import random +import numpy as np +from contextlib import contextmanager + +def set_seed(seed=42): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.npu.is_available(): + torch.npu.manual_seed(seed) + torch.npu.manual_seed_all(seed) + +@contextmanager +def seed_memory_time_context(seed=42): + set_seed(seed) + torch.npu.reset_peak_memory_stats() + torch.npu.empty_cache() + start_time = time.time() + + stats = {} + + try: + yield stats + finally: + exec_time = time.time() - start_time + peak_mem_gb = torch.npu.max_memory_allocated() / (1024 ** 3) + torch.npu.empty_cache() + stats["peak_mem"] = peak_mem_gb + stats["exec_time"] = exec_time + + +class TransformerBlock(nn.Module): + def __init__(self, dim=256, num_heads=4): + super().__init__() + self.dim = dim + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.ffn = nn.Sequential( + nn.Linear(dim, dim * 4), + nn.ReLU(), + nn.Linear(dim * 4, dim) + ) + + def forward(self, x): + attn_out, _ = self.attn(x, x, x) + x = x + attn_out + x = self.norm1(x) + x = x + self.ffn(x) + x = self.norm2(x) + return x + +class SimpleTransformer(nn.Module): + def __init__(self, vocab_size=1000, dim=2048, depth=6): + super().__init__() + self.embed = nn.Embedding(vocab_size, dim) + self.layers = nn.ModuleList([TransformerBlock(dim) for _ in range(depth)]) + self.norm = nn.LayerNorm(dim) + self.head = nn.Linear(dim, vocab_size) + + def forward(self, x): + x = self.embed(x) + for block in self.layers: + x = block(x) + x = self.norm(x) + return self.head(x) + +def prepare_data(batch_size=8, seq_len=512, num_samples=64): + dataset = TensorDataset( + torch.randint(0, 10000, (num_samples, seq_len)), + torch.randint(0, 10000, (num_samples, seq_len)) + ) + dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False) + return dataloader + +def apply_recompute(model, mode): + op_non_recompute = { + torch.ops.aten.matmul.default, + torch.ops.aten.addmm.default, + torch.ops.aten.bmm.default + } + if mode == "none": + return model + elif mode == "recompute": + for i, layer in enumerate(model.layers): + model.layers[i] = checkpoint_wrapper(layer) + elif mode == "swap": + def policy_fn(ctx, op, *args, **kwargs): + if op in op_non_recompute: + return CheckpointPolicy.MUST_SWAP + return CheckpointPolicy.MUST_RECOMPUTE + + for i, layer in enumerate(model.layers): + model.layers[i] = checkpoint_wrapper(layer, policy_fn=policy_fn) + + for i in range(len(model.layers) - 1): + SwapManager().set_forward_prefetch_layer(model.layers[i], model.layers[i + 1]) + elif mode == 'save': + def policy_fn(ctx, op, *args, **kwargs): + if op in op_non_recompute: + return CheckpointPolicy.MUST_SAVE + return CheckpointPolicy.MUST_RECOMPUTE + + for i, layer in enumerate(model.layers): + model.layers[i] = checkpoint_wrapper(layer, policy_fn=policy_fn) + else: + raise ValueError(f"Unknown mode: {mode}") + return model + + +def train_one_mode(mode, dataloader, train_steps=5): + """Run training for one mode and return metrics.""" + with seed_memory_time_context() as stats: + model = SimpleTransformer(vocab_size=32000, dim=2048, depth=16).npu() + model = apply_recompute(model, mode) + + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) + losses = [] + for step, (x, y) in enumerate(dataloader): + if step >= train_steps: + break + x, y = x.npu(), y.npu() + optimizer.zero_grad() + logits = model(x) + loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1)) + loss.backward() + optimizer.step() + losses.append(loss.item()) + + return losses, stats.get("peak_mem"), stats.get("exec_time") + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_ac_memory_comparison(): + """ + Feature: Activation Checkpointing and Swapping Memory Behavior + Description: Compare peak memory usage across four modes: + 'none' (baseline), 'recompute' (full activation checkpointing), + 'save' (partial activation saving), and 'swap' (fine-grained tensor swapping). + Validates that losses are numerically identical at every training step + and that peak memory follows the expected hierarchy: + NONE > SAVE > RECOMPUTE ≈ SWAP. + Expectation: All modes produce consistent losses (within 1e-5 tolerance), + the memory usage trend is satisfied, and no OOM occurs. + """ + print("🚀 Starting memory and time comparison: none vs recompute vs save vs swap") + dataloader = prepare_data() + train_steps=3 + + modes = ["none", "recompute", "save", "swap"] + results = {} + + for mode in modes: + print(f"\n--- Running mode: {mode.upper()} ---") + losses, peak_mem, duration = train_one_mode(mode, dataloader, train_steps) + results[mode] = { + "losses": losses, + "peak_mem_gb": peak_mem, + "time_sec": duration + } + print(f"{mode}: Loss={losses[-1]:.4f}, Peak Mem={peak_mem:.5f} GB, Time={duration:.5f}s") + + print("\n" + "="*70) + print("📊 FINAL COMPARISON") + print("="*70) + print(f"{'Mode':<12} | {'Peak Mem (GB)':<15} | {'Time (s)':<10} | {'Final Loss':<12}") + print("-"*70) + for mode in modes: + r = results[mode] + print(f"{mode.upper():<12} | {r['peak_mem_gb']:<15.5f} | {r['time_sec']:<10.5f} | {r['losses'][-1]:<12.4f}") + + # loss assert + base_losses = results["none"]["losses"] + tol = 1e-5 + for step in range(train_steps): + base_val = base_losses[step] + for mode in ["recompute", "save", "swap"]: + val = results[mode]["losses"][step] + diff = abs(val - base_val) + assert diff < tol, ( + f"Loss mismatch at step {step} in mode '{mode}': " + f"none={base_val:.8f}, {mode}={val:.8f}, diff={diff:.2e}" + ) + print(f"\n✅ All {train_steps} steps: losses are consistent across modes (tol={tol}).") + + # mem assert + mem_none = results["none"]["peak_mem_gb"] + mem_save = results["save"]["peak_mem_gb"] + mem_recompute = results["recompute"]["peak_mem_gb"] + mem_swap = results["swap"]["peak_mem_gb"] + + # none > save > recompute ≈ swap + # none > save + assert mem_none > mem_save, f"Expected NONE ({mem_none:.5f}) > SAVE ({mem_save:.5f})" + print(f"✅ Verified: NONE ({mem_none:.5f}) > SAVE ({mem_save:.5f})") + # save > recompute + assert mem_save > mem_recompute, f"Expected SAVE ({mem_save:.5f}) > RECOMPUTE ({mem_recompute:.5f})" + print(f"✅ Verified: SAVE ({mem_save:.5f}) > RECOMPUTE ({mem_recompute:.5f})") + # recompute ≈ swap + TOL_MEM = 0.15 + assert abs(mem_recompute - mem_swap) < TOL_MEM, \ + f"Expected RECOMPUTE ({mem_recompute:.5f}) ≈ SWAP ({mem_swap:.5f})" + print(f"✅ Verified: RECOMPUTE ({mem_recompute:.5f}) ≈ SWAP ({mem_swap:.5f}) within tolerance ({TOL_MEM:.5f} GB)") -- Gitee From dee1eeac19e8c1a1dcf05e4eee7ee0eed3d63941 Mon Sep 17 00:00:00 2001 From: hwfuchao Date: Mon, 19 Jan 2026 19:43:09 +0800 Subject: [PATCH 2/2] activation swap --- .../torch/activation_checkpoint/__init__.py | 1 + .../activation_checkpoint/activation_swap.py | 193 ++++++++++++++++ tests/torch/activation_swap/test_act_swap.py | 217 ++++++++++++++++++ 3 files changed, 411 insertions(+) create mode 100644 hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py create mode 100644 tests/torch/activation_swap/test_act_swap.py diff --git a/hyper_parallel/platform/torch/activation_checkpoint/__init__.py b/hyper_parallel/platform/torch/activation_checkpoint/__init__.py index fb25bcd..a62af34 100644 --- a/hyper_parallel/platform/torch/activation_checkpoint/__init__.py +++ b/hyper_parallel/platform/torch/activation_checkpoint/__init__.py @@ -15,3 +15,4 @@ from .activation_checkpoint import CheckpointPolicy, checkpoint_wrapper 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 new file mode 100644 index 0000000..c86eff0 --- /dev/null +++ b/hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py @@ -0,0 +1,193 @@ +# 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 Swap implementation for PyTorch.""" +# pylint: disable=W0212, W0613 + +import enum +from abc import ABC, abstractmethod +from collections.abc import Iterator +from typing import Optional, Callable, Any +import torch +from torch import nn +from torch.distributed.utils import _replace_by_prefix +from hyper_parallel.utils import logger +from .swap import Storage, SwapManager, SwapTensor + + +_SWAP_WRAPPED_MODULE = "_swap_wrapped_module" +_SWAP_PREFIX = _SWAP_WRAPPED_MODULE + "." + + +class ActivationPolicy(enum.Enum): + """Enum for activation policies.""" + SAVE = 0 + SWAP = 1 + + +def base_check_fn(tensor) -> bool: + """ + Basic check to determine if a tensor is eligible for offloading. + - Skip Parameters and their views. + - Skip empty storage tensors. + """ + if isinstance(tensor._base, torch.nn.parameter.Parameter) or isinstance(tensor, torch.nn.parameter.Parameter): # pylint: disable=W0212 + return False + if tensor.storage().size() <= 0: + return False + return True + + +class AsyncSaveOnCpu(torch.autograd.graph.saved_tensors_hooks): + """ + Context manager to offload tensors to CPU during forward pass. + """ + def __init__(self, policy_fn=None) -> None: + self.add_to_storage = False + self.storage = Storage() + self.count_idx = 0 + self.pack_count = 0 + self.unpack_count = 0 + self.policy_fn = policy_fn + + def pack_to_cpu(tensor: torch.Tensor): + # skip ineligible tensors + if not base_check_fn(tensor): + return tensor + + if (policy_fn is not None) and (policy_fn(tensor)==ActivationPolicy.SAVE): + return tensor + + if not self.add_to_storage: + group_name = SwapManager().get_current_group_name() + SwapManager().add_storage(group_name, self.storage) + self.add_to_storage = True + logger.info(f"swap add to group is {group_name}") + logger.info(f"enter pack, tensor shape is {tensor.shape}, tensor id {id(tensor)}") + self.storage.swap_storage[self.count_idx].append(SwapTensor(tensor)) + idx = self.count_idx + self.count_idx += 1 + self.pack_count += 1 + return idx + + def unpack_from_cpu(idx) -> torch.Tensor: + if isinstance(idx, torch.Tensor): + return idx + + swap_tensor = self.storage.swap_storage[idx].pop(0) + tensor = swap_tensor.get_val() + logger.info(f"enter unpack, tensor shape is {tensor.shape}, tensor id {id(tensor)}") + self.unpack_count += 1 + if self.unpack_count == self.pack_count: + self.storage = None + return tensor + + super().__init__(pack_to_cpu, unpack_from_cpu) + + +class ActivationWrapper(torch.nn.Module, ABC): + """ + Base class for Activation Swap. + + Not meant to be instantiated directly. + """ + + def __init__(self, module): + super().__init__() + self._swap_wrapped_module = module + # state_dict post hook to remove prefix to allow loading into a + # non-checkpoint wrapped module. + self._register_state_dict_hook(self._post_state_dict_hook) + # load_state_dict pre-hook to allow loading back into + # checkpoint-wrapped module. + self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook) + + @abstractmethod + def forward(self, *args, **kwargs): + raise ValueError("Subclasses should implement forward().") + + def __getattr__(self, name: str) -> Any: + """Forward missing attributes to wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + return getattr(self._swap_wrapped_module, name) + + def __getitem__(self, key: int) -> Any: + """Forward indexing calls in case the module is a nn.Sequential.""" + return self._swap_wrapped_module.__getitem__(key) # type: ignore[operator] + + def named_parameters( + self, + *args, + **kwargs, + ) -> Iterator[tuple[str, torch.nn.Parameter]]: + """ + Override :meth:`named_parameters()` to intercept parameter names. + + remove all occurrences of ``_SWAP_PREFIX``. + """ + for param_name, param in super().named_parameters(*args, **kwargs): + yield param_name.replace(_SWAP_PREFIX, ""), param + + @staticmethod + def _post_state_dict_hook( + module: nn.Module, # pylint: disable=W0613 + state_dict: dict[str, Any], + prefix: str, + *args: Any, # pylint: disable=W0613 + ) -> dict[str, Any]: + """ + _post_state_dict_hook() is called after the state_dict() of this FSDP module is executed. + + For ``swap_wrapper``, it will strip swap-wrapped module prefix, + so that this module can be loaded into non-swaped modules. + It would still be able to be loaded into swap-wrapped modules as this class, + adds the prefix back before loading the state_dict. + """ + _replace_by_prefix(state_dict, f"{prefix}{_SWAP_PREFIX}", prefix) + return state_dict + + @staticmethod + def _pre_load_state_dict_hook( + module: nn.Module, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + ) -> None: + """ + ``_pre_state_dict_hook` is called before ``self._load_from_state_dict()`` is called. + + For ``swap_wrapper``, it will add back the module + prefix so that non-swaped modules can be loaded into + swap_wrapper modules properly. + """ + _replace_by_prefix(state_dict, prefix, prefix + f"{_SWAP_PREFIX}") + + +class SwapWrapper(ActivationWrapper): + """ + Customize an nn.Module wrapper class to add an AsyncSaveOnCpu context manager for the target model. + """ + def __init__(self, mod: nn.Module, policy_fn: Optional[Callable] = None): + super().__init__(mod) + self.policy_fn = policy_fn + + def forward(self, *args, **kwargs): + with AsyncSaveOnCpu(policy_fn=self.policy_fn): + return self._swap_wrapped_module(*args, **kwargs) + + +def swap_wrapper(module: nn.Module, policy_fn: Optional[Callable] = None): + return SwapWrapper(module, policy_fn) diff --git a/tests/torch/activation_swap/test_act_swap.py b/tests/torch/activation_swap/test_act_swap.py new file mode 100644 index 0000000..376a60e --- /dev/null +++ b/tests/torch/activation_swap/test_act_swap.py @@ -0,0 +1,217 @@ +# Copyright 2025 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. +# ============================================================================ +"""Test Activation Swap memory comparison: None vs Swap""" +import time +import random +from contextlib import contextmanager +import numpy as np +import torch +from torch import nn +from torch.utils.data import DataLoader, TensorDataset +from tests.common.mark_utils import arg_mark +from hyper_parallel.platform.torch.activation_checkpoint import swap_wrapper, ActivationPolicy, SwapManager + + +def set_seed(seed=42): + """Set random seed for reproducibility.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.npu.is_available(): + torch.npu.manual_seed(seed) + torch.npu.manual_seed_all(seed) + +@contextmanager +def seed_memory_time_context(seed=42): + """Context manager to set seed, track peak memory and execution time.""" + set_seed(seed) + torch.npu.reset_peak_memory_stats() + torch.npu.empty_cache() + start_time = time.time() + + stats = {} + + try: + yield stats + finally: + exec_time = time.time() - start_time + peak_mem_gb = torch.npu.max_memory_allocated() / (1024 ** 3) + torch.npu.empty_cache() + stats["peak_mem"] = peak_mem_gb + stats["exec_time"] = exec_time + + +class TransformerBlock(nn.Module): + """A simple Transformer block for testing purposes.""" + + def __init__(self, dim=256, num_heads=4): + super().__init__() + self.dim = dim + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.ffn = nn.Sequential( + nn.Linear(dim, dim * 4), + nn.ReLU(), + nn.Linear(dim * 4, dim) + ) + + def forward(self, x): + attn_out, _ = self.attn(x, x, x) + x = x + attn_out + x = self.norm1(x) + x = x + self.ffn(x) + x = self.norm2(x) + return x + +class SimpleTransformer(nn.Module): + """A simple Transformer model for testing purposes.""" + + def __init__(self, vocab_size=1000, dim=2048, depth=6): + super().__init__() + self.embed = nn.Embedding(vocab_size, dim) + self.layers = nn.ModuleList([TransformerBlock(dim) for _ in range(depth)]) + self.norm = nn.LayerNorm(dim) + self.head = nn.Linear(dim, vocab_size) + + def forward(self, x): + x = self.embed(x) + for block in self.layers: + x = block(x) + x = self.norm(x) + return self.head(x) + +def prepare_data(batch_size=8, seq_len=512, num_samples=64): + dataset = TensorDataset( + torch.randint(0, 10000, (num_samples, seq_len)), + torch.randint(0, 10000, (num_samples, seq_len)) + ) + dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False) + return dataloader + +def apply_recompute(model, mode): + """Apply activation swap to the model based on the specified mode.""" + if mode == "none": + return model + + if mode == "swap": + for i, layer in enumerate(model.layers): + model.layers[i].attn = swap_wrapper(layer.attn) + + for i in range(len(model.layers) - 1): + SwapManager().set_forward_prefetch_layer(model.layers[i], model.layers[i + 1]) + elif mode == "swap_with_policy": + def policy_fn(x): + if x.storage().size() <= 32*512*512: + return ActivationPolicy.SAVE + return ActivationPolicy.SWAP + + for i, layer in enumerate(model.layers): + model.layers[i].attn = swap_wrapper(layer.attn, policy_fn) + + for i in range(len(model.layers) - 1): + SwapManager().set_forward_prefetch_layer(model.layers[i], model.layers[i + 1]) + else: + raise ValueError(f"Unknown mode: {mode}") + return model + + +def train_one_mode(mode, dataloader, train_steps=5): + """Run training for one mode and return metrics.""" + with seed_memory_time_context() as stats: + model = SimpleTransformer(vocab_size=32000, dim=2048, depth=16).npu() + model = apply_recompute(model, mode) + + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) + losses = [] + for step, (x, y) in enumerate(dataloader): + if step >= train_steps: + break + x, y = x.npu(), y.npu() + optimizer.zero_grad() + logits = model(x) + loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1)) + loss.backward() + optimizer.step() + losses.append(loss.item()) + + return losses, stats.get("peak_mem"), stats.get("exec_time") + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_act_swap_memory_comparison(): + """ + Feature: Activation Swap Memory Behavior + Description: Compare peak memory usage across four modes: + 'none' (baseline), + 'swap' (swap all tensors), + 'swap_with_policy' (swap part of tensors according to the policy function). + Validates that losses are numerically identical at every training step + and that peak memory follows the expected hierarchy: + Expectation: NONE > SWAP_WITH_POLICY > SWAP, + the memory usage trend is satisfied, and no OOM occurs. + """ + print("🚀 Starting memory and time comparison: none vs swap vs swap_with_policy") + dataloader = prepare_data() + train_steps=3 + + modes = ["none", "swap", "swap_with_policy"] + results = {} + + for mode in modes: + print(f"\n--- Running mode: {mode.upper()} ---") + losses, peak_mem, duration = train_one_mode(mode, dataloader, train_steps) + results[mode] = { + "losses": losses, + "peak_mem_gb": peak_mem, + "time_sec": duration + } + print(f"{mode}: Loss={losses[-1]:.4f}, Peak Mem={peak_mem:.5f} GB, Time={duration:.5f}s") + + print("\n" + "="*70) + print("📊 FINAL COMPARISON") + print("="*70) + print(f"{'Mode':<15} | {'Peak Mem (GB)':<15} | {'Time (s)':<10} | {'Final Loss':<12}") + print("-"*70) + for mode in modes: + r = results[mode] + print(f"{mode.upper():<20} | {r['peak_mem_gb']:<15.5f} | {r['time_sec']:<10.5f} | {r['losses'][-1]:<12.4f}") + + # loss assert + base_losses = results["none"]["losses"] + tol = 1e-5 + for step in range(train_steps): + base_val = base_losses[step] + for mode in ["swap", "swap_with_policy"]: + val = results[mode]["losses"][step] + diff = abs(val - base_val) + assert diff < tol, ( + f"Loss mismatch at step {step} in mode '{mode}': " + f"none={base_val:.8f}, {mode}={val:.8f}, diff={diff:.2e}" + ) + print(f"\n✅ All {train_steps} steps: losses are consistent across modes (tol={tol}).") + + # memory assert + mem_none = results["none"]["peak_mem_gb"] + mem_swap = results["swap"]["peak_mem_gb"] + mem_swap_with_policy = results["swap_with_policy"]["peak_mem_gb"] + + # mem_none > mem_swap_with_policy + assert mem_none > mem_swap_with_policy, \ + f"Expected NONE ({mem_none:.5f}) > SWAP_WITH_POLICY ({mem_swap_with_policy:.5f})" + print(f"✅ Verified: NONE ({mem_none:.5f}) > SWAP_WITH_POLICY ({mem_swap_with_policy:.5f})") + # mem_swap_with_policy > mem_swap + assert mem_swap_with_policy > mem_swap, \ + f"Expected SWAP_WITH_POLICY ({mem_swap_with_policy:.5f}) > SWAP ({mem_swap:.5f})" + print(f"✅ Verified: SWAP_WITH_POLICY ({mem_swap_with_policy:.5f}) > SWAP ({mem_swap:.5f})") -- Gitee