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

193 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-08-22 04:23 +0800

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 

16# https://github.com/pytorch/pytorch/blob/release/2.6/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py 

17# enhanced with activation swap functionality. 

18# ============================================================================ 

19"""Activation Swap implementation for PyTorch.""" 

20# pylint: disable=W0212, W0613 

21 

22from abc import ABC, abstractmethod 

23from collections.abc import Iterator 

24from typing import Optional, Callable, Any, Union 

25import types 

26import warnings 

27import torch 

28from torch import nn 

29from torch.distributed.utils import _replace_by_prefix 

30from hyper_parallel.core.activation_checkpoint.activation_checkpoint import CheckpointPolicy 

31from hyper_parallel.core.activation_checkpoint.swap import SwapManager, SwapTensor, Storage 

32 

33 

34_SWAP_WRAPPED_MODULE = "_swap_wrapped_module" 

35_SWAP_PREFIX = _SWAP_WRAPPED_MODULE + "." 

36 

37 

38def _raise_if_compiling(feature_name: str) -> None: 

39 """Reject HyperParallel activation-swap extensions during Torch capture.""" 

40 if torch.compiler.is_compiling(): 

41 raise ValueError( 

42 f"HyperParallel {feature_name} is not supported in compile mode. " 

43 "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies." 

44 ) 

45 

46 

47class FuncModule(nn.Module): 

48 """ 

49 Thin :class:`~torch.nn.Module` adapter that wraps a plain callable. 

50 

51 Allows ordinary Python functions (or any callable without Module 

52 parameters) to be passed to :func:`swap_wrapper` and 

53 :func:`~hyper_parallel.core.activation_checkpoint.checkpoint_wrapper` 

54 in place of an :class:`~torch.nn.Module`. 

55 The wrapped function is stored as ``_fn`` and invoked in 

56 :meth:`forward`; the module has no trainable parameters. 

57 

58 Args: 

59 fn (callable): The function to wrap. 

60 

61 Example: 

62 >>> wrapped = swap_wrapper(lambda x: x * 2) 

63 """ 

64 

65 def __init__(self, fn: Callable): 

66 super().__init__() 

67 self._fn = fn 

68 

69 def forward(self, *args, **kwargs): 

70 """Invoke the wrapped callable with the given arguments.""" 

71 return self._fn(*args, **kwargs) 

72 

73 

74def _is_callable_exempt_from_overlap_check(callable_obj: Callable) -> bool: 

75 """Return True for callables that cannot be reliably overlap-tracked by object marks.""" 

76 return isinstance(callable_obj, (types.FunctionType, types.BuiltinFunctionType, types.MethodType)) 

77 

78 

79def _iter_wrappable_callable_attrs(module: nn.Module) -> Iterator[tuple[str, Callable]]: 

80 """Yield public per-instance callable attributes not registered as child modules. 

81 

82 Plain functions, builtins and bound methods are skipped: these are stateless 

83 module-level utilities shared by reference across many modules (e.g. 

84 ``self.act = F.gelu`` repeated in every layer). They are never standalone 

85 checkpoint regions, and marking a shared function's ``_is_wrapped`` flag 

86 would both mutate a global object and falsely flag every sibling module that 

87 references the same function as an overlapping wrap. Only per-instance 

88 callables participate in overlap tracking. 

89 """ 

90 for attr_name, attr_value in vars(module).items(): 

91 if attr_name.startswith("_") or isinstance(attr_value, nn.Module): 

92 continue 

93 if _is_callable_exempt_from_overlap_check(attr_value): 

94 continue 

95 if callable(attr_value): 

96 yield attr_name, attr_value 

97 

98 

99def _mark_wrapped(obj: Any) -> None: 

100 try: 

101 obj._is_wrapped = True # pylint: disable=W0212 

102 except (AttributeError, TypeError): 

103 pass 

104 

105 

106def _get_wrapped_callable(module: nn.Module) -> Optional[Callable]: 

107 wrapped_module = getattr(module, _SWAP_WRAPPED_MODULE, None) 

108 if isinstance(wrapped_module, FuncModule): 

109 return getattr(wrapped_module, "_fn", None) 

110 if isinstance(module, FuncModule): 

111 return getattr(module, "_fn", None) 

112 return None 

113 

114 

115def _raise_callable_already_wrapped(callable_obj: Callable) -> None: 

