« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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"""PyTorch wrapper for regions that should be saved instead of recomputed."""
16from collections import defaultdict, deque
17from dataclasses import dataclass
18from functools import lru_cache
19from typing import Any, Callable, cast, Deque, Dict, List, Optional, Tuple
21import torch
23from hyper_parallel.core.activation_checkpoint.recompute_state import get_recompute_state
24from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import ActivationWrapper
27_RECOMPUTE_INPUT_HANDLE_ATTR = "_hyper_parallel_recompute_input_handle"
28_SAVE_OUTPUT_SOURCE_ATTR = "_hyper_parallel_save_output_source"
29_MISSING = object()
30_InputPath = Tuple[Tuple[str, Any], ...]
31_TensorInput = Tuple[_InputPath, Any]
34class _RecomputedInputHandle:
35 """Defer one excluded-region saved input until checkpoint replay."""
37 def __init__(self) -> None:
38 """Initialize an unused and unresolved handle."""
39 self.used = False
40 self._tensor = None
42 def mark_used(self) -> None:
43 """Record that an exclude operation saved this input for backward."""
44 self.used = True
46 def materialize(self, tensor: Any) -> None:
47 """Bind the handle to the matching input produced during replay."""
48 self._tensor = tensor
50 def get_recomputed_tensor(self) -> Any:
51 """Return the replay-produced input for backward."""
52 if self._tensor is None:
53 raise RuntimeError("Checkpoint-excluded input was requested before recomputation")
54 return self._tensor
57@dataclass(frozen=True)
58class _InputBinding:
59 """Map one input path to its deferred saved-tensor handle."""
61 path: _InputPath
62 handle: _RecomputedInputHandle
65@dataclass
66class _ExcludeCacheEntry:
67 """Store one excluded call's replay value and deferred input bindings."""
69 output: Any
70 input_bindings: List[_InputBinding]
71 output_tensor_count: int = 1
74class _ExcludeCache:
75 """Store excluded-region call entries for one checkpoint invocation."""
77 def __init__(self) -> None:
78 """Initialize an empty per-checkpoint output cache."""
79 self._entries: Dict[int, Deque[_ExcludeCacheEntry]] = defaultdict(deque)
81 def save(self, wrapper_id: int, entry: _ExcludeCacheEntry) -> None:
82 """Save one call entry produced by a checkpoint-excluded region."""
83 self._entries[wrapper_id].append(entry)
85 def pop(self, wrapper_id: int) -> _ExcludeCacheEntry:
86 """Return the matching forward call entry during recomputation."""
87 entries = self._entries.get(wrapper_id)
88 if not entries:
89 raise RuntimeError("No cached forward output is available for this checkpoint exclusion wrapper")
90 entry = entries.popleft()
91 if not entries:
92 self._entries.pop(wrapper_id)
93 return entry
95 def clear(self) -> None:
96 """Release outputs not consumed because recomputation stopped early."""
97 self._entries.clear()
100def _pack_saved_tensor(tensor: Any) -> Any:
101 """Return a deferred input handle or detached tensor data."""
102 handle = getattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, None)
103 if isinstance(handle, _RecomputedInputHandle):
104 handle.mark_used()
105 return handle
106 return tensor.detach() if tensor.requires_grad else tensor
109def _unpack_saved_tensor(value: Any) -> Any:
110 """Restore the saved tensor for backward."""
111 if isinstance(value, _RecomputedInputHandle):
112 return value.get_recomputed_tensor()
113 return value
116def _saved_tensors_context() -> Any:
117 """Create an inner hook that stores real tensors instead of outer holders."""
118 return torch.autograd.graph.saved_tensors_hooks(_pack_saved_tensor, _unpack_saved_tensor)
121_EXCLUDE_CACHE_KEY = object()
124def _append_tensor_inputs(
125 value: Any,
126 path: _InputPath,
127 leaves: List[_TensorInput],
128 tensor_type: type,
129) -> None:
130 """Append tensor leaves without creating a self-referential local function."""
131 if isinstance(value, tensor_type):
132 leaves.append((path, value))
133 return
134 if isinstance(value, (tuple, list)):
135 for index, item in enumerate(value):
136 _append_tensor_inputs(item, path + (("index", index),), leaves, tensor_type)
137 return
138 if isinstance(value, dict):
139 for key, item in value.items():
140 _append_tensor_inputs(item, path + (("key", key),), leaves, tensor_type)
143def _collect_tensor_inputs(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> List[_TensorInput]:
144 """Return tensor leaves and self-describing paths from one excluded-region call."""
145 leaves = []
146 for index, arg in enumerate(args):
147 _append_tensor_inputs(arg, (("arg", index),), leaves, torch.Tensor)
148 for key, value in kwargs.items():
149 _append_tensor_inputs(value, (("kwarg", key),), leaves, torch.Tensor)
150 return leaves
153def _mark_recompute_inputs(
154 invocation_id: object,
155 args: Tuple[Any, ...],
156 kwargs: Dict[str, Any],
157) -> Tuple[List[_InputBinding], List[Tuple[Any, Any]]]:
158 """Attach deferred handles to inputs that replay reproduces."""
159 bindings = []
160 previous_handles = []
161 seen_tensor_ids = set()
162 try:
163 for path, tensor in _collect_tensor_inputs(args, kwargs):
164 if (
165 isinstance(tensor, torch.nn.Parameter)
166 or id(tensor) in seen_tensor_ids
167 or getattr(tensor, _SAVE_OUTPUT_SOURCE_ATTR, None) is invocation_id
168 ):
169 continue
170 handle = _RecomputedInputHandle()
171 previous = getattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, _MISSING)
172 previous_handles.append((tensor, previous))
173 setattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, handle)
174 bindings.append(_InputBinding(path, handle))
175 seen_tensor_ids.add(id(tensor))
176 except BaseException:
177 _restore_recompute_inputs(previous_handles)
178 raise
179 return bindings, previous_handles
182def _restore_recompute_inputs(previous_handles: List[Tuple[Any, Any]]) -> None:
183 """Restore Tensor attributes overwritten for one excluded call."""
184 for tensor, previous in previous_handles:
185 if previous is _MISSING:
186 if hasattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR):
187 delattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR)
188 else:
189 setattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, previous)
192def _resolve_input(args: Tuple[Any, ...], kwargs: Dict[str, Any], path: _InputPath) -> Any:
193 """Resolve one replay input from its forward argument path."""
194 root_kind, root_key = path[0]
195 value = args[root_key] if root_kind == "arg" else kwargs[root_key]
196 for _, token_value in path[1:]:
197 value = value[token_value]
198 return value
201def _materialize_recompute_inputs(
202 entry: _ExcludeCacheEntry,
203 args: Tuple[Any, ...],
204 kwargs: Dict[str, Any],
205) -> None:
206 """Bind used input handles to tensors produced during checkpoint replay."""
207 for binding in entry.input_bindings:
208 if not binding.handle.used:
209 continue
210 tensor = _resolve_input(args, kwargs, binding.path)
211 if not isinstance(tensor, torch.Tensor):
212 raise RuntimeError(
213 "Checkpoint replay did not reproduce a tensor input required by a checkpoint-excluded region"
214 )
215 binding.handle.materialize(tensor.detach())
218def _has_used_input(input_bindings: List[_InputBinding]) -> bool:
219 """Return whether the excluded call saved any marked input."""
220 return any(binding.handle.used for binding in input_bindings)
223@lru_cache(maxsize=1)
224def _get_replay_placeholder() -> Any:
225 """Create a zero-element placeholder returned by an elided SAVE replay."""
226 return torch.empty(0, device="cpu")
229def _make_replay_placeholder_output(tensor_count: int) -> Any:
230 """Create one placeholder leaf for each forward output tensor."""
231 if tensor_count == 0:
232 return ()
233 placeholder = _get_replay_placeholder()
234 if tensor_count == 1:
235 return placeholder
236 return (placeholder,) * tensor_count
239@lru_cache(maxsize=1)
240def _get_recompute_trigger() -> Any:
241 """Create the differentiable zero-element input used by recompute boundaries."""
242 return torch.empty(0, device="cpu", requires_grad=True)
245@lru_cache(maxsize=1)
246def _get_recompute_boundary() -> Any:
247 """Create the autograd Function used to trigger outer checkpoint replay."""
248 class _RecomputeBoundary(torch.autograd.Function):
249 """Trigger the outer checkpoint hook before excluded-region backward."""
251 @staticmethod
252 def forward(ctx: Any, tensor: Any, trigger: Any) -> Any:
253 """Save one zero-element outer-hook dependency and return the tensor unchanged."""
254 ctx.save_for_backward(trigger)
255 return tensor
257 @staticmethod
258 def backward(ctx: Any, grad_output: Any) -> Tuple[Any, None]:
259 """Trigger dependency unpack and pass the gradient through."""
260 _ = ctx.saved_tensors
261 return grad_output, None
263 return _RecomputeBoundary
266def _finalize_save_outputs_impl(
267 output: Any,
268 add_recompute_boundary: bool,
269 invocation_id: Optional[object],
270 tensor_leaf_count: Optional[List[int]],
271 tensor_type: type,
272) -> Any:
273 """Recursively finalize SAVE output leaves."""
274 if isinstance(output, tensor_type):
275 if tensor_leaf_count is not None:
276 tensor_leaf_count[0] += 1
277 if add_recompute_boundary:
278 output = _get_recompute_boundary().apply(output, _get_recompute_trigger())
279 if invocation_id is not None:
280 setattr(output, _SAVE_OUTPUT_SOURCE_ATTR, invocation_id)
281 return output
282 if isinstance(output, list):
283 return [
284 _finalize_save_outputs_impl(item, add_recompute_boundary, invocation_id, tensor_leaf_count, tensor_type)
285 for item in output
286 ]
287 if isinstance(output, tuple):
288 items = [
289 _finalize_save_outputs_impl(item, add_recompute_boundary, invocation_id, tensor_leaf_count, tensor_type)
290 for item in output
291 ]
292 if hasattr(output, "_fields"):
293 return type(output)(*items)
294 return tuple(items)
295 if isinstance(output, dict):
296 return type(output)(
297 (
298 key,
299 _finalize_save_outputs_impl(
300 value,
301 add_recompute_boundary,
302 invocation_id,
303 tensor_leaf_count,
304 tensor_type,
305 ),
306 )
307 for key, value in output.items()
308 )
309 return output
312def _finalize_save_outputs(
313 output: Any,
314 add_recompute_boundary: bool,
315 invocation_id: Optional[object],
316 tensor_leaf_count: Optional[List[int]] = None,
317) -> Any:
318 """Apply the required boundary and SAVE provenance to output tensor leaves."""
319 return _finalize_save_outputs_impl(
320 output,
321 add_recompute_boundary,
322 invocation_id,
323 tensor_leaf_count,
324 torch.Tensor,
325 )
328class CheckpointExcludeWrapper(ActivationWrapper):
329 """Exclude a callable region from checkpoint recomputation."""
331 def __init__(self, module: Callable[..., Any], *, save_output: bool = True) -> None:
332 """Initialize a checkpoint exclusion wrapper for a PyTorch module or function."""
333 if not callable(module):
334 raise ValueError("module must be a PyTorch Module or callable")
335 if not isinstance(save_output, bool):
336 raise ValueError(f"save_output must be a bool, got {type(save_output).__name__}")
337 super().__init__(module, track_overlaps=False)
338 self.save_output = save_output
340 def forward(self, *args: Any, **kwargs: Any) -> Any:
341 """Execute normally outside recompute and return the cached output in recompute."""
342 wrapped_module = cast(Callable[..., Any], self._wrapped_module)
343 state = get_recompute_state()
344 if state is None:
345 return wrapped_module(*args, **kwargs)
346 cache = state.get_resource(_EXCLUDE_CACHE_KEY, _ExcludeCache)
347 if state.is_recomputing:
348 entry = cache.pop(id(self))
349 _materialize_recompute_inputs(entry, args, kwargs)
350 output = (
351 entry.output
352 if self.save_output
353 else _make_replay_placeholder_output(entry.output_tensor_count)
354 )
355 return _finalize_save_outputs(output, _has_used_input(entry.input_bindings), None)
357 input_bindings, previous_handles = _mark_recompute_inputs(state.invocation_id, args, kwargs)
358 try:
359 with _saved_tensors_context():
360 output = wrapped_module(*args, **kwargs)
361 finally:
362 _restore_recompute_inputs(previous_handles)
363 needs_recompute_boundary = _has_used_input(input_bindings)
364 tensor_leaf_count = None if self.save_output else [0]
365 finalized_output = _finalize_save_outputs(
366 output,
367 needs_recompute_boundary,
368 state.invocation_id,
369 tensor_leaf_count,
370 )
371 replay_output = output if self.save_output else None
372 output_tensor_count = 1 if tensor_leaf_count is None else tensor_leaf_count[0]
373 cache.save(id(self), _ExcludeCacheEntry(replay_output, input_bindings, output_tensor_count))
374 return finalized_output
377def checkpoint_exclude_wrapper(
378 module: Callable[..., Any],
379 *,
380 save_output: bool = True,
381) -> CheckpointExcludeWrapper:
382 """Wrap a PyTorch module or function so its region is not recomputed.
384 Args:
385 module: PyTorch Module or callable to execute only during the original
386 checkpoint forward pass.
387 save_output: Whether to retain the region output for checkpoint replay.
388 Set this to ``False`` only when the output is passed directly as one
389 argument to another checkpoint exclusion wrapper. Default: ``True``.
391 Returns:
392 A wrapper that saves the callable's autograd tensors and reuses its
393 forward output while replaying a non-reentrant checkpoint.
395 Note:
396 This feature requires eager mode and a surrounding HyperParallel
397 checkpoint configured with ``use_reentrant=False``. Nested checkpoint
398 exclusion wrappers are not supported.
399 """
400 return CheckpointExcludeWrapper(module, save_output=save_output)