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 

71class _CachingMindSporeDispatchMode(MsDispatchMode): 

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

73 self.policy_fn = policy_fn 

74 self.swap_storage = swap_storage 

75 self.storage = storage 

76 self.add_to_storage = False 

77 self.group_swap = group_swap 

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

79 self._swap_manager = SwapManager() 

80 self._group_prefix = "" 

81 

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

83 kwargs = {} if kwargs is None else kwargs 

84 if func.name in SAC_IGNORED_OPS: 

85 return func(*args, **kwargs) 

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

87 func, *args, **kwargs) 

88 

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

90 

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

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

93 platform.tree_map( 

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

95 ) 

96 ) 

97 elif policy == CheckpointPolicy.MUST_SWAP: 

98 if not self.add_to_storage: 

99 group_name = self._swap_manager.get_current_group_name() 

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

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

102 self.add_to_storage = True 

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

104 group_swap = self.group_swap 

105 entries = platform.tree_map( 

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

107 ) 

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

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

110 ) 

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

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

113 ) 

114 elif policy != CheckpointPolicy.MUST_RECOMPUTE: 

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

116 return out 

117 

118 

119class _CachedMindSporeDispatchMode(MsDispatchMode): 

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

121 self.policy_fn = policy_fn 

122 self.swap_storage = swap_storage 

123 self.storage = storage 

124 self.allow_cache_entry_mutation = allow_cache_entry_mutation 

125 self._swap_cleared = False 

126 

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

128 kwargs = {} if kwargs is None else kwargs 

129 if func.name in SAC_IGNORED_OPS: 

130 return func(*args, **kwargs) 

131 

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

133 func, *args, **kwargs) 

134 

135 if not self._swap_cleared: 

136 self.swap_storage.clear() 

137 self._swap_cleared = True 

138 

139 # MUST_SAVE and MUST_SWAP both restore from storage identically. 

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

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

142 if storage is None: 

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

144 if len(storage) == 0: 

145 raise RuntimeError( 

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

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

148 ) 

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

150 else: 

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

152 return out 

153 

154 

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

156 if policy_fn_or_list is None: 

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

158 return CheckpointPolicy.PREFER_RECOMPUTE 

159 elif callable(policy_fn_or_list): 

160 policy_fn = policy_fn_or_list 

161 else: 

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

163 

164 swap_storage = Storage() 

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

166 return ( 

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

168 _CachedMindSporeDispatchMode(policy_fn, swap_storage, storage, allow_cache_entry_mutation) 

169 )