116 warnings.warn( 

117 f"Callable '{callable_obj.__class__.__name__}' is already wrapped. " 

118 "Wrapping overlapping module regions is not allowed." 

119 ) 

120 

121 

122def _check_callable_attr_not_wrapped(owner: nn.Module, attr_name: str, attr_value: Callable) -> None: 

123 del owner, attr_name 

124 if getattr(attr_value, '_is_wrapped', False): 

125 _raise_callable_already_wrapped(attr_value) 

126 

127 

128def _check_and_mark_callable(callable_obj: Callable) -> None: 

129 if _is_callable_exempt_from_overlap_check(callable_obj): 

130 return 

131 if getattr(callable_obj, '_is_wrapped', False): 

132 warnings.warn( 

133 f"Callable '{callable_obj.__class__.__name__}' or one of its ancestors is already wrapped. " 

134 "Wrapping overlapping module regions is not allowed." 

135 ) 

136 _mark_wrapped(callable_obj) 

137 

138 

139def _check_and_mark_wrapped(module: nn.Module) -> None: 

140 """Validate no wrapping overlap, then mark module and all descendants as wrapped.""" 

141 if getattr(module, '_is_wrapped', False): 

142 warnings.warn( 

143 f"Module '{module.__class__.__name__}' or one of its ancestors is already wrapped. " 

144 "Wrapping overlapping module regions is not allowed." 

145 ) 

146 for submodule in module.modules(): 

147 if submodule is module: 

148 continue 

149 wrapped_callable = _get_wrapped_callable(submodule) 

150 if wrapped_callable is not None and _is_callable_exempt_from_overlap_check(wrapped_callable): 

151 continue 

152 if getattr(submodule, '_is_wrapped', False): 

153 if wrapped_callable is not None: 

154 _raise_callable_already_wrapped(wrapped_callable) 

155 warnings.warn( 

156 f"Submodule '{getattr(submodule, '_swap_wrapped_module', submodule).__class__.__name__}' of " 

157 f"'{module.__class__.__name__}' is already wrapped. " 

158 "Wrapping overlapping module regions is not allowed." 

159 ) 

160 for submodule in module.modules(): 

161 for attr_name, attr_value in _iter_wrappable_callable_attrs(submodule): 

162 _check_callable_attr_not_wrapped(submodule, attr_name, attr_value) 

163 for submodule in module.modules(): 

164 _mark_wrapped(submodule) 

165 for _, attr_value in _iter_wrappable_callable_attrs(submodule): 

166 _mark_wrapped(attr_value) 

167 

168 

169def base_check_fn(tensor) -> bool: 

170 """ 

171 Basic check to determine if a tensor is eligible for offloading. 

172 - Skip Parameters and their views. 

173 - Skip empty storage tensors. 

174 """ 

175 if isinstance(tensor._base, torch.nn.parameter.Parameter) or isinstance(tensor, torch.nn.parameter.Parameter): # pylint: disable=W0212 

176 return False 

177 if tensor.untyped_storage().size() == 0: 

178 return False 

179 return True 

180 

181 

182class AsyncSaveOnCpu(torch.autograd.graph.saved_tensors_hooks): 

183 """ 

184 Context manager to offload tensors to CPU during forward pass. 

185 """ 

186 def __init__(self, policy_fn=None, group_swap: bool = False) -> None: 

187 _raise_if_compiling("AsyncSaveOnCpu") 

188 self.add_to_storage = False 

189 self.storage = Storage() 

190 self.count_idx = 0 

191 self.policy_fn = policy_fn 

192 

193 # Cache per-context-manager state once to avoid per-tensor singleton lookups. 

194 swap_manager = SwapManager() 

195 

196 def pack_to_cpu(tensor: torch.Tensor): 

197 if not base_check_fn(tensor): 

198 return tensor.detach() 

199 if policy_fn is not None: 

200 if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE: 

201 return tensor.detach() 

202 if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP: 

203 raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}") 

204 group_name = swap_manager.get_current_group_name() 

205 if not group_name: 

206 return tensor.detach() 

207 if not self.add_to_storage: 

208 swap_manager.add_storage(group_name, self.storage) 

209 self.add_to_storage = True 

210 funcname = f"{group_name}::{tensor.shape}" 

211 detached = tensor.detach() 

