Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / reentrant_exclude.py: 96%
169 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"""Local autograd graphs for excluded regions inside a reentrant checkpoint."""
16import contextlib
17import contextvars
18import weakref
19from typing import Any, Iterator, Optional, Tuple
21import torch
22from torch.utils._pytree import tree_flatten, tree_unflatten
24from . import checkpoint_exclude as recipes
27_CURRENT_FRAME = contextvars.ContextVar("hyper_reentrant_frame", default=None)
30@contextlib.contextmanager
31def execution_frame(frame: Any) -> Iterator[None]:
32 """Expose a frame only while running its forward or replay function."""
33 token = _CURRENT_FRAME.set(frame)
34 try:
35 yield
36 finally:
37 _CURRENT_FRAME.reset(token)
40class _BorrowedInputs:
41 """Pass real activations to a Function without making them autograd inputs."""
43 def __init__(self, tensors: Tuple[torch.Tensor, ...]) -> None:
44 """Borrow Tensor data only until the boundary forward has detached it."""
45 self.tensors = tensors
46 self.slots = []
49class _InputBoundary(torch.autograd.Function):
50 """Connect local gradients to a zero-element root without retaining inputs."""
52 @staticmethod
53 def forward(ctx: Any, borrowed: _BorrowedInputs, dummy: torch.Tensor) -> tuple:
54 """Return detached aliases while storing gradient slots, not activations."""
55 del dummy
56 ctx.set_materialize_grads(False)
57 ctx.slots = [None] * len(borrowed.tensors)
58 borrowed.slots = ctx.slots
59 outputs = tuple(tensor.detach() for tensor in borrowed.tensors)
60 borrowed.tensors = None
61 return outputs
63 @staticmethod
64 def backward(ctx: Any, *grads: torch.Tensor) -> tuple:
65 """Record local input sensitivities for the replay bridge."""
66 ctx.slots[:] = grads
67 return None, None
70class _LocalTape:
71 """Keep the original excluded graph and only lightweight input gradient edges."""
73 def __init__(self, outputs: tuple, input_edges: tuple, slots: list, frame: Any) -> None:
74 """Retain the local graph while holding only gradient edges to its inputs."""
75 self.outputs = outputs
76 self.input_edges = input_edges
77 self.slots = slots
78 self.frame_ref = weakref.ref(frame)
79 self.output_mask = tuple(item.requires_grad for item in outputs)
80 self.param_input_mask = ()
81 self.param_aliases = ()
83 def gradients(self, grad_outputs: tuple) -> tuple:
84 """Compute local gradients according to the owning frame's explicit phase."""
85 frame = self.frame_ref()
86 if frame is None:
87 raise RuntimeError("Reentrant checkpoint frame was released before exclude backward")
88 pairs = [(out, grad) for out, grad in zip(self.outputs, grad_outputs) if out.requires_grad and grad is not None]
89 if not pairs:
90 return (None,) * len(self.input_edges)
91 outputs, grads = zip(*pairs)
92 if frame.phase in {"ordinary", "dw"}:
93 return self._backward(outputs, grads, frame.retain_graph)
94 return self._dx_gradients(outputs, grads)
96 def _backward(self, outputs: tuple, grads: tuple, retain_graph: bool) -> tuple:
97 """Accumulate captured parameters once and return explicit input sensitivities."""
98 self.slots[:] = [None] * len(self.slots)
99 for alias in self.param_aliases:
100 if alias is not None:
101 alias.grad = None
102 torch.autograd.backward(outputs, grads, retain_graph=retain_graph)
103 slot_grads = iter(self.slots)
104 seen_aliases = set()
106 def alias_grad(alias: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
107 """Return a shared explicit proxy gradient only at its first input edge."""
108 if alias is None or id(alias) in seen_aliases:
109 return None
110 seen_aliases.add(id(alias))
111 return alias.grad
113 return tuple(
114 alias_grad(alias) if is_param else next(slot_grads)
115 for is_param, alias in zip(self.param_input_mask, self.param_aliases)
116 )
118 def _dx_gradients(self, outputs: tuple, grads: tuple) -> tuple:
119 """Target only activation boundaries, leaving parameter hooks untouched."""
120 targets = [edge for edge, is_param in zip(self.input_edges, self.param_input_mask)
121 if edge is not None and not is_param]
122 if not targets:
123 return (None,) * len(self.input_edges)
124 result = iter(torch.autograd.grad(outputs, targets, grads, retain_graph=True, allow_unused=True))
125 return tuple(
126 next(result) if edge is not None and not is_param else None
127 for edge, is_param in zip(self.input_edges, self.param_input_mask)
128 )
130 def release(self) -> None:
131 """Drop graph outputs and boundary gradient edges after use."""
132 self.outputs = ()
133 self.input_edges = ()
134 self.slots = []
135 self.param_aliases = ()
136 self.param_input_mask = ()
139class _ReplayBridge(torch.autograd.Function):
140 """Connect the cached local graph to current replay inputs and its trigger."""
142 @staticmethod
143 def forward(ctx: Any, entry: Any, input_count: int, *tensors: torch.Tensor) -> tuple:
144 """Return detached outputs with the recorded differentiability mask."""
145 del tensors
146 ctx.entry = entry
147 ctx.input_count = input_count
148 ctx.set_materialize_grads(False)
149 outputs = tuple(tensor.detach() for tensor in entry.tape.outputs)
150 ctx.mark_non_differentiable(*(
151 output for output, differentiable in zip(outputs, entry.tape.output_mask)
152 if not differentiable
153 ))
154 return outputs
156 @staticmethod
157 def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple:
158 """Propagate through the original local graph using the frame phase."""
159 input_grads = ctx.entry.tape.gradients(grad_outputs)
160 return (None, None, *input_grads, None)
163def _replace_tensor_leaves(value, tensors):
164 leaves, spec = tree_flatten(value)
165 replacements = iter(tensors)
166 return tree_unflatten([next(replacements) if isinstance(item, torch.Tensor) else item for item in leaves], spec)
169def _tensor_leaves(value: Any) -> tuple:
170 """Flatten the supported input/output pytree to its Tensor leaves."""
171 return tuple(item for item in tree_flatten(value)[0] if isinstance(item, torch.Tensor))
174def _replay_exclude(wrapper: Any, args: tuple, kwargs: dict, frame: Any) -> Any:
175 """Bind recipes and bridge cached outputs to the replay graph."""
176 entry = frame.exclude_cache.pop(id(wrapper))
177 frame.entries.append(entry)
178 recipes._restore_rng_state(entry.rng_state) # pylint: disable=protected-access
179 recipes._materialize_recompute_inputs(entry, args, kwargs) # pylint: disable=protected-access
180 current = _tensor_leaves((args, kwargs))
181 if len(current) != len(entry.tape.input_edges):
182 raise RuntimeError("Checkpoint exclusion Tensor input count changed during replay")
183 result = _ReplayBridge.apply(entry, len(current), *current, torch.empty(0, requires_grad=True))
184 tensors = (result,) if isinstance(result, torch.Tensor) else result
185 return _replace_tensor_leaves(entry.output, tensors)
188def _local_inputs(original: tuple, param_ids: set) -> tuple:
189 """Build a local activation boundary and explicit parameter proxies."""
190 borrowed = _BorrowedInputs(tuple(item for item in original if id(item) not in param_ids))
191 boundary = _InputBoundary.apply(borrowed, torch.empty(0, requires_grad=True)) if borrowed.tensors else ()
192 boundary = (boundary,) if isinstance(boundary, torch.Tensor) else boundary
193 replacements = iter(boundary)
194 param_aliases = {}
195 for item in original:
196 if id(item) in param_ids and id(item) not in param_aliases:
197 param_aliases[id(item)] = item.detach().requires_grad_(item.requires_grad)
198 local = (param_aliases[id(item)] if id(item) in param_ids else next(replacements) for item in original)
199 return tuple(local), boundary, borrowed.slots, param_aliases
202def _input_edges(original: tuple, boundary: tuple, param_ids: set, param_aliases: dict) -> tuple:
203 """Expose one gradient edge per original input without retaining activations."""
204 boundary_edges = iter(
205 torch.autograd.graph.get_gradient_edge(tensor) if tensor.requires_grad else None for tensor in boundary
206 )
207 seen_params = set()
208 edges = []
209 for item in original:
210 if id(item) not in param_ids:
211 edges.append(next(boundary_edges))
212 elif item.requires_grad and id(item) not in seen_params:
213 edges.append(torch.autograd.graph.get_gradient_edge(param_aliases[id(item)]))
214 seen_params.add(id(item))
215 else:
216 edges.append(None)
217 return tuple(edges)
220def _record_exclude(wrapper: Any, module: Any, args: tuple, kwargs: dict, frame: Any) -> Any:
221 """Execute an excluded region once under localized saved-tensor hooks."""
222 if frame.recording_exclude:
223 raise ValueError("Nested checkpoint exclusion wrappers are not supported")
224 original = _tensor_leaves((args, kwargs))
225 param_ids = {id(param) for param in frame.params or ()} if frame.split else set()
226 with torch.enable_grad():
227 local, boundary, slots, param_aliases = _local_inputs(original, param_ids)
228 local_args, local_kwargs = _replace_tensor_leaves((args, kwargs), local)
229 infos = recipes._capture_recompute_inputs(local_args, local_kwargs) # pylint: disable=protected-access
230 bindings = []
231 swap_hooks = wrapper._make_exclude_swap_hooks() # pylint: disable=protected-access
232 token = recipes._ACTIVE_PACK_STATE.set( # pylint: disable=protected-access
233 recipes._PackState(infos, bindings, swap_hooks=swap_hooks) # pylint: disable=protected-access
234 )
235 frame.recording_exclude = True
236 try:
237 with recipes._saved_tensors_context(): # pylint: disable=protected-access
238 output = module(*local_args, **local_kwargs)
239 finally:
240 frame.recording_exclude = False
241 recipes._ACTIVE_PACK_STATE.reset(token) # pylint: disable=protected-access
242 if swap_hooks is not None:
243 wrapper._register_exclude_outputs(output) # pylint: disable=protected-access
244 outputs = _tensor_leaves(output)
245 if not outputs:
246 raise ValueError("Reentrant checkpoint exclusion requires at least one Tensor output")
247 tape = _LocalTape(outputs, _input_edges(original, boundary, param_ids, param_aliases), slots, frame)
248 tape.param_input_mask = tuple(id(item) in param_ids for item in original)
249 tape.param_aliases = tuple(param_aliases.get(id(item)) for item in original)
250 entry = recipes._ExcludeCacheEntry( # pylint: disable=protected-access
251 output, bindings, rng_state=recipes._capture_rng_state( # pylint: disable=protected-access
252 args, kwargs, output, frame.preserve_rng_state
253 ),
254 tape=tape,
255 )
256 frame.exclude_cache.save(id(wrapper), entry)
257 detached = tuple(tensor.detach() for tensor in outputs)
258 frame.excluded_output_mask.update(
259 (id(item), (weakref.ref(item), mask)) for item, mask in zip(detached, tape.output_mask)
260 )
261 return _replace_tensor_leaves(output, detached)
264def run_exclude(wrapper: Any, module: Any, args: tuple, kwargs: dict, frame: Any) -> Any:
265 """Record an excluded body once, or bridge its output during replay."""
266 if frame.replaying:
267 return _replay_exclude(wrapper, args, kwargs, frame)
268 return _record_exclude(wrapper, module, args, kwargs, frame)
271def current_frame() -> Optional[Any]:
272 """Return the frame active in the caller's forward/replay thread."""
273 return _CURRENT_FRAME.get()