Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / pipeline_swap.py: 96%
197 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05: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"""Pipeline-parallel activation swap scheduling helpers."""
17from collections import defaultdict
18from contextlib import nullcontext
19from enum import IntEnum
20import itertools
21from typing import Any, ContextManager, FrozenSet, List
23from hyper_parallel.core.activation_checkpoint.swap import SwapManager
24from hyper_parallel.platform import get_platform
26MIN_SWAP_GAP = 4
27platform = get_platform()
30class _BeforeActionPriority(IntEnum):
31 WAIT_LOAD = 10
32 LAUNCH_LOAD = 20
35class _AfterActionPriority(IntEnum):
36 WAIT_OFFLOAD = 10
37 LAUNCH_OFFLOAD = 20
40class PipelineSwapSession:
41 """Own activation-swap groups for one pipeline schedule run."""
43 _generation = itertools.count()
45 def __init__(self, eligible_keys: FrozenSet[tuple[int, int]]) -> None:
46 """Create one run-scoped session from build-time eligible chunk keys."""
47 self._generation_id = next(self._generation)
48 self._manager = SwapManager()
49 self._eligible_keys = eligible_keys
50 self._group_names = {
51 key: f"pp_swap_run{self._generation_id}_s{key[0]}_m{key[1]}"
52 for key in self._eligible_keys
53 }
54 self._forward_context_keys = set()
56 @staticmethod
57 def _key(step: Any) -> tuple[int, int]:
58 """Return the logical chunk key carried by a swap or compute step."""
59 return step.stage_index, step.micro_index
61 def manages(self, step: Any) -> bool:
62 """Return whether this session manages the step's pipeline chunk."""
63 return self._key(step) in self._eligible_keys
65 def group_context(self, step: Any) -> ContextManager[None]:
66 """Enter the run-scoped swap group for one forward leaf."""
67 if not self.manages(step):
68 return nullcontext()
69 key = self._key(step)
70 group_name = self._group_names[key]
71 self._manager.ensure_group(group_name)
72 self._forward_context_keys.add(key)
73 return self._manager.group_context(group_name)
75 def require_forward_context(self, step: Any) -> None:
76 """Require the managed forward leaf to have entered its swap context."""
77 key = self._key(step)
78 if key not in self._forward_context_keys:
79 raise RuntimeError(
80 "Pipeline swap did not observe the matching forward leaf. "
81 "Custom overlap callbacks must call schedule.execute_fwd_leaf()."
82 )
84 def group_name(self, step: Any) -> str:
85 """Return the physical group name for a managed step."""
86 key = self._key(step)
87 if key not in self._group_names:
88 raise RuntimeError(f"Pipeline swap does not manage chunk {key}.")
89 return self._group_names[key]
91 def wait_load(self, step: Any) -> None:
92 """Wait for H2D on the scheduler's current compute stream."""
93 if self.manages(step):
94 self._manager.wait_load(self.group_name(step))
96 def protect_aliases(self, step: Any, tensors: Any) -> None:
97 """Keep pipeline-owned aliases resident for a managed chunk."""
98 if self.manages(step):
99 self._manager.protect_alias_tensors(self.group_name(step), tensors)
101 def close(self) -> None:
102 """Release every group still owned by this run."""
103 for group_name in self._group_names.values():
104 self._manager.abort_group(group_name)
107def _is_compute_step(step) -> bool:
108 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415
110 return step is not None and step.type in (
111 MetaStepType.FWD,
112 MetaStepType.BWD,
113 MetaStepType.BWD_INPUT,
114 MetaStepType.BWD_WEIGHT,
115 )
118def _is_comm_step(step) -> bool:
119 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415
121 return step is not None and step.type in (
122 MetaStepType.FWD_RECV,
123 MetaStepType.FWD_SEND,
124 MetaStepType.BWD_RECV,
125 MetaStepType.BWD_SEND,
126 MetaStepType.BATCH_SEND_RECV,
127 )
130def _is_composite_compute_step(step) -> bool:
131 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415
133 return (
134 step is not None
135 and step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F)
136 and step.sub_steps
137 )
140class _ComputeLeaf:
141 """A real FWD/BWD leaf and the top-level container that owns it."""
143 __slots__ = ("step", "container_index", "compute_index")
145 def __init__(self, step: Any, container_index: int, compute_index: int) -> None:
146 """Record a compute leaf and its physical container positions."""
147 self.step = step
148 self.container_index = container_index
149 self.compute_index = compute_index
152def _iter_compute_leaf_steps(step):
153 """Yield real FWD/BWD steps, expanding composite containers."""
154 if _is_compute_step(step):
155 yield step
156 return
157 if _is_composite_compute_step(step):
158 for sub_step in step.sub_steps:
159 if _is_compute_step(sub_step):
160 yield sub_step
163def _collect_compute_leaves(order):
164 """Collect compute leaves while counting each composite as one slot."""
165 leaves = []
166 container_by_compute_index = {}
167 compute_index = 0
168 for container_index, step in enumerate(order):
169 leaf_steps = list(_iter_compute_leaf_steps(step))
170 if not leaf_steps:
171 continue
172 container_by_compute_index[compute_index] = container_index
173 for leaf_step in leaf_steps:
174 leaves.append(_ComputeLeaf(leaf_step, container_index, compute_index))
175 compute_index += 1
176 return leaves, container_by_compute_index
179def _append_after(after_steps, index, priority, step):
180 after_steps[index].append((priority, step))
183def _append_before(before_steps, index, priority, step):
184 before_steps[index].append((priority, step))
187def _iter_steps_by_priority(priority_steps):
188 """Yield steps from high priority to low priority."""
189 for _, step in sorted(priority_steps, key=lambda item: item[0], reverse=True):
190 yield step
193def _comm_block_anchor(order, index):
194 """Return the last immediately following communication step."""
195 anchor = index
196 for next_index in range(index + 1, len(order)):
197 next_step = order[next_index]
198 if _is_comm_step(next_step):
199 anchor = next_index
200 continue
201 break
202 return anchor
205def _post_compute_anchor(order, index, leaf_step=None):
206 """Return the safe index after which post-compute swap steps may run."""
207 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415
209 step = leaf_step if leaf_step is not None else order[index]
210 fallback_anchor = _comm_block_anchor(order, index)
211 if step.type == MetaStepType.FWD:
212 send_type = MetaStepType.FWD_SEND
213 elif step.type in (MetaStepType.BWD, MetaStepType.BWD_INPUT):
214 send_type = MetaStepType.BWD_SEND
215 else:
216 return fallback_anchor
218 for next_index in range(index + 1, fallback_anchor + 1):
219 next_step = order[next_index]
220 next_valid = (next_step is not None and next_step.type == send_type
221 and next_step.stage_index == step.stage_index and next_step.micro_index == step.micro_index)
222 if next_valid:
223 return next_index
224 if next_step is not None and next_step.type == MetaStepType.BATCH_SEND_RECV:
225 for sub_step in next_step.sub_steps:
226 if (
227 sub_step.type == send_type
228 and sub_step.stage_index == step.stage_index
229 and sub_step.micro_index == step.micro_index):
230 return next_index
231 return fallback_anchor
234def _post_compute_launch_anchor(leaf):
235 """Return the fallback point after compute where D2H may be launched."""
236 return leaf.container_index
239def _load_launch_anchor(
240 order: List[Any], fwd_leaf: _ComputeLeaf, bwd_leaf: _ComputeLeaf,
241 compute_between: List[int]) -> int:
242 """Choose the latest safe H2D launch point for plain or FSDP execution."""
243 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415
245 has_fsdp_steps = any(
246 step is not None and step.type in (
247 MetaStepType.FSDP_UNSHARD,
248 MetaStepType.FSDP_RESHARD,
249 MetaStepType.FSDP_REDUCE_GRAD,
250 )
251 for step in order
252 )
253 if not has_fsdp_steps:
254 return compute_between[-1]
256 for index in range(bwd_leaf.container_index - 1, fwd_leaf.container_index, -1):
257 step = order[index]
258 if (
259 step is not None
260 and step.type == MetaStepType.FSDP_UNSHARD
261 and step.stage_index == bwd_leaf.step.stage_index):
262 return index
263 return bwd_leaf.container_index
266def inject_pipeline_swap_steps(order: List[Any]) -> List[Any]:
267 """Inject asynchronous transfer steps into one rank's pipeline order.
269 Forward collection is executed directly by the forward leaf executor.
270 Transfer launch/wait actions, including the H2D wait before the backward
271 consumer container, appear in the top-level order.
272 """
273 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStep, MetaStepType # pylint: disable=C0415
275 fwd_index = {}
276 bwd_index = {}
277 compute_leaves, container_by_compute_index = _collect_compute_leaves(order)
278 for leaf in compute_leaves:
279 step = leaf.step
280 key = (step.stage_index, step.micro_index)
281 if step.type == MetaStepType.FWD:
282 fwd_index[key] = leaf
283 elif step.type in (MetaStepType.BWD, MetaStepType.BWD_INPUT):
284 # BWD_INPUT (dxdw split) needs activations → place WAIT_LOAD before it.
285 # BWD_WEIGHT is intentionally excluded: it does not consume the
286 # original forward activations restored by swap.
287 bwd_index[key] = leaf
289 before_steps = defaultdict(list)
290 after_steps = defaultdict(list)
291 chunk_gaps = {
292 key: bwd_index[key].compute_index - fwd_leaf.compute_index
293 for key, fwd_leaf in fwd_index.items()
294 if key in bwd_index
295 }
296 for key, fwd_leaf in fwd_index.items():
297 bwd_leaf = bwd_index.get(key)
298 if bwd_leaf is None:
299 continue
300 if chunk_gaps[key] < MIN_SWAP_GAP:
301 continue
302 compute_between = [
303 container_by_compute_index[index]
304 for index in range(fwd_leaf.compute_index + 1, bwd_leaf.compute_index)
305 ]
306 if not compute_between:
307 continue
308 stage_index, micro_index = key
310 first_between_anchor = _post_compute_anchor(order, compute_between[0])
312 # Always launch offload immediately after the FWD container so that
313 # the async D2H starts before any FSDP_RESHARD or FWD_SEND that may
314 # sit between the FWD and the next compute step.
315 fwd_anchor = _post_compute_launch_anchor(fwd_leaf)
316 _append_after(
317 after_steps, fwd_anchor, _AfterActionPriority.LAUNCH_OFFLOAD,
318 MetaStep(micro_index, MetaStepType.SWAP_LAUNCH_OFFLOAD, stage_index),
319 )
321 _append_after(
322 after_steps, first_between_anchor, _AfterActionPriority.WAIT_OFFLOAD,
323 MetaStep(micro_index, MetaStepType.SWAP_WAIT_OFFLOAD, stage_index),
324 )
326 load_launch_anchor = _load_launch_anchor(order, fwd_leaf, bwd_leaf, compute_between)
327 _append_before(
328 before_steps, load_launch_anchor, _BeforeActionPriority.LAUNCH_LOAD,
329 MetaStep(micro_index, MetaStepType.SWAP_LAUNCH_LOAD, stage_index),
330 )
331 _append_before(
332 before_steps, bwd_leaf.container_index, _BeforeActionPriority.WAIT_LOAD,
333 MetaStep(micro_index, MetaStepType.SWAP_WAIT_LOAD, stage_index),
334 )
335 injected = []
336 for index, step in enumerate(order):
337 injected.extend(_iter_steps_by_priority(before_steps[index]))
338 injected.append(step)
339 injected.extend(_iter_steps_by_priority(after_steps[index]))
340 return injected
343def _protect_pipeline_owned_tensors(step, schedule, arg_mbs, kwarg_mbs, group_name: str) -> None:
344 """Keep long-lived module and pipeline-owned tensors alive on device.
346 Swap offload clears the device storage of saved tensors after D2H copy.
347 If a saved tensor aliases a parameter, registered buffer, or pipeline
348 boundary tensor, clearing it would invalidate the long-lived owner. The
349 alias protection below marks those saved tensors as keep-on-device.
350 """
351 stage = schedule._stage_dict[step.stage_index] # pylint: disable=protected-access
352 manager = SwapManager()
354 # Saved-tensor hooks may receive a plain Tensor view of a Parameter. Some
355 # backends do not preserve parameter metadata on that view, so protect by
356 # storage ownership before any group member can be resized.
357 parameters = tuple(param for _, param in platform.parameters_dict(stage.submodule))
358 buffers = tuple(buffer for _, buffer in platform.buffers_dict(stage.submodule))
359 if stage.is_first_stage:
360 # First-stage inputs come from split_microbatches(), outside the
361 # wrapped stage. They are not stage outputs, but they can alias
362 # tensors saved by the first layer and must not have their storage
363 # resized by the swap group.
364 boundary_inputs = (arg_mbs[step.micro_index], kwarg_mbs[step.micro_index])
365 else:
366 recv_infos = stage.args_recv_info.get(step.micro_index, ())
367 boundary_inputs = tuple(info.buffer for info in recv_infos if info.buffer is not None)
369 # Protect parameters, registered buffers, and pipeline-owned inputs in one
370 # pass over the collected saved tensors. Forward outputs are protected
371 # earlier by execute_fwd_leaf(), while the direct return value is still
372 # available.
373 manager.protect_alias_tensors(group_name, (parameters, buffers, boundary_inputs))
376def swap_launch_offload(
377 step: Any,
378 schedule: Any,
379 arg_mbs: List[Any],
380 kwarg_mbs: List[Any],
381 session: PipelineSwapSession) -> None:
382 """Launch D2H for a pipeline swap group."""
383 session.require_forward_context(step)
384 group_name = session.group_name(step)
385 manager = SwapManager()
386 _protect_pipeline_owned_tensors(step, schedule, arg_mbs, kwarg_mbs, group_name)
387 manager.launch_offload(group_name)
390def swap_wait_offload(step: Any, session: PipelineSwapSession) -> None:
391 """Wait for a pipeline swap group's D2H and release device storage."""
392 SwapManager().wait_offload(session.group_name(step))
395def swap_launch_load(step: Any, session: PipelineSwapSession) -> None:
396 """Launch H2D for a pipeline swap group."""
397 SwapManager().launch_load(session.group_name(step))
400def swap_wait_load(step: Any, session: PipelineSwapSession) -> None:
401 """Wait for a pipeline swap group's H2D before its backward container."""
402 session.wait_load(step)