Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / reentrant_checkpoint.py: 94%
234 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""Eager reentrant checkpoint with explicit scheduled replay and local SAVE graphs."""
16import contextlib
17import contextvars
18import threading
19import weakref
20from dataclasses import dataclass
21from typing import Any, Callable, Iterator, Optional, Sequence, Union
23import torch
24from torch.utils._pytree import tree_flatten, tree_unflatten
26from .checkpoint import _ExecutionState, _noop_context_fn, _RECOMPUTE_COLLECTOR, _RECOMPUTE_SESSION
27from .checkpoint_exclude import _ExcludeCache
28from .reentrant_exclude import current_frame, execution_frame
31_SCOPE = contextvars.ContextVar("hyper_reentrant_checkpoint_scope", default=None)
32_SESSIONS = {}
33_SESSIONS_LOCK = threading.RLock()
36def _read_context(context_ref: Callable[[], Any]) -> Any:
37 """Dereference the context without retaining its autograd graph."""
38 return context_ref()
41@dataclass(frozen=True)
42class _InputSlot:
43 """Index of a Tensor in the positional Function inputs."""
45 index: int
48class ReentrantCheckpointHandle:
49 """One checkpoint invocation owned by an explicit stage-forward scope.
51 In a split stage, DX targets activations and must run before DW. DW targets
52 each handle's ``backward_trigger`` as well as the stage parameters. The
53 trigger makes closure-only parameter regions reachable without a synthetic
54 parameter edge that would invoke a leaf hook with ``None``. During DW,
55 captured parameters accumulate naturally in the retained replay graph;
56 explicit parameter inputs return their gradients through their original
57 outer edges. Distinct contributions to a shared parameter may each invoke
58 its hook, as they do in ordinary reentrant backward.
59 """
61 def __init__(self, function: Callable, template: list, spec: Any, params: Any,
62 preserve_rng_state: bool, context_fn: Callable) -> None:
63 """Capture invocation inputs and state without retaining transient activations."""
64 self.function = function
65 self.template = template
66 self.spec = spec
67 self.params = params
68 self.split = params is not None
69 self.preserve_rng_state = preserve_rng_state
70 self.context_fn = context_fn
71 self.exclude_cache = _ExcludeCache()
72 self.entries = []
73 self.replaying = False
74 self.phase = "ordinary"
75 self.replay_outputs = None
76 self.replay_inputs = ()
77 self.state = None
78 self.context_ref: Optional[Callable[[], Any]] = None
79 self.input_count = 0
80 self.param_input_mask = ()
81 self.backward_trigger = None
82 self.released = False
83 self.retain_graph = False
84 self.active_session = None
85 self.output_spec = None
86 self.output_template = ()
87 self.excluded_output_mask = {}
88 self.recording_exclude = False
90 def _call(self, tensors):
91 leaves = [tensors[item.index] if isinstance(item, _InputSlot) else item for item in self.template]
92 args, kwargs = tree_unflatten(leaves, self.spec)
93 with execution_frame(self):
94 return self.function(*args, **kwargs)
96 def _gradients(self, values: tuple, grads: tuple) -> tuple:
97 """Compute replay-input gradients without returning captured parameter edges."""
98 if self.split and self.phase == "dx":
99 targets = [item for item, is_param in zip(self.replay_inputs, self.param_input_mask)
100 if item.requires_grad and not is_param]
101 if not targets:
102 return (None,) * self.input_count
103 calculated = iter(torch.autograd.grad(values, targets, grads, retain_graph=True, allow_unused=True))
104 return tuple(next(calculated) if item.requires_grad and not is_param else None
105 for item, is_param in zip(self.replay_inputs, self.param_input_mask))
107 for item in self.replay_inputs:
108 if item.requires_grad:
109 item.grad = None
110 torch.autograd.backward(values, grads, retain_graph=self.retain_graph)
111 return tuple(item.grad if item.requires_grad else None for item in self.replay_inputs)
113 def recompute(self) -> None:
114 """Replay this frame once; retain its graph until DW or explicit cleanup."""
115 if self.released:
116 raise RuntimeError("Cannot replay a released reentrant checkpoint")
117 if self.replay_outputs is not None:
118 return
119 context_ref = self.context_ref
120 if context_ref is None:
121 raise RuntimeError("Reentrant checkpoint output graph was released before replay")
122 ctx = _read_context(context_ref)
123 if ctx is None:
124 raise RuntimeError("Reentrant checkpoint output graph was released before replay")
125 inputs = ctx.saved_tensors
126 self.replay_inputs = tuple(
127 tensor.detach().requires_grad_(tensor.requires_grad) if tensor.is_floating_point() or tensor.is_complex()
128 else tensor.detach() for tensor in inputs
129 )
130 self.replaying = True
131 try:
132 def replay(*unused: Any) -> None:
133 """Store the differentiable replay output for the owning frame."""
134 del unused
135 self.replay_outputs = self._call(self.replay_inputs)
137 with torch.enable_grad():
138 self.state.recompute_fn(replay, {})
139 finally:
140 self.replaying = False
142 def set_phase(self, phase: str) -> None:
143 """Select the phase used by this frame's engine callbacks (dx or dw)."""
144 if not self.split or phase not in {"dx", "dw"}:
145 raise ValueError("Split reentrant handles require phase='dx' or phase='dw'")
146 if self.released:
147 raise RuntimeError("Cannot set phase on a released reentrant checkpoint")
148 self.phase = phase
150 def release(self) -> None:
151 """Release all retained replay and excluded local graphs."""
152 self.exclude_cache.clear()
153 for entry in self.entries:
154 entry.tape.release()
155 entry.input_bindings.clear()
156 entry.output = None
157 self.entries.clear()
158 self.excluded_output_mask.clear()
159 self.replay_outputs = None
160 self.replay_inputs = ()
161 self.released = True
163 def bind_session(self, session_id: Any) -> None:
164 """Register retained state for cleanup through the unified session API."""
165 activation = _RECOMPUTE_SESSION.get()
166 if activation is not None and activation.session_id != session_id:
167 raise ValueError("Active session does not match the reentrant handle session")
168 with _SESSIONS_LOCK:
169 if self.active_session is not None and self.active_session.session_id != session_id:
170 raise RuntimeError("Concurrent reentrant recompute sessions on one frame are unsupported")
171 _SESSIONS.setdefault(session_id, weakref.WeakSet()).add(self)
172 if activation is not None:
173 self.active_session = activation
174 self.retain_graph = activation.retain_on_unpack
177def clear_session(session_id: Any) -> None:
178 """Release any reentrant handles registered with a recompute session."""
179 with _SESSIONS_LOCK:
180 handles = _SESSIONS.pop(session_id, ())
181 for handle in handles:
182 handle.release()
185def activate_session(activation: Any) -> None:
186 """Bind a session's retention policy before autograd enters worker threads."""
187 with _SESSIONS_LOCK:
188 for handle in _SESSIONS.get(activation.session_id, ()):
189 if handle.active_session is not None and handle.active_session is not activation:
190 raise RuntimeError("Concurrent reentrant recompute sessions on one frame are unsupported")
191 handle.active_session = activation
192 handle.retain_graph = activation.retain_on_unpack
195def deactivate_session(activation: Any) -> None:
196 """Unbind the policy without discarding registered frames or retained graph."""
197 with _SESSIONS_LOCK:
198 for handle in _SESSIONS.get(activation.session_id, ()):
199 if handle.active_session is activation:
200 handle.active_session = None
201 handle.retain_graph = False
204class _Checkpoint(torch.autograd.Function):
205 """Expose input edges and a split-backward trigger to the outer engine."""
207 @staticmethod
208 def forward(ctx: Any, frame: ReentrantCheckpointHandle, dummy: torch.Tensor,
209 *tensors: torch.Tensor) -> tuple:
210 """Record one no-grad invocation and attach its replay frame to outputs."""
211 del dummy
212 ctx.frame = frame
213 ctx.save_for_backward(*tensors[:frame.input_count])
214 frame.context_ref = weakref.ref(ctx)
215 frame.state = _ExecutionState.capture(frame.context_fn, frame.preserve_rng_state, tensors[:frame.input_count])
216 with torch.no_grad(), frame.state.forward_context:
217 output = frame._call(tensors[:frame.input_count]) # pylint: disable=protected-access
218 frame.state.check_device_state_preserved(frame.preserve_rng_state)
219 leaves, frame.output_spec = tree_flatten(output)
220 frame.output_template = tuple(
221 _InputSlot(index) if isinstance(value, torch.Tensor) else value
222 for index, value in enumerate(leaves)
223 )
224 outputs = tuple(value.detach() for value in leaves if isinstance(value, torch.Tensor))
225 if not outputs:
226 raise ValueError("Reentrant checkpoint requires at least one Tensor output")
227 originals = (value for value in leaves if isinstance(value, torch.Tensor))
228 def differentiable(original: torch.Tensor) -> bool:
229 """Match only the same live detached object, never a reused Python id."""
230 matched = frame.excluded_output_mask.get(id(original))
231 if matched is not None and matched[0]() is original:
232 return matched[1]
233 return original.is_floating_point() or original.is_complex()
235 ctx.mark_non_differentiable(*(
236 detached for detached, original in zip(outputs, originals) if not differentiable(original)
237 ))
238 frame.excluded_output_mask.clear()
239 return outputs
241 @staticmethod
242 def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple:
243 """Return target-limited gradients to the outer engine."""
244 try:
245 return _Checkpoint._run_backward(ctx, *grad_outputs)
246 except BaseException:
247 if not ctx.frame.retain_graph:
248 ctx.frame.release()
249 raise
251 @staticmethod
252 def _run_backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple:
253 """Perform the replay and route gradients according to the explicit phase."""
254 frame = ctx.frame
255 if frame.split and frame.phase == "ordinary":
256 raise RuntimeError("Set dx or dw phase on reentrant handles before split backward")
257 frame.recompute()
258 outputs = tuple(value for value in tree_flatten(frame.replay_outputs)[0] if isinstance(value, torch.Tensor))
259 selected = [
260 (value, grad) for value, grad in zip(outputs, grad_outputs)
261 if value.requires_grad and grad is not None
262 ]
263 input_grads = (frame._gradients(*zip(*selected)) # pylint: disable=protected-access
264 if selected else (None,) * frame.input_count)
265 if (not frame.split and not frame.retain_graph) or (frame.split and frame.phase == "dw"):
266 frame.release()
267 return (None, None, *input_grads, None)
270@contextlib.contextmanager
271def reentrant_checkpoint_scope(
272 parameters: Union[Sequence[torch.Tensor], Callable[[], Sequence[torch.Tensor]]],
273) -> Iterator[list]:
274 """Collect split checkpoint frames during stage forward using explicit parameters.
276 Args:
277 parameters: Stable iterable or zero-argument provider of stage parameter Tensors.
278 A provider is only evaluated if a reentrant checkpoint runs in the scope.
280 Yields:
281 List of invocation handles. Set each handle's phase before DX and DW;
282 include every ``backward_trigger`` in DW's autograd ``inputs`` tuple.
283 Release handles on cancellation; DW releases them after consumption.
284 """
285 if _SCOPE.get() is not None:
286 raise ValueError("Nested reentrant checkpoint scopes are not supported")
287 scope = [parameters, [], None]
288 token = _SCOPE.set(scope)
289 try:
290 yield scope[1]
291 finally:
292 _SCOPE.reset(token)
295def _checkpoint_inputs(args: tuple, kwargs: dict) -> tuple:
296 """Index Tensor leaves once while preserving the original pytree structure."""
297 leaves, spec = tree_flatten((args, kwargs))
298 inputs = []
299 template = []
300 for item in leaves:
301 if isinstance(item, torch.Tensor):
302 template.append(_InputSlot(len(inputs)))
303 inputs.append(item)
304 else:
305 template.append(item)
306 return tuple(inputs), template, spec
309def _scope_parameters(scope: Optional[list]) -> Optional[tuple]:
310 """Resolve and deduplicate scoped parameters only on the first checkpoint."""
311 if scope is None:
312 return None
313 if scope[2] is None:
314 supplied = tuple(scope[0]() if callable(scope[0]) else scope[0])
315 if any(not isinstance(param, torch.Tensor) for param in supplied):
316 raise ValueError("Scope parameters must be Tensors")
317 scope[2] = tuple(dict((id(param), param) for param in supplied if param.requires_grad).values())
318 return scope[2]
321def _collect_frame(frame: ReentrantCheckpointHandle, scope: Optional[list]) -> Optional[list]:
322 """Register a frame for the current scope or unified collector."""
323 if scope is not None:
324 scope[1].append(frame)
325 return scope[1]
326 collector = _RECOMPUTE_COLLECTOR.get()
327 if collector is not None:
328 collector.append(frame)
329 return collector
332def checkpoint(function: Callable, *args: Any, preserve_rng_state: bool = True,
333 context_fn: Callable = _noop_context_fn, **kwargs: Any) -> Any:
334 """Run an eager reentrant checkpoint with localized excluded-region graphs.
336 Ordinary backward accumulates captured parameter gradients naturally.
337 Unscoped ``autograd.grad`` does not filter these nested accumulations by
338 its target list; use ``backward`` for ordinary execution and an explicit
339 scope for staged DX/DW. Nested reentrant checkpoints are unsupported.
340 """
341 if current_frame() is not None:
342 raise ValueError("Nested reentrant checkpoints are not supported")
343 if not torch.is_grad_enabled():
344 return function(*args, **kwargs)
345 inputs, template, spec = _checkpoint_inputs(args, kwargs)
346 scope = _SCOPE.get()
347 params = _scope_parameters(scope)
348 frame = ReentrantCheckpointHandle(function, template, spec, params, preserve_rng_state, context_fn)
349 frame.input_count = len(inputs)
350 frame.param_input_mask = tuple(id(item) in {id(param) for param in params or ()} for item in inputs)
351 frame.backward_trigger = torch.empty(0, requires_grad=True) if frame.split else None
352 collector = _collect_frame(frame, scope)
353 try:
354 result = _Checkpoint.apply(frame, torch.empty(0, requires_grad=True), *inputs, frame.backward_trigger)
355 except BaseException:
356 frame.release()
357 if collector is not None:
358 collector.remove(frame)
359 raise
360 tensors = (result,) if isinstance(result, torch.Tensor) else result
361 iterator = iter(tensors)
362 return tree_unflatten(
363 [next(iterator) if isinstance(item, _InputSlot) else item for item in frame.output_template],
364 frame.output_spec,
365 )