Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / execution / base.py: 100%
113 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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"""Base executor definition."""
17from __future__ import annotations
19import logging
20import weakref
21from abc import ABC, abstractmethod
22from collections import defaultdict
23from collections.abc import Callable
24from typing import TYPE_CHECKING, Any
26import torch
27from torch.utils._pytree import tree_flatten, tree_unflatten
29from hyper_parallel.auto_parallel.hyper_offload.execution.tensor import ShadowTensor
31if TYPE_CHECKING:
32 from hyper_parallel.auto_parallel.hyper_offload.runtime.residency import ResidencyManager
34logger = logging.getLogger(__name__)
37class OpaqueRegionStart(torch.autograd.Function):
38 """Autograd function to mark the start of an opaque region."""
40 @staticmethod
41 def forward(ctx: Any, executor: BaseExecutor, func_name: str, dummy: torch.Tensor, *inputs: Any) -> Any:
42 """Forward pass for the start boundary."""
43 ctx.executor = executor
44 ctx.func_name = func_name
45 return (dummy,) + inputs
47 @staticmethod
48 def backward(ctx: Any, grad_dummy: Any, *grad_inputs: Any) -> Any: # pylint: disable=unused-argument
49 """Backward pass for the start boundary."""
50 executor = ctx.executor
51 executor.exit_opaque_region()
53 # Reconstruct the backward op using on_op_end.
54 # Outputs of the backward op are grad_inputs.
55 # func/args/kwargs were cached by on_op_begin in OpaqueRegionEnd.backward.
56 # on_op_end already shadows via apply_shadows and caches the bindings,
57 # so we can use its return value directly.
58 grad_inputs = executor.on_op_end(grad_inputs)
60 return (None, None, None) + tuple(grad_inputs)
63class OpaqueRegionEnd(torch.autograd.Function):
64 """Autograd function to mark the end of an opaque region."""
66 @staticmethod
67 def forward(ctx: Any, executor: BaseExecutor, func_name: str, dummy: torch.Tensor, *outputs: Any) -> Any: # pylint: disable=unused-argument
68 """Forward pass for the end boundary.
70 Wraps raw tensors into :class:`ShadowTensor` **inside** the
71 autograd boundary. This is required for the autograd engine to
72 correctly link tensor-subclass outputs into the computation graph
73 (subclass fixup). The wrapping happens here instead of in
74 :meth:`execute_opaque_op` so that :meth:`on_op_end` always
75 receives raw tensors regardless of execution path.
76 """
77 ctx.executor = executor
78 ctx.func_name = func_name
79 # 'outputs' is a flat tuple of tensors passed via .apply(*flat_res).
80 # on_op_end has already been called and cached output_bindings
81 # via apply_shadows; use them to avoid a redundant traversal.
82 # pylint: disable=protected-access
83 return tuple(executor.apply_shadows(outputs, executor._last_output_bindings))
85 @staticmethod
86 def backward(ctx: Any, *grad_outputs: Any) -> Any:
87 """Backward pass for the end boundary."""
88 executor = ctx.executor
90 def bwd_dummy(*_args: Any, **_kwargs: Any) -> Any:
91 """Dummy backward function that records the op boundary."""
93 bwd_dummy.__name__ = ctx.func_name + "_bwd"
95 executor.on_op_begin(bwd_dummy, grad_outputs, {})
97 executor.enter_opaque_region()
98 return (None, None, torch.zeros(1)) + grad_outputs
101class BaseExecutor(ABC):
102 """Abstract base class for execution phases (warmup or replay).
104 Executors implement phase-specific lifecycle callbacks. Raw PyTorch
105 dispatch mechanics are handled by :class:`ActivationDispatchMode`.
106 """
108 def __init__(
109 self,
110 residency_manager: ResidencyManager,
111 ) -> None:
112 self.residency_manager = residency_manager
113 #: sid -> WeakSet of alive shadows (used only for ``retained_sids``
114 #: computation at the end of warmup.
115 self._alive_shadows: dict[int, weakref.WeakSet[ShadowTensor]] = defaultdict(weakref.WeakSet)
116 self.op_idx: int = -1
117 self._opaque_depth: int = 0
119 # Cached by on_op_begin for use in on_op_end.
120 self._last_func = None
121 self._last_args = None
122 self._last_kwargs = None
123 # Cached by apply_shadows for use in autograd boundaries.
124 self._last_output_bindings: dict[int, int] | None = None
126 @property
127 def in_opaque_region(self) -> bool:
128 """Return True if the executor is currently inside an opaque region."""
129 return self._opaque_depth > 0
131 def enter_opaque_region(self) -> None:
132 """Enter an opaque region where fine-grained tracing is suspended."""
133 self._opaque_depth += 1
135 def exit_opaque_region(self) -> None:
136 """Exit an opaque region."""
137 self._opaque_depth -= 1
139 def execute_opaque_op(self, func_name: str, fn: Callable, args: tuple, kwargs: dict) -> Any:
140 """Execute a function as a single virtual op.
142 Wraps the function execution into a single "virtual op" in the
143 execution trace, while suspending fine-grained tracing for
144 internal operations.
146 The lifecycle hook order is:
148 1. :meth:`on_op_begin` — pre-actions.
149 2. Opaque region (inner ops bypass lifecycle hooks).
150 3. :meth:`on_op_end` — trace recording with **raw** tensors
151 (consistent with :meth:`dispatch`).
152 4. :meth:`OpaqueRegionEnd` — autograd boundary that wraps
153 outputs into :class:`ShadowTensor` **inside** the autograd
154 function (required for correct subclass graph linkage).
156 Steps 3 and 4 happen **inside** the opaque region so that any
157 incidental dispatch triggered by the autograd engine during
158 :meth:`OpaqueRegionEnd.apply` does not invoke lifecycle hooks.
159 """
160 if self.in_opaque_region:
161 return fn(*args, **kwargs)
163 def fwd_dummy(*_a: Any, **_kw: Any) -> Any:
164 """Dummy forward function that records the op boundary."""
166 fwd_dummy.__name__ = func_name + "_fwd"
168 self.enter_opaque_region()
169 try:
170 self.on_op_begin(fwd_dummy, args, kwargs)
172 # Inject a dummy tensor to ensure backward graph continuity
173 dummy = torch.zeros(1, requires_grad=True)
174 flat_args, spec_args = tree_flatten((args, kwargs))
176 # Boundary 1: Wrap inputs to delay backward virtual step exit
177 out_start = OpaqueRegionStart.apply(self, func_name, dummy, *flat_args)
178 dummy_out = out_start[0]
179 flat_args_out = out_start[1:]
181 args_out, kwargs_out = tree_unflatten(flat_args_out, spec_args)
183 result = fn(*args_out, **kwargs_out)
185 # Lifecycle: After op — record trace with RAW tensors.
186 self.on_op_end(result)
188 # Boundary 2: Wrap outputs inside autograd (shadowing must
189 # happen inside the autograd Function for the engine to
190 # correctly link ShadowTensor subclasses into the graph).
191 flat_res, spec_res = tree_flatten(result)
192 out_end = OpaqueRegionEnd.apply(self, func_name, dummy_out, *flat_res)
193 result_unflat = tree_unflatten(out_end, spec_res)
194 finally:
195 self.exit_opaque_region()
197 return result_unflat
199 # ------------------------------------------------------------------
200 # Dispatch (template method)
201 # ------------------------------------------------------------------
203 def dispatch(self, func, args, kwargs):
204 """Dispatch *func* with *args*/*kwargs*.
206 When the executor is inside an opaque region (e.g. executing the
207 internals of a virtual op), *func* is called directly,
208 **skipping** the lifecycle hooks
209 (:meth:`on_op_begin` / :meth:`on_op_end`) since those
210 boundaries are managed by ``OpaqueRegionStart`` / ``OpaqueRegionEnd``.
212 Otherwise the standard slow-path is taken:
214 1. :meth:`on_op_begin` — pre-actions (prefetch, etc.).
215 2. ``func(*args, **kwargs)`` — raw execution.
216 3. :meth:`on_op_end` — trace recording, post-actions **and**
217 :class:`ShadowTensor` wrapping (single pass).
218 """
219 if self.in_opaque_region:
220 return func(*args, **kwargs)
222 self.on_op_begin(func, args, kwargs)
223 result = func(*args, **kwargs)
224 return self.on_op_end(result)
226 # ------------------------------------------------------------------
227 # Lifecycle hooks
228 # ------------------------------------------------------------------
230 def on_op_begin(self, func, args, kwargs) -> None:
231 """Lifecycle callback before the operator is executed.
233 Caches *func*, *args*, *kwargs* so that :meth:`on_op_end`
234 can access them via ``self._last_func`` etc. Subclasses that
235 override this method **must** call ``super().on_op_begin(...)``
236 to maintain the cache.
237 """
238 self._last_func = func
239 self._last_args = args
240 self._last_kwargs = kwargs
241 self.op_idx += 1
243 @abstractmethod
244 def on_op_end(self, result) -> Any:
245 """Lifecycle callback after the operator is executed.
247 Responsible for trace recording, phase-specific post-actions,
248 **and** shadow wrapping. Must return the (possibly shadowed)
249 result tree so that :meth:`dispatch` can return it directly.
251 The op's function and arguments are available via
252 ``self._last_func``, ``self._last_args``, ``self._last_kwargs``
253 (cached by :meth:`on_op_begin`).
255 Subclasses should call :meth:`apply_shadows` with the bindings
256 they have already computed (e.g. from the tracker or the guide)
257 to avoid redundant traversal.
258 """
260 def apply_shadows(self, result: Any, bindings: dict[int, int]) -> Any:
261 """Replace result tensors with ShadowTensor instances per an explicit bindings map.
263 Subclasses that have already computed the ``leaf_index → sid``
264 mapping (e.g. during trace recording) pass it here to avoid a
265 redundant second traversal and SID re-resolution.
267 Caches *bindings* in ``_last_output_bindings`` for use by autograd
268 boundaries (:class:`OpaqueRegionEnd`) that need to shadow outputs
269 after :meth:`on_op_end` has already returned.
271 Args:
272 result: The raw output pytree.
273 bindings: Mapping ``leaf_index → storage_id``. Only leaves
274 whose index appears in the map are shadowed.
276 Returns:
277 A pytree of the same structure as *result* with eligible
278 tensors replaced by :class:`ShadowTensor`.
280 """
281 leaves, tree_spec = tree_flatten(result)
282 shadowed = list(leaves)
283 for idx, leaf in enumerate(leaves):
284 if idx in bindings and isinstance(leaf, torch.Tensor):
285 shadowed[idx] = self.make_shadow(bindings[idx], leaf)
286 result = tree_unflatten(shadowed, tree_spec)
287 self._last_output_bindings = bindings
288 return result
290 # ------------------------------------------------------------------
291 # Orchestration helper (output wrapping)
292 # ------------------------------------------------------------------
294 @property
295 def retained_sids(self) -> set[int]:
296 """Return the set of storage IDs that still have alive shadows."""
297 return {sid for sid, shadows in self._alive_shadows.items() if shadows}
299 def make_shadow(self, storage_id: int, tensor: torch.Tensor) -> Any:
300 """Register physical storage, create shadow, and track logically.
302 Composes:
303 #. :meth:`ResidencyManager.bind` — physical registration.
304 #. :class:`ShadowTensor` construction.
305 #. Tracking in ``_alive_shadows`` for ``retained_sids``.
307 Args:
308 storage_id: Physical storage ID.
309 tensor: The device-resident tensor to shadow.
311 Returns:
312 The original *tensor* replaced by a :class:`ShadowTensor`,
313 or the same shadow updated in-place for mutation.
315 """
316 with torch.utils._python_dispatch._disable_current_modes(): # pylint: disable=protected-access
317 if isinstance(tensor, ShadowTensor):
318 return tensor
320 buffer = self.residency_manager.bind(storage_id, tensor)
321 shadow = ShadowTensor(tensor, buffer, storage_id)
322 self._alive_shadows[storage_id].add(shadow)
323 return shadow
325 def reset(self) -> None:
326 """Reset per-cycle state before a new pass."""
327 self._opaque_depth = 0
328 self.op_idx = -1
329 self._alive_shadows.clear()
330 self.residency_manager.clear_runtime()