Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / checkpoint_exclude.py: 100%

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"""Saved-tensor helpers for regions excluded from checkpoint recomputation.""" 

16from collections import defaultdict, deque 

17from dataclasses import dataclass 

18from functools import lru_cache 

19from typing import Any, Deque, Dict, List, Optional, Tuple 

20 

21import torch 

22 

23_RECOMPUTE_INPUT_HANDLE_ATTR = "_hyper_parallel_recompute_input_handle" 

24_SAVE_OUTPUT_SOURCE_ATTR = "_hyper_parallel_save_output_source" 

25_MISSING = object() 

26_InputPath = Tuple[Tuple[str, Any], ...] 

27_TensorInput = Tuple[_InputPath, Any] 

28 

29 

30class _RecomputedInputHandle: 

31 """Defer one excluded-region saved input until checkpoint replay.""" 

32 

33 def __init__(self) -> None: 

34 """Initialize an unused and unresolved handle.""" 

35 self.used = False 

36 self._tensor = None 

37 

38 def mark_used(self) -> None: 

39 """Record that an exclude operation saved this input for backward.""" 

40 self.used = True 

41 

42 def materialize(self, tensor: Any) -> None: 

43 """Bind the handle to the matching input produced during replay.""" 

44 self._tensor = tensor 

45 

46 def get_recomputed_tensor(self) -> Any: 

47 """Return the replay-produced input for backward.""" 

48 if self._tensor is None: 

49 raise RuntimeError("Checkpoint-excluded input was requested before recomputation") 

50 return self._tensor 

51 

52 

53@dataclass(frozen=True) 

54class _InputBinding: 

55 """Map one input path to its deferred saved-tensor handle.""" 

56 

57 path: _InputPath 

58 handle: _RecomputedInputHandle 

59 

60 

61@dataclass 

62class _ExcludeCacheEntry: 

63 """Store one excluded call's replay value and deferred input bindings.""" 

64 

65 output: Any 

66 input_bindings: List[_InputBinding] 

67 output_tensor_count: int = 1 

68 

69 

70class _ExcludeCache: 

71 """Store excluded-region call entries for one checkpoint invocation.""" 

72 

73 def __init__(self) -> None: 

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

75 self._entries: Dict[int, Deque[_ExcludeCacheEntry]] = defaultdict(deque) 

76 

77 def save(self, wrapper_id: int, entry: _ExcludeCacheEntry) -> None: 

78 """Save one call entry produced by a checkpoint-excluded region.""" 

79 self._entries[wrapper_id].append(entry) 

80 

81 def pop(self, wrapper_id: int) -> _ExcludeCacheEntry: 

82 """Return the matching forward call entry during recomputation.""" 

83 entries = self._entries.get(wrapper_id) 

84 if not entries: 

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

86 entry = entries.popleft() 

87 if not entries: 

88 self._entries.pop(wrapper_id) 

89 return entry 

90 

91 def clear(self) -> None: 

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

93 self._entries.clear() 

94 

95 

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

97 """Return a deferred input handle or detached tensor data.""" 

98 handle = getattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, None) 

99 if isinstance(handle, _RecomputedInputHandle): 

100 handle.mark_used() 

101 return handle 

102 return tensor.detach() if tensor.requires_grad else tensor 

103 

104 

105def _unpack_saved_tensor(value: Any) -> Any: 

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

107 if isinstance(value, _RecomputedInputHandle): 

108 return value.get_recomputed_tensor() 

109 return value 

110 

111 

112def _saved_tensors_context() -> Any: 

113 """Create an inner hook that stores real tensors instead of outer holders.""" 

114 return torch.autograd.graph.saved_tensors_hooks(_pack_saved_tensor, _unpack_saved_tensor) 

115 

116 

117_EXCLUDE_CACHE_KEY = object() 

118 

119 

120def _append_tensor_inputs( 

121 value: Any, 

122 path: _InputPath, 

123 leaves: List[_TensorInput], 

124 tensor_type: type, 

125) -> None: 

