« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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 + "."
38class FuncModule(nn.Module):
39 """
40 Thin :class:`~torch.nn.Module` adapter that wraps a plain callable.
42 Allows ordinary Python functions (or any callable without Module
43 parameters) to be passed to :func:`swap_wrapper` and
44 :func:`~hyper_parallel.core.activation_checkpoint.checkpoint_wrapper`
45 in place of an :class:`~torch.nn.Module`.
46 The wrapped function is stored as ``_fn`` and invoked in
47 :meth:`forward`; the module has no trainable parameters.
49 Args:
50 fn (callable): The function to wrap.
52 Example:
53 >>> wrapped = swap_wrapper(lambda x: x * 2)
54 """
56 def __init__(self, fn: Callable):
57 super().__init__()
58 self._fn = fn
60 def forward(self, *args, **kwargs):
61 """Invoke the wrapped callable with the given arguments."""
62 return self._fn(*args, **kwargs)
65def _is_callable_exempt_from_overlap_check(callable_obj: Callable) -> bool:
66 """Return True for callables that cannot be reliably overlap-tracked by object marks."""
67 return isinstance(callable_obj, (types.FunctionType, types.BuiltinFunctionType, types.MethodType))
70def _iter_wrappable_callable_attrs(module: nn.Module) -> Iterator[tuple[str, Callable]]:
71 """Yield public per-instance callable attributes not registered as child modules.
73 Plain functions, builtins and bound methods are skipped: these are stateless
74 module-level utilities shared by reference across many modules (e.g.
75 ``self.act = F.gelu`` repeated in every layer). They are never standalone
76 checkpoint regions, and marking a shared function's ``_is_wrapped`` flag
77 would both mutate a global object and falsely flag every sibling module that
78 references the same function as an overlapping wrap. Only per-instance
79 callables participate in overlap tracking.
80 """
81 for attr_name, attr_value in vars(module).items():
82 if attr_name.startswith("_") or isinstance(attr_value, nn.Module):
83 continue
84 if _is_callable_exempt_from_overlap_check(attr_value):
85 continue
86 if callable(attr_value):
87 yield attr_name, attr_value
90def _mark_wrapped(obj: Any) -> None:
91 try:
92 obj._is_wrapped = True # pylint: disable=W0212
93 except (AttributeError, TypeError):
94 pass
97def _get_wrapped_callable(module: nn.Module) -> Optional[Callable]:
98 wrapped_module = getattr(module, _SWAP_WRAPPED_MODULE, None)
99 if isinstance(wrapped_module, FuncModule):
100 return getattr(wrapped_module, "_fn", None)
101 if isinstance(module, FuncModule):
102 return getattr(module, "_fn", None)
103 return None
106def _raise_callable_already_wrapped(callable_obj: Callable) -> None:
107 warnings.warn(
108 f"Callable '{callable_obj.__class__.__name__}' is already wrapped. "
109 "Wrapping overlapping module regions is not allowed."
110 )
113def _check_callable_attr_not_wrapped(owner: nn.Module, attr_name: str, attr_value: Callable) -> None:
114 del owner, attr_name
115 if getattr(attr_value, '_is_wrapped', False):
116 _raise_callable_already_wrapped(attr_value)
119def _check_and_mark_callable(callable_obj: Callable) -> None:
120 if _is_callable_exempt_from_overlap_check(callable_obj):
121 return
122 if getattr(callable_obj, '_is_wrapped', False):
123 warnings.warn(
124 f"Callable '{callable_obj.__class__.__name__}' or one of its ancestors is already wrapped. "
125 "Wrapping overlapping module regions is not allowed."
126 )
127 _mark_wrapped(callable_obj)
130def _check_and_mark_wrapped(module: nn.Module) -> None:
131 """Validate no wrapping overlap, then mark module and all descendants as wrapped."""
132 if getattr(module, '_is_wrapped', False):
133 warnings.warn(
134 f"Module '{module.__class__.__name__}' or one of its ancestors is already wrapped. "
135 "Wrapping overlapping module regions is not allowed."
136 )
137 for submodule in module.modules():
138 if submodule is module:
139 continue
140 wrapped_callable = _get_wrapped_callable(submodule)
141 if wrapped_callable is not None and _is_callable_exempt_from_overlap_check(wrapped_callable):
142 continue
143 if getattr(submodule, '_is_wrapped', False):
144 if wrapped_callable is not None:
145 _raise_callable_already_wrapped(wrapped_callable)
146 warnings.warn(
147 f"Submodule '{getattr(submodule, '_swap_wrapped_module', submodule).__class__.__name__}' of "
148 f"'{module.__class__.__name__}' is already wrapped. "
149 "Wrapping overlapping module regions is not allowed."
150 )
151 for submodule in module.modules():
152 for attr_name, attr_value in _iter_wrappable_callable_attrs(submodule):
153 _check_callable_attr_not_wrapped(submodule, attr_name, attr_value)
154 for submodule in module.modules():
155 _mark_wrapped(submodule)
156 for _, attr_value in _iter_wrappable_callable_attrs(submodule):
157 _mark_wrapped(attr_value)
160def base_check_fn(tensor) -> bool:
161 """
162 Basic check to determine if a tensor is eligible for offloading.
163 - Skip Parameters and their views.
164 - Skip empty storage tensors.
165 """
166 if isinstance(tensor._base, torch.nn.parameter.Parameter) or isinstance(tensor, torch.nn.parameter.Parameter): # pylint: disable=W0212
167 return False
168 if tensor.untyped_storage().size() == 0:
169 return False
170 return True
173class AsyncSaveOnCpu(torch.autograd.graph.saved_tensors_hooks):
174 """
175 Context manager to offload tensors to CPU during forward pass.
176 """
177 def __init__(self, policy_fn=None, group_swap: bool = False) -> None:
178 self.add_to_storage = False
179 self.storage = Storage()
180 self.count_idx = 0
181 self.policy_fn = policy_fn
183 # Cache per-context-manager state once to avoid per-tensor singleton lookups.
184 swap_manager = SwapManager()
186 def pack_to_cpu(tensor: torch.Tensor):
187 if not base_check_fn(tensor):
188 return tensor.detach()
189 if policy_fn is not None:
190 if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE:
191 return tensor.detach()
192 if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP:
193 raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}")
194 group_name = swap_manager.get_current_group_name()
195 if not group_name:
196 return tensor.detach()
197 if not self.add_to_storage:
198 swap_manager.add_storage(group_name, self.storage)
199 self.add_to_storage = True
200 funcname = f"{group_name}::{tensor.shape}"
201 detached = tensor.detach()
202 self.storage[self.count_idx].append(
203 SwapTensor(detached, funcname, group_swap=group_swap)
204 )
205 self.count_idx += 1
206 return detached
208 def unpack_from_cpu(tensor) -> torch.Tensor:
209 if self.storage is not None:
210 self.storage.clear()
211 self.storage = None
212 return tensor
214 super().__init__(pack_to_cpu, unpack_from_cpu)
217class ActivationWrapper(torch.nn.Module, ABC):
218 """
219 Base class for Activation Swap.
221 Not meant to be instantiated directly.
222 """
224 def __init__(self, module: Union[nn.Module, Callable], *, track_overlaps: bool = True):
225 """Initialize a wrapper and optionally participate in overlap tracking."""
226 if callable(module) and not isinstance(module, nn.Module):
227 if track_overlaps:
228 _check_and_mark_callable(module)
229 module = FuncModule(module)
230 if track_overlaps:
231 _mark_wrapped(module)
232 elif track_overlaps:
233 _check_and_mark_wrapped(module)
234 super().__init__()
235 self._swap_wrapped_module = module
236 self._is_wrapped = track_overlaps
237 # state_dict post hook to remove prefix to allow loading into a
238 # non-swap wrapped module.
239 self._register_state_dict_hook(self._post_state_dict_hook)
240 # load_state_dict pre-hook to allow loading back into
241 # swap-wrapped module.
242 self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook)
244 @property
245 def _wrapped_module(self):
246 return self._swap_wrapped_module
248 @abstractmethod
249 def forward(self, *args, **kwargs):
250 """Run the wrapped module's forward pass with activation swapping. Must be implemented by subclasses."""
251 raise ValueError("Subclasses should implement forward().")
253 def __getattr__(self, name: str) -> Any:
254 """Forward missing attributes to wrapped module."""
255 try:
256 return super().__getattr__(name) # defer to nn.Module's logic
257 except AttributeError:
258 return getattr(self._swap_wrapped_module, name)
260 def __getitem__(self, key: int) -> Any:
261 """Forward indexing calls in case the module is a nn.Sequential."""
262 return self._swap_wrapped_module.__getitem__(key) # type: ignore[operator]
264 def named_modules(
265 self,
266 memo: Optional[set[nn.Module]] = None,
267 prefix: str = "",
268 remove_duplicate: bool = True,
269 ) -> Iterator[tuple[str, nn.Module]]:
270 """
271 Yield wrapped-module children without exposing the internal wrapper prefix.
273 PyTorch parent modules implement ``named_parameters(recurse=True)`` by
274 iterating ``named_modules()`` and reading each module's direct
275 ``_parameters``. They do not call child modules' ``named_parameters()``
276 overrides. Exposing the wrapped module under the wrapper's own prefix
277 keeps root-module traversals aligned with ``state_dict()`` keys.
279 Args:
280 memo (Optional[set[nn.Module]], optional): A memo set to avoid infinite recursion. Default: ``None``.
281 prefix (str, optional): A prefix to prepend to all module names. Default: ``""``.
282 remove_duplicate (bool, optional): Whether to remove duplicate modules. Default: ``True``.
284 Returns:
285 Iterator[tuple[str, nn.Module]] An iterator of (name, module) pairs.
286 """
287 if memo is None:
288 memo = set()
289 if self not in memo:
290 memo.add(self)
291 yield prefix, self
292 yield from self._swap_wrapped_module.named_modules(
293 memo=memo,
294 prefix=prefix,
295 remove_duplicate=remove_duplicate,
296 )
298 def named_parameters(
299 self,
300 *args,
301 **kwargs,
302 ) -> Iterator[tuple[str, torch.nn.Parameter]]:
303 """
304 Override :meth:`named_parameters()` to intercept parameter names.
306 remove all occurrences of ``_SWAP_PREFIX``.
307 """
308 for param_name, param in super().named_parameters(*args, **kwargs):
309 yield param_name.replace(_SWAP_PREFIX, ""), param
311 @staticmethod
312 def _post_state_dict_hook(
313 module: nn.Module, # pylint: disable=W0613
314 state_dict: dict[str, Any],
315 prefix: str,
316 *args: Any, # pylint: disable=W0613
317 ) -> dict[str, Any]:
318 """
319 _post_state_dict_hook() is called after the state_dict() of this FSDP module is executed.
321 For ``swap_wrapper``, it will strip swap-wrapped module prefix,
322 so that this module can be loaded into non-swapped modules.
323 It would still be able to be loaded into swap-wrapped modules as this class,
324 adds the prefix back before loading the state_dict.
325 """
326 _replace_by_prefix(state_dict, f"{prefix}{_SWAP_PREFIX}", prefix)
327 return state_dict
329 @staticmethod
330 def _pre_load_state_dict_hook(
331 module: nn.Module,
332 state_dict: dict[str, Any],
333 prefix: str,
334 *args: Any,
335 ) -> None:
336 """
337 ``_pre_state_dict_hook` is called before ``self._load_from_state_dict()`` is called.
339 For ``swap_wrapper``, it will add back the module
340 prefix so that non-swapped modules can be loaded into
341 swap_wrapper modules properly.
342 """
343 _replace_by_prefix(state_dict, prefix, prefix + f"{_SWAP_PREFIX}")
346class SwapWrapper(ActivationWrapper):
347 """
348 Customize an nn.Module wrapper class to add an AsyncSaveOnCpu context manager for the target model.
349 """
350 def __init__(
351 self,
352 mod: Union[nn.Module, Callable],
353 policy_fn: Optional[Callable] = None,
354 group_swap: bool = False,
355 ):
356 super().__init__(mod)
357 self.policy_fn = policy_fn
358 self.group_swap = group_swap
360 def forward(self, *args, **kwargs):
361 """Run the wrapped module inside an AsyncSaveOnCpu context for activation swapping."""
362 with AsyncSaveOnCpu(policy_fn=self.policy_fn, group_swap=self.group_swap):
363 return self._swap_wrapped_module(*args, **kwargs)
366def swap_wrapper(
367 module: Union[nn.Module, Callable],
368 policy_fn: Optional[Callable] = None,
369 group_swap: bool = False,
370) -> SwapWrapper:
371 """Wrap a module or callable with activation swap functionality."""
372 return SwapWrapper(module, policy_fn, group_swap)
375def swap_tensor_wrapper(target, tag: Optional[str] = None, group_swap: bool = False):
376 """Register selected tensors into the current swap group.
378 This helper is intended to be used inside a forward path that already
379 participates in the existing swap scheduling managed by ``SwapManager``.
380 It preserves the input structure and returns the original tensors.
381 """
382 swap_manager = SwapManager()
383 group_name = swap_manager.get_current_group_name()
384 if not group_name:
385 warnings.warn(
386 f"Tensor {tag} cannot be swapped, for its group is unregistered."
387 )
388 return target
389 if swap_manager.is_last_group(group_name):
390 return target
392 storage = Storage()
393 count_idx = 0
395 def _register_tensor(tensor):
396 nonlocal count_idx
397 if not base_check_fn(tensor):
398 return tensor
400 tensor_tag = tag or f"{group_name}_swap_tensor"
401 funcname = f"{tensor_tag}::{tuple(tensor.shape)}"
402 storage[count_idx].append(SwapTensor(tensor, funcname, group_swap=group_swap))
403 count_idx += 1
404 return tensor
406 wrapped = torch.utils._pytree.tree_map( # pylint: disable=protected-access
407 lambda x: _register_tensor(x) if isinstance(x, torch.Tensor) else x,
408 target,
409 )
410 if count_idx > 0:
411 swap_manager.add_storage(group_name, storage)
412 return wrapped