Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / activation_checkpoint / sac.py: 96%

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# Adapted from https://github.com/pytorch/pytorch/blob/release/2.6/torch/utils/checkpoint.py 

16# enhanced with selective checkpoint support swap 

17# ============================================================================ 

18"""enhanced with selective checkpoint support swap""" 

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

20from collections import defaultdict 

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

22 

23import torch 

24import torch.fx.traceback as fx_traceback 

25from torch._functorch._aot_autograd.functional_utils import is_fun 

26from torch.utils._pytree import tree_map 

27from torch.utils._python_dispatch import TorchDispatchMode 

28from hyper_parallel.core.activation_checkpoint import CheckpointPolicy # patch code 

29from hyper_parallel.core.activation_checkpoint.swap import ( # patch code 

30 SwapManager, 

31 SwapTensor, 

32 Storage, 

33) 

34 

35 

36def _is_compiling(func, args, kwargs): 

37 # Check if we are under AOTAutograd tracing 

38 # There should probably be a better way to do this... 

39 # NOTE: unify _is_compiling across all compile stacks 

40 for arg in args: 

41 if isinstance(arg, torch.Tensor) and is_fun(arg): 

42 return True 

43 return False 

44 

45 

46class _VersionWrapper: 

47 # Check that cached tensors are not mutated. 

48 def __init__(self, val): 

49 self.val: Union[torch.Tensor, Any] = val 

50 self.version: Optional[int] = ( 

51 val._version if isinstance(val, torch.Tensor) else None 

52 ) 

53 

54 def get_val(self, allow_cache_entry_mutation): 

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

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

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

58 raise RuntimeError( 

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

60 ) 

61 return self.val 

62 

63 

64class _SwapCacheEntry: 

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

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

67 self.save = _VersionWrapper(val) 

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

69 

70 

71def _maybe_detach(x, any_ret_has_alias_info): 

72 # We detach for two separate reasons: 

73 # - For view ops, we need to ensure that when the tensor is returned from 

74 # CachedDispatchMode, as_view sees that the AutogradMeta is nullptr 

75 # - Avoid reference cycles 

76 # For case 1, it is not enough to check whether x has differentiable dtype 

77 # because non-differentiable dtype can have non-nullptr AutogradMeta, e.g. 

78 # when the tensor is a view. 

79 need_detach = (isinstance(x, torch.Tensor) 

80 and (x.is_floating_point() or x.is_complex() or any_ret_has_alias_info)) 

81 if need_detach: 

82 with torch._C._SetExcludeDispatchKeyGuard(torch._C.DispatchKey.ADInplaceOrView, False): 

83 # Ensure that view performed beneath autograd properly propagates 

84 # version counter. TODO: Use reentrant_dispatch instead of 

85 # manually manipulating dispatch keys. Using reentrant_dispatch 

86 # would respect inference_mode, though that is not relevant for 

87 # this case. 

88 x = x.detach() 

89 return x 

90 

91 

92class SelectiveCheckpointContext: 

93 """ 

94 Context passed to policy function during selective checkpointing. 

95 

96 This class is used to pass relevant metadata to the policy function during 

97 selective checkpointing. The metadata includes whether the current invocation 

98 of the policy function is during recomputation or not. 

99 

100 Example: 

101 >>> # xdoctest: +SKIP(stub) 

102 >>> 

103 >>> def policy_fn(ctx, op, *args, **kwargs): 

104 >>> print(ctx.is_recompute) 

105 >>> 

106 >>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) 

107 >>> 

108 >>> out = torch.utils.checkpoint.checkpoint( 

109 >>> fn, x, y, 

110 >>> use_reentrant=False, 

111 >>> context_fn=context_fn, 

112 >>> ) 

113 """ 

114 def __init__(self, *, is_recompute): 

115 self.is_recompute = is_recompute 

116 

117 

118def _policy_from_bool(b): 

119 # For backward compatibility 

