Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_checkpoint / activation_checkpoint.py: 95%
55 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05: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.
14# ============================================================================
15"""Activation checkpointing related interfaces"""
16import contextlib
17import enum
18from functools import partial
19from typing import Any, Callable, Optional, Tuple
21from hyper_parallel.platform import get_platform
22from .recompute_state import create_recompute_contexts
23plat = get_platform()
26class CheckpointPolicy(enum.Enum):
27 """
28 Enum for specifying the policy for checkpointing during backpropagation.
30 This enum extends PyTorch's selective activation checkpointing policies
31 by introducing a SWAP-based strategy, which allows activation tensors
32 to be offloaded during the forward pass and loaded back before backward
33 computation.
35 For PyTorch native policies (SAVE / RECOMPUTE semantics and MUST vs PREFER),
36 see: https://docs.pytorch.org/docs/2.6/checkpoint.html#torch.utils.checkpoint.CheckpointPolicy
38 Additional policy:
40 - ``MUST_SWAP``: The operation's output is offloaded to host memory during the
41 forward pass and loaded back asynchronously before backward computation. The backward
42 pass reuses the loaded activations without recomputation.
44 This policy must be used together with :class:`SwapManager` to coordinate
45 asynchronous offload/load and stream synchronization.
47 .. note::
48 ``MUST_SWAP`` is typically applied to operations that are either
49 computationally expensive or have large memory footprints. Note that
50 swapping very small outputs may introduce additional overhead and
51 reduce the effectiveness of asynchronous copy.
52 """
53 MUST_SAVE = 0
54 PREFER_SAVE = 1
55 MUST_RECOMPUTE = 2
56 PREFER_RECOMPUTE = 3
58 # Offload during forward, reload before backward. Requires SwapManager.
59 MUST_SWAP = 4
62class _StackedCtx:
63 """Compose multiple context managers as one — enter in order, exit reversed."""
65 def __init__(self, ctxs) -> None:
66 self._ctxs = list(ctxs)
67 self._stack = contextlib.ExitStack()
69 def __enter__(self):
70 self._stack.__enter__()
71 try:
72 for ctx in self._ctxs:
73 self._stack.enter_context(ctx)
74 except BaseException as exc:
75 self._stack.__exit__(type(exc), exc, exc.__traceback__)
76 raise
77 return self
79 def __exit__(self, exc_type, exc_val, exc_tb):
80 return self._stack.__exit__(exc_type, exc_val, exc_tb)
83def _compose_context_fns(
84 factories: Tuple[Callable[[], Tuple[object, object]], ...],
85) -> Callable[[], Tuple[_StackedCtx, _StackedCtx]]:
86 """Combine ``(forward_ctx, recompute_ctx)`` factories into one factory.
88 ``ms.recompute`` / ``torch.utils.checkpoint(use_reentrant=False)`` call
89 ``context_fn()`` once per invocation and unpack the result as
90 ``(forward_ctx, recompute_ctx)``. This helper calls each input factory
91 once, then stacks all forward contexts and all recompute contexts into
92 two :class:`_StackedCtx` instances so the composite respects the
93 single-call contract.
94 """
95 def factory() -> Tuple[_StackedCtx, _StackedCtx]:
96 pairs = [fn() for fn in factories]
97 fwd_ctxs = [pair[0] for pair in pairs]
98 rec_ctxs = [pair[1] for pair in pairs]
99 return _StackedCtx(fwd_ctxs), _StackedCtx(rec_ctxs)
101 return factory
104def checkpoint(
105 function,
106 *args,
107 swap_inputs: bool = False,
108 policy_fn: Optional[Callable] = None,
109 context_fn: Optional[Callable[[], Tuple[object, object]]] = None,
110 group_swap: bool = False,
111 **kwargs,
112):
113 """
114 Apply activation checkpointing to a function with optional input swapping.
116 Args:
117 function: The function to apply checkpointing to.
118 *args: Arguments to pass to the function.
119 swap_inputs (bool): Whether to enable input swapping using async_save_on_cpu context.
120 policy_fn (callable, optional): Function that determines checkpoint policy for operations.
121 context_fn (callable, optional): A no-arg factory returning a
122 ``(forward_ctx, recompute_ctx)`` pair, matching the
123 ``context_fn`` contract of ``ms.recompute(use_reentrant=False)``
124 and ``torch.utils.checkpoint(use_reentrant=False)``. Use this
125 to bracket the backward-time forward re-run with custom logic.
126 When ``policy_fn``, ``group_swap`` and ``context_fn`` are
127 supplied together, the resulting factories are composed: their
128 forward and recompute contexts are stacked so all enter in
129 order and exit in reverse.
130 group_swap (bool, optional): Whether MUST_SWAP tensors participate in group copy fusion.
131 Only effective when ``policy_fn`` is provided. Default: ``False``.
132 **kwargs: Additional keyword arguments to pass to the function.
134 Returns:
135 The result of applying the function with checkpointing.
136 """
137 factories: list = [create_recompute_contexts]
138 if policy_fn is not None:
139 factories.append(partial(plat.create_selective_checkpoint_contexts, policy_fn, group_swap=group_swap))
140 if context_fn is not None:
141 factories.append(context_fn)
143 if len(factories) == 1:
144 composed_context_fn = factories[0]
145 else:
146 composed_context_fn = _compose_context_fns(tuple(factories))
148 context = partial(plat.async_save_on_cpu, group_swap=group_swap) if swap_inputs else contextlib.nullcontext
149 with context():
150 return plat.checkpoint(
151 function, *args, context_fn=composed_context_fn, use_reentrant=False, **kwargs
152 )
155def swap(function, *args, policy_fn=None, group_swap=False, **kwargs):
156 """Apply activation swap to a function call.
158 Offloads intermediate activations saved by the autograd engine to CPU
159 during the forward pass and loads them back before the backward pass,
160 trading device memory for host memory bandwidth. Unlike
161 :func:`checkpoint`, no recomputation is performed.
163 Args:
164 function (callable): The function whose activations should be swapped.
165 *args: Positional arguments forwarded to *function*.
166 policy_fn (callable, optional): Per-tensor swap policy. Receives
167 a tensor and returns a :class:`CheckpointPolicy` value. Tensors
168 that return ``CheckpointPolicy.MUST_SAVE`` are kept on device;
169 all other eligible tensors are offloaded. When ``None``, all
170 eligible tensors are offloaded.
171 group_swap (bool, optional): Whether swapped tensors participate in
172 group copy fusion. Default: ``False``.
173 **kwargs: Keyword arguments forwarded to *function*.
175 Returns:
176 The return value of ``function(*args, **kwargs)``.
178 Example:
179 >>> output = swap(layer, x, policy_fn=lambda t: CheckpointPolicy.MUST_SAVE)
180 """
181 with plat.async_save_on_cpu(policy_fn=policy_fn, group_swap=group_swap):
182 return function(*args, **kwargs)
185def checkpoint_exclude_wrapper(module: Any) -> Any:
186 """Wrap a callable whose region is excluded from activation recomputation.
188 Args:
189 module: The module or callable to exclude from recomputation.
191 Returns:
192 The platform-specific checkpoint exclusion wrapper.
193 """
194 return plat.checkpoint_exclude_wrapper(module)
197checkpoint_wrapper = plat.checkpoint_wrapper
198swap_wrapper = plat.swap_wrapper
199swap_tensor_wrapper = plat.swap_tensor_wrapper