Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_checkpoint / activation_checkpoint.py: 79%
78 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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 early_stop: bool = True,
112 **kwargs,
113):
114 """
115 Apply activation checkpointing to a function with optional input swapping.
117 Args:
118 function: The function to apply checkpointing to.
119 *args: Arguments to pass to the function.
120 swap_inputs (bool): Whether to enable input swapping using async_save_on_cpu context.
121 policy_fn (callable, optional): Function that determines checkpoint policy for operations.
122 context_fn (callable, optional): A no-arg factory returning a
123 ``(forward_ctx, recompute_ctx)`` pair, matching the
124 ``context_fn`` contract of ``ms.recompute(use_reentrant=False)``
125 and ``torch.utils.checkpoint(use_reentrant=False)``. Use this
126 to bracket the backward-time forward re-run with custom logic.
127 When ``policy_fn``, ``group_swap`` and ``context_fn`` are
128 supplied together, the resulting factories are composed: their
129 forward and recompute contexts are stacked so all enter in
130 order and exit in reverse.
131 group_swap (bool, optional): Whether MUST_SWAP tensors participate in group copy fusion.
132 Only effective when ``policy_fn`` is provided. Default: ``False``.
133 early_stop (bool, optional): Whether recomputation stops after all tensors needed by
134 backward have been produced. This per-call keyword is the only supported way to
135 configure early stop. Default: ``True``.
136 **kwargs: Additional keyword arguments to pass to the function.
138 Returns:
139 The result of applying the function with checkpointing.
140 """
141 if not isinstance(early_stop, bool):
142 raise ValueError(f"early_stop must be bool, but got {type(early_stop).__name__}.")
144 is_compiling = getattr(plat, "is_compiling", None)
145 checkpoint_compile = callable(is_compiling) and is_compiling() is True # pylint: disable=not-callable
146 if checkpoint_compile:
147 unsupported = []
148 if swap_inputs:
149 unsupported.append("swap_inputs")
150 if group_swap:
151 unsupported.append("group_swap")
152 if context_fn is not None:
153 unsupported.append("custom context_fn")
154 if kwargs.get("use_reentrant", False):
155 unsupported.append("use_reentrant=True")
156 if unsupported:
157 raise ValueError(
158 "HyperParallel checkpoint compile mode does not support: "
159 + ", ".join(unsupported)
160 + ". Use Torch-native non-reentrant checkpointing with optional "
161 "SAVE/RECOMPUTE selective policies."
162 )
163 composed_context_fn = (
164 partial(plat.create_native_selective_checkpoint_contexts, policy_fn)
165 if policy_fn is not None
166 else None
167 )
168 else:
169 factories: list = [create_recompute_contexts]
170 if policy_fn is not None:
171 factories.append(partial(plat.create_selective_checkpoint_contexts, policy_fn, group_swap=group_swap))
172 if context_fn is not None:
173 factories.append(context_fn)
175 if len(factories) == 1:
176 composed_context_fn = factories[0]
177 else:
178 composed_context_fn = _compose_context_fns(tuple(factories))
180 context = partial(plat.async_save_on_cpu, group_swap=group_swap) if swap_inputs else contextlib.nullcontext
181 with context():
182 checkpoint_kwargs = {**kwargs, "use_reentrant": False, "early_stop": early_stop}
183 if composed_context_fn is not None:
184 checkpoint_kwargs["context_fn"] = composed_context_fn
185 return plat.checkpoint(function, *args, **checkpoint_kwargs)
188def swap(function, *args, policy_fn=None, group_swap=False, **kwargs):
189 """Apply activation swap to a function call.
191 Offloads intermediate activations saved by the autograd engine to CPU
192 during the forward pass and loads them back before the backward pass,
193 trading device memory for host memory bandwidth. Unlike
194 :func:`checkpoint`, no recomputation is performed.
196 Args:
197 function (callable): The function whose activations should be swapped.
198 *args: Positional arguments forwarded to *function*.
199 policy_fn (callable, optional): Per-tensor swap policy. Receives
200 a tensor and returns a :class:`CheckpointPolicy` value. Tensors
201 that return ``CheckpointPolicy.MUST_SAVE`` are kept on device;
202 all other eligible tensors are offloaded. When ``None``, all
203 eligible tensors are offloaded.
204 group_swap (bool, optional): Whether swapped tensors participate in
205 group copy fusion. Default: ``False``.
206 **kwargs: Keyword arguments forwarded to *function*.
208 Returns:
209 The return value of ``function(*args, **kwargs)``.
211 Example:
212 >>> output = swap(layer, x, policy_fn=lambda t: CheckpointPolicy.MUST_SAVE)
213 """
214 is_compiling = getattr(plat, "is_compiling", None)
215 if callable(is_compiling) and is_compiling() is True: # pylint: disable=not-callable
216 raise ValueError(
217 "HyperParallel activation swap is not supported in compile mode. "
218 "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
219 )
220 with plat.async_save_on_cpu(policy_fn=policy_fn, group_swap=group_swap):
221 return function(*args, **kwargs)
224def checkpoint_exclude_wrapper(module: Any, *, save_output: bool = True) -> Any:
225 """Wrap a callable whose region is excluded from activation recomputation.
227 Args:
228 module: The module or callable to exclude from recomputation.
229 save_output: Whether to retain the region output for checkpoint replay.
230 Set this to ``False`` only when the output is passed directly as one
231 argument to another excluded region. Default: ``True``.
233 Returns:
234 The platform-specific checkpoint exclusion wrapper.
235 """
236 return plat.checkpoint_exclude_wrapper(module, save_output=save_output)
239checkpoint_wrapper = plat.checkpoint_wrapper
240swap_wrapper = plat.swap_wrapper
241swap_tensor_wrapper = plat.swap_tensor_wrapper