120 return CheckpointPolicy.MUST_SAVE if b else CheckpointPolicy.PREFER_RECOMPUTE 

121 

122 

123SAC_IGNORED_OPS = { 

124 # AC inserts different number of detach during forward and recompute. 

125 torch.ops.aten.detach.default, 

126 # AC's determinism check invokes additional metadata ops during forward. 

127 # With subclasses involved, these metadata ops become dispatchable, this 

128 # can result in incorrectness if these ops are selected cached. 

129 torch.ops.prim.device.default, 

130} | set(torch._subclasses.functional_tensor.FunctionalTensor.metadata_fns) 

131 

132 

133class _CachingTorchDispatchMode(TorchDispatchMode): 

134 # Used together with _CachedTorchDispatchMode to implement SAC. 

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

136 self.policy_fn = policy_fn 

137 self.swap_storage = swap_storage 

138 self.storage = storage 

139 self.add_to_storage = False 

140 self.group_swap = group_swap 

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

142 self._swap_manager = SwapManager() 

143 self._group_prefix = "" 

144 

145 def __torch_dispatch__(self, func, types, args=(), kwargs=None): 

146 if func in SAC_IGNORED_OPS: 

147 return func(*args, **kwargs) 

148 

149 kwargs = {} if kwargs is None else kwargs 

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

151 func, *args, **kwargs) 

152 if isinstance(policy, bool): 

153 policy = _policy_from_bool(policy) 

154 

155 is_compiling = _is_compiling(func, args, kwargs) 

156 

157 if is_compiling: 

158 # Overwrite each node's "recompute" tag to add in the user annotation. 

159 fx_traceback.current_meta["recompute"] = policy 

160 

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

162 

163 has_alias = any(ret.alias_info is not None for ret in func._schema.returns) 

164 

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

166 self.storage[func].append( 

167 tree_map(lambda x: _VersionWrapper(_maybe_detach(x, has_alias)), out) 

168 ) 

169 elif policy == CheckpointPolicy.MUST_SWAP: # patch code 

170 if not self.add_to_storage: 

171 group_name = self._swap_manager.get_current_group_name() 

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

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

174 self.add_to_storage = True 

175 funcname = f"{self._group_prefix}{func}" 

176 group_swap = self.group_swap 

177 entries = tree_map( 

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

179 ) 

180 self.storage[func].append(tree_map(lambda x: x.save, entries)) 

181 self.swap_storage[func].append(tree_map(lambda x: x.swap, entries)) 

182 elif policy != CheckpointPolicy.MUST_RECOMPUTE: 

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

184 return out 

185 

186 

187class _CachedTorchDispatchMode(TorchDispatchMode): 

188 # Used together with _CachingTorchDispatchMode to implement SAC. 

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

190 self.policy_fn = policy_fn 

191 self.swap_storage = swap_storage 

192 self.storage = storage 

193 self.allow_cache_entry_mutation = allow_cache_entry_mutation 

194 self._swap_cleared = False 

195 

196 def __torch_dispatch__(self, func, types, args=(), kwargs=None): 

197 if func in SAC_IGNORED_OPS: 

198 return func(*args, **kwargs) 

199 

200 kwargs = {} if kwargs is None else kwargs 

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

202 func, *args, **kwargs) 

203 if isinstance(policy, bool): 

204 policy = _policy_from_bool(policy) 

205 

206 is_compiling = _is_compiling(func, args, kwargs) 

207 

208 if not self._swap_cleared: 

209 self.swap_storage.clear() 

210 self._swap_cleared = True 

211 

212 # MUST_SAVE, PREFER_SAVE, and MUST_SWAP all restore from storage identically. 

213 if (policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE, CheckpointPolicy.MUST_SWAP) 

214 or is_compiling): 

215 storage = self.storage.get(func) # patch code 

216 if storage is None: 

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

218 if len(storage) == 0: 

219 raise RuntimeError( 

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

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

222 ) 

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

224 else: 

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

226 return out 

227 

228 

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

