Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / wrapper.py: 95%
310 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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.pinned_memory_pool import PinnedMemoryPool
32from hyper_parallel.core.activation_memory.recompute_state import get_recompute_state
33from hyper_parallel.core.activation_memory.swap import SwapManager, SwapTensor, Storage
34from . import checkpoint_exclude as _checkpoint_exclude
35from .reentrant_exclude import current_frame, run_exclude
38_SWAP_WRAPPED_MODULE = "_swap_wrapped_module"
39_SWAP_PREFIX = _SWAP_WRAPPED_MODULE + "."
42def _raise_if_compiling(feature_name: str) -> None:
43 """Reject HyperParallel activation-swap extensions during Torch capture."""
44 if torch.compiler.is_compiling():
45 raise ValueError(
46 f"HyperParallel {feature_name} is not supported in compile mode. "
47 "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
48 )
51class FuncModule(nn.Module):
52 """
53 Thin :class:`~torch.nn.Module` adapter that wraps a plain callable.
55 Allows ordinary Python functions (or any callable without Module
56 parameters) to be passed to :func:`swap_wrapper` and
57 :func:`~hyper_parallel.core.activation_memory.checkpoint_wrapper`
58 in place of an :class:`~torch.nn.Module`.
59 The wrapped function is stored as ``_fn`` and invoked in
60 :meth:`forward`; the module has no trainable parameters.
62 Args:
63 fn (callable): The function to wrap.
65 Example:
66 >>> wrapped = swap_wrapper(lambda x: x * 2)
67 """
69 def __init__(self, fn: Callable):
70 super().__init__()
71 self._fn = fn
73 def forward(self, *args, **kwargs):
74 """Invoke the wrapped callable with the given arguments."""
75 return self._fn(*args, **kwargs)
78def _is_callable_exempt_from_overlap_check(callable_obj: Callable) -> bool:
79 """Return True for callables that cannot be reliably overlap-tracked by object marks."""
80 return isinstance(callable_obj, (types.FunctionType, types.BuiltinFunctionType, types.MethodType))
83def _iter_wrappable_callable_attrs(module: nn.Module) -> Iterator[tuple[str, Callable]]:
84 """Yield public per-instance callable attributes not registered as child modules.
86 Plain functions, builtins and bound methods are skipped: these are stateless
87 module-level utilities shared by reference across many modules (e.g.
88 ``self.act = F.gelu`` repeated in every layer). They are never standalone
89 checkpoint regions, and marking a shared function's ``_is_wrapped`` flag
90 would both mutate a global object and falsely flag every sibling module that
91 references the same function as an overlapping wrap. Only per-instance
92 callables participate in overlap tracking.
93 """
94 for attr_name, attr_value in vars(module).items():
95 if attr_name.startswith("_") or isinstance(attr_value, nn.Module):
96 continue
97 if _is_callable_exempt_from_overlap_check(attr_value):
98 continue
99 if callable(attr_value):
100 yield attr_name, attr_value
103def _mark_wrapped(obj: Any) -> None:
104 """Tag ``obj`` as already wrapped, ignoring objects that reject new attributes."""
105 try:
106 setattr(obj, "_is_wrapped", True)
107 except (AttributeError, TypeError):
108 pass
111def _get_wrapped_callable(module: nn.Module) -> Optional[Callable]:
112 wrapped_module = getattr(module, _SWAP_WRAPPED_MODULE, None)
113 if isinstance(wrapped_module, FuncModule):
114 return getattr(wrapped_module, "_fn", None)
115 if isinstance(module, FuncModule):
116 return getattr(module, "_fn", None)
117 return None
120def _raise_callable_already_wrapped(callable_obj: Callable) -> None:
121 warnings.warn(
122 f"Callable '{callable_obj.__class__.__name__}' is already wrapped. "
123 "Wrapping overlapping module regions is not allowed."
124 )
127def _check_callable_attr_not_wrapped(owner: nn.Module, attr_name: str, attr_value: Callable) -> None:
128 del owner, attr_name
129 if getattr(attr_value, '_is_wrapped', False):
130 _raise_callable_already_wrapped(attr_value)
133def _check_and_mark_callable(callable_obj: Callable) -> None:
134 if _is_callable_exempt_from_overlap_check(callable_obj):
135 return
136 if getattr(callable_obj, '_is_wrapped', False):
137 warnings.warn(
138 f"Callable '{callable_obj.__class__.__name__}' or one of its ancestors is already wrapped. "
139 "Wrapping overlapping module regions is not allowed."
140 )
141 _mark_wrapped(callable_obj)
144def _check_and_mark_wrapped(module: nn.Module) -> None:
145 """Validate no wrapping overlap, then mark module and all descendants as wrapped."""
146 if getattr(module, '_is_wrapped', False):
147 warnings.warn(
148 f"Module '{module.__class__.__name__}' or one of its ancestors is already wrapped. "
149 "Wrapping overlapping module regions is not allowed."
150 )
151 for submodule in module.modules():
152 if submodule is module:
153 continue
154 wrapped_callable = _get_wrapped_callable(submodule)
155 if wrapped_callable is not None and _is_callable_exempt_from_overlap_check(wrapped_callable):
156 continue
157 if getattr(submodule, '_is_wrapped', False):
158 if wrapped_callable is not None:
159 _raise_callable_already_wrapped(wrapped_callable)
160 # A param-free module shared across siblings (one rotary embedding
161 # per decoder layer) is never its own region; skip, do not flag.
162 if next(submodule.parameters(recurse=True), None) is None:
163 continue
164 warnings.warn(
165 f"Submodule '{getattr(submodule, '_swap_wrapped_module', submodule).__class__.__name__}' of "
166 f"'{module.__class__.__name__}' is already wrapped. "
167 "Wrapping overlapping module regions is not allowed."
168 )
169 for submodule in module.modules():
170 for attr_name, attr_value in _iter_wrappable_callable_attrs(submodule):
171 _check_callable_attr_not_wrapped(submodule, attr_name, attr_value)
172 for submodule in module.modules():
173 _mark_wrapped(submodule)
174 for _, attr_value in _iter_wrappable_callable_attrs(submodule):
175 _mark_wrapped(attr_value)
178def base_check_fn(tensor) -> bool:
179 """
180 Basic check to determine if a tensor is eligible for offloading.
181 - Skip Parameters and their views.
182 - Skip empty storage tensors.
183 """
184 if isinstance(tensor._base, torch.nn.parameter.Parameter) or isinstance(tensor, torch.nn.parameter.Parameter): # pylint: disable=W0212
185 return False
186 if tensor.untyped_storage().size() == 0:
187 return False
188 return True
191class AsyncSaveOnCpu(torch.autograd.graph.saved_tensors_hooks):
192 """
193 Context manager to offload tensors to CPU during forward pass.
194 """
195 def __init__(self, policy_fn=None, group_swap: bool = False, cpu_pool=None) -> None:
196 _raise_if_compiling("AsyncSaveOnCpu")
197 self.add_to_storage = False
198 self.storage = Storage()
199 self.count_idx = 0
200 self.policy_fn = policy_fn
201 self.cpu_pool = cpu_pool
203 # Cache per-context-manager state once to avoid per-tensor singleton lookups.
204 swap_manager = SwapManager()
206 def pack_to_cpu(tensor: torch.Tensor):
207 if not base_check_fn(tensor):
208 return tensor.detach()
209 if policy_fn is not None:
210 if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE:
211 return tensor.detach()
212 if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP:
213 raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}")
214 group_name = swap_manager.get_current_group_name()
215 if not group_name:
216 swap_manager.warn_missing_group_once()
217 return tensor.detach()
218 if swap_manager.is_last_group(group_name):
219 return tensor.detach()
220 if not self.add_to_storage:
221 swap_manager.add_storage(group_name, self.storage)
222 self.add_to_storage = True
223 funcname = f"{group_name}::{tensor.shape}"
224 detached = tensor.detach()
225 self.storage[self.count_idx].append(
226 SwapTensor(detached, funcname, group_swap=group_swap, cpu_pool=cpu_pool)
227 )
228 self.count_idx += 1
229 return detached
231 def unpack_from_cpu(tensor) -> torch.Tensor:
232 if self.storage is not None:
233 self.storage.clear()
234 self.storage = None
235 return tensor
237 super().__init__(pack_to_cpu, unpack_from_cpu)
240def _register_checkpoint_exclude_outputs(
241 output: Any,
242 *,
243 group_swap: bool,
244 cpu_pool: Optional[PinnedMemoryPool],
245) -> Any:
246 """Register eligible output leaves without changing their identity or structure."""
247 swap_manager = SwapManager()
248 group_name = swap_manager.get_current_group_name()
249 if not group_name:
250 swap_manager.warn_missing_group_once()
251 return output
252 if swap_manager.is_last_group(group_name):
253 return output
255 storage = Storage()
256 swappable_outputs = []
257 for tensor in _checkpoint_exclude._iter_tensor_leaves(output):
258 if not base_check_fn(tensor):
259 continue
260 funcname = f"{group_name}::{tensor.shape}"
261 storage[len(swappable_outputs)].append(
262 SwapTensor(tensor, funcname, group_swap=group_swap, cpu_pool=cpu_pool)
263 )
264 swappable_outputs.append(tensor)
266 if swappable_outputs:
267 swap_manager.add_storage(group_name, storage)
268 return output
271class ActivationWrapper(torch.nn.Module, ABC):
272 """
273 Base class for Activation Swap.
275 Not meant to be instantiated directly.
276 """
278 def __init__(self, module: Union[nn.Module, Callable], *, track_overlaps: bool = True):
279 """Initialize a wrapper and optionally participate in overlap tracking."""
280 if callable(module) and not isinstance(module, nn.Module):
281 if track_overlaps:
282 _check_and_mark_callable(module)
283 module = FuncModule(module)
284 if track_overlaps:
285 _mark_wrapped(module)
286 elif track_overlaps:
287 _check_and_mark_wrapped(module)
288 super().__init__()
289 self._swap_wrapped_module = module
290 self._is_wrapped = track_overlaps
291 # state_dict post hook to remove prefix to allow loading into a
292 # non-swap wrapped module.
293 self._register_state_dict_hook(self._post_state_dict_hook)
294 # load_state_dict pre-hook to allow loading back into
295 # swap-wrapped module.
296 self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook)
298 @property
299 def _wrapped_module(self):
300 return self._swap_wrapped_module
302 @abstractmethod
303 def forward(self, *args, **kwargs):
304 """Run the wrapped module's forward pass with activation swapping. Must be implemented by subclasses."""
305 raise ValueError("Subclasses should implement forward().")
307 def __getattr__(self, name: str) -> Any:
308 """Forward missing attributes to wrapped module."""
309 try:
310 return super().__getattr__(name) # defer to nn.Module's logic
311 except AttributeError:
312 return getattr(self._swap_wrapped_module, name)
314 def __getitem__(self, key: int) -> Any:
315 """Forward indexing calls in case the module is a nn.Sequential."""
316 return self._swap_wrapped_module.__getitem__(key) # type: ignore[operator]
318 def named_modules(
319 self,
320 memo: Optional[set[nn.Module]] = None,
321 prefix: str = "",
322 remove_duplicate: bool = True,
323 ) -> Iterator[tuple[str, nn.Module]]:
324 """
325 Yield wrapped-module children without exposing the internal wrapper prefix.
327 PyTorch parent modules implement ``named_parameters(recurse=True)`` by
328 iterating ``named_modules()`` and reading each module's direct
329 ``_parameters``. They do not call child modules' ``named_parameters()``
330 overrides. Exposing the wrapped module under the wrapper's own prefix
331 keeps root-module traversals aligned with ``state_dict()`` keys.
333 Args:
334 memo (Optional[set[nn.Module]], optional): A memo set to avoid infinite recursion. Default: ``None``.
335 prefix (str, optional): A prefix to prepend to all module names. Default: ``""``.
336 remove_duplicate (bool, optional): Whether to remove duplicate modules. Default: ``True``.
338 Returns:
339 Iterator[tuple[str, nn.Module]] An iterator of (name, module) pairs.
340 """
341 if memo is None:
342 memo = set()
343 if self not in memo:
344 memo.add(self)
345 yield prefix, self
346 yield from self._swap_wrapped_module.named_modules(
347 memo=memo,
348 prefix=prefix,
349 remove_duplicate=remove_duplicate,
350 )
352 def named_parameters(
353 self,
354 *args,
355 **kwargs,
356 ) -> Iterator[tuple[str, torch.nn.Parameter]]:
357 """
358 Override :meth:`named_parameters()` to intercept parameter names.
360 remove all occurrences of ``_SWAP_PREFIX``.
361 """
362 for param_name, param in super().named_parameters(*args, **kwargs):
363 yield param_name.replace(_SWAP_PREFIX, ""), param
365 def _fqn_modifiers(self) -> dict[str, str]:
366 """Describe the wrapper's state-dict FQN rewrite to PyTorch DCP.
368 Distributed state-dict traversal does not use ``named_modules()`` and
369 therefore cannot infer that this wrapper removes its internal module
370 prefix. Map every direct child state name to the wrapped module so
371 optimizer save and restore use the same canonical FQNs.
372 """
373 wrapped_module = self._swap_wrapped_module
374 exposed_names = set(wrapped_module._modules) # pylint: disable=protected-access
375 exposed_names.update(wrapped_module._parameters) # pylint: disable=protected-access
376 exposed_names.update(wrapped_module._buffers) # pylint: disable=protected-access
377 if (
378 getattr(wrapped_module.__class__, "get_extra_state", nn.Module.get_extra_state)
379 != nn.Module.get_extra_state
380 ):
381 exposed_names.add(nn.modules.module._EXTRA_STATE_KEY_SUFFIX)
382 return {name: _SWAP_WRAPPED_MODULE for name in exposed_names}
384 @staticmethod
385 def _post_state_dict_hook(
386 module: nn.Module, # pylint: disable=W0613
387 state_dict: dict[str, Any],
388 prefix: str,
389 *args: Any, # pylint: disable=W0613
390 ) -> dict[str, Any]:
391 """
392 _post_state_dict_hook() is called after the state_dict() of this FSDP module is executed.
394 For ``swap_wrapper``, it will strip swap-wrapped module prefix,
395 so that this module can be loaded into non-swapped modules.
396 It would still be able to be loaded into swap-wrapped modules as this class,
397 adds the prefix back before loading the state_dict.
398 """
399 _replace_by_prefix(state_dict, f"{prefix}{_SWAP_PREFIX}", prefix)
400 return state_dict
402 @staticmethod
403 def _pre_load_state_dict_hook(
404 module: nn.Module,
405 state_dict: dict[str, Any],
406 prefix: str,
407 *args: Any,
408 ) -> None:
409 """
410 ``_pre_state_dict_hook` is called before ``self._load_from_state_dict()`` is called.
412 For ``swap_wrapper``, it will add back the module
413 prefix so that non-swapped modules can be loaded into
414 swap_wrapper modules properly.
415 """
416 _replace_by_prefix(state_dict, prefix, prefix + f"{_SWAP_PREFIX}")
419class SwapWrapper(ActivationWrapper):
420 """
421 Customize an nn.Module wrapper class to add an AsyncSaveOnCpu context manager for the target model.
422 """
423 def __init__(
424 self,
425 mod: Union[nn.Module, Callable],
426 policy_fn: Optional[Callable] = None,
427 group_swap: bool = False,
428 cpu_pool=None,
429 ):
430 super().__init__(mod)
431 self.policy_fn = policy_fn
432 self.group_swap = group_swap
433 self.cpu_pool = cpu_pool
435 def forward(self, *args, **kwargs):
436 """Run the wrapped module inside an AsyncSaveOnCpu context for activation swapping."""
437 _raise_if_compiling("swap_wrapper")
438 async_kwargs = {
439 "policy_fn": self.policy_fn,
440 "group_swap": self.group_swap,
441 "cpu_pool": self.cpu_pool,
442 }
443 with AsyncSaveOnCpu(**async_kwargs):
444 return self._swap_wrapped_module(*args, **kwargs)
447def swap_wrapper(
448 module: Union[nn.Module, Callable],
449 policy_fn: Optional[Callable] = None,
450 group_swap: bool = False,
451 cpu_pool=None,
452) -> SwapWrapper:
453 """Wrap a module or callable with activation swap functionality.
455 Args:
456 module: Module or callable to wrap.
457 policy_fn: Optional per-tensor swap policy.
458 group_swap: Whether tensors participate in group copy fusion.
459 cpu_pool: Optional pinned host memory pool for swapped tensors.
461 Returns:
462 Configured activation swap wrapper.
463 """
464 return SwapWrapper(module, policy_fn, group_swap, cpu_pool)
467def swap_tensor_wrapper(target, tag: Optional[str] = None, group_swap: bool = False, cpu_pool=None):
468 """Register selected tensors into the current swap group.
470 This helper is intended to be used inside a forward path that already
471 participates in the existing swap scheduling managed by ``SwapManager``.
472 It preserves the input structure and returns the original tensors.
474 Args:
475 target: Tensor or nested tensor structure to register.
476 tag: Optional debug tag.
477 group_swap: Whether tensors participate in group copy fusion.
478 cpu_pool: Optional pinned host memory pool for swapped tensors.
480 Returns:
481 The registered tensor structure.
482 """
483 _raise_if_compiling("swap_tensor_wrapper")
484 swap_manager = SwapManager()
485 group_name = swap_manager.get_current_group_name()
486 if not group_name:
487 warnings.warn(
488 f"Tensor {tag} cannot be swapped, for its group is unregistered."
489 )
490 return target
491 if swap_manager.is_last_group(group_name):
492 return target
494 storage = Storage()
495 count_idx = 0
497 def _register_tensor(tensor):
498 nonlocal count_idx
499 if not base_check_fn(tensor):
500 return tensor
502 tensor_tag = tag or f"{group_name}_swap_tensor"
503 funcname = f"{tensor_tag}::{tuple(tensor.shape)}"
504 storage[count_idx].append(
505 SwapTensor(tensor, funcname, group_swap=group_swap, cpu_pool=cpu_pool)
506 )
507 count_idx += 1
508 return tensor
510 wrapped = torch.utils._pytree.tree_map( # pylint: disable=protected-access
511 lambda x: _register_tensor(x) if isinstance(x, torch.Tensor) else x,
512 target,
513 )
514 if count_idx > 0:
515 swap_manager.add_storage(group_name, storage)
516 return wrapped
519class CheckpointWrapper(ActivationWrapper):
520 """Wrap a module with activation checkpointing."""
522 def __init__(
523 self,
524 mod: Union[nn.Module, Callable],
525 **checkpoint_kwargs: Any,
526 ):
527 super().__init__(mod)
528 self.checkpoint_kwargs = checkpoint_kwargs
530 def _do_checkpoint(self, wrapped_module: Any, *args: Any, **kwargs: Any) -> Any:
531 """Run the wrapped module through the configured checkpoint function."""
532 # Checkpoint may save inputs before the wrapped module's pre-hook runs.
533 group_name = getattr(self, "_swap_group_name", None)
534 if group_name is not None:
535 SwapManager().set_current_group_name(group_name)
536 # Reentrant backward can bypass this wrapper's full backward hook.
537 # This method is called for a new forward, never for its replay.
538 if getattr(self, "_swap_state", None) == "pre_backward":
539 self._swap_state = "backward"
541 from hyper_parallel.core.activation_memory.api import checkpoint # pylint: disable=C0415
542 return checkpoint(
543 wrapped_module,
544 *args,
545 **self.checkpoint_kwargs,
546 **kwargs,
547 )
549 def forward(self, *args: Any, **kwargs: Any) -> Any:
550 return self._do_checkpoint(self._wrapped_module, *args, **kwargs)
553def ckpt_wrapper(module: Union[nn.Module, Callable], **checkpoint_kwargs: Any) -> CheckpointWrapper:
554 """Wrap *module* with activation checkpointing."""
555 return CheckpointWrapper(module, **checkpoint_kwargs)
558class CheckpointExcludeWrapper(ActivationWrapper):
559 """Exclude a callable region from checkpoint recomputation."""
561 def __init__(
562 self,
563 module: Callable[..., Any],
564 *,
565 swap: bool = False,
566 group_swap: bool = False,
567 cpu_pool: Optional[PinnedMemoryPool] = None,
568 ) -> None:
569 """Initialize a checkpoint exclusion wrapper for a PyTorch module or function."""
570 if not callable(module):
571 raise ValueError("module must be a PyTorch Module or callable")
572 if not isinstance(swap, bool):
573 raise ValueError(f"swap must be a bool, got {type(swap).__name__}")
574 if not isinstance(group_swap, bool):
575 raise ValueError(f"group_swap must be a bool, got {type(group_swap).__name__}")
576 if cpu_pool is not None and not isinstance(cpu_pool, PinnedMemoryPool):
577 raise ValueError(f"cpu_pool must be a PinnedMemoryPool or None, got {type(cpu_pool).__name__}")
578 ignored_options = []
579 if not swap:
580 if group_swap:
581 ignored_options.append("group_swap")
582 if cpu_pool is not None:
583 ignored_options.append("cpu_pool")
584 if ignored_options:
585 warnings.warn(
586 f"checkpoint_exclude_wrapper ignores {', '.join(ignored_options)} when swap=False; "
587 "set swap=True to enable these options.",
588 UserWarning,
589 stacklevel=3,
590 )
591 super().__init__(module, track_overlaps=False)
592 self.swap = swap
593 self.group_swap = group_swap
594 self.cpu_pool = cpu_pool
596 def _make_exclude_swap_hooks(self) -> Optional[AsyncSaveOnCpu]:
597 """Create the fallback hooks shared by both exclusion execution modes."""
598 if not self.swap:
599 return None
600 swap_hooks = AsyncSaveOnCpu(group_swap=self.group_swap, cpu_pool=self.cpu_pool)
601 if not callable(getattr(swap_hooks, "pack_hook", None)) or not callable(
602 getattr(swap_hooks, "unpack_hook", None)
603 ):
604 raise RuntimeError(
605 "There is no callable pack_hook and unpack_hook attributes for checkpoint exclusion swap"
606 )
607 return swap_hooks
609 def _register_exclude_outputs(self, output: Any) -> None:
610 """Make cached outputs available to the active swap group."""
611 _register_checkpoint_exclude_outputs(output, group_swap=self.group_swap, cpu_pool=self.cpu_pool)
613 def forward(self, *args: Any, **kwargs: Any) -> Any:
614 """Execute normally outside recompute and return the cached output in recompute."""
615 wrapped_module = cast(Callable[..., Any], self._wrapped_module)
616 frame = current_frame()
617 if frame is not None:
618 return run_exclude(self, wrapped_module, args, kwargs, frame)
619 state = get_recompute_state()
620 if state is None:
621 return wrapped_module(*args, **kwargs)
622 cache = state.get_resource(_checkpoint_exclude._EXCLUDE_CACHE_KEY, _checkpoint_exclude._ExcludeCache)
623 if state.is_recomputing:
624 entry = cache.pop(id(self))
625 _checkpoint_exclude._restore_rng_state(entry.rng_state)
626 _checkpoint_exclude._materialize_recompute_inputs(entry, args, kwargs)
627 return _checkpoint_exclude._finalize_save_outputs(
628 entry.output,
629 _checkpoint_exclude._has_used_input(entry.input_bindings),
630 )
632 swap_hooks = self._make_exclude_swap_hooks()
633 input_infos = _checkpoint_exclude._capture_recompute_inputs(args, kwargs)
634 input_bindings = []
635 pack_token = _checkpoint_exclude._ACTIVE_PACK_STATE.set(
636 _checkpoint_exclude._PackState(input_infos, input_bindings, swap_hooks=swap_hooks)
637 )
638 try:
639 with _checkpoint_exclude._saved_tensors_context():
640 output = wrapped_module(*args, **kwargs)
641 finally:
642 _checkpoint_exclude._ACTIVE_PACK_STATE.reset(pack_token)
643 if swap_hooks is not None:
644 self._register_exclude_outputs(output)
645 rng_state = _checkpoint_exclude._capture_rng_state(
646 args,
647 kwargs,
648 output,
649 state.preserve_rng_state,
650 )
651 needs_recompute_boundary = _checkpoint_exclude._has_used_input(input_bindings)
652 finalized_output = _checkpoint_exclude._finalize_save_outputs(
653 output,
654 needs_recompute_boundary,
655 )
656 cache.save(
657 id(self),
658 _checkpoint_exclude._ExcludeCacheEntry(
659 output,
660 input_bindings,
661 rng_state,
662 ),
663 )
664 return finalized_output
667def checkpoint_exclude_wrapper(
668 module: Callable[..., Any],
669 *,
670 swap: bool = False,
671 group_swap: bool = False,
672 cpu_pool: Optional[PinnedMemoryPool] = None,
673) -> CheckpointExcludeWrapper:
674 """Wrap a PyTorch module or function so its region is not recomputed.
676 Args:
677 module: PyTorch Module or callable to execute only during the original
678 checkpoint forward pass.
679 swap: Whether eligible outputs and non-recomputable saved tensors should
680 use the active activation-swap group. Default: ``False``.
681 group_swap: Whether swapped tensors participate in group copy fusion.
682 Only used when ``swap=True`` inside an active checkpoint.
683 cpu_pool: Optional pinned-memory pool used for swapped tensors. Only used
684 when ``swap=True`` inside an active checkpoint.
686 Returns:
687 A wrapper that saves the callable's autograd tensors and reuses its
688 forward output while replaying a checkpoint.
690 Note:
691 This feature requires eager mode and a surrounding HyperParallel
692 checkpoint. Both modes cache the excluded outputs; nested checkpoint
693 exclusion wrappers are not supported.
694 """
695 return CheckpointExcludeWrapper(
696 module,
697 swap=swap,
698 group_swap=group_swap,
699 cpu_pool=cpu_pool,
700 )