Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / activation_checkpoint / checkpoint_exclude_wrapper.py: 98%

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 typing import Any, Callable, Deque, Dict 

18 

19from hyper_parallel.core.activation_checkpoint.recompute_state import get_recompute_state 

20from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper 

21 

22 

23class _CheckpointExcludeCache: 

24 """Store excluded-region outputs for one checkpoint invocation.""" 

25 

26 def __init__(self) -> None: 

27 """Initialize an empty per-checkpoint output cache.""" 

28 self._outputs: Dict[int, Deque[Any]] = defaultdict(deque) 

29 

30 def save(self, wrapper_id: int, output: Any) -> None: 

31 """Save one output produced by a checkpoint-excluded region.""" 

32 self._outputs[wrapper_id].append(output) 

33 

34 def pop(self, wrapper_id: int) -> Any: 

35 """Return the matching forward output during recomputation.""" 

36 outputs = self._outputs.get(wrapper_id) 

37 if not outputs: 

38 raise RuntimeError("No cached forward output is available for this checkpoint exclusion wrapper") 

39 output = outputs.popleft() 

40 if not outputs: 

41 self._outputs.pop(wrapper_id) 

42 return output 

43 

44 def clear(self) -> None: 

45 """Release outputs not consumed because recomputation stopped early.""" 

46 self._outputs.clear() 

47 

48 

49def _pack_saved_tensor(tensor: Any) -> Any: 

50 """Return the tensor data without retaining the input tensor object.""" 

51 return tensor.data 

52 

53 

54def _unpack_saved_tensor(tensor: Any) -> Any: 

55 """Restore the saved tensor for backward.""" 

56 return tensor 

57 

58 

59def _saved_tensors_context() -> Any: 

60 """Create an inner hook that stores real tensors instead of placeholders.""" 

61 import mindspore as ms # pylint: disable=C0415 

62 return ms.saved_tensors_hooks(_pack_saved_tensor, _unpack_saved_tensor) 

63 

64 

65_CHECKPOINT_EXCLUDE_CACHE_KEY = object() 

66 

67 

68class CheckpointExcludeWrapper(ActivationWrapper): 

69 """Exclude a callable region from checkpoint recomputation.""" 

70 

71 def __init__(self, module: Callable[..., Any]) -> None: 

72 """Initialize a checkpoint exclusion wrapper for a MindSpore Cell or function.""" 

73 if not callable(module): 

74 raise ValueError("module must be a MindSpore Cell or callable") 

75 super().__init__(module, track_overlaps=False) 

76 

77 def construct(self, *args: Any, **kwargs: Any) -> Any: 

78 """Execute normally outside recompute and return the cached output in recompute.""" 

79 state = get_recompute_state() 

80 if state is None: 

81 return self._ckpt_wrapped_module(*args, **kwargs) 

82 cache = state.get_resource(_CHECKPOINT_EXCLUDE_CACHE_KEY, _CheckpointExcludeCache) 

83 if state.is_recomputing: 

84 return cache.pop(id(self)) 

85 

86 with _saved_tensors_context(): 

87 output = self._ckpt_wrapped_module(*args, **kwargs) 

88 cache.save(id(self), output) 

89 return output 

90 

91 

92def checkpoint_exclude_wrapper(module: Callable[..., Any]) -> CheckpointExcludeWrapper: 

93 """Wrap a MindSpore Cell or function so its region is not recomputed. 

94 

95 Args: 

96 module: MindSpore Cell or callable to execute only during the original 

97 checkpoint forward pass. 

98 

99 Returns: 

100 A wrapper that saves the callable's autograd tensors and reuses its 

101 forward output while replaying a non-reentrant checkpoint. 

102 

103 Note: 

104 This feature requires MindSpore PyNative mode and a surrounding 

105 HyperParallel checkpoint configured with ``use_reentrant=False``. 

106 """ 

107 return CheckpointExcludeWrapper(module)