212 self.storage[self.count_idx].append( 

213 SwapTensor(detached, funcname, group_swap=group_swap) 

214 ) 

215 self.count_idx += 1 

216 return detached 

217 

218 def unpack_from_cpu(tensor) -> torch.Tensor: 

219 if self.storage is not None: 

220 self.storage.clear() 

221 self.storage = None 

222 return tensor 

223 

224 super().__init__(pack_to_cpu, unpack_from_cpu) 

225 

226 

227class ActivationWrapper(torch.nn.Module, ABC): 

228 """ 

229 Base class for Activation Swap. 

230 

231 Not meant to be instantiated directly. 

232 """ 

233 

234 def __init__(self, module: Union[nn.Module, Callable], *, track_overlaps: bool = True): 

235 """Initialize a wrapper and optionally participate in overlap tracking.""" 

236 if callable(module) and not isinstance(module, nn.Module): 

237 if track_overlaps: 

238 _check_and_mark_callable(module) 

239 module = FuncModule(module) 

240 if track_overlaps: 

241 _mark_wrapped(module) 

242 elif track_overlaps: 

243 _check_and_mark_wrapped(module) 

244 super().__init__() 

245 self._swap_wrapped_module = module 

246 self._is_wrapped = track_overlaps 

247 # state_dict post hook to remove prefix to allow loading into a 

248 # non-swap wrapped module. 

249 self._register_state_dict_hook(self._post_state_dict_hook) 

250 # load_state_dict pre-hook to allow loading back into 

251 # swap-wrapped module. 

252 self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook) 

253 

254 @property 

255 def _wrapped_module(self): 

256 return self._swap_wrapped_module 

257 

258 @abstractmethod 

259 def forward(self, *args, **kwargs): 

260 """Run the wrapped module's forward pass with activation swapping. Must be implemented by subclasses.""" 

261 raise ValueError("Subclasses should implement forward().") 

262 

263 def __getattr__(self, name: str) -> Any: 

264 """Forward missing attributes to wrapped module.""" 

265 try: 

266 return super().__getattr__(name) # defer to nn.Module's logic 

267 except AttributeError: 

268 return getattr(self._swap_wrapped_module, name) 

269 

270 def __getitem__(self, key: int) -> Any: 

271 """Forward indexing calls in case the module is a nn.Sequential.""" 

272 return self._swap_wrapped_module.__getitem__(key) # type: ignore[operator] 

273 

274 def named_modules( 

275 self, 

276 memo: Optional[set[nn.Module]] = None, 

277 prefix: str = "", 

278 remove_duplicate: bool = True, 

279 ) -> Iterator[tuple[str, nn.Module]]: 

280 """ 

281 Yield wrapped-module children without exposing the internal wrapper prefix. 

282 

283 PyTorch parent modules implement ``named_parameters(recurse=True)`` by 

284 iterating ``named_modules()`` and reading each module's direct 

285 ``_parameters``. They do not call child modules' ``named_parameters()`` 

286 overrides. Exposing the wrapped module under the wrapper's own prefix 

287 keeps root-module traversals aligned with ``state_dict()`` keys. 

288 

289 Args: 

290 memo (Optional[set[nn.Module]], optional): A memo set to avoid infinite recursion. Default: ``None``. 

291 prefix (str, optional): A prefix to prepend to all module names. Default: ``""``. 

292 remove_duplicate (bool, optional): Whether to remove duplicate modules. Default: ``True``. 

293 

294 Returns: 

295 Iterator[tuple[str, nn.Module]] An iterator of (name, module) pairs. 

296 """ 

297 if memo is None: 

298 memo = set() 

299 if self not in memo: 

300 memo.add(self) 

301 yield prefix, self 

302 yield from self._swap_wrapped_module.named_modules( 

303 memo=memo, 

304 prefix=prefix, 

305 remove_duplicate=remove_duplicate, 

306 ) 

307 

308 def named_parameters( 

309 self, 

310 *args, 

311 **kwargs, 

312 ) -> Iterator[tuple[str, torch.nn.Parameter]]: 

313 """ 

314 Override :meth:`named_parameters()` to intercept parameter names. 

315 

316 remove all occurrences of ``_SWAP_PREFIX``. 

317 """ 

318 for param_name, param in super().named_parameters(*args, **kwargs): 

