Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / pipeline_swap.py: 96%
209 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"""Pipeline-parallel activation swap scheduling helpers."""
17import warnings
18from collections import defaultdict
19from contextlib import nullcontext
20from enum import IntEnum
21import itertools
22from typing import Any, ContextManager, FrozenSet, Iterable, List
24from hyper_parallel.core.activation_memory.swap import SwapManager
25from hyper_parallel.core.pipeline_parallel.utils import MetaStep, MetaStepType
27MIN_SWAP_GAP = 4
30class _BeforeActionPriority(IntEnum):
31 WAIT_LOAD = 10
32 LAUNCH_LOAD = 20
35class _AfterActionPriority(IntEnum):
36 WAIT_OFFLOAD = 10
37 LAUNCH_OFFLOAD = 20
40def unregister_layer_swap_hooks(stages: Iterable[Any]) -> int:
41 """Remove layer-level swap hooks from pipeline stage modules.
43 Pipeline schedule swap groups are keyed by stage and micro-batch, while
44 hooks installed by ``SwapManager.set_forward_prefetch_layer`` replace the
45 active group with a layer-level group. When pipeline swap is enabled, the
46 schedule takes precedence and the layer hooks are permanently removed.
48 Args:
49 stages: Pipeline stages whose submodules are managed by the schedule.
51 Returns:
52 Number of removed hook handles.
53 """
54 removed_count = 0
55 visited_modules = set()
56 manager = SwapManager()
57 for stage in stages:
58 for _, module in stage.submodule.named_modules():
59 module_id = id(module)
60 if module_id in visited_modules:
61 continue
62 visited_modules.add(module_id)
63 # Full teardown, not just handle removal: also drops the module's
64 # swap group from the process-wide singleton plus the
65 # _swap_group_name/_swap_group_order/_swap_state attributes, so a
66 # permanently-removed layer does not leave an orphaned SwapGroup
67 # holding pinned host/device storage (see B.cache_no_eviction).
68 removed_count += manager.unregister_forward_prefetch_layer(module)
70 if removed_count:
71 warnings.warn(
72 "Pipeline schedule swap=True takes precedence over layer-level activation swap; "
73 "hooks registered by SwapManager.set_forward_prefetch_layer() were removed.",
74 UserWarning,
75 stacklevel=2,
76 )
77 return removed_count
80class PipelineSwapSession:
81 """Own activation-swap groups for one pipeline schedule run."""
83 _generation = itertools.count()
85 def __init__(self, eligible_keys: FrozenSet[tuple[int, int]]) -> None:
86 """Create one run-scoped session from build-time eligible chunk keys."""
87 self._generation_id = next(self._generation)
88 self._manager = SwapManager()
89 self._eligible_keys = eligible_keys
90 self._group_names = {
91 key: f"pp_swap_run{self._generation_id}_s{key[0]}_m{key[1]}"
92 for key in self._eligible_keys
93 }
94 self._forward_context_keys = set()
96 @staticmethod
97 def _key(step: Any) -> tuple[int, int]:
98 """Return the logical chunk key carried by a swap or compute step."""
99 return step.stage_index, step.micro_index
101 def manages(self, step: Any) -> bool:
102 """Return whether this session manages the step's pipeline chunk."""
103 return self._key(step) in self._eligible_keys
105 def group_context(self, step: Any) -> ContextManager[None]:
106 """Enter the run-scoped swap group for one forward leaf."""
107 if not self.manages(step):
108 return nullcontext()
109 key = self._key(step)
110 group_name = self._group_names[key]
111 self._manager.ensure_group(group_name)
112 self._forward_context_keys.add(key)
113 return self._manager.group_context(group_name)
115 def require_forward_context(self, step: Any) -> None:
116 """Require the managed forward leaf to have entered its swap context."""
117 key = self._key(step)
118 if key not in self._forward_context_keys:
119 raise RuntimeError(
120 "Pipeline swap did not observe the matching forward leaf. "
121 "Custom overlap callbacks must call schedule.execute_fwd_leaf()."
122 )
124 def group_name(self, step: Any) -> str:
125 """Return the physical group name for a managed step."""
126 key = self._key(step)
127 if key not in self._group_names:
128 raise RuntimeError(f"Pipeline swap does not manage chunk {key}.")
129 return self._group_names[key]
131 def wait_load(self, step: Any) -> None:
132 """Wait for H2D on the scheduler's current compute stream."""
133 if self.manages(step):
134 self._manager.wait_load(self.group_name(step))
136 def protect_aliases(self, step: Any, tensors: Any) -> None:
137 """Keep pipeline-owned aliases resident for a managed chunk."""
138 if self.manages(step):
139 self._manager.protect_alias_tensors(self.group_name(step), tensors)
141 def close(self) -> None:
142 """Release every group still owned by this run."""
143 for group_name in self._group_names.values():
144 self._manager.abort_group(group_name)
147def _is_compute_step(step) -> bool:
148 return step is not None and step.type in (
149 MetaStepType.FWD,
150 MetaStepType.BWD,
151 MetaStepType.BWD_INPUT,
152 MetaStepType.BWD_WEIGHT,
153 )
156def _is_comm_step(step) -> bool:
157 return step is not None and step.type in (
158 MetaStepType.FWD_RECV,
159 MetaStepType.FWD_SEND,
160 MetaStepType.BWD_RECV,
161 MetaStepType.BWD_SEND,
162 MetaStepType.BATCH_SEND_RECV,
163 )
166def _is_composite_compute_step(step) -> bool:
167 return (
168 step is not None
169 and step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F)
170 and step.sub_steps
171 )
174class _ComputeLeaf:
175 """A real FWD/BWD leaf and the top-level container that owns it."""
177 __slots__ = ("step", "container_index", "compute_index")
179 def __init__(self, step: Any, container_index: int, compute_index: int) -> None:
180 """Record a compute leaf and its physical container positions."""
181 self.step = step
182 self.container_index = container_index
183 self.compute_index = compute_index
186def _iter_compute_leaf_steps(step):
187 """Yield real FWD/BWD steps, expanding composite containers."""
188 if _is_compute_step(step):
189 yield step
190 return
191 if _is_composite_compute_step(step):
192 for sub_step in step.sub_steps:
193 if _is_compute_step(sub_step):
194 yield sub_step
197def _collect_compute_leaves(order):
198 """Collect compute leaves while counting each composite as one slot."""
199 leaves = []
200 container_by_compute_index = {}
201 compute_index = 0
202 for container_index, step in enumerate(order):
203 leaf_steps = list(_iter_compute_leaf_steps(step))
204 if not leaf_steps:
205 continue
206 container_by_compute_index[compute_index] = container_index
207 for leaf_step in leaf_steps:
208 leaves.append(_ComputeLeaf(leaf_step, container_index, compute_index))
209 compute_index += 1
210 return leaves, container_by_compute_index
213def _iter_steps_by_priority(priority_steps):
214 """Yield steps from high priority to low priority."""
215 for _, step in sorted(priority_steps, key=lambda item: item[0], reverse=True):
216 yield step
219def _comm_block_anchor(order, index):
220 """Return the last immediately following communication step."""
221 anchor = index
222 for next_index in range(index + 1, len(order)):
223 next_step = order[next_index]
224 if _is_comm_step(next_step):
225 anchor = next_index
226 continue
227 break
228 return anchor
231def _post_compute_anchor(order, index, leaf_step=None):
232 """Return the safe index after which post-compute swap steps may run."""
233 step = leaf_step if leaf_step is not None else order[index]
234 fallback_anchor = _comm_block_anchor(order, index)
235 if step.type == MetaStepType.FWD:
236 send_type = MetaStepType.FWD_SEND
237 elif step.type in (MetaStepType.BWD, MetaStepType.BWD_INPUT):
238 send_type = MetaStepType.BWD_SEND
239 else:
240 return fallback_anchor
242 for next_index in range(index + 1, fallback_anchor + 1):
243 next_step = order[next_index]
244 next_valid = (next_step is not None and next_step.type == send_type
245 and next_step.stage_index == step.stage_index and next_step.micro_index == step.micro_index)
246 if next_valid:
247 return next_index
248 if next_step is not None and next_step.type == MetaStepType.BATCH_SEND_RECV:
249 for sub_step in next_step.sub_steps:
250 if (
251 sub_step.type == send_type
252 and sub_step.stage_index == step.stage_index
253 and sub_step.micro_index == step.micro_index):
254 return next_index
255 return fallback_anchor
258def _post_compute_launch_anchor(leaf):
259 """Return the fallback point after compute where D2H may be launched."""
260 return leaf.container_index
263def _load_launch_anchor(
264 order: List[Any], fwd_leaf: _ComputeLeaf, bwd_leaf: _ComputeLeaf,
265 compute_between: List[int]) -> int:
266 """Choose the latest safe H2D launch point for plain or FSDP execution."""
267 has_fsdp_steps = any(
268 step is not None and step.type in (
269 MetaStepType.FSDP_UNSHARD,
270 MetaStepType.FSDP_RESHARD,
271 MetaStepType.FSDP_REDUCE_GRAD,
272 )
273 for step in order
274 )
275 if not has_fsdp_steps:
276 return compute_between[-1]
278 for index in range(bwd_leaf.container_index - 1, fwd_leaf.container_index, -1):
279 step = order[index]
280 if (
281 step is not None
282 and step.type == MetaStepType.FSDP_UNSHARD
283 and step.stage_index == bwd_leaf.step.stage_index):
284 return index
285 return bwd_leaf.container_index
288def _index_compute_leaves(order):
289 """Index compute leaves by chunk key, separating FWD from BWD consumers."""
290 fwd_index = {}
291 bwd_index = {}
292 leaves, container_by_compute_index = _collect_compute_leaves(order)
293 for leaf in leaves:
294 step = leaf.step
295 key = (step.stage_index, step.micro_index)
296 if step.type == MetaStepType.FWD:
297 fwd_index[key] = leaf
298 elif step.type in (MetaStepType.BWD, MetaStepType.BWD_INPUT):
299 # BWD_INPUT (dxdw split) needs activations → place WAIT_LOAD before it.
300 # BWD_WEIGHT is intentionally excluded: it does not consume the
301 # original forward activations restored by swap.
302 bwd_index[key] = leaf
303 return fwd_index, bwd_index, container_by_compute_index
306def _chunk_swap_actions(order, fwd_leaf, bwd_leaf, container_by_compute_index):
307 """Return the swap actions for one chunk as ``(before, after)`` entries.
309 Each entry is ``(anchor_index, priority, meta_step)``. Returns ``None``
310 when the chunk is not swappable.
311 """
312 if bwd_leaf.compute_index - fwd_leaf.compute_index < MIN_SWAP_GAP:
313 return None
314 compute_between = [
315 container_by_compute_index[index]
316 for index in range(fwd_leaf.compute_index + 1, bwd_leaf.compute_index)
317 ]
318 if not compute_between:
319 return None
321 stage_index = fwd_leaf.step.stage_index
322 micro_index = fwd_leaf.step.micro_index
323 # Always launch offload immediately after the FWD container so that the
324 # async D2H starts before any FSDP_RESHARD or FWD_SEND that may sit
325 # between the FWD and the next compute step.
326 after = [
327 (
328 _post_compute_launch_anchor(fwd_leaf),
329 _AfterActionPriority.LAUNCH_OFFLOAD,
330 MetaStep(micro_index, MetaStepType.SWAP_LAUNCH_OFFLOAD, stage_index),
331 ),
332 (
333 _post_compute_anchor(order, compute_between[0]),
334 _AfterActionPriority.WAIT_OFFLOAD,
335 MetaStep(micro_index, MetaStepType.SWAP_WAIT_OFFLOAD, stage_index),
336 ),
337 ]
338 before = [
339 (
340 _load_launch_anchor(order, fwd_leaf, bwd_leaf, compute_between),
341 _BeforeActionPriority.LAUNCH_LOAD,
342 MetaStep(micro_index, MetaStepType.SWAP_LAUNCH_LOAD, stage_index),
343 ),
344 (
345 bwd_leaf.container_index,
346 _BeforeActionPriority.WAIT_LOAD,
347 MetaStep(micro_index, MetaStepType.SWAP_WAIT_LOAD, stage_index),
348 ),
349 ]
350 return before, after
353def _collect_swap_actions(order):
354 """Bucket swap actions by the order index they attach to.
356 Returns a ``(before_pairs, after_pairs)`` tuple of
357 ``index -> [(priority, MetaStep), ...]`` maps. Each map holds the actions
358 to emit immediately before and immediately after the top-level step at
359 that index.
360 """
361 fwd_index, bwd_index, container_by_compute_index = _index_compute_leaves(order)
363 maps = (defaultdict(list), defaultdict(list))
364 for key, fwd_leaf in fwd_index.items():
365 bwd_leaf = bwd_index.get(key)
366 if bwd_leaf is None:
367 continue
368 actions = _chunk_swap_actions(order, fwd_leaf, bwd_leaf, container_by_compute_index)
369 if actions is None:
370 continue
371 for bucket, action_group in enumerate(actions):
372 for anchor, priority, swap_step in action_group:
373 maps[bucket][anchor].append((priority, swap_step))
374 return maps
377def inject_pipeline_swap_steps(order: List[Any]) -> List[Any]:
378 """Inject asynchronous transfer steps into one rank's pipeline order.
380 Forward collection is executed directly by the forward leaf executor.
381 Transfer launch/wait actions, including the H2D wait before the backward
382 consumer container, appear in the top-level order.
383 """
384 before_steps, after_steps = _collect_swap_actions(order)
386 injected = []
387 for index, step in enumerate(order):
388 injected.extend(_iter_steps_by_priority(before_steps[index]))
389 injected.append(step)
390 injected.extend(_iter_steps_by_priority(after_steps[index]))
391 return injected
394def _protect_pipeline_owned_tensors(step, schedule, arg_mbs, kwarg_mbs, group_name: str) -> None:
395 """Keep long-lived module and pipeline-owned tensors alive on device.
397 Swap offload clears the device storage of saved tensors after D2H copy.
398 If a saved tensor aliases a parameter, registered buffer, or pipeline
399 boundary tensor, clearing it would invalidate the long-lived owner. The
400 alias protection below marks those saved tensors as keep-on-device.
401 """
402 stage = schedule._stage_dict[step.stage_index] # pylint: disable=protected-access
403 manager = SwapManager()
405 # Saved-tensor hooks may receive a plain Tensor view of a Parameter. Some
406 # backends do not preserve parameter metadata on that view, so protect by
407 # storage ownership before any group member can be resized.
408 parameters = tuple(stage.submodule.parameters())
409 buffers = tuple(stage.submodule.buffers())
410 if stage.is_first_stage:
411 # First-stage inputs come from split_microbatches(), outside the
412 # wrapped stage. They are not stage outputs, but they can alias
413 # tensors saved by the first layer and must not have their storage
414 # resized by the swap group.
415 boundary_inputs = (arg_mbs[step.micro_index], kwarg_mbs[step.micro_index])
416 else:
417 recv_infos = stage.args_recv_info.get(step.micro_index, ())
418 boundary_inputs = tuple(info.buffer for info in recv_infos if info.buffer is not None)
420 # Protect parameters, registered buffers, and pipeline-owned inputs in one
421 # pass over the collected saved tensors. Forward outputs are protected
422 # earlier by execute_fwd_leaf(), while the direct return value is still
423 # available.
424 manager.protect_alias_tensors(group_name, (parameters, buffers, boundary_inputs))
427def swap_launch_offload(
428 step: Any,
429 schedule: Any,
430 arg_mbs: List[Any],
431 kwarg_mbs: List[Any],
432 session: PipelineSwapSession) -> None:
433 """Launch D2H for a pipeline swap group."""
434 session.require_forward_context(step)
435 group_name = session.group_name(step)
436 manager = SwapManager()
437 _protect_pipeline_owned_tensors(step, schedule, arg_mbs, kwarg_mbs, group_name)
438 manager.launch_offload(group_name)
441def swap_wait_offload(step: Any, session: PipelineSwapSession) -> None:
442 """Wait for a pipeline swap group's D2H and release device storage."""
443 SwapManager().wait_offload(session.group_name(step))
446def swap_launch_load(step: Any, session: PipelineSwapSession) -> None:
447 """Launch H2D for a pipeline swap group."""
448 SwapManager().launch_load(session.group_name(step))
451def swap_wait_load(step: Any, session: PipelineSwapSession) -> None:
452 """Wait for a pipeline swap group's H2D before its backward container."""
453 session.wait_load(step)