Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / api.py: 100%
135 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.
14# ============================================================================
15"""Activation checkpointing related interfaces"""
17__all__ = [
18 "CheckpointPolicy",
19 "checkpoint",
20 "swap",
21 "checkpoint_wrapper",
22 "checkpoint_exclude_wrapper",
23 "swap_wrapper",
24 "swap_tensor_wrapper",
25 "recompute_handle_collector_ctx",
26 "recompute_handle",
27 "recompute_session_ctx",
28 "clear_recompute_session",
29 "ignore_sac_ops",
30 "create_selective_checkpoint_contexts",
31 "create_native_selective_checkpoint_contexts",
32 "async_save_on_cpu",
33 "get_class_activation_wrapper",
34 "noop_context_fn",
35 "is_compiling",
36]
38import contextlib
39from functools import partial
40from typing import Any, Callable, Optional, Tuple
42import torch
44from .policy import CheckpointPolicy
45from .recompute_state import create_recompute_contexts
48class _StackedCtx:
49 """Compose multiple context managers as one — enter in order, exit reversed."""
51 def __init__(self, ctxs) -> None:
52 self._ctxs = list(ctxs)
53 self._stack = contextlib.ExitStack()
55 def __enter__(self):
56 self._stack.__enter__()
57 try:
58 for ctx in self._ctxs:
59 self._stack.enter_context(ctx)
60 except BaseException as exc:
61 self._stack.__exit__(type(exc), exc, exc.__traceback__)
62 raise
63 return self
65 def __exit__(self, exc_type, exc_val, exc_tb):
66 return self._stack.__exit__(exc_type, exc_val, exc_tb)
69def _compose_context_fns(
70 factories: Tuple[Callable[[], Tuple[object, object]], ...],
71) -> Callable[[], Tuple[_StackedCtx, _StackedCtx]]:
72 """Combine ``(forward_ctx, recompute_ctx)`` factories into one factory.
74 ``ms.recompute`` / ``torch.utils.checkpoint(use_reentrant=False)`` call
75 ``context_fn()`` once per invocation and unpack the result as
76 ``(forward_ctx, recompute_ctx)``. This helper calls each input factory
77 once, then stacks all forward contexts and all recompute contexts into
78 two :class:`_StackedCtx` instances so the composite respects the
79 single-call contract.
80 """
81 def factory() -> Tuple[_StackedCtx, _StackedCtx]:
82 pairs = [fn() for fn in factories]
83 fwd_ctxs = [pair[0] for pair in pairs]
84 rec_ctxs = [pair[1] for pair in pairs]
85 return _StackedCtx(fwd_ctxs), _StackedCtx(rec_ctxs)
87 return factory
90def is_compiling() -> bool:
91 """Return whether the caller runs inside a ``torch.compile`` capture."""
92 return torch.compiler.is_compiling()
95def noop_context_fn():
96 """The Torch-native no-op ``context_fn``, for callers that need an explicit one."""
97 # pylint: disable=C0415
98 from torch.utils.checkpoint import noop_context_fn as _torch_noop_context_fn
99 return _torch_noop_context_fn
102def checkpoint_exclude_wrapper(module: Any, *, save_output: bool = True) -> Any:
103 """Wrap a callable whose region is excluded from activation recomputation.
105 Args:
106 module: The module or callable to exclude from recomputation.
107 save_output: Whether to retain the region output for checkpoint replay.
108 Set this to ``False`` only when the output is passed directly as one
109 argument to another excluded region. Default: ``True``.
111 Returns:
112 The checkpoint exclusion wrapper.
113 """
114 # pylint: disable=C0415
115 from .wrapper import checkpoint_exclude_wrapper as _wrapper
116 return _wrapper(module, save_output=save_output)
119def checkpoint_wrapper(module, **checkpoint_kwargs):
120 """Wrap *module* with HyperParallel activation checkpointing."""
121 # pylint: disable=C0415
122 from .wrapper import ckpt_wrapper
123 return ckpt_wrapper(module, **checkpoint_kwargs)
126def swap_wrapper(module, policy_fn=None, group_swap=False, cpu_pool=None):
127 """Wrap *module* so its activations are offloaded to host memory.
129 Args:
130 module: Module or callable to wrap.
131 policy_fn: Optional per-tensor swap policy.
132 group_swap: Whether tensors participate in group copy fusion.
133 cpu_pool: Optional pinned host memory pool for the swapped tensors.
134 Only supported by the Torch backend, so it is forwarded only when set.
136 Returns:
137 The configured swap wrapper.
138 """
139 # pylint: disable=C0415
140 from .wrapper import swap_wrapper as _swap_wrapper
141 kwargs = {}
142 if cpu_pool is not None:
143 kwargs["cpu_pool"] = cpu_pool
144 return _swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap, **kwargs)
147def swap_tensor_wrapper(target, tag=None, group_swap=False, cpu_pool=None):
148 """Wrap a single tensor or callable for activation swap.
150 Args:
151 target: Tensor or nested tensor structure to register.
152 tag: Optional debug tag.
153 group_swap: Whether tensors participate in group copy fusion.
154 cpu_pool: Optional pinned host memory pool for the swapped tensors.
155 Only supported by the Torch backend, so it is forwarded only when set.
157 Returns:
158 The registered tensor structure.
159 """
160 # pylint: disable=C0415
161 from .wrapper import swap_tensor_wrapper as _swap_tensor_wrapper
162 kwargs = {}
163 if cpu_pool is not None:
164 kwargs["cpu_pool"] = cpu_pool
165 return _swap_tensor_wrapper(target, tag=tag, group_swap=group_swap, **kwargs)
168def get_class_activation_wrapper():
169 """Return the :class:`ActivationWrapper` base class of this backend."""
170 # pylint: disable=C0415
171 from .wrapper import ActivationWrapper
172 return ActivationWrapper
175def ignore_sac_ops(ignore_ops: list) -> None:
176 """Exclude backend operators from selective-AC replay accounting."""
177 # pylint: disable=C0415
178 from .sac import ignore_sac_ops as _ignore_sac_ops
179 _ignore_sac_ops(ignore_ops)
182def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False,
183 group_swap=False, cpu_pool=None):
184 """Create HyperParallel selective-checkpoint ``(forward_ctx, recompute_ctx)`` contexts."""
185 # pylint: disable=C0415
186 from .sac import create_selective_checkpoint_contexts as _create_contexts
187 return _create_contexts(
188 policy_fn_or_list,
189 allow_cache_entry_mutation=allow_cache_entry_mutation,
190 group_swap=group_swap,
191 cpu_pool=cpu_pool,
192 )
195def create_native_selective_checkpoint_contexts(policy_fn):
196 """Create Torch-native selective-checkpoint contexts for compile capture."""
197 # pylint: disable=C0415
198 from .compile_adapter import create_native_selective_checkpoint_contexts as _create_contexts
199 return _create_contexts(policy_fn)
202def async_save_on_cpu(policy_fn=None, group_swap: bool = False, cpu_pool=None):
203 """Return the async host-offload context used by :func:`swap` and ``swap_inputs``."""
204 # pylint: disable=C0415
205 from .wrapper import AsyncSaveOnCpu
206 return AsyncSaveOnCpu(policy_fn=policy_fn, group_swap=group_swap, cpu_pool=cpu_pool)
209def recompute_handle_collector_ctx():
210 """Return the context that collects recompute handles for a checkpoint region."""
211 # pylint: disable=C0415
212 from .checkpoint import recompute_handle_collector_ctx as _collector_ctx
213 return _collector_ctx()
216def recompute_handle(handle, session_id):
217 """Recompute the region behind *handle* under *session_id*."""
218 # pylint: disable=C0415
219 from .checkpoint import recompute_handle as _recompute_handle
220 return _recompute_handle(handle, session_id)
223def recompute_session_ctx(session_id, retain_on_unpack=False):
224 """Open a recompute session so replay can unpack saved activations."""
225 if session_id is None:
226 raise ValueError("session_id must not be None.")
227 # pylint: disable=C0415
228 from .checkpoint import recompute_session_ctx as _session_ctx
229 return _session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack)
232def clear_recompute_session(session_id):
233 """Release all state held for *session_id*."""
234 # pylint: disable=C0415
235 from .checkpoint import clear_recompute_session as _clear
236 return _clear(session_id)
239def _validate_compile_checkpoint_options(
240 swap_inputs: bool,
241 group_swap: bool,
242 cpu_pool: Any,
243 context_fn: Optional[Callable],
244 use_reentrant: bool,
245) -> None:
246 """Reject eager-only checkpoint options before compile capture."""
247 unsupported = []
248 if swap_inputs:
249 unsupported.append("swap_inputs")
250 if group_swap:
251 unsupported.append("group_swap")
252 if cpu_pool is not None:
253 unsupported.append("cpu_pool")
254 if context_fn is not None:
255 unsupported.append("custom context_fn")
256 if use_reentrant:
257 unsupported.append("use_reentrant=True")
258 if unsupported:
259 raise ValueError(
260 "HyperParallel checkpoint compile mode does not support: "
261 + ", ".join(unsupported)
262 + ". Use Torch-native non-reentrant checkpointing with optional "
263 "SAVE/RECOMPUTE selective policies."
264 )
267def _make_checkpoint_context_fn(
268 policy_fn: Optional[Callable],
269 context_fn: Optional[Callable],
270 group_swap: bool,
271 cpu_pool: Any,
272) -> Callable:
273 """Compose eager recompute, selective, and caller-provided contexts."""
274 factories: list = [create_recompute_contexts]
275 if policy_fn is not None:
276 selective_kwargs = {"group_swap": group_swap}
277 if cpu_pool is not None:
278 selective_kwargs["cpu_pool"] = cpu_pool
279 factories.append(partial(create_selective_checkpoint_contexts, policy_fn, **selective_kwargs))
280 if context_fn is not None:
281 factories.append(context_fn)
282 if len(factories) == 1:
283 return factories[0]
284 return _compose_context_fns(tuple(factories))
287def _checkpoint_input_context(swap_inputs: bool, group_swap: bool, cpu_pool: Any) -> Any:
288 """Create the optional context that offloads checkpoint inputs."""
289 if not swap_inputs:
290 return contextlib.nullcontext()
291 async_kwargs = {"group_swap": group_swap}
292 if cpu_pool is not None:
293 async_kwargs["cpu_pool"] = cpu_pool
294 return async_save_on_cpu(**async_kwargs)
297def checkpoint(
298 function,
299 *args,
300 swap_inputs: bool = False,
301 policy_fn: Optional[Callable] = None,
302 context_fn: Optional[Callable[[], Tuple[object, object]]] = None,
303 group_swap: bool = False,
304 cpu_pool=None,
305 early_stop: bool = True,
306 **kwargs,
307):
308 """
309 Apply activation checkpointing to a function with optional input swapping.
311 Args:
312 function: The function to apply checkpointing to.
313 *args: Arguments to pass to the function.
314 swap_inputs (bool): Whether to enable input swapping using async_save_on_cpu context.
315 policy_fn (callable, optional): Function that determines checkpoint policy for operations.
316 context_fn (callable, optional): A no-arg factory returning a
317 ``(forward_ctx, recompute_ctx)`` pair, matching the
318 ``context_fn`` contract of ``ms.recompute(use_reentrant=False)``
319 and ``torch.utils.checkpoint(use_reentrant=False)``. Use this
320 to bracket the backward-time forward re-run with custom logic.
321 When ``policy_fn``, ``group_swap`` and ``context_fn`` are
322 supplied together, the resulting factories are composed: their
323 forward and recompute contexts are stacked so all enter in
324 order and exit in reverse.
325 group_swap (bool, optional): Whether MUST_SWAP tensors participate in group copy fusion.
326 Only effective when ``policy_fn`` is provided. Default: ``False``.
327 cpu_pool (PinnedMemoryPool, optional): Explicit pinned host memory pool used by tensors
328 selected for swapping. Currently supported only by the Torch backend. Default: ``None``.
329 early_stop (bool, optional): Whether recomputation stops after all tensors needed by
330 backward have been produced. This per-call keyword is the only supported way to
331 configure early stop. Default: ``True``.
332 **kwargs: Additional keyword arguments to pass to the function.
334 Returns:
335 The result of applying the function with checkpointing.
336 """
337 if not isinstance(early_stop, bool):
338 raise ValueError(f"early_stop must be bool, but got {type(early_stop).__name__}.")
340 if is_compiling():
341 _validate_compile_checkpoint_options(
342 swap_inputs, group_swap, cpu_pool, context_fn, kwargs.get("use_reentrant", False)
343 )
344 composed_context_fn = (
345 partial(create_native_selective_checkpoint_contexts, policy_fn)
346 if policy_fn is not None
347 else None
348 )
349 else:
350 composed_context_fn = _make_checkpoint_context_fn(policy_fn, context_fn, group_swap, cpu_pool)
352 with _checkpoint_input_context(swap_inputs, group_swap, cpu_pool):
353 checkpoint_kwargs = {**kwargs, "use_reentrant": False, "early_stop": early_stop}
354 if composed_context_fn is not None:
355 checkpoint_kwargs["context_fn"] = composed_context_fn
356 # pylint: disable=C0415
357 from .checkpoint import checkpoint as hyper_checkpoint
358 return hyper_checkpoint(function, *args, **checkpoint_kwargs)
361def swap(function, *args, policy_fn=None, group_swap=False, cpu_pool=None, **kwargs):
362 """Apply activation swap to a function call.
364 Offloads intermediate activations saved by the autograd engine to CPU
365 during the forward pass and loads them back before the backward pass,
366 trading device memory for host memory bandwidth. Unlike
367 :func:`checkpoint`, no recomputation is performed.
369 Args:
370 function (callable): The function whose activations should be swapped.
371 *args: Positional arguments forwarded to *function*.
372 policy_fn (callable, optional): Per-tensor swap policy. Receives
373 a tensor and returns a :class:`CheckpointPolicy` value. Tensors
374 that return ``CheckpointPolicy.MUST_SAVE`` are kept on device;
375 all other eligible tensors are offloaded. When ``None``, all
376 eligible tensors are offloaded.
377 group_swap (bool, optional): Whether swapped tensors participate in
378 group copy fusion. Default: ``False``.
379 cpu_pool (PinnedMemoryPool, optional): Explicit pinned host memory pool used by swapped
380 tensors. Currently supported only by the Torch backend. Default: ``None``.
381 **kwargs: Keyword arguments forwarded to *function*.
383 Returns:
384 The return value of ``function(*args, **kwargs)``.
386 Example:
387 >>> output = swap(layer, x, policy_fn=lambda t: CheckpointPolicy.MUST_SAVE)
388 """
389 if is_compiling():
390 raise ValueError(
391 "HyperParallel activation swap is not supported in compile mode. "
392 "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
393 )
394 async_kwargs = {"policy_fn": policy_fn, "group_swap": group_swap}
395 if cpu_pool is not None:
396 async_kwargs["cpu_pool"] = cpu_pool
397 with async_save_on_cpu(**async_kwargs):
398 return function(*args, **kwargs)