319 yield param_name.replace(_SWAP_PREFIX, ""), param 

320 

321 @staticmethod 

322 def _post_state_dict_hook( 

323 module: nn.Module, # pylint: disable=W0613 

324 state_dict: dict[str, Any], 

325 prefix: str, 

326 *args: Any, # pylint: disable=W0613 

327 ) -> dict[str, Any]: 

328 """ 

329 _post_state_dict_hook() is called after the state_dict() of this FSDP module is executed. 

330 

331 For ``swap_wrapper``, it will strip swap-wrapped module prefix, 

332 so that this module can be loaded into non-swapped modules. 

333 It would still be able to be loaded into swap-wrapped modules as this class, 

334 adds the prefix back before loading the state_dict. 

335 """ 

336 _replace_by_prefix(state_dict, f"{prefix}{_SWAP_PREFIX}", prefix) 

337 return state_dict 

338 

339 @staticmethod 

340 def _pre_load_state_dict_hook( 

341 module: nn.Module, 

342 state_dict: dict[str, Any], 

343 prefix: str, 

344 *args: Any, 

345 ) -> None: 

346 """ 

347 ``_pre_state_dict_hook` is called before ``self._load_from_state_dict()`` is called. 

348 

349 For ``swap_wrapper``, it will add back the module 

350 prefix so that non-swapped modules can be loaded into 

351 swap_wrapper modules properly. 

352 """ 

353 _replace_by_prefix(state_dict, prefix, prefix + f"{_SWAP_PREFIX}") 

354 

355 

356class SwapWrapper(ActivationWrapper): 

357 """ 

358 Customize an nn.Module wrapper class to add an AsyncSaveOnCpu context manager for the target model. 

359 """ 

360 def __init__( 

361 self, 

362 mod: Union[nn.Module, Callable], 

363 policy_fn: Optional[Callable] = None, 

364 group_swap: bool = False, 

365 ): 

366 super().__init__(mod) 

367 self.policy_fn = policy_fn 

368 self.group_swap = group_swap 

369 

370 def forward(self, *args, **kwargs): 

371 """Run the wrapped module inside an AsyncSaveOnCpu context for activation swapping.""" 

372 _raise_if_compiling("swap_wrapper") 

373 with AsyncSaveOnCpu(policy_fn=self.policy_fn, group_swap=self.group_swap): 

374 return self._swap_wrapped_module(*args, **kwargs) 

375 

376 

377def swap_wrapper( 

378 module: Union[nn.Module, Callable], 

379 policy_fn: Optional[Callable] = None, 

380 group_swap: bool = False, 

381) -> SwapWrapper: 

382 """Wrap a module or callable with activation swap functionality.""" 

383 return SwapWrapper(module, policy_fn, group_swap) 

384 

385 

386def swap_tensor_wrapper(target, tag: Optional[str] = None, group_swap: bool = False): 

387 """Register selected tensors into the current swap group. 

388 

389 This helper is intended to be used inside a forward path that already 

390 participates in the existing swap scheduling managed by ``SwapManager``. 

391 It preserves the input structure and returns the original tensors. 

392 """ 

393 _raise_if_compiling("swap_tensor_wrapper") 

394 swap_manager = SwapManager() 

395 group_name = swap_manager.get_current_group_name() 

396 if not group_name: 

397 warnings.warn( 

398 f"Tensor {tag} cannot be swapped, for its group is unregistered." 

399 ) 

400 return target 

401 if swap_manager.is_last_group(group_name): 

402 return target 

403 

404 storage = Storage() 

405 count_idx = 0 

406 

407 def _register_tensor(tensor): 

408 nonlocal count_idx 

409 if not base_check_fn(tensor): 

410 return tensor 

411 

412 tensor_tag = tag or f"{group_name}_swap_tensor" 

413 funcname = f"{tensor_tag}::{tuple(tensor.shape)}" 

414 storage[count_idx].append(SwapTensor(tensor, funcname, group_swap=group_swap)) 

415 count_idx += 1 

416 return tensor 

417 

418 wrapped = torch.utils._pytree.tree_map( # pylint: disable=protected-access 

419 lambda x: _register_tensor(x) if isinstance(x, torch.Tensor) else x, 

420 target, 

421 ) 

422 if count_idx > 0: 

423 swap_manager.add_storage(group_name, storage) 

424 return wrapped