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

187 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-08-21 04:29 +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 

38class FuncModule(nn.Module): 

39 """ 

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

41 

42 Allows ordinary Python functions (or any callable without Module 

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

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

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

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

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

48 

49 Args: 

50 fn (callable): The function to wrap. 

51 

52 Example: 

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

54 """ 

55 

56 def __init__(self, fn: Callable): 

57 super().__init__() 

58 self._fn = fn 

59 

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

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

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

63 

64 

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

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

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

68 

69 

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

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

72 

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

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

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

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

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

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

79 callables participate in overlap tracking. 

80 """ 

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

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

83 continue 

84 if _is_callable_exempt_from_overlap_check(attr_value): 

85 continue 

86 if callable(attr_value): 

87 yield attr_name, attr_value 

88 

89 

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

91 try: 

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

93 except (AttributeError, TypeError): 

94 pass 

95 

96 

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

98 wrapped_module = getattr(module, _SWAP_WRAPPED_MODULE, None) 

99 if isinstance(wrapped_module, FuncModule): 

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

101 if isinstance(module, FuncModule): 

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

103 return None 

104 

105 

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

107 warnings.warn( 

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

109 "Wrapping overlapping module regions is not allowed." 

110 ) 

111 

112 

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

114 del owner, attr_name 

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

116 _raise_callable_already_wrapped(attr_value) 

117 

118 

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

120 if _is_callable_exempt_from_overlap_check(callable_obj): 

121 return 

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

123 warnings.warn( 

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

125 "Wrapping overlapping module regions is not allowed." 

126 ) 

127 _mark_wrapped(callable_obj) 

128 

129 

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

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

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

133 warnings.warn( 

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

135 "Wrapping overlapping module regions is not allowed." 

136 ) 

137 for submodule in module.modules(): 

138 if submodule is module: 

139 continue 

140 wrapped_callable = _get_wrapped_callable(submodule) 

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

142 continue 

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

144 if wrapped_callable is not None: 

145 _raise_callable_already_wrapped(wrapped_callable) 

146 warnings.warn( 

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

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

149 "Wrapping overlapping module regions is not allowed." 

150 ) 

151 for submodule in module.modules(): 

152 for attr_name, attr_value in _iter_wrappable_callable_attrs(submodule): 

153 _check_callable_attr_not_wrapped(submodule, attr_name, attr_value) 

154 for submodule in module.modules(): 

155 _mark_wrapped(submodule) 

156 for _, attr_value in _iter_wrappable_callable_attrs(submodule): 

157 _mark_wrapped(attr_value) 

158 

159 

160def base_check_fn(tensor) -> bool: 

161 """ 

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

163 - Skip Parameters and their views. 

164 - Skip empty storage tensors. 

165 """ 

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

167 return False 

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

169 return False 

170 return True 

171 

172 

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

174 """ 

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

176 """ 

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

178 self.add_to_storage = False 

179 self.storage = Storage() 

180 self.count_idx = 0 

181 self.policy_fn = policy_fn 

182 

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

184 swap_manager = SwapManager() 

185 

186 def pack_to_cpu(tensor: torch.Tensor): 

187 if not base_check_fn(tensor): 

188 return tensor.detach() 

189 if policy_fn is not None: 

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

191 return tensor.detach() 

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

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

194 group_name = swap_manager.get_current_group_name() 

195 if not group_name: 

196 return tensor.detach() 

197 if not self.add_to_storage: 

198 swap_manager.add_storage(group_name, self.storage) 

199 self.add_to_storage = True 

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

201 detached = tensor.detach() 

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

203 SwapTensor(detached, funcname, group_swap=group_swap) 

204 ) 

205 self.count_idx += 1 

206 return detached 

207 

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

209 if self.storage is not None: 

210 self.storage.clear() 

211 self.storage = None 

212 return tensor 

213 

214 super().__init__(pack_to_cpu, unpack_from_cpu) 

215 

216 

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

218 """ 

