Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_checkpoint / activation_checkpoint.py: 95%

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"""Activation checkpointing related interfaces""" 

16import contextlib 

17import enum 

18from functools import partial 

19from typing import Any, Callable, Optional, Tuple 

20 

21from hyper_parallel.platform import get_platform 

22from .recompute_state import create_recompute_contexts 

23plat = get_platform() 

24 

25 

26class CheckpointPolicy(enum.Enum): 

27 """ 

28 Enum for specifying the policy for checkpointing during backpropagation. 

29 

30 This enum extends PyTorch's selective activation checkpointing policies 

31 by introducing a SWAP-based strategy, which allows activation tensors 

32 to be offloaded during the forward pass and loaded back before backward 

33 computation. 

34 

35 For PyTorch native policies (SAVE / RECOMPUTE semantics and MUST vs PREFER), 

36 see: https://docs.pytorch.org/docs/2.6/checkpoint.html#torch.utils.checkpoint.CheckpointPolicy 

37 

38 Additional policy: 

39 

40 - ``MUST_SWAP``: The operation's output is offloaded to host memory during the 

41 forward pass and loaded back asynchronously before backward computation. The backward 

42 pass reuses the loaded activations without recomputation. 

43 

44 This policy must be used together with :class:`SwapManager` to coordinate 

45 asynchronous offload/load and stream synchronization. 

46 

47 .. note:: 

48 ``MUST_SWAP`` is typically applied to operations that are either 

49 computationally expensive or have large memory footprints. Note that 

50 swapping very small outputs may introduce additional overhead and 

51 reduce the effectiveness of asynchronous copy. 

52 """ 

53 MUST_SAVE = 0 

54 PREFER_SAVE = 1 

55 MUST_RECOMPUTE = 2 

56 PREFER_RECOMPUTE = 3 

57 

58 # Offload during forward, reload before backward. Requires SwapManager. 

59 MUST_SWAP = 4 

60 

61 

62class _StackedCtx: 

63 """Compose multiple context managers as one — enter in order, exit reversed.""" 

64 

65 def __init__(self, ctxs) -> None: 

66 self._ctxs = list(ctxs) 

67 self._stack = contextlib.ExitStack() 

68 

69 def __enter__(self): 

70 self._stack.__enter__() 

71 try: 

72 for ctx in self._ctxs: 

73 self._stack.enter_context(ctx) 

74 except BaseException as exc: 

75 self._stack.__exit__(type(exc), exc, exc.__traceback__) 

76 raise 

77 return self 

78 

79 def __exit__(self, exc_type, exc_val, exc_tb): 

80 return self._stack.__exit__(exc_type, exc_val, exc_tb) 

81 

82 

83def _compose_context_fns( 

84 factories: Tuple[Callable[[], Tuple[object, object]], ...], 

85) -> Callable[[], Tuple[_StackedCtx, _StackedCtx]]: 

86 """Combine ``(forward_ctx, recompute_ctx)`` factories into one factory. 

87 

88 ``ms.recompute`` / ``torch.utils.checkpoint(use_reentrant=False)`` call 

89 ``context_fn()`` once per invocation and unpack the result as 

90 ``(forward_ctx, recompute_ctx)``. This helper calls each input factory 

91 once, then stacks all forward contexts and all recompute contexts into 

92 two :class:`_StackedCtx` instances so the composite respects the 

93 single-call contract. 

94 """ 

95 def factory() -> Tuple[_StackedCtx, _StackedCtx]: 

96 pairs = [fn() for fn in factories] 

97 fwd_ctxs = [pair[0] for pair in pairs] 

98 rec_ctxs = [pair[1] for pair in pairs] 

99 return _StackedCtx(fwd_ctxs), _StackedCtx(rec_ctxs) 

100 

101 return factory 

102 

103 

104def checkpoint( 

105 function, 

106 *args, 

107 swap_inputs: bool = False, 

108 policy_fn: Optional[Callable] = None, 

109 context_fn: Optional[Callable[[], Tuple[object, object]]] = None, 

110 group_swap: bool = False, 

111 early_stop: bool = True, 

112 **kwargs, 

113): 