126 """Append tensor leaves without creating a self-referential local function.""" 

127 if isinstance(value, tensor_type): 

128 leaves.append((path, value)) 

129 return 

130 if isinstance(value, (tuple, list)): 

131 for index, item in enumerate(value): 

132 _append_tensor_inputs(item, path + (("index", index),), leaves, tensor_type) 

133 return 

134 if isinstance(value, dict): 

135 for key, item in value.items(): 

136 _append_tensor_inputs(item, path + (("key", key),), leaves, tensor_type) 

137 

138 

139def _collect_tensor_inputs(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> List[_TensorInput]: 

140 """Return tensor leaves and self-describing paths from one excluded-region call.""" 

141 leaves = [] 

142 for index, arg in enumerate(args): 

143 _append_tensor_inputs(arg, (("arg", index),), leaves, torch.Tensor) 

144 for key, value in kwargs.items(): 

145 _append_tensor_inputs(value, (("kwarg", key),), leaves, torch.Tensor) 

146 return leaves 

147 

148 

149def _mark_recompute_inputs( 

150 invocation_id: object, 

151 args: Tuple[Any, ...], 

152 kwargs: Dict[str, Any], 

153) -> Tuple[List[_InputBinding], List[Tuple[Any, Any]]]: 

154 """Attach deferred handles to inputs that replay reproduces.""" 

155 bindings = [] 

156 previous_handles = [] 

157 seen_tensor_ids = set() 

158 try: 

159 for path, tensor in _collect_tensor_inputs(args, kwargs): 

160 if ( 

161 isinstance(tensor, torch.nn.Parameter) 

162 or id(tensor) in seen_tensor_ids 

163 or getattr(tensor, _SAVE_OUTPUT_SOURCE_ATTR, None) is invocation_id 

164 ): 

165 continue 

166 handle = _RecomputedInputHandle() 

167 previous = getattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, _MISSING) 

168 previous_handles.append((tensor, previous)) 

169 setattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, handle) 

170 bindings.append(_InputBinding(path, handle)) 

171 seen_tensor_ids.add(id(tensor)) 

172 except BaseException: 

173 _restore_recompute_inputs(previous_handles) 

174 raise 

175 return bindings, previous_handles 

176 

177 

178def _restore_recompute_inputs(previous_handles: List[Tuple[Any, Any]]) -> None: 

179 """Restore Tensor attributes overwritten for one excluded call.""" 

180 for tensor, previous in previous_handles: 

181 if previous is _MISSING: 

182 if hasattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR): 

183 delattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR) 

184 else: 

185 setattr(tensor, _RECOMPUTE_INPUT_HANDLE_ATTR, previous) 

186 

187 

188def _resolve_input(args: Tuple[Any, ...], kwargs: Dict[str, Any], path: _InputPath) -> Any: 

189 """Resolve one replay input from its forward argument path.""" 

190 root_kind, root_key = path[0] 

191 value = args[root_key] if root_kind == "arg" else kwargs[root_key] 

192 for _, token_value in path[1:]: 

193 value = value[token_value] 

194 return value 

195 

196 

197def _materialize_recompute_inputs( 

198 entry: _ExcludeCacheEntry, 

199 args: Tuple[Any, ...], 

200 kwargs: Dict[str, Any], 

201) -> None: 

202 """Bind used input handles to tensors produced during checkpoint replay.""" 

203 for binding in entry.input_bindings: 

204 if not binding.handle.used: 

205 continue 

206 tensor = _resolve_input(args, kwargs, binding.path) 

207 if not isinstance(tensor, torch.Tensor): 

208 raise RuntimeError( 

209 "Checkpoint replay did not reproduce a tensor input required by a checkpoint-excluded region" 

210 ) 

211 binding.handle.materialize(tensor.detach()) 

212 

213 

214def _has_used_input(input_bindings: List[_InputBinding]) -> bool: 

215 """Return whether the excluded call saved any marked input.""" 

