Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / sac.py: 98%
130 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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.
15# Adapted from https://github.com/pytorch/pytorch/blob/release/2.6/torch/utils/checkpoint.py
16# enhanced with selective checkpoint support swap
17# ============================================================================
18"""enhanced with selective checkpoint support swap"""
19# pylint: disable=W0212, W0613, C0115, C0116, C0103, R1705
20from collections import defaultdict
21from functools import partial
22from typing import Any, Dict, List, Optional, Union
24import torch
25import torch.fx.traceback as fx_traceback
26from torch._functorch._aot_autograd.functional_utils import is_fun
27from torch.utils._pytree import tree_map
28from torch.utils._python_dispatch import TorchDispatchMode
29from hyper_parallel.core.activation_memory.policy import CheckpointPolicy
30from hyper_parallel.core.activation_memory.swap import ( # patch code
31 SwapManager,
32 SwapTensor,
33 Storage,
34)
37def _is_compiling(func, args, kwargs):
38 # Check if we are under AOTAutograd tracing
39 # There should probably be a better way to do this...
40 # NOTE: unify _is_compiling across all compile stacks
41 for arg in args:
42 if isinstance(arg, torch.Tensor) and is_fun(arg):
43 return True
44 return False
47class _VersionWrapper:
48 # Check that cached tensors are not mutated.
49 def __init__(self, val):
50 self.val: Union[torch.Tensor, Any] = val
51 self.version: Optional[int] = (
52 val._version if isinstance(val, torch.Tensor) else None
53 )
55 def get_val(self, allow_cache_entry_mutation):
56 if self.version is not None and not allow_cache_entry_mutation:
57 if self.val._version != self.version:
58 # Can we give user a stack trace of where the mutation happened?
59 raise RuntimeError(
60 "Tensor cached during selective activation checkpoint has been mutated"
61 )
62 return self.val
65class _SwapCacheEntry:
66 """Pair the recompute cache and swap record around the same tensor object."""
67 def __init__(self, val, funcname, group_swap=False, cpu_pool=None):
68 self.save = _VersionWrapper(val)
69 self.swap = SwapTensor(val, funcname, group_swap=group_swap, cpu_pool=cpu_pool)
72def _maybe_detach(x, any_ret_has_alias_info):
73 # We detach for two separate reasons:
74 # - For view ops, we need to ensure that when the tensor is returned from
75 # CachedDispatchMode, as_view sees that the AutogradMeta is nullptr
76 # - Avoid reference cycles
77 # For case 1, it is not enough to check whether x has differentiable dtype
78 # because non-differentiable dtype can have non-nullptr AutogradMeta, e.g.
79 # when the tensor is a view.
80 need_detach = (isinstance(x, torch.Tensor)
81 and (x.is_floating_point() or x.is_complex() or any_ret_has_alias_info))
82 if need_detach:
83 with torch._C._SetExcludeDispatchKeyGuard(torch._C.DispatchKey.ADInplaceOrView, False):
84 # Ensure that view performed beneath autograd properly propagates
85 # version counter. TODO: Use reentrant_dispatch instead of
86 # manually manipulating dispatch keys. Using reentrant_dispatch
87 # would respect inference_mode, though that is not relevant for
88 # this case.
89 x = x.detach()
90 return x
93class SelectiveCheckpointContext:
94 """
95 Context passed to policy function during selective checkpointing.
97 This class is used to pass relevant metadata to the policy function during
98 selective checkpointing. The metadata includes whether the current invocation
99 of the policy function is during recomputation or not.
101 Example:
102 >>> # xdoctest: +SKIP(stub)
103 >>>
104 >>> def policy_fn(ctx, op, *args, **kwargs):
105 >>> print(ctx.is_recompute)
106 >>>
107 >>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn)
108 >>>
109 >>> out = torch.utils.checkpoint.checkpoint(
110 >>> fn, x, y,
111 >>> use_reentrant=False,
112 >>> context_fn=context_fn,
113 >>> )
114 """
115 def __init__(self, *, is_recompute):
116 self.is_recompute = is_recompute
119def _policy_from_bool(b):
120 # For backward compatibility
121 return CheckpointPolicy.MUST_SAVE if b else CheckpointPolicy.PREFER_RECOMPUTE
124def _make_swap_entry(x, *, has_alias, funcname, group_swap, cpu_pool):
125 return _SwapCacheEntry(
126 _maybe_detach(x, has_alias), funcname, group_swap=group_swap, cpu_pool=cpu_pool
127 )
130SAC_IGNORED_OPS = {
131 # AC inserts different number of detach during forward and recompute.
132 torch.ops.aten.detach.default,
133 # AC's determinism check invokes additional metadata ops during forward.
134 # With subclasses involved, these metadata ops become dispatchable, this
135 # can result in incorrectness if these ops are selected cached.
136 torch.ops.prim.device.default,
137} | set(torch._subclasses.functional_tensor.FunctionalTensor.metadata_fns)
140def ignore_sac_ops(ignore_ops: List[Optional[object]]) -> None:
141 """Add available operators to the selective-checkpoint ignore set.
143 Args:
144 ops (List[Optional[object]]): Operators to execute without selective-AC replay accounting.
145 ``None`` entries are ignored for optional-version compatibility.
146 """
147 SAC_IGNORED_OPS.update(op for op in ignore_ops if op is not None)
150class _CachingTorchDispatchMode(TorchDispatchMode):
151 # Used together with _CachedTorchDispatchMode to implement SAC.
152 def __init__(self, policy_fn, swap_storage, storage, group_swap=False, cpu_pool=None):
153 self.policy_fn = policy_fn
154 self.swap_storage = swap_storage
155 self.storage = storage
156 self.add_to_storage = False
157 self.group_swap = group_swap
158 self.cpu_pool = cpu_pool
159 # Cache context and singleton to avoid per-dispatch allocation / lookup.
160 self._swap_manager = SwapManager()
161 self._group_prefix = ""
163 def __torch_dispatch__(self, func, types, args=(), kwargs=None):
164 if func in SAC_IGNORED_OPS:
165 return func(*args, **kwargs)
167 kwargs = {} if kwargs is None else kwargs
168 policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False),
169 func, *args, **kwargs)
170 if isinstance(policy, bool):
171 policy = _policy_from_bool(policy)
173 is_compiling = _is_compiling(func, args, kwargs)
175 if is_compiling:
176 # Overwrite each node's "recompute" tag to add in the user annotation.
177 fx_traceback.current_meta["recompute"] = policy
179 out = func(*args, **kwargs)
181 has_alias = any(ret.alias_info is not None for ret in func._schema.returns)
183 if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE):
184 self.storage[func].append(
185 tree_map(lambda x: _VersionWrapper(_maybe_detach(x, has_alias)), out)
186 )
187 elif policy == CheckpointPolicy.MUST_SWAP: # patch code
188 group_name = self._swap_manager.get_current_group_name()
189 if not group_name:
190 # No swap group is active for this region, which is the normal
191 # case for pipeline chunks the schedule did not mark swappable
192 # (forward/backward distance below MIN_SWAP_GAP). Degrade to
193 # MUST_SAVE: cache the activation on device and register it
194 # nowhere, so the empty-named global group never accumulates
195 # storage that is neither offloaded nor released.
196 self._swap_manager.warn_missing_group_once()
197 self.storage[func].append(
198 tree_map(lambda x: _VersionWrapper(_maybe_detach(x, has_alias)), out)
199 )
200 else:
201 if not self.add_to_storage:
202 self._group_prefix = f"{group_name}::"
203 self._swap_manager.add_storage(group_name, self.swap_storage)
204 self.add_to_storage = True
205 funcname = f"{self._group_prefix}{func}"
206 entries = tree_map(
207 partial(
208 _make_swap_entry,
209 has_alias=has_alias,
210 funcname=funcname,
211 group_swap=self.group_swap,
212 cpu_pool=self.cpu_pool,
213 ),
214 out,
215 )
216 self.storage[func].append(tree_map(lambda x: x.save, entries))
217 self.swap_storage[func].append(tree_map(lambda x: x.swap, entries))
218 elif policy != CheckpointPolicy.MUST_RECOMPUTE:
219 raise RuntimeError(f"Checkpoint Activation: {func} encountered an invalid policy {policy}")
220 return out
223class _CachedTorchDispatchMode(TorchDispatchMode):
224 # Used together with _CachingTorchDispatchMode to implement SAC.
225 def __init__(self, policy_fn, swap_storage, storage, allow_cache_entry_mutation):
226 self.policy_fn = policy_fn
227 self.swap_storage = swap_storage
228 self.storage = storage
229 self.allow_cache_entry_mutation = allow_cache_entry_mutation
230 self._swap_cleared = False
232 def __torch_dispatch__(self, func, types, args=(), kwargs=None):
233 if func in SAC_IGNORED_OPS:
234 return func(*args, **kwargs)
236 kwargs = {} if kwargs is None else kwargs
237 policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=True),
238 func, *args, **kwargs)
239 if isinstance(policy, bool):
240 policy = _policy_from_bool(policy)
242 is_compiling = _is_compiling(func, args, kwargs)
244 if not self._swap_cleared:
245 self.swap_storage.clear()
246 self._swap_cleared = True
248 # MUST_SAVE, PREFER_SAVE, and MUST_SWAP all restore from storage identically.
249 if (policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE, CheckpointPolicy.MUST_SWAP)
250 or is_compiling):
251 storage = self.storage.get(func) # patch code
252 if storage is None:
253 raise RuntimeError(f"{func} encountered during backward, but not found in storage")
254 if len(storage) == 0:
255 raise RuntimeError(
256 "Trying to backward an extra time. You are only allowed to backward once "
257 "on any region computed under selective activation checkpoint."
258 )
259 out = tree_map(lambda x: x.get_val(self.allow_cache_entry_mutation), storage.pop(0))
260 else:
261 out = func(*args, **kwargs)
262 return out
265def create_selective_checkpoint_contexts(
266 policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False, cpu_pool=None
267):
268 """
269 Helper to avoid recomputing certain ops during activation checkpointing.
271 Use this with `torch.utils.checkpoint.checkpoint` to control which
272 operations are recomputed during the backward pass.
274 Args:
275 policy_fn_or_list (Callable or List):
276 - If a policy function is provided, it should accept a
277 :class:`SelectiveCheckpointContext`, the :class:`OpOverload`, args and
278 kwargs to the op, and return a :class:`CheckpointPolicy` enum value
279 indicating whether the execution of the op should be recomputed or not.
280 - If a list of operations is provided, it is equivalent to a policy
281 returning `CheckpointPolicy.MUST_SAVE` for the specified
282 operations and `CheckpointPolicy.PREFER_RECOMPUTE` for all other
283 operations.
284 allow_cache_entry_mutation (bool, optional): By default, an error is
285 raised if any tensors cached by selective activation checkpoint are
286 mutated in order to ensure correctness. If set to `True`, this check
287 is disabled.
288 Returns:
289 A tuple of two context managers.
291 Example:
292 >>> # xdoctest: +REQUIRES(LINUX)
293 >>> import functools
294 >>>
295 >>> x = torch.rand(10, 10, requires_grad=True)
296 >>> y = torch.rand(10, 10, requires_grad=True)
297 >>>
298 >>> ops_to_save = [
299 >>> torch.ops.aten.mm.default,
300 >>> ]
301 >>>
302 >>> def policy_fn(ctx, op, *args, **kwargs):
303 >>> if op in ops_to_save:
304 >>> return CheckpointPolicy.MUST_SAVE
305 >>> else:
306 >>> return CheckpointPolicy.PREFER_RECOMPUTE
307 >>>
308 >>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn)
309 >>>
310 >>> # or equivalently
311 >>> context_fn = functools.partial(create_selective_checkpoint_contexts, ops_to_save)
312 >>>
313 >>> def fn(x, y):
314 >>> return torch.sigmoid(torch.matmul(torch.matmul(x, y), y)) * y
315 >>>
316 >>> out = torch.utils.checkpoint.checkpoint(
317 >>> fn, x, y,
318 >>> use_reentrant=False,
319 >>> context_fn=context_fn,
320 >>> )
321 """
322 # NB: If grad_mode is disabled, checkpoint would not run forward under
323 # context_fn anyway, so proceed as usual.
324 if policy_fn_or_list is None:
325 def policy_fn(_ctx, _op, *_args, **_kwargs):
326 return CheckpointPolicy.PREFER_RECOMPUTE
327 elif isinstance(policy_fn_or_list, list):
328 for op in policy_fn_or_list:
329 if not isinstance(op, torch._ops.OpOverload):
330 _extra_msg = (
331 "Please update the OpOverloadPacket to a specific OpOverload."
332 "For example, if you have `torch.ops.aten.mm`, change it to `torch.ops.aten.mm.default`."
333 ) if isinstance(op, torch._ops.OpOverloadPacket) else ""
334 raise ValueError(
335 f"Expected op in `op_list` to be an OpOverload but got: {op} "
336 f"of type {type(op)}. {_extra_msg}"
337 )
339 def policy_fn(ctx, op, *args, **kwargs):
340 if op in policy_fn_or_list:
341 return CheckpointPolicy.MUST_SAVE
342 else:
343 return CheckpointPolicy.PREFER_RECOMPUTE
344 elif callable(policy_fn_or_list):
345 policy_fn = policy_fn_or_list
346 else:
347 raise TypeError("policy_fn_or_list must be either a function or a list of ops.")
349 swap_storage = Storage() # patch code
350 storage: Dict[Any, List[Any]] = defaultdict(list)
351 return (
352 _CachingTorchDispatchMode(
353 policy_fn, swap_storage, storage, group_swap=group_swap, cpu_pool=cpu_pool
354 ),
355 _CachedTorchDispatchMode(policy_fn, swap_storage, storage, allow_cache_entry_mutation),
356 )