« prev ^ index » next coverage.py v7.13.1, created at 2026-08-22 04:23 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15# Adapted from
16# https://github.com/pytorch/pytorch/blob/release/2.6/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py
17# enhanced with activation swap functionality.
18# ============================================================================
19"""Activation Swap implementation for PyTorch."""
20# pylint: disable=W0212, W0613
22from abc import ABC, abstractmethod
23from collections.abc import Iterator
24from typing import Optional, Callable, Any, Union
25import types
26import warnings
27import torch
28from torch import nn
29from torch.distributed.utils import _replace_by_prefix
30from hyper_parallel.core.activation_checkpoint.activation_checkpoint import CheckpointPolicy
31from hyper_parallel.core.activation_checkpoint.swap import SwapManager, SwapTensor, Storage
34_SWAP_WRAPPED_MODULE = "_swap_wrapped_module"
35_SWAP_PREFIX = _SWAP_WRAPPED_MODULE + "."
38def _raise_if_compiling(feature_name: str) -> None:
39 """Reject HyperParallel activation-swap extensions during Torch capture."""
40 if torch.compiler.is_compiling():
41 raise ValueError(
42 f"HyperParallel {feature_name} is not supported in compile mode. "
43 "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
44 )
47class FuncModule(nn.Module):
48 """
49 Thin :class:`~torch.nn.Module` adapter that wraps a plain callable.
51 Allows ordinary Python functions (or any callable without Module
52 parameters) to be passed to :func:`swap_wrapper` and
53 :func:`~hyper_parallel.core.activation_checkpoint.checkpoint_wrapper`
54 in place of an :class:`~torch.nn.Module`.
55 The wrapped function is stored as ``_fn`` and invoked in
56 :meth:`forward`; the module has no trainable parameters.
58 Args:
59 fn (callable): The function to wrap.
61 Example:
62 >>> wrapped = swap_wrapper(lambda x: x * 2)
63 """
65 def __init__(self, fn: Callable):
66 super().__init__()
67 self._fn = fn
69 def forward(self, *args, **kwargs):
70 """Invoke the wrapped callable with the given arguments."""
71 return self._fn(*args, **kwargs)
74def _is_callable_exempt_from_overlap_check(callable_obj: Callable) -> bool:
75 """Return True for callables that cannot be reliably overlap-tracked by object marks."""
76 return isinstance(callable_obj, (types.FunctionType, types.BuiltinFunctionType, types.MethodType))
79def _iter_wrappable_callable_attrs(module: nn.Module) -> Iterator[tuple[str, Callable]]:
80 """Yield public per-instance callable attributes not registered as child modules.
82 Plain functions, builtins and bound methods are skipped: these are stateless
83 module-level utilities shared by reference across many modules (e.g.
84 ``self.act = F.gelu`` repeated in every layer). They are never standalone
85 checkpoint regions, and marking a shared function's ``_is_wrapped`` flag
86 would both mutate a global object and falsely flag every sibling module that
87 references the same function as an overlapping wrap. Only per-instance
88 callables participate in overlap tracking.
89 """
90 for attr_name, attr_value in vars(module).items():
91 if attr_name.startswith("_") or isinstance(attr_value, nn.Module):
92 continue
93 if _is_callable_exempt_from_overlap_check(attr_value):
94 continue
95 if callable(attr_value):
96 yield attr_name, attr_value
99def _mark_wrapped(obj: Any) -> None:
100 try:
101 obj._is_wrapped = True # pylint: disable=W0212
102 except (AttributeError, TypeError):
103 pass
106def _get_wrapped_callable(module: nn.Module) -> Optional[Callable]:
107 wrapped_module = getattr(module, _SWAP_WRAPPED_MODULE, None)
108 if isinstance(wrapped_module, FuncModule):
109 return getattr(wrapped_module, "_fn", None)
110 if isinstance(module, FuncModule):
111 return getattr(module, "_fn", None)
112 return None
115def _raise_callable_already_wrapped(callable_obj: Callable) -> None:
116 warnings.warn(
117 f"Callable '{callable_obj.__class__.__name__}' is already wrapped. "
118 "Wrapping overlapping module regions is not allowed."
119 )
122def _check_callable_attr_not_wrapped(owner: nn.Module, attr_name: str, attr_value: Callable) -> None:
123 del owner, attr_name
124 if getattr(attr_value, '_is_wrapped', False):
125 _raise_callable_already_wrapped(attr_value)
128def _check_and_mark_callable(callable_obj: Callable) -> None:
129 if _is_callable_exempt_from_overlap_check(callable_obj):
130 return
131 if getattr(callable_obj, '_is_wrapped', False):
132 warnings.warn(
133 f"Callable '{callable_obj.__class__.__name__}' or one of its ancestors is already wrapped. "
134 "Wrapping overlapping module regions is not allowed."
135 )
136 _mark_wrapped(callable_obj)
139def _check_and_mark_wrapped(module: nn.Module) -> None:
140 """Validate no wrapping overlap, then mark module and all descendants as wrapped."""
141 if getattr(module, '_is_wrapped', False):
142 warnings.warn(
143 f"Module '{module.__class__.__name__}' or one of its ancestors is already wrapped. "
144 "Wrapping overlapping module regions is not allowed."
145 )
146 for submodule in module.modules():
147 if submodule is module:
148 continue
149 wrapped_callable = _get_wrapped_callable(submodule)
150 if wrapped_callable is not None and _is_callable_exempt_from_overlap_check(wrapped_callable):
151 continue
152 if getattr(submodule, '_is_wrapped', False):
153 if wrapped_callable is not None:
154 _raise_callable_already_wrapped(wrapped_callable)
155 warnings.warn(
156 f"Submodule '{getattr(submodule, '_swap_wrapped_module', submodule).__class__.__name__}' of "
157 f"'{module.__class__.__name__}' is already wrapped. "
158 "Wrapping overlapping module regions is not allowed."
159 )
160 for submodule in module.modules():
161 for attr_name, attr_value in _iter_wrappable_callable_attrs(submodule):
162 _check_callable_attr_not_wrapped(submodule, attr_name, attr_value)
163 for submodule in module.modules():
164 _mark_wrapped(submodule)
165 for _, attr_value in _iter_wrappable_callable_attrs(submodule):
166 _mark_wrapped(attr_value)
169def base_check_fn(tensor) -> bool:
170 """
171 Basic check to determine if a tensor is eligible for offloading.
172 - Skip Parameters and their views.
173 - Skip empty storage tensors.
174 """
175 if isinstance(tensor._base, torch.nn.parameter.Parameter) or isinstance(tensor, torch.nn.parameter.Parameter): # pylint: disable=W0212
176 return False
177 if tensor.untyped_storage().size() == 0:
178 return False
179 return True
182class AsyncSaveOnCpu(torch.autograd.graph.saved_tensors_hooks):
183 """
184 Context manager to offload tensors to CPU during forward pass.
185 """
186 def __init__(self, policy_fn=None, group_swap: bool = False) -> None:
187 _raise_if_compiling("AsyncSaveOnCpu")
188 self.add_to_storage = False
189 self.storage = Storage()
190 self.count_idx = 0
191 self.policy_fn = policy_fn
193 # Cache per-context-manager state once to avoid per-tensor singleton lookups.
194 swap_manager = SwapManager()
196 def pack_to_cpu(tensor: torch.Tensor):
197 if not base_check_fn(tensor):
198 return tensor.detach()
199 if policy_fn is not None:
200 if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE:
201 return tensor.detach()
202 if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP:
203 raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}")
204 group_name = swap_manager.get_current_group_name()
205 if not group_name:
206 return tensor.detach()
207 if not self.add_to_storage:
208 swap_manager.add_storage(group_name, self.storage)
209 self.add_to_storage = True
210 funcname = f"{group_name}::{tensor.shape}"
211 detached = tensor.detach()
212 self.storage[self.count_idx].append(
213 SwapTensor(detached, funcname, group_swap=group_swap)
214 )
215 self.count_idx += 1
216 return detached
218 def unpack_from_cpu(tensor) -> torch.Tensor:
219 if self.storage is not None:
220 self.storage.clear()
221 self.storage = None
222 return tensor
224 super().__init__(pack_to_cpu, unpack_from_cpu)
227class ActivationWrapper(torch.nn.Module, ABC):
228 """
229 Base class for Activation Swap.
231 Not meant to be instantiated directly.
232 """
234 def __init__(self, module: Union[nn.Module, Callable], *, track_overlaps: bool = True):
235 """Initialize a wrapper and optionally participate in overlap tracking."""
236 if callable(module) and not isinstance(module, nn.Module):
237 if track_overlaps:
238 _check_and_mark_callable(module)
239 module = FuncModule(module)
240 if track_overlaps:
241 _mark_wrapped(module)
242 elif track_overlaps:
243 _check_and_mark_wrapped(module)
244 super().__init__()
245 self._swap_wrapped_module = module
246 self._is_wrapped = track_overlaps
247 # state_dict post hook to remove prefix to allow loading into a
248 # non-swap wrapped module.
249 self._register_state_dict_hook(self._post_state_dict_hook)
250 # load_state_dict pre-hook to allow loading back into
251 # swap-wrapped module.
252 self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook)
254 @property
255 def _wrapped_module(self):
256 return self._swap_wrapped_module
258 @abstractmethod
259 def forward(self, *args, **kwargs):
260 """Run the wrapped module's forward pass with activation swapping. Must be implemented by subclasses."""
261 raise ValueError("Subclasses should implement forward().")
263 def __getattr__(self, name: str) -> Any:
264 """Forward missing attributes to wrapped module."""
265 try:
266 return super().__getattr__(name) # defer to nn.Module's logic
267 except AttributeError:
268 return getattr(self._swap_wrapped_module, name)
270 def __getitem__(self, key: int) -> Any:
271 """Forward indexing calls in case the module is a nn.Sequential."""
272 return self._swap_wrapped_module.__getitem__(key) # type: ignore[operator]
274 def named_modules(
275 self,
276 memo: Optional[set[nn.Module]] = None,
277 prefix: str = "",
278 remove_duplicate: bool = True,
279 ) -> Iterator[tuple[str, nn.Module]]:
280 """
281 Yield wrapped-module children without exposing the internal wrapper prefix.
283 PyTorch parent modules implement ``named_parameters(recurse=True)`` by
284 iterating ``named_modules()`` and reading each module's direct
285 ``_parameters``. They do not call child modules' ``named_parameters()``
286 overrides. Exposing the wrapped module under the wrapper's own prefix
287 keeps root-module traversals aligned with ``state_dict()`` keys.
289 Args:
290 memo (Optional[set[nn.Module]], optional): A memo set to avoid infinite recursion. Default: ``None``.
291 prefix (str, optional): A prefix to prepend to all module names. Default: ``""``.
292 remove_duplicate (bool, optional): Whether to remove duplicate modules. Default: ``True``.
294 Returns:
295 Iterator[tuple[str, nn.Module]] An iterator of (name, module) pairs.
296 """
297 if memo is None:
298 memo = set()
299 if self not in memo:
300 memo.add(self)
301 yield prefix, self
302 yield from self._swap_wrapped_module.named_modules(
303 memo=memo,
304 prefix=prefix,
305 remove_duplicate=remove_duplicate,
306 )
308 def named_parameters(
309 self,
310 *args,
311 **kwargs,
312 ) -> Iterator[tuple[str, torch.nn.Parameter]]:
313 """
314 Override :meth:`named_parameters()` to intercept parameter names.
316 remove all occurrences of ``_SWAP_PREFIX``.
317 """
318 for param_name, param in super().named_parameters(*args, **kwargs):
319 yield param_name.replace(_SWAP_PREFIX, ""), param
321 @staticmethod
322 def _post_state_dict_hook(
323 module: nn.Module, # pylint: disable=W0613
324 state_dict: dict[str, Any],
325 prefix: str,
326 *args: Any, # pylint: disable=W0613
327 ) -> dict[str, Any]:
328 """
329 _post_state_dict_hook() is called after the state_dict() of this FSDP module is executed.
331 For ``swap_wrapper``, it will strip swap-wrapped module prefix,
332 so that this module can be loaded into non-swapped modules.
333 It would still be able to be loaded into swap-wrapped modules as this class,
334 adds the prefix back before loading the state_dict.
335 """
336 _replace_by_prefix(state_dict, f"{prefix}{_SWAP_PREFIX}", prefix)
337 return state_dict
339 @staticmethod
340 def _pre_load_state_dict_hook(
341 module: nn.Module,
342 state_dict: dict[str, Any],
343 prefix: str,
344 *args: Any,
345 ) -> None:
346 """
347 ``_pre_state_dict_hook` is called before ``self._load_from_state_dict()`` is called.
349 For ``swap_wrapper``, it will add back the module
350 prefix so that non-swapped modules can be loaded into
351 swap_wrapper modules properly.
352 """
353 _replace_by_prefix(state_dict, prefix, prefix + f"{_SWAP_PREFIX}")
356class SwapWrapper(ActivationWrapper):
357 """
358 Customize an nn.Module wrapper class to add an AsyncSaveOnCpu context manager for the target model.
359 """
360 def __init__(
361 self,
362 mod: Union[nn.Module, Callable],
363 policy_fn: Optional[Callable] = None,
364 group_swap: bool = False,
365 ):
366 super().__init__(mod)
367 self.policy_fn = policy_fn
368 self.group_swap = group_swap
370 def forward(self, *args, **kwargs):
371 """Run the wrapped module inside an AsyncSaveOnCpu context for activation swapping."""
372 _raise_if_compiling("swap_wrapper")
373 with AsyncSaveOnCpu(policy_fn=self.policy_fn, group_swap=self.group_swap):
374 return self._swap_wrapped_module(*args, **kwargs)
377def swap_wrapper(
378 module: Union[nn.Module, Callable],
379 policy_fn: Optional[Callable] = None,
380 group_swap: bool = False,
381) -> SwapWrapper:
382 """Wrap a module or callable with activation swap functionality."""
383 return SwapWrapper(module, policy_fn, group_swap)
386def swap_tensor_wrapper(target, tag: Optional[str] = None, group_swap: bool = False):
387 """Register selected tensors into the current swap group.
389 This helper is intended to be used inside a forward path that already
390 participates in the existing swap scheduling managed by ``SwapManager``.
391 It preserves the input structure and returns the original tensors.
392 """
393 _raise_if_compiling("swap_tensor_wrapper")
394 swap_manager = SwapManager()
395 group_name = swap_manager.get_current_group_name()
396 if not group_name:
397 warnings.warn(
398 f"Tensor {tag} cannot be swapped, for its group is unregistered."
399 )
400 return target
401 if swap_manager.is_last_group(group_name):
402 return target
404 storage = Storage()
405 count_idx = 0
407 def _register_tensor(tensor):
408 nonlocal count_idx
409 if not base_check_fn(tensor):
410 return tensor
412 tensor_tag = tag or f"{group_name}_swap_tensor"
413 funcname = f"{tensor_tag}::{tuple(tensor.shape)}"
414 storage[count_idx].append(SwapTensor(tensor, funcname, group_swap=group_swap))
415 count_idx += 1
416 return tensor
418 wrapped = torch.utils._pytree.tree_map( # pylint: disable=protected-access
419 lambda x: _register_tensor(x) if isinstance(x, torch.Tensor) else x,
420 target,
421 )
422 if count_idx > 0:
423 swap_manager.add_storage(group_name, storage)
424 return wrapped