« 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"""MindSpore 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, Deque, Dict, List, Optional, Tuple
21import mindspore as ms
22from mindspore.common._grad_function import _Function
24from hyper_parallel.core.activation_checkpoint.recompute_state import get_recompute_state
25from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper
28_RECOMPUTE_INPUT_HANDLE_KEY = "hyper_parallel_recompute_input_handle"
29_SAVE_OUTPUT_SOURCE_KEY = "hyper_parallel_save_output_source"
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 = tensor._get_user_data(_RECOMPUTE_INPUT_HANDLE_KEY) # pylint: disable=protected-access
103 if isinstance(handle, _RecomputedInputHandle):
104 handle.mark_used()
105 return handle
106 return tensor.data
109def _unpack_saved_tensor(tensor: Any) -> Any:
110 """Restore the saved tensor for backward."""
111 if isinstance(tensor, _RecomputedInputHandle):
112 return tensor.get_recomputed_tensor()
113 return tensor
116def _saved_tensors_context() -> Any:
117 """Create an inner hook that stores real tensors instead of placeholders."""
118 return ms.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) -> None:
129 """Append tensor leaves without creating a self-referential local function."""
130 if isinstance(value, ms.Tensor):
131 leaves.append((path, value))
132 return
133 if isinstance(value, (tuple, list)):
134 for index, item in enumerate(value):
135 _append_tensor_inputs(item, path + (("index", index),), leaves)
136 return
137 if isinstance(value, dict):
138 for key, item in value.items():
139 _append_tensor_inputs(item, path + (("key", key),), leaves)
142def _collect_tensor_inputs(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> List[_TensorInput]:
143 """Return tensor leaves and self-describing paths from one excluded-region call."""
144 leaves = []
146 for index, arg in enumerate(args):
147 _append_tensor_inputs(arg, (("arg", index),), leaves)
148 for key, value in kwargs.items():
149 _append_tensor_inputs(value, (("kwarg", key),), leaves)
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, ms.Parameter)
166 or id(tensor) in seen_tensor_ids
167 or tensor._get_user_data(_SAVE_OUTPUT_SOURCE_KEY) is invocation_id # pylint: disable=protected-access
168 ):
169 continue
170 handle = _RecomputedInputHandle()
171 previous = tensor._get_user_data(_RECOMPUTE_INPUT_HANDLE_KEY) # pylint: disable=protected-access
172 previous_handles.append((tensor, previous))
173 tensor._set_user_data(_RECOMPUTE_INPUT_HANDLE_KEY, handle) # pylint: disable=protected-access
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 user data overwritten for one excluded call."""
184 for tensor, previous in previous_handles:
185 tensor._set_user_data(_RECOMPUTE_INPUT_HANDLE_KEY, previous) # pylint: disable=protected-access
188def _resolve_input(args: Tuple[Any, ...], kwargs: Dict[str, Any], path: _InputPath) -> Any:
189 """Resolve one replay input from its forward argument path."""
190 root_kind, root_key = path[0]
191 value = args[root_key] if root_kind == "arg" else kwargs[root_key]
192 for _, token_value in path[1:]:
193 value = value[token_value]
194 return value
197def _materialize_recompute_inputs(
198 entry: _ExcludeCacheEntry,
199 args: Tuple[Any, ...],
200 kwargs: Dict[str, Any],
201) -> None:
202 """Bind used input handles to tensors produced during checkpoint replay."""
203 for binding in entry.input_bindings:
204 if not binding.handle.used:
205 continue
206 tensor = _resolve_input(args, kwargs, binding.path)
207 if not isinstance(tensor, ms.Tensor):
208 raise RuntimeError(
209 "Checkpoint replay did not reproduce a tensor input required by a checkpoint-excluded region"
210 )
211 binding.handle.materialize(tensor.data)
214def _has_used_input(input_bindings: List[_InputBinding]) -> bool:
215 """Return whether the excluded call saved any marked input."""
216 return any(binding.handle.used for binding in input_bindings)
219@lru_cache(maxsize=1)
220def _get_replay_placeholder() -> Any:
221 """Create a zero-element placeholder returned by an elided SAVE replay."""
222 return ms.Tensor([], dtype=ms.float32)
225def _make_replay_placeholder_output(tensor_count: int) -> Any:
226 """Create one placeholder leaf for each forward output tensor."""
227 if tensor_count == 0:
228 return ()
229 placeholder = _get_replay_placeholder()
230 if tensor_count == 1:
231 return placeholder
232 return (placeholder,) * tensor_count
235@lru_cache(maxsize=1)
236def _get_recompute_trigger() -> Any:
237 """Create the differentiable zero-element input used by recompute boundaries."""
238 tensor = ms.Tensor([], dtype=ms.float32)
239 tensor.requires_grad_()
240 return tensor
243class _RecomputeBoundary(_Function):
244 """Trigger the outer checkpoint hook before excluded-region backward."""
246 @staticmethod
247 def forward(ctx: Any, tensor: Any, trigger: Any) -> Any:
248 """Save one zero-element outer-hook dependency and return the tensor unchanged."""
249 ctx.save_for_backward(trigger)
250 return tensor
252 @staticmethod
253 def backward(ctx: Any, grad_output: Any) -> Tuple[Any, None]:
254 """Trigger dependency unpack and pass the gradient through."""
255 _ = ctx.saved_tensors
256 return grad_output, None
259def _finalize_save_outputs(
260 output: Any,
261 add_recompute_boundary: bool,
262 invocation_id: Optional[object],
263 tensor_leaf_count: Optional[List[int]] = None,
264) -> Any:
265 """Apply the required boundary and SAVE provenance to output tensor leaves."""
266 if isinstance(output, ms.Tensor):
267 if tensor_leaf_count is not None:
268 tensor_leaf_count[0] += 1
269 if add_recompute_boundary:
270 output = _RecomputeBoundary.apply(output, _get_recompute_trigger())
271 if invocation_id is not None:
272 output._set_user_data(_SAVE_OUTPUT_SOURCE_KEY, invocation_id) # pylint: disable=protected-access
273 return output
274 if isinstance(output, list):
275 return [
276 _finalize_save_outputs(item, add_recompute_boundary, invocation_id, tensor_leaf_count)
277 for item in output
278 ]
279 if isinstance(output, tuple):
280 items = [
281 _finalize_save_outputs(item, add_recompute_boundary, invocation_id, tensor_leaf_count)
282 for item in output
283 ]
284 if hasattr(output, "_fields"):
285 return type(output)(*items)
286 return tuple(items)
287 if isinstance(output, dict):
288 return type(output)(
289 (key, _finalize_save_outputs(value, add_recompute_boundary, invocation_id, tensor_leaf_count))
290 for key, value in output.items()
291 )
292 return output
295class CheckpointExcludeWrapper(ActivationWrapper):
296 """Exclude a callable region from checkpoint recomputation."""
298 def __init__(self, module: Callable[..., Any], *, save_output: bool = True) -> None:
299 """Initialize a checkpoint exclusion wrapper for a MindSpore Cell or function."""
300 if not callable(module):
301 raise ValueError("module must be a MindSpore Cell or callable")
302 if not isinstance(save_output, bool):
303 raise ValueError(f"save_output must be a bool, got {type(save_output).__name__}")
304 super().__init__(module, track_overlaps=False)
305 self.save_output = save_output
307 def construct(self, *args: Any, **kwargs: Any) -> Any:
308 """Execute normally outside recompute and return the cached output in recompute."""
309 state = get_recompute_state()
310 if state is None:
311 return self._ckpt_wrapped_module(*args, **kwargs)
312 cache = state.get_resource(_EXCLUDE_CACHE_KEY, _ExcludeCache)
313 if state.is_recomputing:
314 entry = cache.pop(id(self))
315 _materialize_recompute_inputs(entry, args, kwargs)
316 output = (
317 entry.output
318 if self.save_output
319 else _make_replay_placeholder_output(entry.output_tensor_count)
320 )
321 return _finalize_save_outputs(output, _has_used_input(entry.input_bindings), None)
323 input_bindings, previous_handles = _mark_recompute_inputs(state.invocation_id, args, kwargs)
324 try:
325 with _saved_tensors_context():
326 output = self._ckpt_wrapped_module(*args, **kwargs)
327 finally:
328 _restore_recompute_inputs(previous_handles)
329 needs_recompute_boundary = _has_used_input(input_bindings)
330 tensor_leaf_count = None if self.save_output else [0]
331 finalized_output = _finalize_save_outputs(
332 output,
333 needs_recompute_boundary,
334 state.invocation_id,
335 tensor_leaf_count,
336 )
337 replay_output = output if self.save_output else None
338 output_tensor_count = 1 if tensor_leaf_count is None else tensor_leaf_count[0]
339 cache.save(id(self), _ExcludeCacheEntry(replay_output, input_bindings, output_tensor_count))
340 return finalized_output
343def checkpoint_exclude_wrapper(
344 module: Callable[..., Any],
345 *,
346 save_output: bool = True,
347) -> CheckpointExcludeWrapper:
348 """Wrap a MindSpore Cell or function so its region is not recomputed.
350 Args:
351 module: MindSpore Cell or callable to execute only during the original
352 checkpoint forward pass.
353 save_output: Whether to retain the region output for checkpoint replay.
354 Set this to ``False`` only when the output is passed directly as one
355 argument to another checkpoint exclusion wrapper. Default: ``True``.
357 Returns:
358 A wrapper that saves the callable's autograd tensors and reuses its
359 forward output while replaying a non-reentrant checkpoint.
361 Note:
362 This feature requires MindSpore PyNative mode and a surrounding
363 HyperParallel checkpoint configured with ``use_reentrant=False``.
364 Nested checkpoint exclusion wrappers are not supported.
365 """
366 return CheckpointExcludeWrapper(module, save_output=save_output)