Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / comm_compute_overlap.py: 92%
79 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.
14# ============================================================================
15"""Two-thread comm/compute overlap orchestrator.
17This module provides :class:`CommComputeOverlap`, a helper that wraps
18MoE-style dispatch / combine phases with four synchronization hooks
19(``A``, ``B``, ``C``, ``D``) and drives a forward + backward pass on two
20threads with deterministic comm-first dispatch ordering via
21:class:`HookCoordinator`.
23The mechanism is independent of any specific pipeline schedule. It is
24typically driven by the ``OVERLAP_B_F`` callback registered on a
25schedule (e.g. ``ScheduleInterleaved1F1B(overlap_b_f=True)``), but the
26same orchestrator could be reused by other concurrent-dispatch overlap
27scenarios (TP+CP, FSDP prefetch, etc.) without modification.
29Every rendezvous is a strict COMM + COMPUTE pair — including layer
30boundaries — so the NCCL kernel is always enqueued before the paired
31compute kernel::
33 [A] ─► dispatch ─► [B] ─► module ─► [C] ─► combine ─► [D] ─► (Attention) ─► [A_next]
35At layer boundaries the D / A hooks coordinate combine (COMM) with the
36other thread's Attention (COMPUTE), preserving overlap across layers.
38Typical integration::
40 overlap = CommComputeOverlap()
42 # Wrap the expert-parallel dispatch / combine callables:
43 wrapped_dispatch = overlap.wrap_dispatch(original_dispatch)
44 wrapped_combine = overlap.wrap_combine(original_combine)
46 # At schedule time, run forward and backward in parallel:
47 overlap.run(
48 fwd_fn=lambda: fwd_stage.forward_one_chunk(mb, *args),
49 bwd_fn=lambda: bwd_stage.backward_one_chunk(mb, loss=loss),
50 )
51"""
52import threading
53from dataclasses import dataclass
54from typing import Callable
56from hyper_parallel.core.pipeline_parallel._sync_hook import _SyncHookFunction
57from hyper_parallel.core.pipeline_parallel.hook_coordinator import HookCoordinator
60@dataclass
61class _BwdTask:
62 """One backward task running in an overlap window."""
64 bwd_fn: Callable[[], None]
65 done: threading.Event
66 exc_box: list[Exception]
69class CommComputeOverlap:
70 """Orchestrator for two-thread comm/compute overlap.
72 Manages a :class:`HookCoordinator` and provides helpers to insert the
73 four synchronization hooks (``A``, ``B``, ``C``, ``D``) around MoE
74 dispatch / combine phases and to run forward + backward concurrently
75 with deterministic comm-first kernel launch ordering. Each overlap window
76 runs backward in its own thread.
78 Example:
79 >>> overlap = CommComputeOverlap()
80 >>> wrapped_dispatch = overlap.wrap_dispatch(ep_dispatch_fn)
81 >>> wrapped_combine = overlap.wrap_combine(ep_combine_fn, is_last_layer=is_last)
82 >>> overlap.run(fwd_fn, bwd_fn) # doctest: +SKIP
83 """
85 def __init__(self) -> None:
86 """Initialize the forward/backward coordinator."""
87 self._coordinator = HookCoordinator()
89 @property
90 def coordinator(self) -> HookCoordinator:
91 """The underlying :class:`HookCoordinator` instance."""
92 return self._coordinator
94 # ------------------------------------------------------------------
95 # Wrapping helpers
96 # ------------------------------------------------------------------
98 def wrap_dispatch(self, dispatch_fn: Callable) -> Callable:
99 """Return a wrapped version of ``dispatch_fn`` bracketed by hooks A/B.
101 The returned callable inserts synchronization hooks on the **first
102 positional tensor argument** before and after the call::
104 A ─► dispatch_fn ─► B
106 Args:
107 dispatch_fn: The original dispatch callable.
109 Returns:
110 A new callable with the same signature.
111 """
112 coordinator = self._coordinator
114 def _wrapped(*args, **kwargs):
115 first, rest = args[0], args[1:]
116 first = _SyncHookFunction.apply(first, "A", coordinator)
117 result = dispatch_fn(first, *rest, **kwargs)
118 if isinstance(result, tuple):
119 hooked = _SyncHookFunction.apply(result[0], "B", coordinator)
120 return (hooked,) + result[1:]
121 return _SyncHookFunction.apply(result, "B", coordinator)
123 return _wrapped
125 def wrap_combine(self, combine_fn: Callable, is_last_layer: bool = False) -> Callable:
126 """Return a wrapped version of ``combine_fn`` bracketed by hooks C/D.
128 The returned callable inserts synchronization hooks on the **first
129 positional tensor argument** before and after the call::
131 C ─► combine_fn ─► D
133 Args:
134 combine_fn: The original combine callable.
135 is_last_layer: If ``True``, the closing D hook is tagged
136 ``"D_LAST"`` so the rendezvous is skipped both in
137 forward (no Attention follows the last layer) and in
138 backward (this is the first BWD hook to fire and
139 combine.bwd has already dispatched freely). Tagging
140 this hook statically replaces the old runtime cycle
141 counter and BWD-D-skip mechanisms.
143 Returns:
144 A new callable with the same signature.
145 """
146 coordinator = self._coordinator
147 d_hook = "D_LAST" if is_last_layer else "D"
149 def _wrapped(*args, **kwargs):
150 first, rest = args[0], args[1:]
151 first = _SyncHookFunction.apply(first, "C", coordinator)
152 result = combine_fn(first, *rest, **kwargs)
153 if isinstance(result, tuple):
154 hooked = _SyncHookFunction.apply(result[0], d_hook, coordinator)
155 return (hooked,) + result[1:]
156 return _SyncHookFunction.apply(result, d_hook, coordinator)
158 return _wrapped
160 # ------------------------------------------------------------------
161 # Execution
162 # ------------------------------------------------------------------
164 def _start_bwd_task(self, task: _BwdTask) -> threading.Thread:
165 """Start the backward thread for one overlap window."""
166 thread = threading.Thread(target=self._run_bwd_task, args=(task,), daemon=True)
167 thread.start()
168 return thread
170 def _run_bwd_task(self, task: _BwdTask) -> None:
171 """Execute one backward task and preserve the existing cleanup contract."""
172 coordinator = self._coordinator
173 try:
174 try:
175 task.bwd_fn()
176 except Exception as exc: # pylint: disable=W0718
177 task.exc_box.append(exc)
178 # BWD died — disable the coordinator so any FWD rendezvous
179 # waiting on a barrier/event unblocks immediately. Without
180 # this the FWD thread hangs forever at the very first hook
181 # it reaches and the outer ``finally`` never runs.
182 coordinator.disable()
183 finally:
184 # Graceful one-party-left: this BWD chunk is done. If the
185 # paired FWD chunk has MORE hooks (e.g. more layers) it would
186 # otherwise block forever on the 2-party barrier waiting for
187 # a partner that has exited. ``depart`` aborts the barrier
188 # and flags the coordinator so FWD's remaining hooks run
189 # solo. Required for correctness, not just on error.
190 coordinator.depart()
191 finally:
192 task.done.set()
194 def _run(
195 self,
196 fwd_fn: Callable[[], None],
197 bwd_fn: Callable[[], None],
198 ) -> None:
199 """Run one overlap window and wait for both threads to finish."""
200 self._coordinator.enable()
201 done = threading.Event()
202 exc_box: list[Exception] = []
203 try:
204 thread = self._start_bwd_task(_BwdTask(bwd_fn, done, exc_box))
205 except Exception:
206 self._coordinator.disable()
207 raise
209 fwd_exc: list[Exception] = []
210 try:
211 fwd_fn()
212 except Exception as exc: # pylint: disable=W0718
213 fwd_exc.append(exc)
214 # Symmetric: if FWD dies, unblock BWD so it can exit.
215 self._coordinator.disable()
216 finally:
217 # Graceful one-party-left, mirroring the worker task: if the
218 # paired BWD chunk has MORE hooks, ``depart`` lets it drain its
219 # remaining rendezvous solo instead of hanging on the barrier.
220 # Must precede waiting on the task so a still-running BWD is
221 # released.
222 self._coordinator.depart()
223 done.wait()
224 if thread is not None:
225 thread.join()
226 # Full reset after both sides are done (idempotent with any
227 # earlier disable on the FWD error path).
228 self._coordinator.disable()
230 if exc_box:
231 raise RuntimeError(
232 "Exception in backward thread during dual-pipe overlap"
233 ) from exc_box[0]
234 if fwd_exc:
235 raise fwd_exc[0]
237 def run(
238 self,
239 fwd_fn: Callable[[], None],
240 bwd_fn: Callable[[], None],
241 ) -> None:
242 """Run ``fwd_fn`` and ``bwd_fn`` with comm/compute overlap.
244 Args:
245 fwd_fn: Callable that executes the forward pass.
246 bwd_fn: Callable that executes the backward pass. If it needs
247 a specific device stream, wrap that logic inside
248 ``bwd_fn``.
250 Raises:
251 RuntimeError: If the backward thread raises an exception, it
252 is re-raised on the main thread after the task finishes.
253 """
254 self._run(fwd_fn, bwd_fn)