230 """ 

231 Helper to avoid recomputing certain ops during activation checkpointing. 

232 

233 Use this with `torch.utils.checkpoint.checkpoint` to control which 

234 operations are recomputed during the backward pass. 

235 

236 Args: 

237 policy_fn_or_list (Callable or List): 

238 - If a policy function is provided, it should accept a 

239 :class:`SelectiveCheckpointContext`, the :class:`OpOverload`, args and 

240 kwargs to the op, and return a :class:`CheckpointPolicy` enum value 

241 indicating whether the execution of the op should be recomputed or not. 

242 - If a list of operations is provided, it is equivalent to a policy 

243 returning `CheckpointPolicy.MUST_SAVE` for the specified 

244 operations and `CheckpointPolicy.PREFER_RECOMPUTE` for all other 

245 operations. 

246 allow_cache_entry_mutation (bool, optional): By default, an error is 

247 raised if any tensors cached by selective activation checkpoint are 

248 mutated in order to ensure correctness. If set to `True`, this check 

249 is disabled. 

250 Returns: 

251 A tuple of two context managers. 

252 

253 Example: 

254 >>> # xdoctest: +REQUIRES(LINUX) 

255 >>> import functools 

256 >>> 

257 >>> x = torch.rand(10, 10, requires_grad=True) 

258 >>> y = torch.rand(10, 10, requires_grad=True) 

259 >>> 

260 >>> ops_to_save = [ 

261 >>> torch.ops.aten.mm.default, 

262 >>> ] 

263 >>> 

264 >>> def policy_fn(ctx, op, *args, **kwargs): 

265 >>> if op in ops_to_save: 

266 >>> return CheckpointPolicy.MUST_SAVE 

267 >>> else: 

268 >>> return CheckpointPolicy.PREFER_RECOMPUTE 

269 >>> 

270 >>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) 

271 >>> 

272 >>> # or equivalently 

273 >>> context_fn = functools.partial(create_selective_checkpoint_contexts, ops_to_save) 

274 >>> 

275 >>> def fn(x, y): 

276 >>> return torch.sigmoid(torch.matmul(torch.matmul(x, y), y)) * y 

277 >>> 

278 >>> out = torch.utils.checkpoint.checkpoint( 

279 >>> fn, x, y, 

280 >>> use_reentrant=False, 

281 >>> context_fn=context_fn, 

282 >>> ) 

283 """ 

284 # NB: If grad_mode is disabled, checkpoint would not run forward under 

285 # context_fn anyway, so proceed as usual. 

286 if policy_fn_or_list is None: 

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

288 return CheckpointPolicy.PREFER_RECOMPUTE 

289 elif isinstance(policy_fn_or_list, list): 

290 for op in policy_fn_or_list: 

291 if not isinstance(op, torch._ops.OpOverload): 

292 _extra_msg = ( 

293 "Please update the OpOverloadPacket to a specific OpOverload." 

294 "For example, if you have `torch.ops.aten.mm`, change it to `torch.ops.aten.mm.default`." 

295 ) if isinstance(op, torch._ops.OpOverloadPacket) else "" 

296 raise ValueError( 

297 f"Expected op in `op_list` to be an OpOverload but got: {op} " 

298 f"of type {type(op)}. {_extra_msg}" 

299 ) 

300 

301 def policy_fn(ctx, op, *args, **kwargs): 

302 if op in policy_fn_or_list: 

303 return CheckpointPolicy.MUST_SAVE 

304 else: 

305 return CheckpointPolicy.PREFER_RECOMPUTE 

306 elif callable(policy_fn_or_list): 

307 policy_fn = policy_fn_or_list 

308 else: 

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

310 

311 swap_storage = Storage() # patch code 

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

313 return ( 

314 _CachingTorchDispatchMode(policy_fn, swap_storage, storage, group_swap=group_swap), 

315 _CachedTorchDispatchMode(policy_fn, swap_storage, storage, allow_cache_entry_mutation), 

316 )