Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / activation_checkpoint / sac.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"""enhanced with selective checkpoint support swap""" 

16# pylint: disable=W0212, W0613, C0115, C0116, C0103, R1705 

17from collections import defaultdict 

18from typing import Any, Dict, List, Optional, Union 

19 

20import mindspore as ms 

21from mindspore import MsDispatchMode 

22from hyper_parallel.core.activation_checkpoint.swap import ( 

23 SwapManager, 

24 Storage, 

25 SwapTensor, 

26) 

27from hyper_parallel.core.activation_checkpoint.activation_checkpoint import CheckpointPolicy 

28from hyper_parallel.platform import get_platform 

29 

30platform = get_platform() 

31 

32 

33class _VersionWrapper: 

34 # Check that cached tensors are not mutated. 

35 def __init__(self, val): 

36 self.val: Union[ms.Tensor, Any] = val 

37 self.version: Optional[int] = ( 

38 val._version if isinstance(val, ms.Tensor) else None 

39 ) 

40 

41 def get_val(self, allow_cache_entry_mutation): 

42 if self.version is not None and not allow_cache_entry_mutation: 

43 if self.val._version != self.version: 

44 # Can we give user a stack trace of where the mutation happened? 

45 raise RuntimeError( 

46 "Tensor cached during selective activation checkpoint has been mutated" 

47 ) 

48 return self.val 

49 

50 

51class _SwapCacheEntry: 

52 """Pair the recompute cache and swap record around the same tensor object.""" 

53 def __init__(self, val, funcname, group_swap=False): 

54 self.save = _VersionWrapper(val) 

55 self.swap = SwapTensor(val, funcname, group_swap=group_swap) 

56 

57 

58def _maybe_detach(x): 

59 if isinstance(x, ms.Tensor) and (x.is_floating_point() or x.is_complex()): 

60 x = x.detach() 

61 return x 

62 

63 

64class SelectiveCheckpointContext: 

65 def __init__(self, *, is_recompute): 

66 self.is_recompute = is_recompute 

67 

68SAC_IGNORED_OPS = {"StopGradient"} 

69 

70 

71def ignore_sac_ops(ignore_ops: List[Optional[object]]) -> None: 

72 """Add available operator names to the selective-checkpoint ignore set. 

73 

74 Args: 

75 ops (List[Optional[object]]): MindSpore operator names to execute without selective-AC replay 

76 accounting. ``None`` entries are ignored. 

77 """ 

78 SAC_IGNORED_OPS.update(op for op in ignore_ops if op is not None) 

79 

80 

81class _CachingMindSporeDispatchMode(MsDispatchMode): 

82 def __init__(self, policy_fn, swap_storage, storage, group_swap=False): 

83 self.policy_fn = policy_fn 

84 self.swap_storage = swap_storage 

85 self.storage = storage 

86 self.add_to_storage = False 

87 self.group_swap = group_swap 

88 # Cache context and singleton to avoid per-dispatch allocation / lookup. 

89 self._swap_manager = SwapManager() 

90 self._group_prefix = "" 

91 

92 def __ms_dispatch__(self, func, args=(), kwargs=None): 

93 kwargs = {} if kwargs is None else kwargs 

94 if func.name in SAC_IGNORED_OPS: 

95 return func(*args, **kwargs) 

96 policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False), 

97 func, *args, **kwargs) 

98 

99 out = func(*args, **kwargs) 

100 

101 if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE): 

102 self.storage[func.name].append( 

103 platform.tree_map( 

104 lambda x: _VersionWrapper(_maybe_detach(x)), out 

105 ) 

106 ) 

107 elif policy == CheckpointPolicy.MUST_SWAP: 

108 if not self.add_to_storage: 

109 group_name = self._swap_manager.get_current_group_name() 

110 self._group_prefix = f"{group_name}::" 

111 self._swap_manager.add_storage(group_name, self.swap_storage) 

112 self.add_to_storage = True 

113 funcname = f"{self._group_prefix}{func.name}" 

114 group_swap = self.group_swap 

115 entries = platform.tree_map( 

116 lambda x: _SwapCacheEntry(_maybe_detach(x), funcname, group_swap=group_swap), out 

117 ) 

118 self.storage[func.name].append( 

119 platform.tree_map(lambda x: x.save, entries) 

120 ) 

121 self.swap_storage[func.name].append( 

122 platform.tree_map(lambda x: x.swap, entries) 

123 ) 

124 elif policy != CheckpointPolicy.MUST_RECOMPUTE: 

125 raise RuntimeError(f"Checkpoint Activation: {func.name} encountered an invalid policy {policy}") 

126 return out 

127 

128 

129class _CachedMindSporeDispatchMode(MsDispatchMode): 

130 def __init__(self, policy_fn, swap_storage, storage, allow_cache_entry_mutation): 

131 self.policy_fn = policy_fn 

132 self.swap_storage = swap_storage 

133 self.storage = storage 

134 self.allow_cache_entry_mutation = allow_cache_entry_mutation 

135 self._swap_cleared = False 

136 

137 def __ms_dispatch__(self, func, args=(), kwargs=None): 

138 kwargs = {} if kwargs is None else kwargs 

139 if func.name in SAC_IGNORED_OPS: 

140 return func(*args, **kwargs) 

141 

142 policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=True), 

143 func, *args, **kwargs) 

144 

145 if not self._swap_cleared: 

146 self.swap_storage.clear() 

147 self._swap_cleared = True 

148 

149 # MUST_SAVE and MUST_SWAP both restore from storage identically. 

150 if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE, CheckpointPolicy.MUST_SWAP): 

151 storage = self.storage.get(func.name) 

152 if storage is None: 

153 raise RuntimeError(f"{func} encountered during backward, but not found in storage") 

154 if len(storage) == 0: 

155 raise RuntimeError( 

156 "Trying to backward an extra time. You are only allowed to backward once " 

157 "on any region computed under selective activation checkpoint." 

158 ) 

159 out = platform.tree_map(lambda x: x.get_val(self.allow_cache_entry_mutation), storage.pop(0)) 

160 else: 

161 out = func(*args, **kwargs) 

162 return out 

163 

164 

165def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False): 

166 if policy_fn_or_list is None: 

167 def policy_fn(_ctx, _op, *_args, **_kwargs): 

168 return CheckpointPolicy.PREFER_RECOMPUTE 

169 elif callable(policy_fn_or_list): 

170 policy_fn = policy_fn_or_list 

171 else: 

172 raise TypeError("policy_fn_or_list must be either a function or a list of ops.") 

173 

174 swap_storage = Storage() 

175 storage: Dict[Any, List[Any]] = defaultdict(list) 

176 return ( 

177 _CachingMindSporeDispatchMode(policy_fn, swap_storage, storage, group_swap=group_swap), 

178 _CachedMindSporeDispatchMode(policy_fn, swap_storage, storage, allow_cache_entry_mutation) 

179 )