Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / checkpoint_exclude.py: 88%
392 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"""Saved-tensor helpers for regions excluded from checkpoint recomputation."""
16import weakref
17from collections import defaultdict, deque
18from contextvars import ContextVar
19from dataclasses import dataclass
20from functools import cached_property, lru_cache
21from typing import Any, Deque, Dict, List, Optional, Tuple
23import torch
26_InputPath = Tuple[Tuple[str, Any], ...]
27_TensorInput = Tuple[_InputPath, Any]
30@dataclass(frozen=True)
31class _TensorMetadata:
32 """Alias-relevant metadata captured while an excluded call is running."""
34 tensor_id: int
35 storage_ref: Any
36 storage_nbytes: int
37 dtype: torch.dtype
38 shape: Tuple[int, ...]
39 stride: Tuple[int, ...]
40 storage_offset: int
41 version: int
42 itemsize: int
43 numel: int
44 is_contiguous: bool
46 @cached_property
47 def layout_key(self) -> Tuple[Any, ...]:
48 """Return the dtype/layout identity used to validate a rebuilt alias."""
49 return (self.dtype, self.shape, self.stride, self.version, self.is_contiguous)
52@dataclass(frozen=True)
53class _InputInfo:
54 """Describe one canonical local input without retaining its tensor/storage."""
56 path: _InputPath
57 metadata: _TensorMetadata
60@dataclass(frozen=True)
61class _ViewRecipe:
62 """Rebuild one saved input alias from its matching replay input."""
64 shape: Tuple[int, ...]
65 stride: Tuple[int, ...]
66 relative_offset: int
67 dtype: torch.dtype
68 base_shape: Tuple[int, ...]
69 base_stride: Tuple[int, ...]
70 base_version: int
71 base_is_contiguous: bool
72 exact_input: bool
74 @cached_property
75 def base_layout_key(self) -> Tuple[Any, ...]:
76 """Return the replay-input layout identity this recipe was built from."""
77 return (self.dtype, self.base_shape, self.base_stride, self.base_version, self.base_is_contiguous)
80class _RecomputedInputHandle:
81 """Defer one excluded-region saved alias until checkpoint replay."""
83 def __init__(self, recipe: _ViewRecipe) -> None:
84 """Initialize an unused and unresolved handle."""
85 self.recipe = recipe
86 self._tensor = None
88 def materialize(self, tensor: Any) -> None:
89 """Bind the handle to the matching input produced during replay."""
90 self._tensor = tensor
92 def get_recomputed_tensor(self) -> Any:
93 """Return the replay-produced input for backward."""
94 if self._tensor is None:
95 raise RuntimeError("Checkpoint-excluded input was requested before recomputation")
96 if self._tensor._version != self.recipe.base_version: # pylint: disable=protected-access
97 raise RuntimeError(
98 "Checkpoint-excluded saved tensor was modified in-place after replay materialization"
99 )
100 return self._tensor
103@dataclass(frozen=True)
104class _SavedTensorData:
105 """Hold one detached saved tensor with its save-time version counter."""
107 tensor: torch.Tensor
108 version: int
109 swap_hooks: Optional[Any] = None
112@dataclass(frozen=True)
113class _InputBinding:
114 """Map one replay input path to a deferred saved-alias handle."""
116 path: _InputPath
117 handle: _RecomputedInputHandle
120@dataclass
121class _PackState:
122 """Inputs and matched saved aliases for the active excluded call."""
124 inputs: Tuple[_InputInfo, ...]
125 bindings: List[_InputBinding]
126 swap_hooks: Optional[Any] = None
129_ACTIVE_PACK_STATE: ContextVar[Any] = ContextVar(
130 "hyper_parallel_checkpoint_exclude_pack_state",
131 default=None,
132)
135@dataclass(frozen=True)
136class _DeviceRNGState:
137 """One accelerator device's RNG state captured at a SAVE boundary."""
139 device: torch.device
140 state: Any
143@dataclass(frozen=True)
144class _RNGStateSnapshot:
145 """CPU and accelerator RNG states captured after one excluded call."""
147 cpu_state: Any
148 device_states: Tuple[_DeviceRNGState, ...]
151@dataclass
152class _ExcludeCacheEntry:
153 """Store one excluded call's replay value, deferred inputs, and RNG state."""
155 output: Any
156 input_bindings: List[_InputBinding]
157 rng_state: Optional[_RNGStateSnapshot] = None
158 tape: Any = None
161class _ExcludeCache:
162 """Store excluded-region call entries for one checkpoint invocation."""
164 def __init__(self) -> None:
165 """Initialize an empty per-checkpoint output cache."""
166 self._entries: Dict[int, Deque[_ExcludeCacheEntry]] = defaultdict(deque)
168 def save(self, wrapper_id: int, entry: _ExcludeCacheEntry) -> None:
169 """Save one call entry produced by a checkpoint-excluded region."""
170 self._entries[wrapper_id].append(entry)
172 def pop(self, wrapper_id: int) -> _ExcludeCacheEntry:
173 """Return the matching forward call entry during recomputation."""
174 entries = self._entries.get(wrapper_id)
175 if not entries:
176 raise RuntimeError("No cached forward output is available for this checkpoint exclusion wrapper")
177 entry = entries.popleft()
178 if not entries:
179 self._entries.pop(wrapper_id)
180 return entry
182 def clear(self) -> None:
183 """Release outputs not consumed because recomputation stopped early."""
184 for entries in self._entries.values():
185 for entry in entries:
186 if entry.tape is not None:
187 entry.tape.release()
188 entry.output = None
189 self._entries.clear()
192@lru_cache(maxsize=1)
193def _get_dtensor_type() -> Optional[type]:
194 """Return the local DTensor type without importing it during module import."""
195 try:
196 from hyper_parallel.core.dtensor.dtensor_base import DTensorBase # pylint: disable=C0415
197 except ImportError:
198 return None
199 return DTensorBase
202def _is_dtensor(tensor: Any) -> bool:
203 """Return whether a value is a distributed tensor wrapper."""
204 dtensor_type = _get_dtensor_type()
205 return dtensor_type is not None and isinstance(tensor, dtensor_type)
208def _canonical_input_tensor(tensor: Any) -> Any:
209 """Return the local tensor carried by a DTensor wrapper."""
210 if _is_dtensor(tensor):
211 return tensor.to_local()
212 return tensor
215def _tensor_metadata(tensor: Any) -> Optional[_TensorMetadata]:
216 """Return storage/layout metadata, or ``None`` when aliasing is unavailable."""
217 if not isinstance(tensor, torch.Tensor) or tensor.numel() == 0:
218 return None
219 try:
220 storage = tensor.untyped_storage()
221 storage_nbytes = storage.size()
222 itemsize = tensor.element_size()
223 if storage_nbytes == 0 or itemsize <= 0:
224 return None
225 return _TensorMetadata(
226 tensor_id=id(tensor),
227 storage_ref=weakref.ref(storage),
228 storage_nbytes=storage_nbytes,
229 dtype=tensor.dtype,
230 shape=tuple(tensor.shape),
231 stride=tuple(tensor.stride()),
232 storage_offset=tensor.storage_offset(),
233 version=tensor._version, # pylint: disable=protected-access
234 itemsize=itemsize,
235 numel=tensor.numel(),
236 is_contiguous=tensor.is_contiguous(),
237 )
238 except (AttributeError, RuntimeError, TypeError, ValueError):
239 return None
242def _storage_span(metadata: _TensorMetadata) -> Optional[Tuple[int, int]]:
243 """Return inclusive element offsets touched by a non-negative-stride tensor."""
244 if metadata.numel == 0 or any(step < 0 for step in metadata.stride):
245 return None
246 max_offset = metadata.storage_offset
247 for size, step in zip(metadata.shape, metadata.stride):
248 if size == 0:
249 return None
250 max_offset += (size - 1) * step
251 if metadata.storage_offset < 0 or (max_offset + 1) * metadata.itemsize > metadata.storage_nbytes:
252 return None
253 return metadata.storage_offset, max_offset
256def _exact_layout(saved: _TensorMetadata, base: _TensorMetadata) -> bool:
257 """Return whether two tensors have identical dtype and logical layout."""
258 return (
259 saved.dtype == base.dtype
260 and saved.shape == base.shape
261 and saved.stride == base.stride
262 and saved.storage_offset == base.storage_offset
263 )
266def _is_non_overlapping_and_dense(metadata: _TensorMetadata) -> bool:
267 """Return whether a non-contiguous layout covers every element exactly once."""
268 expected_stride = 1
269 for stride, size in sorted(
270 (stride, size)
271 for size, stride in zip(metadata.shape, metadata.stride)
272 if size > 1
273 ):
274 if stride != expected_stride:
275 return False
276 expected_stride *= size
277 return True
280def _make_exact_recipe(base: _TensorMetadata) -> _ViewRecipe:
281 """Build a replay recipe for an unchanged exact input."""
282 return _ViewRecipe(
283 shape=base.shape,
284 stride=base.stride,
285 relative_offset=0,
286 dtype=base.dtype,
287 base_shape=base.shape,
288 base_stride=base.stride,
289 base_version=base.version,
290 base_is_contiguous=base.is_contiguous,
291 exact_input=True,
292 )
295def _make_view_recipe(saved: _TensorMetadata, base: _TensorMetadata) -> Optional[_ViewRecipe]:
296 """Return a safe exact-input/view recipe, or ``None`` for a real save."""
297 base_storage = base.storage_ref()
298 if base_storage is None or base_storage is not saved.storage_ref():
299 return None
300 if saved.dtype != base.dtype or saved.version != base.version:
301 return None
303 exact_input = _exact_layout(saved, base)
304 if not exact_input:
305 if not base.is_contiguous and not _is_non_overlapping_and_dense(base):
306 return None
307 saved_span = _storage_span(saved)
308 base_span = _storage_span(base)
309 if saved_span is None or base_span is None:
310 return None
311 if saved_span[0] < base_span[0] or saved_span[1] > base_span[1]:
312 return None
314 return _ViewRecipe(
315 shape=saved.shape,
316 stride=saved.stride,
317 relative_offset=saved.storage_offset - base.storage_offset,
318 dtype=saved.dtype,
319 base_shape=base.shape,
320 base_stride=base.stride,
321 base_version=base.version,
322 base_is_contiguous=base.is_contiguous,
323 exact_input=exact_input,
324 )
327def _match_recompute_input(
328 tensor: torch.Tensor,
329 inputs: Tuple[_InputInfo, ...],
330) -> Optional[Tuple[_InputInfo, _ViewRecipe]]:
331 """Match a saved tensor to an excluded-call input by storage and layout."""
332 if _is_dtensor(tensor):
333 return None
334 for input_info in inputs:
335 base = input_info.metadata
336 if id(tensor) != base.tensor_id:
337 continue
338 if (
339 tensor.dtype != base.dtype
340 or tuple(tensor.shape) != base.shape
341 or tuple(tensor.stride()) != base.stride
342 or tensor.storage_offset() != base.storage_offset
343 or tensor._version != base.version # pylint: disable=protected-access
344 ):
345 return None
346 return input_info, _make_exact_recipe(base)
348 saved = _tensor_metadata(tensor)
349 if saved is None:
350 return None
352 exact_identity = []
353 exact_layout = []
354 views = []
355 for input_info in inputs:
356 recipe = _make_view_recipe(saved, input_info.metadata)
357 if recipe is None:
358 continue
359 candidate = (input_info, recipe)
360 if recipe.exact_input and saved.tensor_id == input_info.metadata.tensor_id:
361 exact_identity.append(candidate)
362 elif recipe.exact_input:
363 exact_layout.append(candidate)
364 else:
365 views.append(candidate)
367 if exact_identity:
368 return exact_identity[0]
369 if exact_layout:
370 return exact_layout[0]
371 if views:
372 return min(views, key=lambda item: item[0].metadata.numel)
373 return None
376def _pack_saved_tensor(tensor: Any) -> Any:
377 """Return a replayable input-alias handle or detached tensor data."""
378 pack_state = _ACTIVE_PACK_STATE.get()
379 if pack_state is not None:
380 if (
381 isinstance(tensor, torch.nn.Parameter)
382 or isinstance(tensor._base, torch.nn.Parameter) # pylint: disable=protected-access
383 ):
384 return _SavedTensorData(tensor.detach(), tensor._version) # pylint: disable=protected-access
385 matched = _match_recompute_input(tensor, pack_state.inputs)
386 if matched is not None:
387 input_info, recipe = matched
388 handle = _RecomputedInputHandle(recipe)
389 pack_state.bindings.append(_InputBinding(input_info.path, handle))
390 return handle
391 if pack_state.swap_hooks is not None:
392 detached = pack_state.swap_hooks.pack_hook(tensor)
393 return _SavedTensorData(
394 detached,
395 tensor._version, # pylint: disable=protected-access
396 pack_state.swap_hooks,
397 )
398 return _SavedTensorData(tensor.detach(), tensor._version) # pylint: disable=protected-access
401def _unpack_saved_tensor(value: Any) -> Any:
402 """Restore the saved tensor for backward."""
403 if isinstance(value, _RecomputedInputHandle):
404 return value.get_recomputed_tensor()
405 if isinstance(value, _SavedTensorData):
406 tensor = value.tensor
407 if value.swap_hooks is not None:
408 tensor = value.swap_hooks.unpack_hook(tensor)
409 if tensor._version != value.version: # pylint: disable=protected-access
410 raise RuntimeError("Checkpoint-excluded saved tensor was modified in-place after it was saved")
411 return tensor
412 return value
415def _saved_tensors_context() -> Any:
416 """Create an inner hook that stores real tensors instead of outer holders."""
417 return torch.autograd.graph.saved_tensors_hooks(_pack_saved_tensor, _unpack_saved_tensor)
420_EXCLUDE_CACHE_KEY = object()
423def _append_tensor_inputs(value: Any, path: _InputPath, leaves: List[_TensorInput]) -> None:
424 """Append tensor leaves without creating a self-referential local function."""
425 if isinstance(value, torch.Tensor):
426 leaves.append((path, value))
427 return
428 if isinstance(value, (tuple, list)):
429 for index, item in enumerate(value):
430 _append_tensor_inputs(item, path + (("index", index),), leaves)
431 return
432 if isinstance(value, dict):
433 for key, item in value.items():
434 _append_tensor_inputs(item, path + (("key", key),), leaves)
437def _collect_tensor_inputs(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> List[_TensorInput]:
438 """Return tensor leaves and self-describing paths from one excluded-region call."""
439 leaves = []
440 for index, arg in enumerate(args):
441 _append_tensor_inputs(arg, (("arg", index),), leaves)
442 for key, value in kwargs.items():
443 _append_tensor_inputs(value, (("kwarg", key),), leaves)
444 return leaves
447def _iter_tensor_leaves(value: Any) -> List[torch.Tensor]:
448 """Return tensor leaves from one structured excluded-region value."""
449 if isinstance(value, torch.Tensor):
450 return [value]
451 if isinstance(value, (tuple, list)):
452 leaves = []
453 for item in value:
454 leaves.extend(_iter_tensor_leaves(item))
455 return leaves
456 if isinstance(value, dict):
457 leaves = []
458 for item in value.values():
459 leaves.extend(_iter_tensor_leaves(item))
460 return leaves
461 return []
464def _capture_recompute_inputs(
465 args: Tuple[Any, ...],
466 kwargs: Dict[str, Any],
467) -> Tuple[_InputInfo, ...]:
468 """Capture canonical local input storage/layout without retaining tensors."""
469 inputs = []
470 seen_tensor_ids = set()
471 for path, tensor in _collect_tensor_inputs(args, kwargs):
472 if isinstance(tensor, torch.nn.Parameter):
473 continue
474 canonical = _canonical_input_tensor(tensor)
475 if id(canonical) in seen_tensor_ids:
476 continue
477 metadata = _tensor_metadata(canonical)
478 if metadata is None:
479 continue
480 inputs.append(_InputInfo(path, metadata))
481 seen_tensor_ids.add(id(canonical))
482 return tuple(inputs)
485def _resolve_input(args: Tuple[Any, ...], kwargs: Dict[str, Any], path: _InputPath) -> Any:
486 """Resolve one replay input from its forward argument path."""
487 root_kind, root_key = path[0]
488 value = args[root_key] if root_kind == "arg" else kwargs[root_key]
489 for _, token_value in path[1:]:
490 value = value[token_value]
491 return value
494def _rebuild_saved_alias(tensor: Any, recipe: _ViewRecipe) -> Any:
495 """Rebuild a detached saved input/view from one replay-produced base."""
496 canonical = _canonical_input_tensor(tensor)
497 metadata = _tensor_metadata(canonical)
498 if metadata is None:
499 raise RuntimeError("Checkpoint replay input does not expose usable storage metadata")
500 if metadata.layout_key != recipe.base_layout_key:
501 raise RuntimeError(
502 "Checkpoint replay input layout/version changed for a checkpoint-excluded saved tensor"
503 )
505 if recipe.exact_input:
506 return canonical.detach()
508 view_offset = metadata.storage_offset + recipe.relative_offset
509 max_offset = view_offset
510 numel = 1
511 for size, step in zip(recipe.shape, recipe.stride):
512 if size == 0:
513 raise RuntimeError("Checkpoint-excluded empty saved views should have been saved normally")
514 numel *= size
515 max_offset += (size - 1) * step
516 if numel == 0 or view_offset < 0 or (max_offset + 1) * metadata.itemsize > metadata.storage_nbytes:
517 raise RuntimeError("Checkpoint-excluded saved view is outside the replay input storage")
519 return canonical.as_strided(recipe.shape, recipe.stride, view_offset).detach()
522def _materialize_recompute_inputs(
523 entry: _ExcludeCacheEntry,
524 args: Tuple[Any, ...],
525 kwargs: Dict[str, Any],
526) -> None:
527 """Bind saved-alias handles to tensors rebuilt from checkpoint replay inputs."""
528 for binding in entry.input_bindings:
529 tensor = _resolve_input(args, kwargs, binding.path)
530 if not isinstance(tensor, torch.Tensor):
531 raise RuntimeError(
532 "Checkpoint replay did not reproduce a tensor input required by a checkpoint-excluded region"
533 )
534 binding.handle.materialize(_rebuild_saved_alias(tensor, binding.handle.recipe))
537def _has_used_input(input_bindings: List[_InputBinding]) -> bool:
538 """Return whether the excluded call saved any replayable input alias."""
539 return bool(input_bindings)
542def _collect_rng_devices(args: Tuple[Any, ...], kwargs: Dict[str, Any], output: Any) -> Tuple[torch.device, ...]:
543 """Return distinct accelerator devices touched by one excluded call."""
544 devices = []
545 seen = set()
546 tensors = _iter_tensor_leaves((args, kwargs))
547 tensors.extend(_iter_tensor_leaves(output))
548 for tensor in tensors:
549 device = _canonical_input_tensor(tensor).device
550 if device.type in {"cpu", "meta"}:
551 continue
552 key = (device.type, device.index)
553 if key in seen:
554 continue
555 devices.append(device)
556 seen.add(key)
557 return tuple(devices)
560def _get_device_rng_state(device: torch.device) -> Any:
561 """Read one accelerator RNG state when its backend exposes the API."""
562 device_module = getattr(torch, device.type, None)
563 get_rng_state = getattr(device_module, "get_rng_state", None)
564 if not callable(get_rng_state):
565 return None
566 try:
567 return get_rng_state(device)
568 except TypeError:
569 return get_rng_state()
572def _set_device_rng_state(device: torch.device, state: Any) -> None:
573 """Restore one accelerator RNG state when its backend exposes the API."""
574 device_module = getattr(torch, device.type, None)
575 set_rng_state = getattr(device_module, "set_rng_state", None)
576 if not callable(set_rng_state):
577 return
578 try:
579 set_rng_state(state, device)
580 except TypeError:
581 set_rng_state(state)
584def _capture_rng_state(
585 args: Tuple[Any, ...],
586 kwargs: Dict[str, Any],
587 output: Any,
588 enabled: bool,
589) -> Optional[_RNGStateSnapshot]:
590 """Snapshot RNG at a SAVE exit so replay can reproduce its state advancement."""
591 if not enabled:
592 return None
593 device_states = []
594 for device in _collect_rng_devices(args, kwargs, output):
595 state = _get_device_rng_state(device)
596 if state is not None:
597 device_states.append(_DeviceRNGState(device, state))
598 return _RNGStateSnapshot(torch.get_rng_state(), tuple(device_states))
601def _restore_rng_state(snapshot: Optional[_RNGStateSnapshot]) -> None:
602 """Restore the RNG state that the skipped SAVE body would have produced."""
603 if snapshot is None:
604 return
605 torch.set_rng_state(snapshot.cpu_state)
606 for device_state in snapshot.device_states:
607 _set_device_rng_state(device_state.device, device_state.state)
610@lru_cache(maxsize=1)
611def _get_recompute_trigger() -> Any:
612 """Create the differentiable zero-element input used by recompute boundaries."""
613 return torch.empty(0, device="cpu", requires_grad=True)
616@lru_cache(maxsize=1)
617def _get_recompute_boundary() -> Any:
618 """Create the autograd Function used to trigger outer checkpoint replay."""
619 class _RecomputeBoundary(torch.autograd.Function):
620 """Trigger the outer checkpoint hook before excluded-region backward."""
622 @staticmethod
623 def forward(ctx: Any, tensor: Any, trigger: Any) -> Any:
624 """Save one zero-element outer-hook dependency and return the tensor unchanged."""
625 ctx.save_for_backward(trigger)
626 return tensor
628 @staticmethod
629 def backward(ctx: Any, grad_output: Any) -> Tuple[Any, None]:
630 """Trigger dependency unpack and pass the gradient through."""
631 _ = ctx.saved_tensors
632 return grad_output, None
634 return _RecomputeBoundary
637def _finalize_save_outputs_impl(
638 output: Any,
639 add_recompute_boundary: bool,
640 tensor_type: type,
641) -> Any:
642 """Recursively finalize SAVE output leaves."""
643 if isinstance(output, tensor_type):
644 if add_recompute_boundary:
645 output = _get_recompute_boundary().apply(output, _get_recompute_trigger())
646 return output
647 if isinstance(output, list):
648 return [
649 _finalize_save_outputs_impl(item, add_recompute_boundary, tensor_type)
650 for item in output
651 ]
652 if isinstance(output, tuple):
653 items = [
654 _finalize_save_outputs_impl(item, add_recompute_boundary, tensor_type)
655 for item in output
656 ]
657 if hasattr(output, "_fields"):
658 return type(output)(*items)
659 return tuple(items)
660 if isinstance(output, dict):
661 return type(output)(
662 (
663 key,
664 _finalize_save_outputs_impl(value, add_recompute_boundary, tensor_type),
665 )
666 for key, value in output.items()
667 )
668 return output
671def _finalize_save_outputs(
672 output: Any,
673 add_recompute_boundary: bool,
674) -> Any:
675 """Apply the required replay boundary to tensor output leaves."""
676 return _finalize_save_outputs_impl(
677 output,
678 add_recompute_boundary,
679 torch.Tensor,
680 )