Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / wrapper.py: 85%
258 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 and checkpoint wrappers for PyTorch."""
20# pylint: disable=W0212, W0613
22from abc import ABC, abstractmethod
23from collections.abc import Iterator
24from typing import Optional, Callable, Any, Union, cast
25import types
26import warnings
27import torch
28from torch import nn
29from torch.distributed.utils import _replace_by_prefix
30from hyper_parallel.core.activation_memory.policy import CheckpointPolicy
31from hyper_parallel.core.activation_memory.recompute_state import get_recompute_state
32from hyper_parallel.core.activation_memory.swap import SwapManager, SwapTensor, Storage
33from . import checkpoint_exclude as _checkpoint_exclude
36_SWAP_WRAPPED_MODULE = "_swap_wrapped_module"
37_SWAP_PREFIX = _SWAP_WRAPPED_MODULE + "."
40def _raise_if_compiling(feature_name: str) -> None:
41 """Reject HyperParallel activation-swap extensions during Torch capture."""
42 if torch.compiler.is_compiling():
43 raise ValueError(
44 f"HyperParallel {feature_name} is not supported in compile mode. "
45 "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
46 )
49class FuncModule(nn.Module):
50 """
51 Thin :class:`~torch.nn.Module` adapter that wraps a plain callable.
53 Allows ordinary Python functions (or any callable without Module
54 parameters) to be passed to :func:`swap_wrapper` and
55 :func:`~hyper_parallel.core.activation_memory.checkpoint_wrapper`
56 in place of an :class:`~torch.nn.Module`.
57 The wrapped function is stored as ``_fn`` and invoked in
58 :meth:`forward`; the module has no trainable parameters.
60 Args:
61 fn (callable): The function to wrap.
63 Example:
64 >>> wrapped = swap_wrapper(lambda x: x * 2)
65 """
67 def __init__(self, fn: Callable):
68 super().__init__()
69 self._fn = fn
71 def forward(self, *args, **kwargs):
72 """Invoke the wrapped callable with the given arguments."""
73 return self._fn(*args, **kwargs)
76def _is_callable_exempt_from_overlap_check(callable_obj: Callable) -> bool:
77 """Return True for callables that cannot be reliably overlap-tracked by object marks."""
78 return isinstance(callable_obj, (types.FunctionType, types.BuiltinFunctionType, types.MethodType))
81def _iter_wrappable_callable_attrs(module: nn.Module) -> Iterator[tuple[str, Callable]]:
82 """Yield public per-instance callable attributes not registered as child modules.
84 Plain functions, builtins and bound methods are skipped: these are stateless
85 module-level utilities shared by reference across many modules (e.g.
86 ``self.act = F.gelu`` repeated in every layer). They are never standalone
87 checkpoint regions, and marking a shared function's ``_is_wrapped`` flag
88 would both mutate a global object and falsely flag every sibling module that
89 references the same function as an overlapping wrap. Only per-instance
90 callables participate in overlap tracking.
91 """
92 for attr_name, attr_value in vars(module).items():
93 if attr_name.startswith("_") or isinstance(attr_value, nn.Module):
94 continue
95 if _is_callable_exempt_from_overlap_check(attr_value):
96 continue
97 if callable(attr_value):
98 yield attr_name, attr_value
101def _mark_wrapped(obj: Any) -> None:
102 """Tag ``obj`` as already wrapped, ignoring objects that reject new attributes."""
103 try:
104 setattr(obj, "_is_wrapped", True)
105 except (AttributeError, TypeError):
106 pass
109def _get_wrapped_callable(module: nn.Module) -> Optional[Callable]:
110 wrapped_module = getattr(module, _SWAP_WRAPPED_MODULE, None)
111 if isinstance(wrapped_module, FuncModule):
112 return getattr(wrapped_module, "_fn", None)
113 if isinstance(module, FuncModule):
114 return getattr(module, "_fn", None)
115 return None
118def _raise_callable_already_wrapped(callable_obj: Callable) -> None:
119 warnings.warn(
120 f"Callable '{callable_obj.__class__.__name__}' is already wrapped. "
121 "Wrapping overlapping module regions is not allowed."
122 )
125def _check_callable_attr_not_wrapped(owner: nn.Module, attr_name: str, attr_value: Callable) -> None:
126 del owner, attr_name
127 if getattr(attr_value, '_is_wrapped', False):
128 _raise_callable_already_wrapped(attr_value)
131def _check_and_mark_callable(callable_obj: Callable) -> None:
132 if _is_callable_exempt_from_overlap_check(callable_obj):
133 return
134 if getattr(callable_obj, '_is_wrapped', False):
135 warnings.warn(
136 f"Callable '{callable_obj.__class__.__name__}' or one of its ancestors is already wrapped. "
137 "Wrapping overlapping module regions is not allowed."
138 )
139 _mark_wrapped(callable_obj)
142def _check_and_mark_wrapped(module: nn.Module) -> None:
143 """Validate no wrapping overlap, then mark module and all descendants as wrapped."""
144 if getattr(module, '_is_wrapped', False):
145 warnings.warn(
146 f"Module '{module.__class__.__name__}' or one of its ancestors is already wrapped. "
147 "Wrapping overlapping module regions is not allowed."
148 )
149 for submodule in module.modules():
150 if submodule is module:
151 continue
152 wrapped_callable = _get_wrapped_callable(submodule)
153 if wrapped_callable is not None and _is_callable_exempt_from_overlap_check(wrapped_callable):
154 continue
155 if getattr(submodule, '_is_wrapped', False):
156 if wrapped_callable is not None:
157 _raise_callable_already_wrapped(wrapped_callable)
158 # A param-free module shared across siblings (one rotary embedding
159 # per decoder layer) is never its own region; skip, do not flag.
160 if next(submodule.parameters(recurse=True), None) is None:
161 continue
162 warnings.warn(
163 f"Submodule '{getattr(submodule, '_swap_wrapped_module', submodule).__class__.__name__}' of "
164 f"'{module.__class__.__name__}' is already wrapped. "
165 "Wrapping overlapping module regions is not allowed."
166 )
167 for submodule in module.modules():
168 for attr_name, attr_value in _iter_wrappable_callable_attrs(submodule):
169 _check_callable_attr_not_wrapped(submodule, attr_name, attr_value)
170 for submodule in module.modules():
171 _mark_wrapped(submodule)
172 for _, attr_value in _iter_wrappable_callable_attrs(submodule):
173 _mark_wrapped(attr_value)
176def base_check_fn(tensor) -> bool:
177 """
178 Basic check to determine if a tensor is eligible for offloading.
179 - Skip Parameters and their views.
180 - Skip empty storage tensors.
181 """
182 if isinstance(tensor._base, torch.nn.parameter.Parameter) or isinstance(tensor, torch.nn.parameter.Parameter): # pylint: disable=W0212
183 return False
184 if tensor.untyped_storage().size() == 0:
185 return False
186 return True
189class AsyncSaveOnCpu(torch.autograd.graph.saved_tensors_hooks):
190 """
191 Context manager to offload tensors to CPU during forward pass.
192 """
193 def __init__(self, policy_fn=None, group_swap: bool = False, cpu_pool=None) -> None:
194 _raise_if_compiling("AsyncSaveOnCpu")
195 self.add_to_storage = False
196 self.storage = Storage()
197 self.count_idx = 0
198 self.policy_fn = policy_fn
199 self.cpu_pool = cpu_pool
201 # Cache per-context-manager state once to avoid per-tensor singleton lookups.
202 swap_manager = SwapManager()
204 def pack_to_cpu(tensor: torch.Tensor):
205 if not base_check_fn(tensor):
206 return tensor.detach()
207 if policy_fn is not None:
208 if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE:
209 return tensor.detach()
210 if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP:
211 raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}")
212 group_name = swap_manager.get_current_group_name()
213 if not group_name:
214 swap_manager.warn_missing_group_once()
215 return tensor.detach()
216 if swap_manager.is_last_group(group_name):
217 return tensor.detach()
218 if not self.add_to_storage:
219 swap_manager.add_storage(group_name, self.storage)
220 self.add_to_storage = True
221 funcname = f"{group_name}::{tensor.shape}"
222 detached = tensor.detach()
223 self.storage[self.count_idx].append(
224 SwapTensor(detached, funcname, group_swap=group_swap, cpu_pool=cpu_pool)
225 )
226 self.count_idx += 1
227 return detached
229 def unpack_from_cpu(tensor) -> torch.Tensor:
230 if self.storage is not None:
231 self.storage.clear()
232 self.storage = None
233 return tensor
235 super().__init__(pack_to_cpu, unpack_from_cpu)
238class ActivationWrapper(torch.nn.Module, ABC):
239 """
240 Base class for Activation Swap.
242 Not meant to be instantiated directly.
243 """
245 def __init__(self, module: Union[nn.Module, Callable], *, track_overlaps: bool = True):
246 """Initialize a wrapper and optionally participate in overlap tracking."""
247 if callable(module) and not isinstance(module, nn.Module):
248 if track_overlaps:
249 _check_and_mark_callable(module)
250 module = FuncModule(module)
251 if track_overlaps:
252 _mark_wrapped(module)
253 elif track_overlaps:
254 _check_and_mark_wrapped(module)
255 super().__init__()
256 self._swap_wrapped_module = module
257 self._is_wrapped = track_overlaps
258 # state_dict post hook to remove prefix to allow loading into a
259 # non-swap wrapped module.
260 self._register_state_dict_hook(self._post_state_dict_hook)
261 # load_state_dict pre-hook to allow loading back into
262 # swap-wrapped module.
263 self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook)
265 @property
266 def _wrapped_module(self):
267 return self._swap_wrapped_module
269 @abstractmethod
270 def forward(self, *args, **kwargs):
271 """Run the wrapped module's forward pass with activation swapping. Must be implemented by subclasses."""
272 raise ValueError("Subclasses should implement forward().")
274 def __getattr__(self, name: str) -> Any:
275 """Forward missing attributes to wrapped module."""
276 try:
277 return super().__getattr__(name) # defer to nn.Module's logic
278 except AttributeError:
279 return getattr(self._swap_wrapped_module, name)
281 def __getitem__(self, key: int) -> Any:
282 """Forward indexing calls in case the module is a nn.Sequential."""
283 return self._swap_wrapped_module.__getitem__(key) # type: ignore[operator]
285 def named_modules(
286 self,
287 memo: Optional[set[nn.Module]] = None,
288 prefix: str = "",
289 remove_duplicate: bool = True,
290 ) -> Iterator[tuple[str, nn.Module]]:
291 """
292 Yield wrapped-module children without exposing the internal wrapper prefix.
294 PyTorch parent modules implement ``named_parameters(recurse=True)`` by
295 iterating ``named_modules()`` and reading each module's direct
296 ``_parameters``. They do not call child modules' ``named_parameters()``
297 overrides. Exposing the wrapped module under the wrapper's own prefix
298 keeps root-module traversals aligned with ``state_dict()`` keys.
300 Args:
301 memo (Optional[set[nn.Module]], optional): A memo set to avoid infinite recursion. Default: ``None``.
302 prefix (str, optional): A prefix to prepend to all module names. Default: ``""``.
303 remove_duplicate (bool, optional): Whether to remove duplicate modules. Default: ``True``.
305 Returns:
306 Iterator[tuple[str, nn.Module]] An iterator of (name, module) pairs.
307 """
308 if memo is None:
309 memo = set()
310 if self not in memo:
311 memo.add(self)
312 yield prefix, self
313 yield from self._swap_wrapped_module.named_modules(
314 memo=memo,
315 prefix=prefix,
316 remove_duplicate=remove_duplicate,
317 )
319 def named_parameters(
320 self,
321 *args,
322 **kwargs,
323 ) -> Iterator[tuple[str, torch.nn.Parameter]]:
324 """
325 Override :meth:`named_parameters()` to intercept parameter names.
327 remove all occurrences of ``_SWAP_PREFIX``.
328 """
329 for param_name, param in super().named_parameters(*args, **kwargs):
330 yield param_name.replace(_SWAP_PREFIX, ""), param
332 def _fqn_modifiers(self) -> dict[str, str]:
333 """Describe the wrapper's state-dict FQN rewrite to PyTorch DCP.
335 Distributed state-dict traversal does not use ``named_modules()`` and
336 therefore cannot infer that this wrapper removes its internal module
337 prefix. Map every direct child state name to the wrapped module so
338 optimizer save and restore use the same canonical FQNs.
339 """
340 wrapped_module = self._swap_wrapped_module
341 exposed_names = set(wrapped_module._modules) # pylint: disable=protected-access
342 exposed_names.update(wrapped_module._parameters) # pylint: disable=protected-access
343 exposed_names.update(wrapped_module._buffers) # pylint: disable=protected-access
344 if (
345 getattr(wrapped_module.__class__, "get_extra_state", nn.Module.get_extra_state)
346 != nn.Module.get_extra_state
347 ):
348 exposed_names.add(nn.modules.module._EXTRA_STATE_KEY_SUFFIX)
349 return {name: _SWAP_WRAPPED_MODULE for name in exposed_names}
351 @staticmethod
352 def _post_state_dict_hook(
353 module: nn.Module, # pylint: disable=W0613
354 state_dict: dict[str, Any],
355 prefix: str,
356 *args: Any, # pylint: disable=W0613
357 ) -> dict[str, Any]:
358 """
359 _post_state_dict_hook() is called after the state_dict() of this FSDP module is executed.
361 For ``swap_wrapper``, it will strip swap-wrapped module prefix,
362 so that this module can be loaded into non-swapped modules.
363 It would still be able to be loaded into swap-wrapped modules as this class,
364 adds the prefix back before loading the state_dict.
365 """
366 _replace_by_prefix(state_dict, f"{prefix}{_SWAP_PREFIX}", prefix)
367 return state_dict
369 @staticmethod
370 def _pre_load_state_dict_hook(
371 module: nn.Module,
372 state_dict: dict[str, Any],
373 prefix: str,
374 *args: Any,
375 ) -> None:
376 """
377 ``_pre_state_dict_hook` is called before ``self._load_from_state_dict()`` is called.
379 For ``swap_wrapper``, it will add back the module
380 prefix so that non-swapped modules can be loaded into
381 swap_wrapper modules properly.
382 """
383 _replace_by_prefix(state_dict, prefix, prefix + f"{_SWAP_PREFIX}")
386class SwapWrapper(ActivationWrapper):
387 """
388 Customize an nn.Module wrapper class to add an AsyncSaveOnCpu context manager for the target model.
389 """
390 def __init__(
391 self,
392 mod: Union[nn.Module, Callable],
393 policy_fn: Optional[Callable] = None,
394 group_swap: bool = False,
395 cpu_pool=None,
396 ):
397 super().__init__(mod)
398 self.policy_fn = policy_fn
399 self.group_swap = group_swap
400 self.cpu_pool = cpu_pool
402 def forward(self, *args, **kwargs):
403 """Run the wrapped module inside an AsyncSaveOnCpu context for activation swapping."""
404 _raise_if_compiling("swap_wrapper")
405 async_kwargs = {
406 "policy_fn": self.policy_fn,
407 "group_swap": self.group_swap,
408 "cpu_pool": self.cpu_pool,
409 }
410 with AsyncSaveOnCpu(**async_kwargs):
411 return self._swap_wrapped_module(*args, **kwargs)
414def swap_wrapper(
415 module: Union[nn.Module, Callable],
416 policy_fn: Optional[Callable] = None,
417 group_swap: bool = False,
418 cpu_pool=None,
419) -> SwapWrapper:
420 """Wrap a module or callable with activation swap functionality.
422 Args:
423 module: Module or callable to wrap.
424 policy_fn: Optional per-tensor swap policy.
425 group_swap: Whether tensors participate in group copy fusion.
426 cpu_pool: Optional pinned host memory pool for swapped tensors.
428 Returns:
429 Configured activation swap wrapper.
430 """
431 return SwapWrapper(module, policy_fn, group_swap, cpu_pool)
434def swap_tensor_wrapper(target, tag: Optional[str] = None, group_swap: bool = False, cpu_pool=None):
435 """Register selected tensors into the current swap group.
437 This helper is intended to be used inside a forward path that already
438 participates in the existing swap scheduling managed by ``SwapManager``.
439 It preserves the input structure and returns the original tensors.
441 Args:
442 target: Tensor or nested tensor structure to register.
443 tag: Optional debug tag.
444 group_swap: Whether tensors participate in group copy fusion.
445 cpu_pool: Optional pinned host memory pool for swapped tensors.
447 Returns:
448 The registered tensor structure.
449 """
450 _raise_if_compiling("swap_tensor_wrapper")
451 swap_manager = SwapManager()
452 group_name = swap_manager.get_current_group_name()
453 if not group_name:
454 warnings.warn(
455 f"Tensor {tag} cannot be swapped, for its group is unregistered."
456 )
457 return target
458 if swap_manager.is_last_group(group_name):
459 return target
461 storage = Storage()
462 count_idx = 0
464 def _register_tensor(tensor):
465 nonlocal count_idx
466 if not base_check_fn(tensor):
467 return tensor
469 tensor_tag = tag or f"{group_name}_swap_tensor"
470 funcname = f"{tensor_tag}::{tuple(tensor.shape)}"
471 storage[count_idx].append(
472 SwapTensor(tensor, funcname, group_swap=group_swap, cpu_pool=cpu_pool)
473 )
474 count_idx += 1
475 return tensor
477 wrapped = torch.utils._pytree.tree_map( # pylint: disable=protected-access
478 lambda x: _register_tensor(x) if isinstance(x, torch.Tensor) else x,
479 target,
480 )
481 if count_idx > 0:
482 swap_manager.add_storage(group_name, storage)
483 return wrapped
486class CheckpointWrapper(ActivationWrapper):
487 """Wrap a module with activation checkpointing."""
489 def __init__(
490 self,
491 mod: Union[nn.Module, Callable],
492 **checkpoint_kwargs: Any,
493 ):
494 super().__init__(mod)
495 self.checkpoint_kwargs = checkpoint_kwargs
497 def _do_checkpoint(self, wrapped_module: Any, *args: Any, **kwargs: Any) -> Any:
498 # Checkpoint may save inputs before the wrapped module's pre-hook runs.
499 group_name = getattr(self, "_swap_group_name", None)
500 if group_name is not None:
501 SwapManager().set_current_group_name(group_name)
503 from hyper_parallel.core.activation_memory.api import checkpoint # pylint: disable=C0415
504 return checkpoint(
505 wrapped_module,
506 *args,
507 **self.checkpoint_kwargs,
508 **kwargs,
509 )
511 def forward(self, *args: Any, **kwargs: Any) -> Any:
512 return self._do_checkpoint(self._wrapped_module, *args, **kwargs)
515def ckpt_wrapper(module: Union[nn.Module, Callable], **checkpoint_kwargs: Any) -> CheckpointWrapper:
516 """Wrap *module* with activation checkpointing."""
517 return CheckpointWrapper(module, **checkpoint_kwargs)
520class CheckpointExcludeWrapper(ActivationWrapper):
521 """Exclude a callable region from checkpoint recomputation."""
523 def __init__(self, module: Callable[..., Any], *, save_output: bool = True) -> None:
524 """Initialize a checkpoint exclusion wrapper for a PyTorch module or function."""
525 if not callable(module):
526 raise ValueError("module must be a PyTorch Module or callable")
527 if not isinstance(save_output, bool):
528 raise ValueError(f"save_output must be a bool, got {type(save_output).__name__}")
529 super().__init__(module, track_overlaps=False)
530 self.save_output = save_output
532 def forward(self, *args: Any, **kwargs: Any) -> Any:
533 """Execute normally outside recompute and return the cached output in recompute."""
534 wrapped_module = cast(Callable[..., Any], self._wrapped_module)
535 state = get_recompute_state()
536 if state is None:
537 return wrapped_module(*args, **kwargs)
538 cache = state.get_resource(_checkpoint_exclude._EXCLUDE_CACHE_KEY, _checkpoint_exclude._ExcludeCache)
539 if state.is_recomputing:
540 entry = cache.pop(id(self))
541 _checkpoint_exclude._materialize_recompute_inputs(entry, args, kwargs)
542 output = (
543 entry.output
544 if self.save_output
545 else _checkpoint_exclude._make_replay_placeholder_output(entry.output_tensor_count)
546 )
547 return _checkpoint_exclude._finalize_save_outputs(
548 output,
549 _checkpoint_exclude._has_used_input(entry.input_bindings),
550 None,
551 )
553 input_bindings, previous_handles = _checkpoint_exclude._mark_recompute_inputs(
554 state.invocation_id,
555 args,
556 kwargs,
557 )
558 try:
559 with _checkpoint_exclude._saved_tensors_context():
560 output = wrapped_module(*args, **kwargs)
561 finally:
562 _checkpoint_exclude._restore_recompute_inputs(previous_handles)
563 needs_recompute_boundary = _checkpoint_exclude._has_used_input(input_bindings)
564 tensor_leaf_count = None if self.save_output else [0]
565 finalized_output = _checkpoint_exclude._finalize_save_outputs(
566 output,
567 needs_recompute_boundary,
568 state.invocation_id,
569 tensor_leaf_count,
570 )
571 replay_output = output if self.save_output else None
572 output_tensor_count = 1 if tensor_leaf_count is None else tensor_leaf_count[0]
573 cache.save(
574 id(self),
575 _checkpoint_exclude._ExcludeCacheEntry(replay_output, input_bindings, output_tensor_count),
576 )
577 return finalized_output
580def checkpoint_exclude_wrapper(
581 module: Callable[..., Any],
582 *,
583 save_output: bool = True,
584) -> CheckpointExcludeWrapper:
585 """Wrap a PyTorch module or function so its region is not recomputed.
587 Args:
588 module: PyTorch Module or callable to execute only during the original
589 checkpoint forward pass.
590 save_output: Whether to retain the region output for checkpoint replay.
591 Set this to ``False`` only when the output is passed directly as one
592 argument to another checkpoint exclusion wrapper. Default: ``True``.
594 Returns:
595 A wrapper that saves the callable's autograd tensors and reuses its
596 forward output while replaying a non-reentrant checkpoint.
598 Note:
599 This feature requires eager mode and a surrounding HyperParallel
600 checkpoint configured with ``use_reentrant=False``. Nested checkpoint
601 exclusion wrappers are not supported.
602 """
603 return CheckpointExcludeWrapper(module, save_output=save_output)