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

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 

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

134 """Add available operators to the selective-checkpoint ignore set. 

135 

136 Args: 

137 ops (List[Optional[object]]): Operators to execute without selective-AC replay accounting. 

138 ``None`` entries are ignored for optional-version compatibility. 

139 """ 

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

141 

142 

143class _CachingTorchDispatchMode(TorchDispatchMode): 

144 # Used together with _CachedTorchDispatchMode to implement SAC. 

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

146 self.policy_fn = policy_fn 

147 self.swap_storage = swap_storage 

148 self.storage = storage 

149 self.add_to_storage = False 

150 self.group_swap = group_swap 

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

152 self._swap_manager = SwapManager() 

153 self._group_prefix = "" 

154 

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

156 if func in SAC_IGNORED_OPS: 

157 return func(*args, **kwargs) 

158 

159 kwargs = {} if kwargs is None else kwargs 

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

161 func, *args, **kwargs) 

162 if isinstance(policy, bool): 

163 policy = _policy_from_bool(policy) 

164 

165 is_compiling = _is_compiling(func, args, kwargs) 

166 

167 if is_compiling: 

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

169 fx_traceback.current_meta["recompute"] = policy 

170 

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

172 

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

174 

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

176 self.storage[func].append( 

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

178 ) 

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

180 if not self.add_to_storage: 

181 group_name = self._swap_manager.get_current_group_name() 

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

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

184 self.add_to_storage = True 

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

186 group_swap = self.group_swap 

187 entries = tree_map( 

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

189 ) 

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

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

192 elif policy != CheckpointPolicy.MUST_RECOMPUTE: 

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

194 return out 

195 

196 

197class _CachedTorchDispatchMode(TorchDispatchMode): 

198 # Used together with _CachingTorchDispatchMode to implement SAC. 

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

200 self.policy_fn = policy_fn 

201 self.swap_storage = swap_storage 

202 self.storage = storage 

203 self.allow_cache_entry_mutation = allow_cache_entry_mutation 

204 self._swap_cleared = False 

205 

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

207 if func in SAC_IGNORED_OPS: 

208 return func(*args, **kwargs) 

209 

210 kwargs = {} if kwargs is None else kwargs 

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

212 func, *args, **kwargs) 

213 if isinstance(policy, bool): 

214 policy = _policy_from_bool(policy) 

215 

216 is_compiling = _is_compiling(func, args, kwargs) 

217 

218 if not self._swap_cleared: 

219 self.swap_storage.clear() 

220 self._swap_cleared = True 

221 

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

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

224 or is_compiling): 

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

226 if storage is None: 

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

228 if len(storage) == 0: 

229 raise RuntimeError( 

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

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

232 ) 

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

234 else: 

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

236 return out 

237 

238 

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

240 """ 

241 Helper to avoid recomputing certain ops during activation checkpointing. 

242 

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

244 operations are recomputed during the backward pass. 

245 

246 Args: 

247 policy_fn_or_list (Callable or List): 

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

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

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

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

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

253 returning `CheckpointPolicy.MUST_SAVE` for the specified 

254 operations and `CheckpointPolicy.PREFER_RECOMPUTE` for all other 

255 operations. 

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

257 raised if any tensors cached by selective activation checkpoint are 

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

259 is disabled. 

260 Returns: 

261 A tuple of two context managers. 

262 

263 Example: 

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

265 >>> import functools 

266 >>> 

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

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

269 >>> 

270 >>> ops_to_save = [ 

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

272 >>> ] 

273 >>> 

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

275 >>> if op in ops_to_save: 

276 >>> return CheckpointPolicy.MUST_SAVE 

277 >>> else: 

278 >>> return CheckpointPolicy.PREFER_RECOMPUTE 

279 >>> 

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

281 >>> 

282 >>> # or equivalently 

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

284 >>> 

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

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

287 >>> 

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

289 >>> fn, x, y, 

290 >>> use_reentrant=False, 

291 >>> context_fn=context_fn, 

292 >>> ) 

293 """ 

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

295 # context_fn anyway, so proceed as usual. 

296 if policy_fn_or_list is None: 

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

298 return CheckpointPolicy.PREFER_RECOMPUTE 

299 elif isinstance(policy_fn_or_list, list): 

300 for op in policy_fn_or_list: 

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

302 _extra_msg = ( 

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

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

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

306 raise ValueError( 

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

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

309 ) 

310 

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

312 if op in policy_fn_or_list: 

313 return CheckpointPolicy.MUST_SAVE 

314 else: 

315 return CheckpointPolicy.PREFER_RECOMPUTE 

316 elif callable(policy_fn_or_list): 

317 policy_fn = policy_fn_or_list 

318 else: 

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

320 

321 swap_storage = Storage() # patch code 

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

323 return ( 

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

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

326 )