216 return any(binding.handle.used for binding in input_bindings) 

217 

218 

219@lru_cache(maxsize=1) 

220def _get_replay_placeholder() -> Any: 

221 """Create a zero-element placeholder returned by an elided SAVE replay.""" 

222 return torch.empty(0, device="cpu") 

223 

224 

225def _make_replay_placeholder_output(tensor_count: int) -> Any: 

226 """Create one placeholder leaf for each forward output tensor.""" 

227 if tensor_count == 0: 

228 return () 

229 placeholder = _get_replay_placeholder() 

230 if tensor_count == 1: 

231 return placeholder 

232 return (placeholder,) * tensor_count 

233 

234 

235@lru_cache(maxsize=1) 

236def _get_recompute_trigger() -> Any: 

237 """Create the differentiable zero-element input used by recompute boundaries.""" 

238 return torch.empty(0, device="cpu", requires_grad=True) 

239 

240 

241@lru_cache(maxsize=1) 

242def _get_recompute_boundary() -> Any: 

243 """Create the autograd Function used to trigger outer checkpoint replay.""" 

244 class _RecomputeBoundary(torch.autograd.Function): 

245 """Trigger the outer checkpoint hook before excluded-region backward.""" 

246 

247 @staticmethod 

248 def forward(ctx: Any, tensor: Any, trigger: Any) -> Any: 

249 """Save one zero-element outer-hook dependency and return the tensor unchanged.""" 

250 ctx.save_for_backward(trigger) 

251 return tensor 

252 

253 @staticmethod 

254 def backward(ctx: Any, grad_output: Any) -> Tuple[Any, None]: 

255 """Trigger dependency unpack and pass the gradient through.""" 

256 _ = ctx.saved_tensors 

257 return grad_output, None 

258 

259 return _RecomputeBoundary 

260 

261 

262def _finalize_save_outputs_impl( 

263 output: Any, 

264 add_recompute_boundary: bool, 

265 invocation_id: Optional[object], 

266 tensor_leaf_count: Optional[List[int]], 

267 tensor_type: type, 

268) -> Any: 

269 """Recursively finalize SAVE output leaves.""" 

270 if isinstance(output, tensor_type): 

271 if tensor_leaf_count is not None: 

272 tensor_leaf_count[0] += 1 

273 if add_recompute_boundary: 

274 output = _get_recompute_boundary().apply(output, _get_recompute_trigger()) 

275 if invocation_id is not None: 

276 setattr(output, _SAVE_OUTPUT_SOURCE_ATTR, invocation_id) 

277 return output 

278 if isinstance(output, list): 

279 return [ 

280 _finalize_save_outputs_impl(item, add_recompute_boundary, invocation_id, tensor_leaf_count, tensor_type) 

281 for item in output 

282 ] 

283 if isinstance(output, tuple): 

284 items = [ 

285 _finalize_save_outputs_impl(item, add_recompute_boundary, invocation_id, tensor_leaf_count, tensor_type) 

286 for item in output 

287 ] 

288 if hasattr(output, "_fields"): 

289 return type(output)(*items) 

290 return tuple(items) 

291 if isinstance(output, dict): 

292 return type(output)( 

293 ( 

294 key, 

295 _finalize_save_outputs_impl( 

296 value, 

297 add_recompute_boundary, 

298 invocation_id, 

299 tensor_leaf_count, 

300 tensor_type, 

301 ), 

302 ) 

303 for key, value in output.items() 

304 ) 

305 return output 

306 

307 

308def _finalize_save_outputs( 

309 output: Any, 

310 add_recompute_boundary: bool, 

311 invocation_id: Optional[object], 

312 tensor_leaf_count: Optional[List[int]] = None, 

313) -> Any: 

314 """Apply the required boundary and SAVE provenance to output tensor leaves.""" 

315 return _finalize_save_outputs_impl( 

316 output, 

317 add_recompute_boundary, 

318 invocation_id, 

319 tensor_leaf_count, 

320 torch.Tensor, 

321 )