Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / api.py: 95%
194 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.
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 "reentrant_checkpoint_scope",
37]
39import contextlib
40from contextlib import AbstractContextManager
41from functools import partial
42from typing import Any, Callable, Optional, Tuple
44import torch
46from .policy import CheckpointPolicy
47from .recompute_state import create_recompute_contexts
50class _StackedCtx:
51 """Compose multiple context managers as one — enter in order, exit reversed."""
53 def __init__(self, ctxs) -> None:
54 self._ctxs = list(ctxs)
55 self._stack = contextlib.ExitStack()
57 def __enter__(self):
58 self._stack.__enter__()
59 try:
60 for ctx in self._ctxs:
61 self._stack.enter_context(ctx)
62 except BaseException as exc:
63 self._stack.__exit__(type(exc), exc, exc.__traceback__)
64 raise
65 return self
67 def __exit__(self, exc_type, exc_val, exc_tb):
68 return self._stack.__exit__(exc_type, exc_val, exc_tb)
71def _compose_context_fns(
72 factories: Tuple[Callable[[], Tuple[object, object]], ...],
73) -> Callable[[], Tuple[_StackedCtx, _StackedCtx]]:
74 """Combine ``(forward_ctx, recompute_ctx)`` factories into one factory.
76 ``ms.recompute`` / ``torch.utils.checkpoint(use_reentrant=False)`` call
77 ``context_fn()`` once per invocation and unpack the result as
78 ``(forward_ctx, recompute_ctx)``. This helper calls each input factory
79 once, then stacks all forward contexts and all recompute contexts into
80 two :class:`_StackedCtx` instances so the composite respects the
81 single-call contract.
82 """
83 def factory() -> Tuple[_StackedCtx, _StackedCtx]:
84 pairs = [fn() for fn in factories]
85 fwd_ctxs = [pair[0] for pair in pairs]
86 rec_ctxs = [pair[1] for pair in pairs]
87 return _StackedCtx(fwd_ctxs), _StackedCtx(rec_ctxs)
89 return factory
92def is_compiling() -> bool:
93 """Return whether the caller runs inside a ``torch.compile`` capture."""
94 return torch.compiler.is_compiling()
97def noop_context_fn():
98 """The Torch-native no-op ``context_fn``, for callers that need an explicit one."""
99 # pylint: disable=C0415
100 from torch.utils.checkpoint import noop_context_fn as _torch_noop_context_fn
101 return _torch_noop_context_fn
104def checkpoint_exclude_wrapper( # pylint: disable=redefined-outer-name
105 module: Any,
106 *,
107 swap: bool = False,
108 group_swap: bool = False,
109 cpu_pool: Any = None,
110) -> Any:
111 """Wrap a callable whose region is excluded from activation recomputation.
113 Args:
114 module: The module or callable to exclude from recomputation.
115 swap: Whether eligible outputs and non-recomputable saved tensors should
116 use the active activation-swap group. Default: ``False``.
117 group_swap: Whether swapped tensors participate in group copy fusion.
118 cpu_pool: Optional pinned-memory pool used for swapped tensors.
120 Returns:
121 The checkpoint exclusion wrapper.
122 """
123 # pylint: disable=C0415
124 from .wrapper import checkpoint_exclude_wrapper as _wrapper
125 return _wrapper(
126 module,
127 swap=swap,
128 group_swap=group_swap,
129 cpu_pool=cpu_pool,
130 )
133def checkpoint_wrapper(module, **checkpoint_kwargs):
134 """Wrap *module* with HyperParallel activation checkpointing."""
135 # pylint: disable=C0415
136 from .wrapper import ckpt_wrapper
137 return ckpt_wrapper(module, **checkpoint_kwargs)
140def reentrant_checkpoint_scope(parameters: Any) -> AbstractContextManager:
141 """Collect reentrant frames with explicit stage parameters for split backward.
143 Args:
144 parameters: Stage parameters or lazy zero-argument provider. The provider
145 is evaluated only when a reentrant checkpoint is invoked.
147 Returns:
148 A context manager yielding handles with ``set_phase``, ``release``, and
149 ``backward_trigger``. For split DX, set phase to ``dx`` and target stage
150 input activations. For DW, set phase to ``dw`` and target stage parameters
151 plus every handle's ``backward_trigger``. Captured parameter hooks run
152 only during DW; shared parameters can receive separate contributions.
153 """
154 from .reentrant_checkpoint import reentrant_checkpoint_scope as _scope # pylint: disable=C0415
155 return _scope(parameters)
158def swap_wrapper(module, policy_fn=None, group_swap=False, cpu_pool=None):
159 """Wrap *module* so its activations are offloaded to host memory.
161 Args:
162 module: Module or callable to wrap.
163 policy_fn: Optional per-tensor swap policy.
164 group_swap: Whether tensors participate in group copy fusion.
165 cpu_pool: Optional pinned host memory pool for the swapped tensors.
166 Only supported by the Torch backend, so it is forwarded only when set.
168 Returns:
169 The configured swap wrapper.
170 """
171 # pylint: disable=C0415
172 from .wrapper import swap_wrapper as _swap_wrapper
173 kwargs = {}
174 if cpu_pool is not None:
175 kwargs["cpu_pool"] = cpu_pool
176 return _swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap, **kwargs)
179def swap_tensor_wrapper(target, tag=None, group_swap=False, cpu_pool=None):
180 """Wrap a single tensor or callable for activation swap.
182 Args:
183 target: Tensor or nested tensor structure to register.
184 tag: Optional debug tag.
185 group_swap: Whether tensors participate in group copy fusion.
186 cpu_pool: Optional pinned host memory pool for the swapped tensors.
187 Only supported by the Torch backend, so it is forwarded only when set.
189 Returns:
190 The registered tensor structure.
191 """
192 # pylint: disable=C0415
193 from .wrapper import swap_tensor_wrapper as _swap_tensor_wrapper
194 kwargs = {}
195 if cpu_pool is not None:
196 kwargs["cpu_pool"] = cpu_pool
197 return _swap_tensor_wrapper(target, tag=tag, group_swap=group_swap, **kwargs)
200def get_class_activation_wrapper():
201 """Return the :class:`ActivationWrapper` base class of this backend."""
202 # pylint: disable=C0415
203 from .wrapper import ActivationWrapper
204 return ActivationWrapper
207def ignore_sac_ops(ignore_ops: list) -> None:
208 """Exclude backend operators from selective-AC replay accounting."""
209 # pylint: disable=C0415
210 from .sac import ignore_sac_ops as _ignore_sac_ops
211 _ignore_sac_ops(ignore_ops)
214def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False,
215 group_swap=False, cpu_pool=None):
216 """Create HyperParallel selective-checkpoint ``(forward_ctx, recompute_ctx)`` contexts."""
217 # pylint: disable=C0415
218 from .sac import create_selective_checkpoint_contexts as _create_contexts
219 return _create_contexts(
220 policy_fn_or_list,
221 allow_cache_entry_mutation=allow_cache_entry_mutation,
222 group_swap=group_swap,
223 cpu_pool=cpu_pool,
224 )
227def create_native_selective_checkpoint_contexts(policy_fn):
228 """Create Torch-native selective-checkpoint contexts for compile capture."""
229 # pylint: disable=C0415
230 from .compile_adapter import create_native_selective_checkpoint_contexts as _create_contexts
231 return _create_contexts(policy_fn)
234def async_save_on_cpu(policy_fn=None, group_swap: bool = False, cpu_pool=None):
235 """Return the async host-offload context used by :func:`swap` and ``swap_inputs``."""
236 # pylint: disable=C0415
237 from .wrapper import AsyncSaveOnCpu
238 return AsyncSaveOnCpu(policy_fn=policy_fn, group_swap=group_swap, cpu_pool=cpu_pool)
241def recompute_handle_collector_ctx() -> AbstractContextManager:
242 """Return the context that collects recompute handles for a checkpoint region."""
243 # pylint: disable=C0415
244 from .checkpoint import recompute_handle_collector_ctx as _collector_ctx # pylint: disable=C0415
245 return _collector_ctx()
248def recompute_handle(handle: Any, session_id: Any) -> Any:
249 """Recompute the region behind *handle* under *session_id*."""
250 from .checkpoint import _validate_session_id # pylint: disable=C0415
251 _validate_session_id(session_id)
252 from .reentrant_checkpoint import ReentrantCheckpointHandle # pylint: disable=C0415
253 if isinstance(handle, ReentrantCheckpointHandle):
254 handle.bind_session(session_id)
255 return handle.recompute()
256 # pylint: disable=C0415
257 from .checkpoint import recompute_handle as _recompute_handle
258 return _recompute_handle(handle, session_id)
261def recompute_session_ctx(session_id: Any, retain_on_unpack: bool = False) -> AbstractContextManager:
262 """Open a recompute session so replay can unpack saved activations."""
263 if session_id is None:
264 raise ValueError("session_id must not be None.")
265 # pylint: disable=C0415
266 from .checkpoint import recompute_session_ctx as _session_ctx
267 from .reentrant_checkpoint import activate_session, deactivate_session # pylint: disable=C0415
268 from .checkpoint import _RECOMPUTE_SESSION # pylint: disable=C0415
270 @contextlib.contextmanager
271 def combined_session() -> Any:
272 """Bind replay retention to registered frames for this session scope."""
273 with _session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack):
274 activation = _RECOMPUTE_SESSION.get()
275 activate_session(activation)
276 try:
277 yield session_id
278 finally:
279 deactivate_session(activation)
281 return combined_session()
284def clear_recompute_session(session_id: Any) -> Any:
285 """Release all state held for *session_id*."""
286 from .checkpoint import _validate_session_id # pylint: disable=C0415
287 _validate_session_id(session_id)
288 from .reentrant_checkpoint import clear_session as _clear_reentrant # pylint: disable=C0415
289 _clear_reentrant(session_id)
290 # pylint: disable=C0415
291 from .checkpoint import clear_recompute_session as _clear
292 return _clear(session_id)
295def _validate_compile_checkpoint_options(
296 swap_inputs: bool,
297 group_swap: bool,
298 cpu_pool: Any,
299 context_fn: Optional[Callable],
300 use_reentrant: bool,
301) -> None:
302 """Reject eager-only checkpoint options before compile capture."""
303 unsupported = []
304 if swap_inputs:
305 unsupported.append("swap_inputs")
306 if group_swap:
307 unsupported.append("group_swap")
308 if cpu_pool is not None:
309 unsupported.append("cpu_pool")
310 if context_fn is not None:
311 unsupported.append("custom context_fn")
312 if use_reentrant:
313 unsupported.append("use_reentrant=True")
314 if unsupported:
315 raise ValueError(
316 "HyperParallel checkpoint compile mode does not support: "
317 + ", ".join(unsupported)
318 + ". Use Torch-native non-reentrant checkpointing with optional "
319 "SAVE/RECOMPUTE selective policies."
320 )
323def _make_checkpoint_context_fn(
324 policy_fn: Optional[Callable],
325 context_fn: Optional[Callable],
326 group_swap: bool,
327 cpu_pool: Any,
328 preserve_rng_state: bool,
329) -> Callable:
330 """Compose eager recompute, selective, and caller-provided contexts."""
331 factories: list = (
332 [create_recompute_contexts]
333 if preserve_rng_state is True
334 else [partial(create_recompute_contexts, preserve_rng_state=preserve_rng_state)]
335 )
336 if policy_fn is not None:
337 selective_kwargs = {"group_swap": group_swap}
338 if cpu_pool is not None:
339 selective_kwargs["cpu_pool"] = cpu_pool
340 factories.append(partial(create_selective_checkpoint_contexts, policy_fn, **selective_kwargs))
341 if context_fn is not None:
342 factories.append(context_fn)
343 if len(factories) == 1:
344 return factories[0]
345 return _compose_context_fns(tuple(factories))
348def _checkpoint_input_context(swap_inputs: bool, group_swap: bool, cpu_pool: Any) -> Any:
349 """Create the optional context that offloads checkpoint inputs."""
350 if not swap_inputs:
351 return contextlib.nullcontext()
352 async_kwargs = {"group_swap": group_swap}
353 if cpu_pool is not None:
354 async_kwargs["cpu_pool"] = cpu_pool
355 return async_save_on_cpu(**async_kwargs)
358def _checkpoint_reentrant(
359 function: Callable,
360 args: tuple,
361 context_fn: Optional[Callable[[], Tuple[object, object]]],
362 kwargs: dict,
363) -> Any:
364 """Apply the eager reentrant checkpoint implementation."""
365 from .reentrant_checkpoint import checkpoint as _reentrant # pylint: disable=C0415
367 preserve_rng_state = kwargs.pop("preserve_rng_state", True)
368 if not isinstance(preserve_rng_state, bool):
369 raise ValueError("preserve_rng_state must be bool")
370 if context_fn is not None and not callable(context_fn):
371 raise ValueError("context_fn must be callable")
372 state_context = partial(create_recompute_contexts, preserve_rng_state=preserve_rng_state)
373 reentrant_context = (
374 _compose_context_fns((state_context, context_fn)) if context_fn is not None else state_context
375 )
376 reentrant_options = {"preserve_rng_state": preserve_rng_state,
377 "context_fn": reentrant_context, **kwargs}
378 return _reentrant(function, *args, **reentrant_options)
381def _validate_reentrant_options(
382 swap_inputs: bool,
383 policy_fn: Optional[Callable],
384 group_swap: bool,
385 cpu_pool: Any,
386 early_stop: bool,
387 kwargs: dict,
388) -> None:
389 """Validate options unsupported by eager reentrant checkpointing."""
390 unsupported = []
391 if swap_inputs:
392 unsupported.append("swap_inputs")
393 if policy_fn is not None:
394 unsupported.append("policy_fn")
395 if group_swap:
396 unsupported.append("group_swap")
397 if cpu_pool is not None:
398 unsupported.append("cpu_pool")
399 if early_stop is not True:
400 unsupported.append("early_stop")
401 for name in ("debug", "determinism_check"):
402 if name in kwargs:
403 unsupported.append(name)
404 if is_compiling():
405 unsupported.append("use_reentrant=True under torch.compile")
406 if unsupported:
407 raise ValueError("Reentrant checkpoint does not support: " + ", ".join(unsupported))
410def checkpoint(
411 function: Callable,
412 *args: Any,
413 swap_inputs: bool = False,
414 policy_fn: Optional[Callable] = None,
415 context_fn: Optional[Callable[[], Tuple[object, object]]] = None,
416 group_swap: bool = False,
417 cpu_pool: Any = None,
418 early_stop: bool = True,
419 use_reentrant: bool = False,
420 **kwargs: Any,
421) -> Any:
422 """
423 Apply activation checkpointing to a function with optional input swapping.
425 Args:
426 function: The function to apply checkpointing to.
427 *args: Arguments to pass to the function.
428 swap_inputs (bool): Whether to enable input swapping using async_save_on_cpu context.
429 policy_fn (callable, optional): Function that determines checkpoint policy for operations.
430 context_fn (callable, optional): A no-arg factory returning a
431 ``(forward_ctx, recompute_ctx)`` pair, matching the
432 ``context_fn`` contract of ``ms.recompute(use_reentrant=False)``
433 and ``torch.utils.checkpoint(use_reentrant=False)``. Use this
434 to bracket the backward-time forward re-run with custom logic.
435 When ``policy_fn``, ``group_swap`` and ``context_fn`` are
436 supplied together, the resulting factories are composed: their
437 forward and recompute contexts are stacked so all enter in
438 order and exit in reverse.
439 group_swap (bool, optional): Whether MUST_SWAP tensors participate in group copy fusion.
440 Only effective when ``policy_fn`` is provided. Default: ``False``.
441 cpu_pool (PinnedMemoryPool, optional): Explicit pinned host memory pool used by tensors
442 selected for swapping. Currently supported only by the Torch backend. Default: ``None``.
443 early_stop (bool, optional): Whether recomputation stops after all tensors needed by
444 backward have been produced. This per-call keyword is the only supported way to
445 configure early stop. Default: ``True``.
446 use_reentrant (bool, optional): Use eager reentrant checkpointing with
447 localized exclusion graphs. Ordinary backward can accumulate
448 captured parameters even when ``autograd.grad`` targets only an
449 activation; use a reentrant scope for split DX/DW execution.
450 **kwargs: Additional keyword arguments to pass to the function.
452 Returns:
453 The result of applying the function with checkpointing.
454 """
455 if not isinstance(early_stop, bool):
456 raise ValueError(f"early_stop must be bool, but got {type(early_stop).__name__}.")
457 if not isinstance(use_reentrant, bool):
458 raise ValueError("use_reentrant must be bool")
459 if use_reentrant:
460 _validate_reentrant_options(
461 swap_inputs, policy_fn, group_swap, cpu_pool, early_stop, kwargs
462 )
463 return _checkpoint_reentrant(function, args, context_fn, kwargs)
465 if is_compiling():
466 _validate_compile_checkpoint_options(
467 swap_inputs, group_swap, cpu_pool, context_fn, kwargs.get("use_reentrant", False)
468 )
469 composed_context_fn = (
470 partial(create_native_selective_checkpoint_contexts, policy_fn)
471 if policy_fn is not None
472 else None
473 )
474 else:
475 composed_context_fn = _make_checkpoint_context_fn(
476 policy_fn,
477 context_fn,
478 group_swap,
479 cpu_pool,
480 kwargs.get("preserve_rng_state", True),
481 )
483 with _checkpoint_input_context(swap_inputs, group_swap, cpu_pool):
484 checkpoint_kwargs = {**kwargs, "use_reentrant": False, "early_stop": early_stop}
485 if composed_context_fn is not None:
486 checkpoint_kwargs["context_fn"] = composed_context_fn
487 # pylint: disable=C0415
488 from .checkpoint import checkpoint as hyper_checkpoint
489 return hyper_checkpoint(function, *args, **checkpoint_kwargs)
492def swap(function, *args, policy_fn=None, group_swap=False, cpu_pool=None, **kwargs):
493 """Apply activation swap to a function call.
495 Offloads intermediate activations saved by the autograd engine to CPU
496 during the forward pass and loads them back before the backward pass,
497 trading device memory for host memory bandwidth. Unlike
498 :func:`checkpoint`, no recomputation is performed.
500 Args:
501 function (callable): The function whose activations should be swapped.
502 *args: Positional arguments forwarded to *function*.
503 policy_fn (callable, optional): Per-tensor swap policy. Receives
504 a tensor and returns a :class:`CheckpointPolicy` value. Tensors
505 that return ``CheckpointPolicy.MUST_SAVE`` are kept on device;
506 all other eligible tensors are offloaded. When ``None``, all
507 eligible tensors are offloaded.
508 group_swap (bool, optional): Whether swapped tensors participate in
509 group copy fusion. Default: ``False``.
510 cpu_pool (PinnedMemoryPool, optional): Explicit pinned host memory pool used by swapped
511 tensors. Currently supported only by the Torch backend. Default: ``None``.
512 **kwargs: Keyword arguments forwarded to *function*.
514 Returns:
515 The return value of ``function(*args, **kwargs)``.
517 Example:
518 >>> output = swap(layer, x, policy_fn=lambda t: CheckpointPolicy.MUST_SAVE)
519 """
520 if is_compiling():
521 raise ValueError(
522 "HyperParallel activation swap is not supported in compile mode. "
523 "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
524 )
525 async_kwargs = {"policy_fn": policy_fn, "group_swap": group_swap}
526 if cpu_pool is not None:
527 async_kwargs["cpu_pool"] = cpu_pool
528 with async_save_on_cpu(**async_kwargs):
529 return function(*args, **kwargs)