114 """ 

115 Apply activation checkpointing to a function with optional input swapping. 

116 

117 Args: 

118 function: The function to apply checkpointing to. 

119 *args: Arguments to pass to the function. 

120 swap_inputs (bool): Whether to enable input swapping using async_save_on_cpu context. 

121 policy_fn (callable, optional): Function that determines checkpoint policy for operations. 

122 context_fn (callable, optional): A no-arg factory returning a 

123 ``(forward_ctx, recompute_ctx)`` pair, matching the 

124 ``context_fn`` contract of ``ms.recompute(use_reentrant=False)`` 

125 and ``torch.utils.checkpoint(use_reentrant=False)``. Use this 

126 to bracket the backward-time forward re-run with custom logic. 

127 When ``policy_fn``, ``group_swap`` and ``context_fn`` are 

128 supplied together, the resulting factories are composed: their 

129 forward and recompute contexts are stacked so all enter in 

130 order and exit in reverse. 

131 group_swap (bool, optional): Whether MUST_SWAP tensors participate in group copy fusion. 

132 Only effective when ``policy_fn`` is provided. Default: ``False``. 

133 early_stop (bool, optional): Whether recomputation stops after all tensors needed by 

134 backward have been produced. This per-call keyword is the only supported way to 

135 configure early stop. Default: ``True``. 

136 **kwargs: Additional keyword arguments to pass to the function. 

137 

138 Returns: 

139 The result of applying the function with checkpointing. 

140 """ 

141 if not isinstance(early_stop, bool): 

142 raise ValueError(f"early_stop must be bool, but got {type(early_stop).__name__}.") 

143 

144 factories: list = [create_recompute_contexts] 

145 if policy_fn is not None: 

146 factories.append(partial(plat.create_selective_checkpoint_contexts, policy_fn, group_swap=group_swap)) 

147 if context_fn is not None: 

148 factories.append(context_fn) 

149 

150 if len(factories) == 1: 

151 composed_context_fn = factories[0] 

152 else: 

153 composed_context_fn = _compose_context_fns(tuple(factories)) 

154 

155 context = partial(plat.async_save_on_cpu, group_swap=group_swap) if swap_inputs else contextlib.nullcontext 

156 with context(): 

157 return plat.checkpoint( 

158 function, 

159 *args, 

160 context_fn=composed_context_fn, 

161 use_reentrant=False, 

162 early_stop=early_stop, 

163 **kwargs, 

164 ) 

165 

166 

167def swap(function, *args, policy_fn=None, group_swap=False, **kwargs): 

168 """Apply activation swap to a function call. 

169 

170 Offloads intermediate activations saved by the autograd engine to CPU 

171 during the forward pass and loads them back before the backward pass, 

172 trading device memory for host memory bandwidth. Unlike 

173 :func:`checkpoint`, no recomputation is performed. 

174 

175 Args: 

176 function (callable): The function whose activations should be swapped. 

177 *args: Positional arguments forwarded to *function*. 

178 policy_fn (callable, optional): Per-tensor swap policy. Receives 

179 a tensor and returns a :class:`CheckpointPolicy` value. Tensors 

180 that return ``CheckpointPolicy.MUST_SAVE`` are kept on device; 

181 all other eligible tensors are offloaded. When ``None``, all 

182 eligible tensors are offloaded. 

183 group_swap (bool, optional): Whether swapped tensors participate in 

184 group copy fusion. Default: ``False``. 

185 **kwargs: Keyword arguments forwarded to *function*. 

186 

187 Returns: 

188 The return value of ``function(*args, **kwargs)``. 

189 

190 Example: 

191 >>> output = swap(layer, x, policy_fn=lambda t: CheckpointPolicy.MUST_SAVE) 

192 """ 

193 with plat.async_save_on_cpu(policy_fn=policy_fn, group_swap=group_swap): 

194 return function(*args, **kwargs) 

195 

196 

197def checkpoint_exclude_wrapper(module: Any, *, save_output: bool = True) -> Any: 

198 """Wrap a callable whose region is excluded from activation recomputation. 

199 

200 Args: 

201 module: The module or callable to exclude from recomputation. 

202 save_output: Whether to retain the region output for checkpoint replay. 

203 Set this to ``False`` only when the output is passed directly as one 

204 argument to another excluded region. Default: ``True``. 

205 

206 Returns: 

207 The platform-specific checkpoint exclusion wrapper. 

208 """ 

209 return plat.checkpoint_exclude_wrapper(module, save_output=save_output) 

210 

211 

212checkpoint_wrapper = plat.checkpoint_wrapper 

213swap_wrapper = plat.swap_wrapper 

214swap_tensor_wrapper = plat.swap_tensor_wrapper