219 Base class for Activation Swap. 

220 

221 Not meant to be instantiated directly. 

222 """ 

223 

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

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

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

227 if track_overlaps: 

228 _check_and_mark_callable(module) 

229 module = FuncModule(module) 

230 if track_overlaps: 

231 _mark_wrapped(module) 

232 elif track_overlaps: 

233 _check_and_mark_wrapped(module) 

234 super().__init__() 

235 self._swap_wrapped_module = module 

236 self._is_wrapped = track_overlaps 

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

238 # non-swap wrapped module. 

239 self._register_state_dict_hook(self._post_state_dict_hook) 

240 # load_state_dict pre-hook to allow loading back into 

241 # swap-wrapped module. 

242 self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook) 

243 

244 @property 

245 def _wrapped_module(self): 

246 return self._swap_wrapped_module 

247 

248 @abstractmethod 

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

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

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

252 

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

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

255 try: 

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

257 except AttributeError: 

258 return getattr(self._swap_wrapped_module, name) 

259 

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

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

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

263 

264 def named_modules( 

265 self, 

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

267 prefix: str = "", 

268 remove_duplicate: bool = True, 

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

270 """ 

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

272 

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

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

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

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

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

278 

279 Args: 

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

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

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

283 

284 Returns: 

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

286 """ 

287 if memo is None: 

288 memo = set() 

289 if self not in memo: 

290 memo.add(self) 

291 yield prefix, self 

292 yield from self._swap_wrapped_module.named_modules( 

293 memo=memo, 

294 prefix=prefix, 

295 remove_duplicate=remove_duplicate, 

296 ) 

297 

298 def named_parameters( 

299 self, 

300 *args, 

301 **kwargs, 

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

303 """ 

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

305 

306 remove all occurrences of ``_SWAP_PREFIX``. 

307 """ 

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

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

310 

311 @staticmethod 

312 def _post_state_dict_hook( 

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

314 state_dict: dict[str, Any], 

315 prefix: str, 

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

317 ) -> dict[str, Any]: 

318 """ 

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

320 

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

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

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

324 adds the prefix back before loading the state_dict. 

325 """ 

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

327 return state_dict 

328 

329 @staticmethod 

330 def _pre_load_state_dict_hook( 

331 module: nn.Module, 

332 state_dict: dict[str, Any], 

333 prefix: str, 

334 *args: Any, 

335 ) -> None: 

336 """ 

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

338 

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

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

341 swap_wrapper modules properly. 

342 """ 

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

344 

345 

346class SwapWrapper(ActivationWrapper): 

347 """ 

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

349 """ 

350 def __init__( 

351 self, 

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

353 policy_fn: Optional[Callable] = None, 

354 group_swap: bool = False, 

355 ): 

356 super().__init__(mod) 

357 self.policy_fn = policy_fn 

358 self.group_swap = group_swap 

359 

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

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

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

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

364 

365 

366def swap_wrapper( 

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

368 policy_fn: Optional[Callable] = None, 

369 group_swap: bool = False, 

370) -> SwapWrapper: 

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

372 return SwapWrapper(module, policy_fn, group_swap) 

373 

374 

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

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

377 

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

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

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

381 """ 

382 swap_manager = SwapManager() 

383 group_name = swap_manager.get_current_group_name() 

384 if not group_name: 

385 warnings.warn( 

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

387 ) 

388 return target 

389 if swap_manager.is_last_group(group_name): 

390 return target 

391 

392 storage = Storage() 

393 count_idx = 0 

394 

395 def _register_tensor(tensor): 

396 nonlocal count_idx 

397 if not base_check_fn(tensor): 

398 return tensor 

399 

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

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

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

403 count_idx += 1 

404 return tensor 

405 

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

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

408 target, 

409 ) 

410 if count_idx > 0: 

411 swap_manager.add_storage(group_name, storage) 

412 return wrapped