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

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 Callable, Optional, Tuple 

20 

21from hyper_parallel.platform import get_platform 

22plat = get_platform() 

23 

24 

25class CheckpointPolicy(enum.Enum): 

26 """ 

27 Enum for specifying the policy for checkpointing during backpropagation. 

28 

29 This enum extends PyTorch's selective activation checkpointing policies 

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

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

32 computation. 

33 

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

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

36 

37 Additional policy: 

38 

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

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

41 pass reuses the loaded activations without recomputation. 

42 

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

44 asynchronous offload/load and stream synchronization. 

45 

46 .. note:: 

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

48 computationally expensive or have large memory footprints. Note that 

49 swapping very small outputs may introduce additional overhead and 

50 reduce the effectiveness of asynchronous copy. 

51 """ 

52 MUST_SAVE = 0 

53 PREFER_SAVE = 1 

54 MUST_RECOMPUTE = 2 

55 PREFER_RECOMPUTE = 3 

56 

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

58 MUST_SWAP = 4 

59 

60 

61class _StackedCtx: 

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

63 

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

65 self._ctxs = list(ctxs) 

66 self._stack = contextlib.ExitStack() 

67 

68 def __enter__(self): 

69 self._stack.__enter__() 

70 for ctx in self._ctxs: 

71 self._stack.enter_context(ctx) 

72 return self 

73 

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

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

76 

77 

78def _compose_context_fns( 

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

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

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

82 

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

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

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

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

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

88 single-call contract. 

89 """ 

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

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

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

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

94 return _StackedCtx(fwd_ctxs), _StackedCtx(rec_ctxs) 

95 

96 return factory 

97 

98 

99def checkpoint( 

100 function, 

101 *args, 

102 swap_inputs: bool = False, 

103 policy_fn: Optional[Callable] = None, 

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

105 group_swap: bool = False, 

106 **kwargs, 

107): 

108 """ 

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

110 

111 Args: 

112 function: The function to apply checkpointing to. 

113 *args: Arguments to pass to the function. 

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

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

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

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

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

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

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

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

122 supplied together, the resulting factories are composed: their 

123 forward and recompute contexts are stacked so all enter in 

124 order and exit in reverse. 

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

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

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

128 

129 Returns: 

130 The result of applying the function with checkpointing. 

131 """ 

132 factories: list = [] 

133 if policy_fn is not None: 

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

135 if context_fn is not None: 

136 factories.append(context_fn) 

137 

138 if not factories: 

139 composed_context_fn = plat.noop_context_fn 

140 elif len(factories) == 1: 

141 composed_context_fn = factories[0] 

142 else: 

143 composed_context_fn = _compose_context_fns(tuple(factories)) 

144 

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

146 with context(): 

147 return plat.checkpoint( 

148 function, *args, context_fn=composed_context_fn, use_reentrant=False, **kwargs 

149 ) 

150 

151 

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

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

154 

155 Offloads intermediate activations saved by the autograd engine to CPU 

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

157 trading device memory for host memory bandwidth. Unlike 

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

159 

160 Args: 

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

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

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

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

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

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

167 eligible tensors are offloaded. 

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

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

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

171 

172 Returns: 

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

174 

175 Example: 

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

177 """ 

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

179 return function(*args, **kwargs) 

180 

181 

182checkpoint_wrapper = plat.checkpoint_wrapper 

183swap_wrapper = plat.swap_wrapper 

184swap_tensor_wrapper = plat.swap_tensor_wrapper