Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / scheduler.py: 68%
1153 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 schedule"""
16from abc import ABC, abstractmethod
17from contextlib import nullcontext
18from enum import Enum, auto
19from collections import defaultdict
20import itertools
21import bisect
22import logging
23import re
24from typing import Any
26import hyper_parallel
27from hyper_parallel.platform import get_platform
28from hyper_parallel.core.fully_shard.api import HSDPModule
29from hyper_parallel.core.pipeline_parallel.pipeline_swap import (
30 PipelineSwapSession,
31 inject_pipeline_swap_steps,
32 swap_launch_load,
33 swap_launch_offload,
34 swap_wait_load,
35 swap_wait_offload,
36)
37from hyper_parallel.core.pipeline_parallel.utils import BatchDimSpec
38platform = get_platform()
39logger = logging.getLogger(__name__)
42class MetaStepType(Enum):
43 """Specify the enumeration type for MetaStep."""
44 FWD = auto()
45 BWD = auto()
46 BWD_INPUT = auto()
47 BWD_WEIGHT = auto()
48 FWD_RECV = auto()
49 FWD_SEND = auto()
50 BWD_RECV = auto()
51 BWD_SEND = auto()
52 # Composite P2P: a contiguous run of FWD_SEND/FWD_RECV/BWD_SEND/BWD_RECV
53 # coalesced by ``coalesce_p2p`` into one step whose ``sub_steps`` the runtime
54 # groups by peer and issues as ``batch_isend_irecv`` (same-peer send+recv ->
55 # duplex). Only produced under ``p2p_transport="batch"``.
56 BATCH_SEND_RECV = auto()
57 OVERLAP_F_B = auto()
58 OVERLAP_B_F = auto()
59 FSDP_UNSHARD = auto()
60 FSDP_RESHARD = auto()
61 FSDP_REDUCE_GRAD = auto()
62 SWAP_LAUNCH_OFFLOAD = auto()
63 SWAP_WAIT_OFFLOAD = auto()
64 SWAP_LAUNCH_LOAD = auto()
65 SWAP_WAIT_LOAD = auto()
68class MetaStep:
69 """
70 Meta step of PipelineSchedule.
71 An execution list composed of MetaStep can be constructed
72 and fed into the PipelineSchedule for execution.
74 Args:
75 micro_index (int | None): The index of micro-batch. ``None`` for
76 composite types (``OVERLAP_F_B`` / ``OVERLAP_B_F``) whose real
77 micro index lives in each ``sub_steps`` entry.
78 type (MetaStepType): Specify the type of current step.
79 stage_index (int | None): Stage index of current step. ``None``
80 for composite types; use ``sub_steps`` to get each direction's
81 stage.
82 sub_steps (tuple[MetaStep, MetaStep] | None): For composite types
83 only: ``(fwd, bwd)`` for ``OVERLAP_F_B``, ``(bwd, fwd)`` for
84 ``OVERLAP_B_F``.
85 boundary_p2p (tuple[MetaStep, ...] | None): For ``OVERLAP_B_F`` under
86 the ``"boundary"`` P2P transport only: P2P steps to issue at the
87 fwd/bwd boundary inside the overlap (the forward's ``FWD_SEND``
88 plus the next slot's recvs), hoisted out of the following gap by
89 ``attach_fwd_boundary_p2p``. Issued via
90 :meth:`PipelineScheduleRuntime.exec_boundary_p2p`.
91 """
92 def __init__(self, micro_index, meta_type, stage_index, sub_steps=None,
93 boundary_p2p=None):
94 self._type = meta_type
95 self._micro_index = micro_index
96 self._stage_index = stage_index
97 self._sub_steps = sub_steps
98 self._boundary_p2p = boundary_p2p
100 @property
101 def micro_index(self):
102 """Return the micro-batch index of this step."""
103 return self._micro_index
105 @property
106 def stage_index(self):
107 """Return the stage index of this step."""
108 return self._stage_index
110 @property
111 def type(self):
112 """Return the MetaStepType of this step."""
113 return self._type
115 @property
116 def sub_steps(self):
117 """Sub-steps for composite types: ``(fwd, bwd)`` for OVERLAP_F_B,
118 ``(bwd, fwd)`` for OVERLAP_B_F, or ``None``."""
119 return self._sub_steps
121 @property
122 def boundary_p2p(self):
123 """P2P steps to issue at the overlap's fwd/bwd boundary, or ``None``."""
124 return self._boundary_p2p
126 def __eq__(self, value):
127 if not isinstance(value, MetaStep):
128 return NotImplemented
129 return (self.type == value.type
130 and self.micro_index == value.micro_index
131 and self.stage_index == value.stage_index
132 and self.sub_steps == value.sub_steps)
134 def __ne__(self, value):
135 if not isinstance(value, MetaStep):
136 return NotImplemented
137 return not self.__eq__(value)
139 def __hash__(self):
140 return hash((self.type, self.micro_index, self.stage_index))
142 def __str__(self):
143 if self.sub_steps:
144 sub = ", ".join(str(s) for s in self.sub_steps)
145 return (f"MetaStep(type={self.type}, micro_index={self.micro_index}, "
146 f"stage_index={self.stage_index}, sub_steps=[{sub}])")
147 return f"MetaStep(type={self.type}, micro_index={self.micro_index}, stage_index={self.stage_index})"
149 def __repr__(self):
150 return self.__str__()
152 @staticmethod
153 def from_str(step_str):
154 """Parse a MetaStep from its string representation."""
155 pass
158def generate_stage_to_rank_mapping(real_stage_num, stage_num, style='loop'):
159 """Generate stage to rank mapping for loop or V schedules."""
160 if style == 'loop':
161 return {stage_index: stage_index % real_stage_num for stage_index in range(stage_num)}
162 if style == 'v':
163 if stage_num % real_stage_num != 0:
164 raise ValueError(
165 f"stage_num {stage_num} must be evenly divisible by real_stage_num {real_stage_num} for V schedules."
166 )
167 mapping = {}
168 rank_index = 0
169 for stage_index in range(stage_num):
170 mapping[stage_index] = rank_index
171 if (stage_index + 1) % real_stage_num == 0:
172 continue
173 if (stage_index // real_stage_num) % 2 == 0:
174 rank_index += 1
175 else:
176 rank_index -= 1
177 return mapping
178 raise ValueError(f"Unsupported stage rank mapping style: {style}")
181def generate_rank_to_stage_mapping(real_stage_num, stage_num, style='loop'):
182 """Invert the stage to rank mapping."""
183 stage_to_rank = generate_stage_to_rank_mapping(real_stage_num, stage_num, style)
184 rank_to_stages = defaultdict(list)
185 for stage_index, rank in stage_to_rank.items():
186 rank_to_stages[rank].append(stage_index)
188 for stages in rank_to_stages.values():
189 stages.sort()
190 return dict(rank_to_stages)
193def iter_leaf_meta_steps(step):
194 """Yield leaf MetaSteps, recursively expanding OVERLAP containers.
196 Both ``OVERLAP_F_B`` and ``OVERLAP_B_F`` carry their real FWD/BWD work in
197 ``sub_steps``; the FSDP unshard/reshard injection pass relies on this to
198 see the FWD/BWD buried inside an overlap. Missing ``OVERLAP_B_F`` here let
199 an overlapped FWD run against a resharded stage → "expected HSDPModule
200 parameters in unsharded state". Mirror the other expansion sites
201 (run/_expand/add_fsdp_*) which already handle both composite types.
202 """
203 if step is None:
204 return
205 if step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F) and step.sub_steps:
206 for sub_step in step.sub_steps:
208 yield from iter_leaf_meta_steps(sub_step)
209 return
210 yield step
213class PipelineContext:
214 """Per-run state handed to a custom execution function (see
215 :meth:`PipelineScheduleRuntime.register_custom_function`).
217 A plain data carrier for one :meth:`PipelineScheduleRuntime.run_microbatches`
218 call. The P2P helpers (``wait_fwd_recv`` / ``wait_bwd_recv`` / ``send_fwd``
219 / ``send_bwd``) and the ``enable_dxdw_split`` flag live on the schedule, so a
220 callback reaches them through :attr:`schedule`, e.g.
221 ``ctx.schedule.send_bwd(stage, micro_index)``.
223 Attributes:
224 schedule: The owning :class:`PipelineScheduleRuntime`.
225 arg_mbs: Per-micro-batch positional args.
226 kwarg_mbs: Per-micro-batch keyword args.
227 losses: Mutable list collecting per-step losses.
228 """
230 def __init__(self, schedule: "PipelineScheduleRuntime", arg_mbs: list,
231 kwarg_mbs: list, losses: list) -> None:
232 """Bundle the active schedule with one run's micro-batch inputs and losses."""
233 self.schedule = schedule
234 self.arg_mbs = arg_mbs
235 self.kwarg_mbs = kwarg_mbs
236 self.losses = losses
239def _exec_fsdp_unshard(stage):
240 """Unshard every HSDPModule in the stage's submodule tree."""
241 for _, module in platform.get_cells_and_names(stage.submodule):
242 if isinstance(module, HSDPModule):
243 module.unshard()
246def _exec_fsdp_reshard(stage):
247 """Reshard every HSDPModule in the stage's submodule tree."""
248 for _, module in platform.get_cells_and_names(stage.submodule):
249 if isinstance(module, HSDPModule):
250 module.reshard()
253def _exec_fsdp_reduce_grad(stage):
254 """Run the stage's FSDP post-backward gradient reduction."""
255 stage.execute_reduce_grad()
258# FSDP control MetaStep -> handler(stage). Membership also marks which
259# MetaStepTypes are FSDP control steps, so the runtime loop dispatches with a
260# single table lookup instead of re-switching on the step type.
261_FSDP_STEP_HANDLERS = {
262 MetaStepType.FSDP_UNSHARD: _exec_fsdp_unshard,
263 MetaStepType.FSDP_RESHARD: _exec_fsdp_reshard,
264 MetaStepType.FSDP_REDUCE_GRAD: _exec_fsdp_reduce_grad,
265}
268class PipelineScheduleRuntime(ABC):
269 """
270 Base class for pipeline schedule.
271 Implements the `split_microbatches` and `run_microbatches` method.
272 Derived classes should implement `run_microbatches` method and `run` method.
274 Supports registering **custom execution functions** for any
275 :class:`MetaStepType` via :meth:`register_custom_function`. When
276 ``run_microbatches`` encounters a step whose type has a registered
277 handler, it creates a :class:`PipelineContext` and delegates execution
278 to the handler instead of using the built-in logic.
280 Args:
281 stages (list[PipelineStage], PipelineStage): PipelineStage used to run_microbatches.
282 micro_batch_num (int): The number of micro-batch.
283 args_batch_dim (int | BatchDimSpec | list | tuple, optional): Per
284 positional-arg batch dim, indexed by arg position. Entries may be
285 plain ``int`` (or ``None`` to keep the default); a single-input
286 model may pass a bare ``int``/``BatchDimSpec`` instead of a
287 one-element list (wrapped automatically). Default ``None``.
288 kwargs_batch_dim (dict, optional): Per keyword-arg batch dim, mapping
289 arg name to a plain ``int`` or ``BatchDimSpec``. Default ``None``.
290 swap (bool, optional): Whether to inject pipeline activation swap
291 control steps. Supported by ``ScheduleGPipe``, ``Schedule1F1B``,
292 and ``ScheduleInterleaved1F1B``. Default ``False``.
293 p2p_transport (str, optional): How pipeline send/recv are issued.
294 ``"auto"`` (default) — gap-time duplex batching on overlap_b_f
295 schedules (``coalesce_p2p``: same-peer send+recv as one
296 ``batch_isend_irecv``, TX||RX; hardware-validated and measured a
297 net win on real workloads), plain per-op ``isend``/``irecv``
298 everywhere else. ``"plain"`` — force per-op ``isend``/``irecv``
299 (escape hatch for transports or topologies where batching
300 misbehaves). ``"batch"`` — duplex batching explicitly (what
301 ``"auto"`` picks under overlap_b_f). ``"boundary"`` —
302 EXPERIMENTAL fwd-boundary batching: each overlap's ``F_SEND`` +
303 the next slot's recvs go out mid-overlap, right after the forward,
304 as per-op solo batches; only ``B_SEND`` waits for the backward.
305 Avoids the duplex handle's send-coupling (a2a-friendly) and posts
306 the activation send ~half a slot early, but is not yet
307 hardware-validated — opt in deliberately. Must be set identically
308 on every rank — HCCL cannot match a batched op against a plain one
309 (EI0005). On backends that require full-group participation before
310 the first batched P2P call, batch-backed modes prepare the PP group
311 once at their first run boundary.
312 """
314 _P2P_TRANSPORTS = ("auto", "plain", "batch", "boundary")
316 def __init__(self,
317 stages,
318 micro_batch_num,
319 args_batch_dim=None,
320 kwargs_batch_dim=None,
321 output_concat_dim=None,
322 overlap_p2p=False,
323 swap=False,
324 p2p_transport="auto"):
325 if p2p_transport not in self._P2P_TRANSPORTS:
326 raise ValueError(
327 f"p2p_transport must be one of {self._P2P_TRANSPORTS}, got "
328 f"{p2p_transport!r}"
329 )
330 self.stages = self._check_stages(stages)
331 self.micro_batch_num = micro_batch_num
332 self._args_batch_dim = self._normalize_args_batch_dim(args_batch_dim)
333 self._kwargs_batch_dim = self._normalize_kwargs_batch_dim(kwargs_batch_dim)
334 self._output_concat_dim = output_concat_dim
335 self.split_micro_batch = platform.micro_batch(self.micro_batch_num,
336 self._args_batch_dim, self._kwargs_batch_dim)
337 self.n_local_stages = len(self.stages)
338 self._stage_dict = self.convert_stages_dict()
339 self.real_stage_num = self.stages[0].stage_num // self.n_local_stages
340 self._stage_num = self.stages[0].stage_num
341 self._stage_to_rank_index = None
342 self._overlap_p2p = overlap_p2p
343 self.exec_order = {}
344 self._init_stages()
345 self._build_stage_to_rank_index()
346 self.fwd_handle_cache = {}
347 self.bwd_handle_cache = {}
348 self._custom_fn_map = {}
349 self._pp_swap_enabled = swap
350 self._swap_keys = frozenset()
351 self._swap_session = None
352 # Outstanding async send handle groups for the in-flight
353 # ``run_microbatches`` call; reset per run and drained at its end.
354 self._send_handles = []
355 # ``p2p_transport`` resolves in ``build_exec_order`` (it needs the
356 # subclass's ``_overlap_b_f``) to one of:
357 #
358 # * ``"batch"`` (the ``"auto"`` default on overlap_b_f schedules) —
359 # gap-time duplex via ``coalesce_p2p``: same-peer send+recv as one
360 # ``batch_isend_irecv`` (TX||RX on the full-duplex link).
361 # Hardware-validated and MEASURED a net win on real workloads — the
362 # duplex saving outweighs its known cost (MS's single handle couples
363 # the riding send into the compute-gating recv wait, which can
364 # shave EP a2a overlap).
365 # * ``"plain"`` (the ``"auto"`` default elsewhere) — per-op
366 # ``isend``/``irecv``, the upstream-original path.
367 # * ``"boundary"`` (EXPERIMENTAL, explicit opt-in only) — fwd-boundary
368 # batching. ``attach_fwd_boundary_p2p`` hangs each steady gap's
369 # ``F_SEND`` (its data is ready when the overlap's forward finishes
370 # — the backward, ~2x FLOPs, is the long pole) plus the next slot's
371 # recvs on the OVERLAP_B_F step; the stage's after-forward hook
372 # fires ``exec_boundary_p2p`` at the fwd/bwd boundary, while the
373 # backward is still running. The send leaves roughly half a slot
374 # early — the only mode that moves the SENDER's post time — and the
375 # recv handles carry no send (a2a-friendly). Every op is a per-op
376 # solo batch and ``coalesce_p2p`` is NOT run: a hoisted recv cannot
377 # stay duplexed with a send that is only ready later, and
378 # asymmetric shapes (one end duplex, other end split) hang —
379 # all-solo keeps per-pair batch sequences complementary
380 # ([S,R] vs [R,S]), safe under both candidate HCCL pairing
381 # semantics. Promote to the auto default only after it earns both
382 # a hardware accuracy pass and a perf win over "batch".
383 #
384 # Two transport invariants for any future rewrite: plain per-pair
385 # streams need the gap's recv-first/send-first complementarity (a
386 # send-crossing hoist made both ends recv-first -> rendezvous deadlock,
387 # 2026-06), and batch pairing needs per-pair shape mirroring (a
388 # one-sided split made 2 solos face 1 duplex -> hang, 2026-06).
389 self._p2p_transport = p2p_transport
390 # Effective mode + per-op batch gating; set by ``build_exec_order``.
391 self._p2p_mode = None
392 self._batch_p2p = False
393 self._batch_p2p_group = None
394 self._batch_p2p_group_initialized = False
395 # OVERLAP steps whose boundary_p2p was already issued this run (the
396 # stage after-forward hook and the post-step safety net are both
397 # allowed to call exec_boundary_p2p; reset per run_microbatches call).
398 self._boundary_issued = set()
399 # (fwd stage_index, micro_index) -> armed OVERLAP step, consumed by the
400 # stage after-forward hook to fire the boundary issue mid-overlap.
401 self._pending_boundary = {}
403 def register_custom_function(self, step_type: MetaStepType, fn) -> None:
404 """Register a custom execution function for the given step type.
406 When :meth:`run_microbatches` encounters a :class:`MetaStep` whose
407 ``type`` matches ``step_type``, it calls ``fn(step, ctx)`` instead
408 of the built-in logic.
410 Args:
411 step_type: The :class:`MetaStepType` to intercept.
412 fn: A callable with signature ``(step: MetaStep, ctx: PipelineContext) -> None``.
414 When pipeline activation swap is enabled, callbacks that
415 execute FWD leaves must call ``ctx.schedule.execute_fwd_leaf``
416 so swap collection runs in the matching chunk context.
418 Example:
419 >>> def my_overlap_callback(step, ctx):
420 ... fwd_step, bwd_step = step.sub_steps
421 ... # custom parallel execution logic
422 >>> schedule.register_custom_function(MetaStepType.OVERLAP_F_B, my_overlap_callback)
423 """
424 self._custom_fn_map[step_type] = fn
426 def _inject_local_fsdp_actions(self):
427 """Annotate the local rank schedule with optional FSDP control actions."""
428 current_rank = self._stage_to_rank_index[self.stages[0].stage_index]
429 managed_stage_indices = {
430 stage.stage_index
431 for stage in self.stages
432 if any(
433 isinstance(module, HSDPModule)
434 for _, module in platform.get_cells_and_names(stage.submodule)
435 )
436 }
437 if not managed_stage_indices:
438 return
439 if len(managed_stage_indices) != len(self.stages):
440 raise RuntimeError(
441 "When injecting fsdp_action, expect all stages to be HSDPModule. "
442 "Check whether all separated modules are wrapped with 'fully_shard'."
443 )
444 rank_actions = add_fsdp_unshard_reshard(self.exec_order[current_rank], managed_stage_indices)
445 self.exec_order[current_rank] = add_fsdp_reduce_grad(
446 rank_actions,
447 managed_stage_indices,
448 self.micro_batch_num,
449 )
451 def _inject_local_pp_swap_actions(self):
452 """Annotate the local rank schedule with pipeline activation-swap actions."""
453 self._swap_keys = frozenset()
454 if not self._pp_swap_enabled:
455 return
456 current_rank = self._stage_to_rank_index[self.stages[0].stage_index]
457 self.exec_order[current_rank] = inject_pipeline_swap_steps(self.exec_order[current_rank])
458 self._swap_keys = frozenset(
459 (step.stage_index, step.micro_index)
460 for step in self.exec_order[current_rank]
461 if step is not None and step.type == MetaStepType.SWAP_LAUNCH_OFFLOAD
462 )
464 @abstractmethod
465 def _build_stage_to_rank_index(self) -> None:
466 """
467 Build attribute of _stage_to_rank_index.
468 Each subclass constructs it according to its own schedule style.
469 """
471 @abstractmethod
472 def construct_exec_order(self) -> None:
473 """Build exec order, PP cmopute and PP comms(Send/Recv)"""
475 def build_exec_order(self) -> None:
476 """Build the execution order and inject optional FSDP / PP-swap actions.
478 Meta-Step Ordering Contract
479 --------------------------
480 The per-rank execution schedule is assembled in layers:
482 1. ``construct_exec_order()``
483 Pure compute (FWD/BWD) + P2P communication (SEND/RECV).
485 2. ``_inject_local_fsdp_actions()``
486 FSDP parameter management: UNSHARD before compute, RESHARD after,
487 REDUCE_GRAD after the last backward of each stage. This layer must
488 run *before* swap injection so that swap can see the already-placed
489 FSDP steps and optimise its placement relative to them:
491 * After FWD: ``SWAP_LAUNCH_OFFLOAD`` before ``FSDP_RESHARD``
492 (start D2H early, freeing parameter memory is fast).
493 * Before BWD: if a direct lookahead ``FSDP_UNSHARD`` exists,
494 ``SWAP_LAUNCH_LOAD`` precedes it so H2D overlaps all-gather.
495 If parameters are already unsharded, load is delayed until the
496 BWD container to avoid extending activation residency.
498 3. ``coalesce_p2p()`` or ``attach_fwd_boundary_p2p()``
499 P2P batching — runs before swap injection so P2P containers are
500 finalized first.
502 4. ``_inject_local_pp_swap_actions()``
503 Activation swap: leaf-local collection around ``FWD``, followed by
504 ``LAUNCH_OFFLOAD → WAIT_OFFLOAD → LAUNCH_LOAD → WAIT_LOAD``.
505 ``WAIT_LOAD`` executes on the scheduler thread immediately before
506 the backward consumer container.
507 When ``enable_dxdw_split`` is active, ``BWD_INPUT`` is used as
508 the backward anchor because it is the activation consumer;
509 ``BWD_WEIGHT`` receives no swap control steps.
511 Swap steps are inserted only into before/after slots around
512 top-level containers and never inside ``BATCH_SEND_RECV.sub_steps``.
513 This preserves symmetric, paired fusion-block members and avoids
514 deadlocks caused by splitting a fused P2P block.
516 Also resolves ``p2p_transport``: ``"auto"`` becomes ``"batch"`` (the
517 measured-beneficial duplex) on schedules running with ``overlap_b_f``
518 and ``"plain"`` everywhere else. The matching P2P order-rewrite pass
519 runs after FSDP injection and before swap injection, so swap sees the
520 finalized top-level P2P containers.
521 """
522 mode = self._p2p_transport
523 if mode == "auto":
524 mode = "batch" if getattr(self, "_overlap_b_f", False) else "plain"
525 self._p2p_mode = mode
526 self._batch_p2p = mode != "plain"
527 self._batch_p2p_group = self._resolve_batch_p2p_group() if self._batch_p2p else None
528 self.construct_exec_order()
529 # FSDP must inject before PP-swap so that swap can optimise
530 # placement relative to FSDP steps (LAUNCH_OFFLOAD before
531 # RESHARD after FWD; LAUNCH_LOAD before UNSHARD for BWD).
532 self._inject_local_fsdp_actions()
533 if mode == "boundary":
534 # fwd-boundary mode: hang the forward's F_SEND + next slot's recvs
535 # on the OVERLAP step (issued mid-overlap, right after the
536 # forward). Everything stays per-op solo batches — deliberately
537 # NO coalesce_p2p (see __init__).
538 self.exec_order = attach_fwd_boundary_p2p(self.exec_order)
539 elif mode == "batch":
540 # Coalesce contiguous P2P runs into BATCH_SEND_RECV so the runtime
541 # issues same-peer send+recv as one duplex batch. NOTE: couples
542 # the riding send into the compute-gating recv wait — see
543 # __init__.
544 self.exec_order = coalesce_p2p(self.exec_order)
545 self._inject_local_pp_swap_actions()
547 def convert_stages_dict(self):
548 """convert stages to dict."""
549 stage_dict = {}
550 for stage in self.stages:
551 stage_dict[stage.stage_index] = stage
552 return stage_dict
554 def split_microbatches(self, args, kwargs):
555 """split_microbatches."""
556 if args or kwargs:
557 args_split, kwargs_split = self.split_micro_batch(args, kwargs)
558 return args_split, kwargs_split
559 return [[] for _ in range(self.micro_batch_num)], [{} for _ in range(self.micro_batch_num)]
561 @staticmethod
562 def _to_spec(elem):
563 """Normalize one batch-dim entry: ``int`` -> ``BatchDimSpec``.
565 ``None`` and ``BatchDimSpec`` pass through unchanged. ``bool`` is
566 rejected even though it is an ``int`` subclass, so ``True``/``False``
567 are not silently read as dims 1/0.
568 """
569 if elem is None or isinstance(elem, BatchDimSpec):
570 return elem
571 if isinstance(elem, int) and not isinstance(elem, bool):
572 return BatchDimSpec(elem)
573 raise TypeError(
574 f"batch-dim entry must be int, BatchDimSpec or None, but got {type(elem)}.")
576 @staticmethod
577 def _normalize_args_batch_dim(args_batch_dim):
578 """Accept a plain ``int``/``BatchDimSpec`` or a ``list``/``tuple`` of them.
580 ``args_batch_dim`` is a per-arg spec indexed by positional-arg position
581 (see ``_MicroBatch``). A single-input model can pass a bare ``int`` /
582 ``BatchDimSpec`` instead of the awkward one-element
583 ``BatchDimSpec.from_tuple((0,))``; elements may be plain ``int`` (or
584 ``None`` to keep the default). Always returns ``None`` or a
585 ``tuple[BatchDimSpec | None]`` so downstream per-arg indexing is
586 unchanged.
587 """
588 if args_batch_dim is None:
589 return None
590 if isinstance(args_batch_dim, BatchDimSpec) or \
591 (isinstance(args_batch_dim, int) and not isinstance(args_batch_dim, bool)):
592 args_batch_dim = (args_batch_dim,)
593 if isinstance(args_batch_dim, (list, tuple)):
594 return tuple(PipelineScheduleRuntime._to_spec(e) for e in args_batch_dim)
595 raise TypeError(
596 f"args_batch_dim must be int, BatchDimSpec or a list/tuple of them, "
597 f"but got {type(args_batch_dim)}.")
599 @staticmethod
600 def _normalize_kwargs_batch_dim(kwargs_batch_dim):
601 """Accept plain ``int`` dict values: ``{\"x\": 0}`` -> ``{\"x\": BatchDimSpec(0)}``.
603 ``kwargs_batch_dim`` maps each keyword-arg name to its batch dim.
604 Returns ``None`` or a ``dict[str, BatchDimSpec | None]`` so downstream
605 per-key indexing is unchanged.
606 """
607 if kwargs_batch_dim is None:
608 return None
609 if not isinstance(kwargs_batch_dim, dict):
610 raise TypeError(
611 f"kwargs_batch_dim must be a dict[str, int | BatchDimSpec], "
612 f"but got {type(kwargs_batch_dim)}.")
613 return {k: PipelineScheduleRuntime._to_spec(v) for k, v in kwargs_batch_dim.items()}
615 def _check_stages(self, stages):
616 """check stages type."""
617 if isinstance(stages, hyper_parallel.PipelineStage):
618 return [stages]
619 if isinstance(stages, (list, tuple)):
620 for stage in stages:
621 if not isinstance(stage, hyper_parallel.PipelineStage):
622 raise TypeError(f"Argument 'stages' must be type of PipelineStage, \
623 list or tuple of PipelineStage, but got list or tuple of {type(stage)}.")
624 return stages
625 raise TypeError(f"Argument 'stages' must be type of PipelineStage, \
626 list or tuple of PipelineStage, but got type of {type(stages)}.")
628 def _init_stages(self):
629 """init stages."""
630 for stage in self.stages:
631 stage.init(self.n_local_stages)
632 # After-forward hook: lets the schedule issue fwd-boundary P2P the
633 # moment a forward chunk completes (no-op unless an OVERLAP step
634 # with boundary_p2p was armed for that (stage, micro)).
635 stage._after_forward_chunk = self._on_forward_chunk_done # pylint: disable=W0212
637 def _on_forward_chunk_done(self, stage_index, micro_index):
638 """Stage after-forward hook: fire the armed boundary P2P, if any.
640 Runs on the thread executing the forward (the overlap callback's
641 ``fwd_fn`` / the main thread), at the fwd/bwd boundary — the paired
642 backward is still running, so the boundary ops overlap it. Keyed by
643 the overlap's forward ``(stage_index, micro_index)``; unrelated
644 forwards (warm-up steps, recompute re-runs of past micros) miss the
645 key and no-op.
646 """
647 step = self._pending_boundary.pop((stage_index, micro_index), None)
648 if step is not None:
649 self.exec_boundary_p2p(step)
651 def run(self, *args: Any, **kwargs: Any) -> list:
652 """schedule run."""
653 losses = []
654 try:
655 if self._pp_swap_enabled:
656 self._swap_session = PipelineSwapSession(self._swap_keys)
657 split_args, split_kwargs = self.split_microbatches(args, kwargs)
658 self.run_microbatches(split_args, split_kwargs, losses)
659 finally:
660 self._drain_inflight_p2p()
661 if self._swap_session is not None:
662 self._swap_session.close()
663 self._swap_session = None
664 return losses
666 def sync_shared_parameters_grad(self):
667 """sync_shared_parameters_grad."""
668 for stage in self.stages:
669 stage.sync_shared_parameters_grad()
671 def update_losses(self, stage, loss, losses):
672 """update_losses."""
673 if stage.is_last_stage:
674 losses.append(loss)
676 @property
677 def enable_dxdw_split(self) -> bool:
678 """Whether this schedule splits ``OVERLAP_B_F`` backward into dx/dw."""
679 return getattr(self, "_enable_dxdw_split", False)
681 def _wait_p2p(self, handles):
682 for handle in handles:
683 if handle is not None:
684 handle.wait()
686 def _resolve_batch_p2p_group(self):
687 """Return the common PP group required by every local stage."""
688 first_stage = self.stages[0]
689 group = first_stage.pp_group
690 for stage in self.stages[1:]:
691 if stage.pp_group is group or stage.pp_group == group:
692 continue
693 raise ValueError(
694 "Batch P2P requires all local stages to use the same PP process group, "
695 f"but stages {first_stage.stage_index} and {stage.stage_index} use different groups."
696 )
697 return group
699 def _ensure_batch_p2p_group_initialized(self) -> None:
700 """Run backend-specific PP-group preparation once before batched P2P."""
701 if (not self._batch_p2p or self._batch_p2p_group_initialized
702 or self.real_stage_num <= 1):
703 return
704 platform.prepare_batch_p2p_group(self._batch_p2p_group)
705 self._batch_p2p_group_initialized = True
707 def _drain_inflight_p2p(self):
708 """Wait every P2P handle still in flight — error-path cleanup.
710 run_microbatches waits its deferred sends only in the end-of-iteration
711 drain; an exception mid-iteration unwinds past that drain, leaving issued
712 isend/irecv handles un-waited in ``_send_handles`` and the recv caches
713 (the ``CommHandle destroyed without calling wait()`` warning). run()'s
714 finally calls this so every handle is still ``wait()``-ed — honoring the
715 comm contract — on the error path too, and pops them so a later run()
716 does not re-wait stale handles (the recv caches are never reset per run).
717 No-op on the normal path: the drain already emptied ``_send_handles`` and
718 every cached recv was consumed.
719 """
720 while self._send_handles:
721 self._wait_p2p(self._send_handles.pop())
722 while self.fwd_handle_cache:
723 self._wait_p2p(self.fwd_handle_cache.popitem()[1])
724 while self.bwd_handle_cache:
725 self._wait_p2p(self.bwd_handle_cache.popitem()[1])
727 def _batched_issue(self, specs):
728 """Launch same-peer P2P ``specs`` as one ``batch_isend_irecv`` group.
730 ``specs`` are ``(op_type, tensor, peer_global_rank)`` from the stage's
731 ``*_specs`` builders (which carry the meta/bookkeeping side effects).
732 Returns ``[handle]`` (the single batch handle) or ``[]`` — shaped like
733 the per-op ``exec_*_ops`` return so the cache / drain paths are
734 unchanged. Only the launch is coalesced; matching stays per-peer FIFO.
735 """
736 if not specs:
737 return []
738 ops = [
739 platform.p2p_op(op_type, tensor, peer, group=self._batch_p2p_group)
740 for op_type, tensor, peer in specs
741 ]
742 handle = platform.batch_isend_irecv(ops)
743 return [handle] if handle is not None else []
745 # --- P2P step primitives ------------------------------------------------
746 # One method per cross-rank comm action, used both by the runtime loop
747 # (``_exec_step``) and by OVERLAP callbacks (via ``ctx.schedule``). With
748 # ``overlap_p2p=True`` comm is decoupled from its compute: a recv caches its
749 # handles for the consuming step to ``wait_*`` later, and a send defers its
750 # handles to the end-of-iteration drain. With ``overlap_p2p=False`` every
751 # op waits inline.
753 def recv_fwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
754 """Post the FWD recv for ``micro_index``; cache it (overlap_p2p) or wait now."""
755 handles = (self._batched_issue(stage.fwd_recv_specs(micro_index))
756 if self._batch_p2p else stage.exec_fwd_recv_ops(micro_index))
757 if self._overlap_p2p:
758 self.fwd_handle_cache[(stage.stage_index, micro_index)] = handles
759 else:
760 self._wait_p2p(handles)
762 def recv_bwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
763 """Post the BWD recv for ``micro_index``; cache it (overlap_p2p) or wait now."""
764 handles = (self._batched_issue(stage.bwd_recv_specs(micro_index))
765 if self._batch_p2p else stage.exec_bwd_recv_ops(micro_index))
766 if self._overlap_p2p:
767 self.bwd_handle_cache[(stage.stage_index, micro_index)] = handles
768 else:
769 self._wait_p2p(handles)
771 def wait_fwd_recv(self, stage_index: int, micro_index: int) -> None:
772 """Wait the FWD recv cached by :meth:`recv_fwd`; no-op if nothing is cached."""
773 handles = self.fwd_handle_cache.pop((stage_index, micro_index), None)
774 if handles:
775 self._wait_p2p(handles)
777 def wait_bwd_recv(self, stage_index: int, micro_index: int) -> None:
778 """Wait the BWD recv cached by :meth:`recv_bwd`; no-op if nothing is cached."""
779 handles = self.bwd_handle_cache.pop((stage_index, micro_index), None)
780 if handles:
781 self._wait_p2p(handles)
783 def send_fwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
784 """Send this stage's forward output for ``micro_index`` to the next stage."""
785 handles = (self._batched_issue(stage.fwd_send_specs(micro_index))
786 if self._batch_p2p else stage.exec_fwd_send_ops(micro_index)) or []
787 if self._overlap_p2p:
788 # Append the whole handle group: run_microbatches drains _send_handles
789 # group by group, so a bare handle would be wrongly iterated as a list.
790 self._send_handles.append(handles)
791 else:
792 self._wait_p2p(handles)
794 def send_bwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
795 """Send this stage's input-gradient for ``micro_index`` to the previous stage.
797 Driven by the scheduler's ``BWD_SEND`` step. It pops the input grad that
798 the backward (unified ``backward_one_chunk`` or, under
799 ``enable_dxdw_split=True``, ``backward_input_one_chunk``) wrote to the
800 stage's ``bwd_cache``. Calling it manually in addition to the scheduled
801 ``BWD_SEND`` would double-send the gradient.
802 """
803 handles = (self._batched_issue(stage.bwd_send_specs(micro_index))
804 if self._batch_p2p else stage.exec_bwd_send_ops(micro_index)) or []
805 if self._overlap_p2p:
806 self._send_handles.append(handles)
807 else:
808 self._wait_p2p(handles)
810 def _arm_boundary(self, step):
811 """Register ``step`` for the stage after-forward hook; return its key.
813 No-op (returns ``None``) unless ``step`` carries ``boundary_p2p``. The
814 key is the overlap's forward ``(stage_index, micro_index)`` — exactly
815 what the hook receives when that forward chunk completes.
816 """
817 if not getattr(step, "boundary_p2p", None) or not step.sub_steps:
818 return None
819 fwd_sub = next((s for s in step.sub_steps
820 if s.type == MetaStepType.FWD), None)
821 if fwd_sub is None:
822 return None
823 key = (fwd_sub.stage_index, fwd_sub.micro_index)
824 self._pending_boundary[key] = step
825 return key
827 def _finish_boundary(self, step, armed_key) -> None:
828 """Post-step safety net: issue any boundary P2P the hook missed."""
829 self.exec_boundary_p2p(step)
830 if armed_key is not None:
831 self._pending_boundary.pop(armed_key, None)
833 def exec_boundary_p2p(self, step) -> None:
834 """Issue ``step.boundary_p2p`` (fwd-boundary P2P) once per run.
836 Fired by the stage after-forward hook (``_on_forward_chunk_done``) the
837 moment the overlap's forward chunk completes — the backward is still
838 running on its own thread, so the F_SEND leaves ~half a slot early and
839 the next slot's recvs are already posted when the peers' sends arrive.
840 Idempotent per ``run_microbatches`` call: the post-step safety net
841 (``_finish_boundary``) also invokes it, so an overlap whose forward
842 never went through ``forward_one_chunk`` degrades to gap-time issue
843 order instead of dropping the ops.
845 Dispatches through the existing per-op helpers (``send_fwd`` /
846 ``recv_fwd`` / ``recv_bwd``), so batching, handle caching for
847 ``wait_*_recv`` and deferred-send bookkeeping behave exactly like the
848 scheduled steps they replace. No-op for steps without
849 ``boundary_p2p``.
850 """
851 ops = getattr(step, "boundary_p2p", None)
852 if not ops or id(step) in self._boundary_issued:
853 return
854 self._boundary_issued.add(id(step))
855 for sub in ops:
856 stage = self._stage_dict[sub.stage_index]
857 if sub.type == MetaStepType.FWD_SEND:
858 self.send_fwd(stage, sub.micro_index)
859 elif sub.type == MetaStepType.FWD_RECV:
860 self.recv_fwd(stage, sub.micro_index)
861 elif sub.type == MetaStepType.BWD_RECV:
862 self.recv_bwd(stage, sub.micro_index)
863 elif sub.type == MetaStepType.BWD_SEND:
864 # attach_fwd_boundary_p2p never hoists BWD_SEND (its grad is
865 # produced by the backward still in flight); defensive only.
866 self.send_bwd(stage, sub.micro_index)
868 # ``op_type -> (specs builder name, route kind)`` for a coalesced sub-step.
869 # ``route`` is the recv-cache kind (so wait_*_recv finds the handle), or
870 # ``None`` for a send (no local consumer).
871 _BATCH_SUB_DISPATCH = {
872 MetaStepType.FWD_RECV: ("fwd_recv_specs", "fwd"),
873 MetaStepType.BWD_RECV: ("bwd_recv_specs", "bwd"),
874 MetaStepType.FWD_SEND: ("fwd_send_specs", None),
875 MetaStepType.BWD_SEND: ("bwd_send_specs", None),
876 }
878 def _exec_batch_send_recv(self, step) -> None:
879 """Execute a coalesced P2P run: one ``batch_isend_irecv`` per peer.
881 Builds each sub-step's specs (same meta/bookkeeping side effects as the
882 per-step ``recv_fwd`` / ``send_fwd`` / ...), groups every op by peer
883 global rank, and issues one batch per peer so a same-peer send+recv runs
884 duplex. Handle routing mirrors the per-step path: under
885 ``overlap_p2p`` a batch carrying a recv is cached for ``wait_*_recv``
886 (its send rides along), a send-only batch defers to ``_send_handles``;
887 without ``overlap_p2p`` every batch waits inline.
888 """
889 # (op_type, tensor, peer, route) per op; route = (kind, stage, micro)
890 # for a recv, else None.
891 tagged = []
892 for sub in step.sub_steps:
893 builder_name, kind = self._BATCH_SUB_DISPATCH[sub.type]
894 stage = self._stage_dict[sub.stage_index]
895 specs = getattr(stage, builder_name)(sub.micro_index)
896 route = (kind, sub.stage_index, sub.micro_index) if kind is not None else None
897 for op_type, tensor, peer in specs:
898 tagged.append((op_type, tensor, peer, route))
900 by_peer = {}
901 for item in tagged:
902 by_peer.setdefault(item[2], []).append(item)
904 for items in by_peer.values():
905 ops = [
906 platform.p2p_op(op_type, tensor, peer, group=self._batch_p2p_group)
907 for op_type, tensor, peer, _ in items
908 ]
909 handle = platform.batch_isend_irecv(ops)
910 if handle is None:
911 continue
912 if not self._overlap_p2p:
913 self._wait_p2p([handle])
914 continue
915 recv_routes = [route for *_, route in items if route is not None]
916 if recv_routes:
917 for kind, si, mi in recv_routes:
918 cache = self.fwd_handle_cache if kind == "fwd" else self.bwd_handle_cache
919 cache[(si, mi)] = [handle]
920 else:
921 self._send_handles.append([handle])
923 def _assert_in_unshard_if_needed(self, stage, check_step):
924 if not isinstance(stage.submodule, HSDPModule):
925 return
926 submodule_hsdp_scheduler = stage.submodule.hsdp_scheduler
927 scheduler_state = submodule_hsdp_scheduler.hsdp_state
928 if scheduler_state.is_shard:
929 raise RuntimeError(
930 f"Executing MetaStep: {check_step}, expected HSDPModule parameters in unsharded "
931 f"state, but got sharded parameters."
932 )
934 def _exec_step(self, cur_step, arg_mbs, kwarg_mbs, losses):
935 """Execute one built-in step (non-custom, non-composite).
937 Each comm step dispatches to a single P2P primitive; each compute step
938 first waits its cached recv (a no-op under ``overlap_p2p=False``) and
939 then runs.
940 """
941 stage = self._stage_dict[cur_step.stage_index]
942 micro_index = cur_step.micro_index
943 step_type = cur_step.type
945 if step_type in (
946 MetaStepType.SWAP_LAUNCH_OFFLOAD,
947 MetaStepType.SWAP_WAIT_OFFLOAD,
948 MetaStepType.SWAP_LAUNCH_LOAD,
949 MetaStepType.SWAP_WAIT_LOAD,
950 ):
951 self._exec_pipeline_swap_step(cur_step, arg_mbs, kwarg_mbs)
953 elif step_type == MetaStepType.FWD_RECV:
954 self.recv_fwd(stage, micro_index)
956 elif step_type == MetaStepType.FWD:
957 self.execute_fwd_leaf(cur_step, arg_mbs, kwarg_mbs, losses)
959 elif step_type == MetaStepType.FWD_SEND:
960 self.send_fwd(stage, micro_index)
962 elif step_type == MetaStepType.BWD_RECV:
963 self.recv_bwd(stage, micro_index)
965 elif step_type == MetaStepType.BWD_INPUT:
966 self._assert_in_unshard_if_needed(stage, cur_step)
967 self.wait_bwd_recv(stage.stage_index, micro_index)
968 stage.backward_input_one_chunk(micro_index)
970 elif step_type == MetaStepType.BWD_WEIGHT:
971 self._assert_in_unshard_if_needed(stage, cur_step)
972 self.wait_bwd_recv(stage.stage_index, micro_index)
973 stage.backward_weight_one_chunk(micro_index)
975 elif step_type == MetaStepType.BWD:
976 self._assert_in_unshard_if_needed(stage, cur_step)
977 self.wait_bwd_recv(stage.stage_index, micro_index)
978 stage.backward_one_chunk(micro_index)
980 elif step_type == MetaStepType.BWD_SEND:
981 self.send_bwd(stage, micro_index)
983 else:
984 # FSDP control steps dispatch via the handler table; any other type
985 # is a no-op here (composite/custom types are handled upstream).
986 fsdp_handler = _FSDP_STEP_HANDLERS.get(step_type)
987 if fsdp_handler is not None:
988 fsdp_handler(stage)
990 def execute_fwd_leaf(self, step: MetaStep, arg_mbs: list, kwarg_mbs: list, losses: list) -> None:
991 """Execute a forward leaf within its run-scoped swap group context."""
992 stage = self._stage_dict[step.stage_index]
993 micro_index = step.micro_index
994 self._assert_in_unshard_if_needed(stage, step)
995 self.wait_fwd_recv(stage.stage_index, micro_index)
996 swap_managed = self._swap_session is not None and self._swap_session.manages(step)
997 group_context = self._swap_session.group_context(step) if swap_managed else nullcontext()
998 with group_context:
999 out = stage.forward_one_chunk(micro_index, arg_mbs[micro_index], kwarg_mbs[micro_index])
1000 if swap_managed:
1001 # Boundary transport may pop the stage output cache from the
1002 # after-forward hook before this leaf returns. Protect the
1003 # local output while it is still directly available.
1004 self._swap_session.protect_aliases(step, out)
1005 self.update_losses(stage, out, losses)
1007 def _exec_pipeline_swap_step(self, cur_step, arg_mbs, kwarg_mbs):
1008 """Execute a pipeline activation-swap control step."""
1009 if self._swap_session is None:
1010 raise RuntimeError("Pipeline swap step executed without an active run session.")
1011 if cur_step.type == MetaStepType.SWAP_LAUNCH_OFFLOAD:
1012 swap_launch_offload(cur_step, self, arg_mbs, kwarg_mbs, self._swap_session)
1013 elif cur_step.type == MetaStepType.SWAP_WAIT_OFFLOAD:
1014 swap_wait_offload(cur_step, self._swap_session)
1015 elif cur_step.type == MetaStepType.SWAP_LAUNCH_LOAD:
1016 swap_launch_load(cur_step, self._swap_session)
1017 elif cur_step.type == MetaStepType.SWAP_WAIT_LOAD:
1018 swap_wait_load(cur_step, self._swap_session)
1020 def run_microbatches(self, arg_mbs: list, kwarg_mbs: list, losses: list) -> None:
1021 """Execute the schedule step by step.
1023 Steps whose :attr:`MetaStep.type` has a registered custom function
1024 are delegated to that function with a :class:`PipelineContext`.
1025 Composite ``OVERLAP_F_B`` / ``OVERLAP_B_F`` steps without a
1026 registered handler fall back to executing their ``sub_steps``
1027 sequentially via :meth:`_exec_step` — correct but without
1028 comm/compute overlap. All other steps are executed by
1029 :meth:`_exec_step`.
1031 Logs one ``DEBUG`` line per non-bubble step showing the rank's
1032 progress: ``rank=<r> step=<i>/<n> <MetaStep>``. Enable with
1033 ``logging.getLogger('hyper_parallel.core.pipeline_parallel.scheduler')
1034 .setLevel(logging.DEBUG)`` to trace per-rank schedule advancement
1035 (handy when diagnosing deadlocks or callback ordering issues).
1036 """
1037 self._ensure_batch_p2p_group_initialized()
1038 real_stage_index = self.stages[0].stage_index % self.real_stage_num
1039 self._send_handles = []
1040 self._boundary_issued = set()
1041 self._pending_boundary = {}
1042 ctx = None # lazily created
1044 ordered = self.exec_order[real_stage_index]
1045 total_steps = len(ordered)
1046 logger.debug(
1047 "run_microbatches start: rank=%d total_steps=%d micro_batch_num=%d",
1048 real_stage_index, total_steps, self.micro_batch_num,
1049 )
1051 for step_idx, cur_step in enumerate(ordered):
1052 if cur_step is None:
1053 continue
1055 logger.debug(
1056 "rank=%d step=%d/%d %s",
1057 real_stage_index, step_idx, total_steps, cur_step,
1058 )
1060 # Arm the fwd-boundary hook: when this step carries boundary_p2p,
1061 # the stage's after-forward hook fires exec_boundary_p2p the moment
1062 # the overlap's forward chunk completes (works for any callback —
1063 # no callback cooperation needed).
1064 armed_key = self._arm_boundary(cur_step)
1066 # Check for registered custom function
1067 custom_fn = self._custom_fn_map.get(cur_step.type)
1068 if custom_fn is not None:
1069 if ctx is None:
1070 ctx = PipelineContext(self, arg_mbs, kwarg_mbs, losses)
1071 custom_fn(cur_step, ctx)
1072 # Safety net: if the forward hook never fired (custom fwd path),
1073 # issue the boundary ops now (gap-time order) instead of
1074 # dropping them. Idempotent.
1075 self._finish_boundary(cur_step, armed_key)
1076 continue
1078 # Coalesced P2P block: group sub-steps by peer, issue one
1079 # batch_isend_irecv per peer (same-peer send+recv -> duplex).
1080 if cur_step.type == MetaStepType.BATCH_SEND_RECV:
1081 self._exec_batch_send_recv(cur_step)
1082 continue
1084 # Default for composite OVERLAP steps: run sub_steps sequentially.
1085 # P2P send/recv around these steps are already laid out in two
1086 # virtual slots by ``add_send_recv``, so sequential execution is
1087 # semantically equivalent to non-overlapped 1F1B.
1088 if (cur_step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F)
1089 and cur_step.sub_steps):
1090 for sub in cur_step.sub_steps:
1091 self._exec_step(sub, arg_mbs, kwarg_mbs, losses)
1092 self._finish_boundary(cur_step, armed_key)
1093 continue
1095 self._exec_step(cur_step, arg_mbs, kwarg_mbs, losses)
1097 logger.debug(
1098 "run_microbatches end: rank=%d pending_send_handles=%d",
1099 real_stage_index, len(self._send_handles),
1100 )
1101 self.sync_shared_parameters_grad()
1102 while self._send_handles:
1103 self._wait_p2p(self._send_handles.pop())
1106class _OverlapPhantom:
1107 """Internal marker used by :func:`add_send_recv` to expand an
1108 ``OVERLAP_F_B`` or ``OVERLAP_B_F`` step into two virtual time slots.
1110 An overlap step composes two sub-steps (``B + F`` or ``F + B``) that
1111 execute concurrently on the GPU but occupy **two** logical time slots
1112 in the column-scan sender timeline — the sender can only finish
1113 emitting the second sub-step's output after the first sub-step has
1114 completed. Treating an overlap step as a single slot places the RECV
1115 triggered by the second sub-step too early on the receiver.
1117 Each overlap step is expanded into two phantoms:
1118 * ``is_first_half=True`` — represents the first sub-step's emission
1119 slot; the original overlap step is emitted into the output
1120 schedule here (only once).
1121 * ``is_first_half=False`` — represents the second sub-step's emission
1122 slot; only its send/recv comms are inserted.
1123 """
1125 __slots__ = ('obf_step', 'sub_step', 'is_first_half')
1127 def __init__(self, obf_step, sub_step, is_first_half: bool):
1128 self.obf_step = obf_step
1129 self.sub_step = sub_step
1130 self.is_first_half = is_first_half
1133def _expand_overlap_slots(scheduler, real_stage_num):
1134 """Expand OVERLAP steps in a per-rank schedule into 2 virtual time slots.
1136 Returns a new ``{rank: [MetaStep | _OverlapPhantom | None, ...]}`` dict
1137 where each OVERLAP step is replaced by a pair of phantoms. Non-OVERLAP
1138 entries pass through unchanged.
1139 """
1140 expanded = {}
1141 for rank in range(real_stage_num):
1142 order = scheduler[rank]
1143 exp = []
1144 for op in order:
1145 if (op is not None
1146 and op.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F)
1147 and op.sub_steps):
1148 exp.append(_OverlapPhantom(op, op.sub_steps[0], is_first_half=True))
1149 exp.append(_OverlapPhantom(op, op.sub_steps[1], is_first_half=False))
1150 else:
1151 exp.append(op)
1152 expanded[rank] = exp
1153 return expanded
1156def _process_rank_items(real_stage_num, current_items, insert_step_comms, new_schedule):
1157 """Run ``insert_step_comms`` for each rank's current item, even ranks first.
1159 Even-before-odd ordering avoids P2P deadlocks between adjacent ranks.
1160 """
1161 for rank in range(0, real_stage_num, 2):
1162 item = current_items.get(rank)
1163 if item is not None:
1164 sub = item.sub_step if isinstance(item, _OverlapPhantom) else item
1165 insert_step_comms(sub, rank, new_schedule)
1166 for rank in range(1, real_stage_num, 2):
1167 item = current_items.get(rank)
1168 if item is not None:
1169 sub = item.sub_step if isinstance(item, _OverlapPhantom) else item
1170 insert_step_comms(sub, rank, new_schedule)
1173def _column_scan_insert_comms(expanded, real_stage_num, insert_step_comms):
1174 """Column-scan over an OVERLAP-expanded schedule to insert SEND/RECV.
1176 Processes ``expanded`` one time slot at a time. Emits the original
1177 overlap step into ``new_schedule`` only once (at the first-half
1178 phantom). Delegates comm insertion to ``insert_step_comms`` for each
1179 plain step or phantom's underlying sub-step.
1181 Even ranks are processed before odd ranks at each time step to avoid
1182 P2P deadlocks between adjacent ranks.
1184 Args:
1185 expanded: Result of :func:`_expand_overlap_slots`.
1186 real_stage_num: Number of physical ranks.
1187 insert_step_comms: Callable ``(step, rank, new_schedule) -> None``
1188 that inserts SEND/RECV for a single FWD/BWD step.
1190 Returns:
1191 ``{rank: [MetaStep, ...]}`` final schedule.
1192 """
1193 max_length = max(len(order) for order in expanded.values())
1194 new_schedule = {rank: [] for rank in range(real_stage_num)}
1196 for time_step in range(max_length):
1197 current_items = {}
1198 for rank in range(real_stage_num):
1199 if time_step < len(expanded[rank]):
1200 item = expanded[rank][time_step]
1201 current_items[rank] = item
1202 if item is None:
1203 # Preserve bubble slots to keep per-rank time-step
1204 # indexing aligned with the column scan. The runtime
1205 # loop skips ``None`` entries, so this is execution-
1206 # semantics-neutral.
1207 new_schedule[rank].append(None)
1208 continue
1209 if isinstance(item, _OverlapPhantom):
1210 # Emit the overlap step only once, at the first-half slot.
1211 if item.is_first_half:
1212 new_schedule[rank].append(item.obf_step)
1213 else:
1214 new_schedule[rank].append(item)
1215 else:
1216 current_items[rank] = None
1218 _process_rank_items(
1219 real_stage_num, current_items, insert_step_comms, new_schedule,
1220 )
1222 return new_schedule
1225_P2P_STEP_TYPES = frozenset({
1226 MetaStepType.FWD_SEND, MetaStepType.FWD_RECV,
1227 MetaStepType.BWD_SEND, MetaStepType.BWD_RECV,
1228})
1231def coalesce_p2p(exec_order):
1232 """Coalesce maximal contiguous runs of >=2 P2P steps into BATCH_SEND_RECV.
1234 A *run* is a maximal sequence of consecutive ``FWD_SEND`` / ``FWD_RECV`` /
1235 ``BWD_SEND`` / ``BWD_RECV`` steps with no compute / overlap / bubble (``None``)
1236 step between them — so no recv in the run is consumed before the batch is
1237 issued, and all sends' data is already produced. Each such run is replaced
1238 by a single :class:`MetaStep` of type ``BATCH_SEND_RECV`` carrying the run as
1239 ``sub_steps`` (order preserved, so per-direction FIFO is kept); the runtime
1240 groups those sub-steps by peer and issues one ``batch_isend_irecv`` per peer
1241 (same-peer send+recv -> duplex). Runs of length 1 are left untouched (the
1242 per-op batched path still batches them, so every transfer is still
1243 batch-vs-batch). Pure ``exec_order -> exec_order`` transform.
1245 Args:
1246 exec_order: ``{rank: [MetaStep | None, ...]}``.
1248 Returns:
1249 A new ``{rank: [...]}`` with contiguous P2P runs coalesced.
1250 """
1251 def _flush(run, new):
1252 if len(run) >= 2:
1253 new.append(MetaStep(None, MetaStepType.BATCH_SEND_RECV, None, sub_steps=tuple(run)))
1254 else:
1255 new.extend(run)
1257 out = {}
1258 for rank, order in exec_order.items():
1259 new = []
1260 run = []
1261 for step in order:
1262 if step is not None and step.type in _P2P_STEP_TYPES:
1263 run.append(step)
1264 continue
1265 _flush(run, new)
1266 run = []
1267 new.append(step)
1268 _flush(run, new)
1269 out[rank] = new
1270 return out
1273_RECV_STEP_TYPES = frozenset({MetaStepType.FWD_RECV, MetaStepType.BWD_RECV})
1276def attach_fwd_boundary_p2p(exec_order):
1277 """Hang each overlap gap's boundary-safe P2P on the OVERLAP_B_F step.
1279 For every ``OVERLAP_B_F`` step, the contiguous P2P run right after it is
1280 split by data readiness at the overlap's fwd/bwd boundary (forward is the
1281 short side; backward, ~2x FLOPs, is the long pole):
1283 * the forward's own ``FWD_SEND`` — its payload exists the moment the
1284 forward sub-step finishes, no need to wait out the backward;
1285 * every ``FWD_RECV`` / ``BWD_RECV`` — no data dependency at all;
1287 are removed from the gap and attached to the OVERLAP step as
1288 ``boundary_p2p`` (order: ``F_SEND`` first, then the recvs in original
1289 order), to be issued by :meth:`PipelineScheduleRuntime.exec_boundary_p2p`
1290 at the boundary, while the backward is still running. ``BWD_SEND`` (its
1291 grad is produced by that backward) and any send not produced by this
1292 overlap's forward stay in the gap.
1294 Pairing shape (the reason this composition is safe where naive
1295 hoist+coalesce hung): with every op issued as a per-op solo batch, each
1296 pair's per-pair batch sequence per slot is ``[F_SEND, B_RECV]`` on the
1297 prev end and ``[F_RECV, B_SEND]`` on the next end — complementary at every
1298 position and equal in count, so it matches under both per-direction FIFO
1299 and per-pair shape-mirroring semantics. Per-direction FIFO data order is
1300 preserved (each direction's ops keep their relative order; they all shift
1301 by the same amount). Pure ``exec_order -> exec_order`` transform.
1303 Args:
1304 exec_order: ``{rank: [MetaStep | None, ...]}``.
1306 Returns:
1307 A new ``{rank: [...]}`` with boundary P2P attached to OVERLAP steps.
1308 """
1309 out = {}
1310 for rank, order in exec_order.items():
1311 new = []
1312 i = 0
1313 while i < len(order):
1314 step = order[i]
1315 if (step is None or step.type != MetaStepType.OVERLAP_B_F
1316 or not step.sub_steps):
1317 new.append(step)
1318 i += 1
1319 continue
1320 run, j = _p2p_run_after(order, i + 1)
1321 boundary, leftover = _split_boundary_run(step, run)
1322 if not boundary:
1323 new.append(step)
1324 i += 1
1325 continue
1326 new.append(MetaStep(step.micro_index, step.type, step.stage_index,
1327 sub_steps=step.sub_steps, boundary_p2p=boundary))
1328 new.extend(leftover)
1329 i = j
1330 out[rank] = new
1331 return out
1334def _p2p_run_after(order, start):
1335 """Collect the contiguous P2P run starting at ``start``.
1337 Returns ``(run, end)`` where ``end`` is the index of the first step past
1338 the run (a compute step, ``None`` bubble, or end of order).
1339 """
1340 run = []
1341 j = start
1342 while j < len(order) and order[j] is not None and order[j].type in _P2P_STEP_TYPES:
1343 run.append(order[j])
1344 j += 1
1345 return run, j
1348def _split_boundary_run(step, run):
1349 """Split an overlap's trailing P2P run by fwd/bwd-boundary data readiness.
1351 Returns ``(boundary, leftover)``: ``boundary`` holds the overlap's own
1352 forward ``FWD_SEND`` (payload ready at the boundary) first, then every
1353 recv (no data dependency), keeping original order; ``leftover`` keeps the
1354 sends produced by the still-running backward, in place.
1355 """
1356 fwd_sub = next((s for s in step.sub_steps
1357 if s.type == MetaStepType.FWD), None)
1359 def _is_own_fwd_send(s):
1360 return (fwd_sub is not None
1361 and s.type == MetaStepType.FWD_SEND
1362 and s.stage_index == fwd_sub.stage_index
1363 and s.micro_index == fwd_sub.micro_index)
1365 boundary = ([s for s in run if _is_own_fwd_send(s)]
1366 + [s for s in run if s.type in _RECV_STEP_TYPES])
1367 taken = {id(s) for s in boundary}
1368 leftover = [s for s in run if id(s) not in taken]
1369 return tuple(boundary), leftover
1372def split_overlap_dxdw(exec_order: dict) -> dict:
1373 """Split each OVERLAP_B_F backward into dx (in the pair) + dw (after the gap).
1375 Rewrites ``(BWD, FWD)`` sub_steps to ``(BWD_INPUT, FWD)`` and inserts the
1376 matching ``BWD_WEIGHT`` after the contiguous P2P run that follows the
1377 overlap. The overlap then joins at ``max(dx, fwd)`` instead of
1378 ``max(dx + dw, fwd)``, so the gap's ``BWD_SEND`` (dx already wrote its
1379 grad to ``bwd_cache``) and the next slot's recvs are issued a dw earlier;
1380 under ``overlap_p2p`` they are async and dw computes while they fly.
1382 Comm placement is untouched: the pass runs after ``add_send_recv`` and
1383 only moves local compute, so the cross-rank matching order is identical,
1384 and the P2P run stays contiguous (dw lands after it, not inside) so
1385 ``coalesce_p2p`` / ``attach_fwd_boundary_p2p`` see the same gap shape.
1387 First-stage (``stage_index == 0``) backwards stay unified: their dx is a
1388 no-op (no input grad to compute or send), so splitting would only move
1389 the whole backward out of the overlap and lose its fwd overlap.
1391 Args:
1392 exec_order: ``{rank: [MetaStep | None, ...]}``.
1394 Returns:
1395 A new ``{rank: [...]}`` with overlap backwards split into dx/dw.
1396 """
1397 out = {}
1398 for rank, order in exec_order.items():
1399 new = []
1400 i = 0
1401 n = len(order)
1402 while i < n:
1403 step = order[i]
1404 i += 1
1405 if (step is None or step.type != MetaStepType.OVERLAP_B_F
1406 or not step.sub_steps):
1407 new.append(step)
1408 continue
1409 bwd_sub, fwd_sub = step.sub_steps
1410 if bwd_sub.type != MetaStepType.BWD or bwd_sub.stage_index == 0:
1411 new.append(step)
1412 continue
1413 dx = MetaStep(bwd_sub.micro_index, MetaStepType.BWD_INPUT, bwd_sub.stage_index)
1414 new.append(MetaStep(step.micro_index, step.type, step.stage_index,
1415 sub_steps=(dx, fwd_sub)))
1416 run, i = _p2p_run_after(order, i)
1417 new.extend(run)
1418 new.append(MetaStep(bwd_sub.micro_index, MetaStepType.BWD_WEIGHT,
1419 bwd_sub.stage_index))
1420 out[rank] = new
1421 return out
1424def add_send_recv(scheduler, stage_num, real_stage_num, style='loop'):
1425 """Insert P2P send/recv operations into a per-rank compute schedule.
1427 For each FWD or BWD step that requires cross-rank communication, a
1428 ``FWD_SEND`` / ``BWD_SEND`` is appended to the sender's schedule and a
1429 ``FWD_RECV`` / ``BWD_RECV`` is appended to the receiver's schedule.
1431 ``OVERLAP_F_B`` / ``OVERLAP_B_F`` composite steps are expanded into
1432 **two** virtual time slots during the column scan so that the RECV
1433 triggered by the **second** sub-step lands in the receiver's schedule
1434 one slot later — matching the fact that the sender can only finish
1435 emitting the second sub-step's output after the first completes.
1437 Even ranks are processed before odd ranks at each time step to avoid
1438 P2P deadlocks between adjacent ranks.
1440 The resulting per-gap op order (steady state:
1441 ``[B_RECV, B_SEND, F_RECV, F_SEND]``) is LOAD-BEARING, not cosmetic.
1442 Each adjacent rank pair shares one comm (HCCL split-by-group) on which
1443 plain send/recv execute in queue order, and this layout makes one end of
1444 every pair recv-first while the other is send-first, so the two queue
1445 heads always match (recv<->send), then the tails match (send<->recv).
1446 Any later pass that reorders P2P ops relative to EACH OTHER breaks this:
1447 a (since removed) hoist variant that moved recvs across sends made both
1448 ends recv-first and deadlocked on hardware (2026-06).
1449 ``attach_fwd_boundary_p2p`` (the ``"boundary"`` transport) is safe: it keeps
1450 every per-direction FIFO and runs on the batch transport with per-op solo
1451 batches, whose per-pair sequences stay complementary.
1453 Args:
1454 scheduler: ``{rank: [MetaStep | None, ...]}`` — compute schedule
1455 with ``None`` for bubble slots.
1456 stage_num: Total number of virtual pipeline stages.
1457 real_stage_num: Number of physical ranks.
1458 style: Topology mapping — ``'loop'`` or ``'v'``.
1460 Returns:
1461 ``{rank: [MetaStep, ...]}`` — schedule with communication ops inserted.
1462 """
1464 def stage_to_rank(stage_index: int) -> int:
1465 """Map a virtual stage index to its physical rank."""
1466 if style == 'loop':
1467 return stage_index % real_stage_num
1468 if style == 'v':
1469 if stage_index < real_stage_num:
1470 return stage_index
1471 return stage_num - 1 - stage_index
1472 raise ValueError(f"Argument 'style' must be 'loop' or 'v', but got {style!r}.")
1474 def _fwd_peer(stage_index: int):
1475 """Return the rank that receives this stage's forward output, or None."""
1476 if stage_index >= stage_num - 1:
1477 return None
1478 peer = stage_to_rank(stage_index + 1)
1479 return peer if peer != stage_to_rank(stage_index) else None
1481 def _bwd_peer(stage_index: int):
1482 """Return the rank that receives this stage's backward gradient, or None."""
1483 if stage_index <= 0:
1484 return None
1485 peer = stage_to_rank(stage_index - 1)
1486 return peer if peer != stage_to_rank(stage_index) else None
1488 def _insert_comms_for_step(step, rank, new_schedule):
1489 """Insert send/recv for a single FWD, BWD, or composite OVERLAP step."""
1490 if step is None:
1491 return
1493 if step.type == MetaStepType.FWD:
1494 peer = _fwd_peer(step.stage_index)
1495 if peer is not None:
1496 new_schedule[rank].append(
1497 MetaStep(step.micro_index, MetaStepType.FWD_SEND, step.stage_index))
1498 new_schedule[peer].append(
1499 MetaStep(step.micro_index, MetaStepType.FWD_RECV, step.stage_index + 1))
1501 elif step.type == MetaStepType.BWD:
1502 peer = _bwd_peer(step.stage_index)
1503 if peer is not None:
1504 new_schedule[rank].append(
1505 MetaStep(step.micro_index, MetaStepType.BWD_SEND, step.stage_index))
1506 new_schedule[peer].append(
1507 MetaStep(step.micro_index, MetaStepType.BWD_RECV, step.stage_index - 1))
1509 elif step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F) and step.sub_steps:
1510 for sub in step.sub_steps:
1511 _insert_comms_for_step(sub, rank, new_schedule)
1513 # --- Main logic: expand OVERLAP steps into 2 virtual slots, then scan ---
1514 expanded = _expand_overlap_slots(scheduler, real_stage_num)
1515 return _column_scan_insert_comms(expanded, real_stage_num, _insert_comms_for_step)
1518_ALIGN_PAD = object()
1519"""Sentinel marking a forced 1F1B-boundary bubble produced during alignment."""
1522def _step_dep_ready(step, rank, t, done, stage_num, stage_to_rank):
1523 """Cross-rank data dependency check used by the alignment simulator.
1525 A FWD step at stage ``s`` depends on FWD at stage ``s-1`` (on a
1526 different rank); BWD at stage ``s`` depends on BWD at stage ``s+1``.
1527 Steps at boundaries or whose producer lives on the same rank are
1528 always ready.
1529 """
1530 si, mi = step.stage_index, step.micro_index
1531 if step.type == MetaStepType.FWD:
1532 if si == 0 or stage_to_rank(si - 1) == rank:
1533 return True
1534 key = (MetaStepType.FWD, si - 1, mi)
1535 return key in done and done[key] < t
1536 if step.type == MetaStepType.BWD:
1537 if si == stage_num - 1 or stage_to_rank(si + 1) == rank:
1538 return True
1539 key = (MetaStepType.BWD, si + 1, mi)
1540 return key in done and done[key] < t
1541 return True
1544def _simulate_aligned_schedule(padded, stage_num, real_stage_num, stage_to_rank):
1545 """Simulate execution time-step by time-step, inserting bubbles where
1546 a step is not yet ready (cross-rank dep) or where the cooldown
1547 rhythm requires it.
1549 Args:
1550 padded: ``{rank: [step | _ALIGN_PAD | None, ...]}`` after
1551 1F1B-boundary padding.
1552 stage_num: Total number of virtual pipeline stages.
1553 real_stage_num: Number of physical ranks.
1554 stage_to_rank: Topology mapping from stage to rank.
1556 Returns:
1557 ``{rank: [step | None, ...]}`` ready for the column-scan SEND/RECV
1558 insertion phase.
1559 """
1560 remaining_fwd = {
1561 rank: sum(
1562 1 for s in padded[rank]
1563 if s is not _ALIGN_PAD and s is not None and s.type == MetaStepType.FWD
1564 )
1565 for rank in range(real_stage_num)
1566 }
1567 cursors = {r: 0 for r in range(real_stage_num)}
1568 aligned = {r: [] for r in range(real_stage_num)}
1569 done = {}
1570 last_was_cooldown_bwd = {r: False for r in range(real_stage_num)}
1571 max_t = sum(len(v) for v in padded.values()) + real_stage_num * 20
1573 def _emit_bubble(rank):
1574 aligned[rank].append(None)
1575 last_was_cooldown_bwd[rank] = False
1577 def _emit_step(rank, step, t, in_cooldown):
1578 aligned[rank].append(step)
1579 done[(step.type, step.stage_index, step.micro_index)] = t
1580 cursors[rank] += 1
1581 if step.type == MetaStepType.FWD:
1582 remaining_fwd[rank] -= 1
1583 last_was_cooldown_bwd[rank] = in_cooldown and step.type == MetaStepType.BWD
1585 def _step_rank_at(t, rank):
1586 if cursors[rank] >= len(padded[rank]):
1587 return
1588 item = padded[rank][cursors[rank]]
1589 if item is _ALIGN_PAD:
1590 _emit_bubble(rank)
1591 cursors[rank] += 1
1592 return
1593 in_cooldown = remaining_fwd[rank] == 0
1594 # Cooldown rhythm: alternate None / BWD in pure-BWD phase.
1595 cooldown_skip = (
1596 in_cooldown
1597 and item.type == MetaStepType.BWD
1598 and last_was_cooldown_bwd[rank]
1599 )
1600 if cooldown_skip:
1601 _emit_bubble(rank)
1602 return
1603 if not _step_dep_ready(item, rank, t, done, stage_num, stage_to_rank):
1604 _emit_bubble(rank)
1605 return
1606 _emit_step(rank, item, t, in_cooldown)
1608 for t in range(max_t):
1609 if all(cursors[r] >= len(padded[r]) for r in range(real_stage_num)):
1610 break
1611 for rank in range(real_stage_num):
1612 _step_rank_at(t, rank)
1613 return aligned
1616def auto_align_and_add_send_recv(scheduler, stage_num, real_stage_num, style='loop'):
1617 """Auto-insert bubble alignment and P2P send/recv into a pure-compute schedule.
1619 Unlike :func:`add_send_recv` which requires the caller to pre-insert
1620 ``None`` bubble slots for time-step alignment, this function accepts a
1621 **pure compute order** (``FWD`` / ``BWD`` only, no ``None`` needed) and
1622 automatically determines bubble placement via execution simulation.
1624 Three constraints are enforced:
1626 1. **Data dependency** — a ``FWD(stage_k)`` cannot execute until
1627 ``FWD(stage_{k-1})`` on its source rank has completed (and
1628 analogously for ``BWD``).
1629 2. **1F1B transition alignment** — ``real_stage_num - 1 - rank`` padding
1630 slots are inserted at the warmup → 1F1B boundary (detected as the
1631 first ``FWD`` immediately followed by a ``BWD`` in the compute order)
1632 so that all ranks enter the 1F1B steady state in lockstep.
1633 3. **Cooldown rhythm** — once a rank exhausts its ``FWD`` ops and enters
1634 pure-``BWD`` cooldown, consecutive ``BWD`` steps are separated by a
1635 ``None`` slot, maintaining the column-phase-sync property (no rank
1636 does ``BWD`` while another does ``FWD`` at the same time step).
1638 After alignment, a column-scan pass inserts ``FWD_SEND`` / ``FWD_RECV``
1639 and ``BWD_SEND`` / ``BWD_RECV`` with the same prefetch semantics as
1640 :func:`add_send_recv`.
1642 Args:
1643 scheduler: ``{rank: [MetaStep, ...]}`` — pure compute schedule.
1644 ``None`` entries are silently stripped before processing.
1645 stage_num: Total number of virtual pipeline stages.
1646 real_stage_num: Number of physical ranks.
1647 style: Topology mapping — ``'loop'`` or ``'v'``.
1649 Returns:
1650 ``{rank: [MetaStep, ...]}`` — fully aligned schedule with bubbles
1651 and communication ops inserted.
1652 """
1654 # ---- topology helpers (shared with column-scan phase) ----
1656 def stage_to_rank(stage_index: int) -> int:
1657 if style == 'loop':
1658 return stage_index % real_stage_num
1659 if style == 'v':
1660 if stage_index < real_stage_num:
1661 return stage_index
1662 return stage_num - 1 - stage_index
1663 raise ValueError(f"Argument 'style' must be 'loop' or 'v', but got {style!r}.")
1665 def _fwd_peer(stage_index: int):
1666 if stage_index >= stage_num - 1:
1667 return None
1668 peer = stage_to_rank(stage_index + 1)
1669 return peer if peer != stage_to_rank(stage_index) else None
1671 def _bwd_peer(stage_index: int):
1672 if stage_index <= 0:
1673 return None
1674 peer = stage_to_rank(stage_index - 1)
1675 return peer if peer != stage_to_rank(stage_index) else None
1677 # ---- Phase 1: strip None, detect 1F1B boundary, insert transition padding ----
1679 def _find_1f1b_boundary(order):
1680 """Index of the first FWD followed by BWD; ``len(order)`` if absent."""
1681 for i in range(len(order) - 1):
1682 if (order[i].type == MetaStepType.FWD
1683 and order[i + 1].type == MetaStepType.BWD):
1684 return i
1685 return len(order)
1687 padded = {}
1688 for rank in range(real_stage_num):
1689 order = [s for s in scheduler[rank] if s is not None]
1690 boundary = _find_1f1b_boundary(order)
1691 pad_count = real_stage_num - 1 - rank
1692 padded[rank] = order[:boundary] + [_ALIGN_PAD] * pad_count + order[boundary:]
1694 # ---- Phase 2: simulate execution with data deps + cooldown rhythm ----
1696 aligned = _simulate_aligned_schedule(padded, stage_num, real_stage_num, stage_to_rank)
1698 # ---- Phase 3: column-scan SEND/RECV insertion (same as add_send_recv) ----
1700 def _insert_comms_for_step(step, rank, new_schedule):
1701 if step is None:
1702 return
1703 if step.type == MetaStepType.FWD:
1704 peer = _fwd_peer(step.stage_index)
1705 if peer is not None:
1706 new_schedule[rank].append(
1707 MetaStep(step.micro_index, MetaStepType.FWD_SEND, step.stage_index))
1708 new_schedule[peer].append(
1709 MetaStep(step.micro_index, MetaStepType.FWD_RECV, step.stage_index + 1))
1710 elif step.type == MetaStepType.BWD:
1711 peer = _bwd_peer(step.stage_index)
1712 if peer is not None:
1713 new_schedule[rank].append(
1714 MetaStep(step.micro_index, MetaStepType.BWD_SEND, step.stage_index))
1715 new_schedule[peer].append(
1716 MetaStep(step.micro_index, MetaStepType.BWD_RECV, step.stage_index - 1))
1717 elif step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F) and step.sub_steps:
1718 for sub in step.sub_steps:
1719 _insert_comms_for_step(sub, rank, new_schedule)
1721 # Expand OVERLAP steps into 2 virtual slots before the column scan so
1722 # the RECV triggered by an overlap's second sub-step lands one slot
1723 # later on the receiver — matching the fact that the sender can only
1724 # finish emitting the second sub-step after the first completes.
1725 expanded = _expand_overlap_slots(aligned, real_stage_num)
1726 return _column_scan_insert_comms(expanded, real_stage_num, _insert_comms_for_step)
1729class ScheduleGPipe(PipelineScheduleRuntime):
1730 """
1731 The Gpipe schedule.
1732 It first executes all forward micro batches and then execute all backward micro batches.
1733 """
1734 def __init__(self,
1735 stages,
1736 micro_batch_num,
1737 args_batch_dim=None,
1738 kwargs_batch_dim=None,
1739 output_concat_dim=None,
1740 swap=False):
1741 super().__init__(stages,
1742 micro_batch_num,
1743 args_batch_dim=args_batch_dim,
1744 kwargs_batch_dim=kwargs_batch_dim,
1745 output_concat_dim=output_concat_dim,
1746 swap=swap)
1747 self.build_exec_order()
1749 def _build_stage_to_rank_index(self) -> None:
1750 self._stage_to_rank_index = generate_stage_to_rank_mapping(
1751 self.real_stage_num, self._stage_num, style='loop'
1752 )
1754 def construct_exec_order(self):
1755 """construct_exec_order of Gpipe."""
1756 for stage_index in range(self.real_stage_num):
1757 order_list = []
1758 for mb_index in range(self.micro_batch_num):
1759 if stage_index != 0:
1760 order_list.append(MetaStep(mb_index, MetaStepType.FWD_RECV, stage_index))
1761 order_list.append(MetaStep(mb_index, MetaStepType.FWD, stage_index))
1762 if stage_index != self.real_stage_num - 1:
1763 order_list.append(MetaStep(mb_index, MetaStepType.FWD_SEND, stage_index))
1764 for mb_index in range(self.micro_batch_num):
1765 if stage_index != self.real_stage_num - 1:
1766 order_list.append(MetaStep(mb_index, MetaStepType.BWD_RECV, stage_index))
1767 order_list.append(MetaStep(mb_index, MetaStepType.BWD, stage_index))
1768 if stage_index != 0:
1769 order_list.append(MetaStep(mb_index, MetaStepType.BWD_SEND, stage_index))
1770 self.exec_order[stage_index] = order_list
1773class Schedule1F1B(PipelineScheduleRuntime):
1774 """
1775 The 1F1B schedule.
1776 It will perform one forward and one backward on the micro batches in steady state.
1777 """
1778 def __init__(self,
1779 stages,
1780 micro_batch_num,
1781 args_batch_dim=None,
1782 kwargs_batch_dim=None,
1783 output_concat_dim=None,
1784 swap=False):
1785 super().__init__(stages,
1786 micro_batch_num,
1787 args_batch_dim=args_batch_dim,
1788 kwargs_batch_dim=kwargs_batch_dim,
1789 output_concat_dim=output_concat_dim,
1790 swap=swap)
1791 self.build_exec_order()
1793 def _build_stage_to_rank_index(self) -> None:
1794 self._stage_to_rank_index = generate_stage_to_rank_mapping(
1795 self.real_stage_num, self._stage_num, style='loop'
1796 )
1798 def construct_exec_order(self):
1799 """construct_exec_order of 1F1B."""
1800 for stage_index in range(self.real_stage_num):
1801 order_list = []
1802 fwd_index = 0
1803 bwd_index = 0
1804 # warmup phase
1805 warmup_micro_batches = min(self.real_stage_num - stage_index, self.micro_batch_num)
1806 for _ in range(warmup_micro_batches):
1807 if stage_index != 0:
1808 order_list.append(MetaStep(fwd_index, MetaStepType.FWD_RECV, stage_index))
1809 if stage_index % 2 == 0:
1810 order_list.append(MetaStep(fwd_index, MetaStepType.FWD, stage_index))
1811 if fwd_index != warmup_micro_batches - 1:
1812 order_list.append(MetaStep(fwd_index, MetaStepType.FWD_SEND, stage_index))
1813 else:
1814 if fwd_index > 0:
1815 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
1816 order_list.append(MetaStep(fwd_index, MetaStepType.FWD, stage_index))
1817 fwd_index += 1
1819 # if warmup phase cannot filled up, then we need to execute fwd send in advance
1820 if self.real_stage_num - stage_index > self.micro_batch_num:
1821 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
1822 fwd_index += 1
1823 # steady phase
1824 steady_micro_batches = self.micro_batch_num - warmup_micro_batches
1825 for _ in range(steady_micro_batches):
1826 if stage_index != self.real_stage_num - 1:
1827 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_RECV, stage_index))
1828 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
1829 order_list.append(MetaStep(bwd_index, MetaStepType.BWD, stage_index))
1831 if stage_index != 0:
1832 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_SEND, stage_index))
1833 order_list.append(MetaStep(fwd_index, MetaStepType.FWD_RECV, stage_index))
1834 order_list.append(MetaStep(fwd_index, MetaStepType.FWD, stage_index))
1835 fwd_index += 1
1836 bwd_index += 1
1838 # cooldown phase
1839 cooldown_micro_batches = warmup_micro_batches
1840 for _ in range(cooldown_micro_batches):
1841 if stage_index != self.real_stage_num - 1:
1842 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_RECV, stage_index))
1843 if bwd_index == self.micro_batch_num - warmup_micro_batches and fwd_index <= self.micro_batch_num:
1844 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
1845 order_list.append(MetaStep(bwd_index, MetaStepType.BWD, stage_index))
1847 if stage_index != 0:
1848 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_SEND, stage_index))
1849 bwd_index += 1
1850 self.exec_order[stage_index] = order_list
1853class ScheduleInterleaved1F1B(PipelineScheduleRuntime):
1854 """The Interleaved 1F1B schedule.
1856 Supports multiple stages per rank. In steady state, performs one
1857 forward followed by one backward on each micro-batch. Handles the
1858 cases where ``micro_batch_num`` is less than, equal to, or greater
1859 than the stage count, including non-evenly-divisible micro counts.
1861 Two orthogonal overlap modes can be enabled via constructor flags:
1863 * ``overlap_p2p=True``: defer P2P recv ``handle.wait()`` until the
1864 consuming FWD/BWD step (or the OVERLAP_B_F callback when
1865 ``overlap_b_f=True``), letting recv overlap with prior compute.
1866 * ``overlap_b_f=True``: in the 1F1B steady state, pair consecutive
1867 ``(B_i, F_{i+1})`` steps into ``OVERLAP_B_F`` composite steps so
1868 a registered callback can drive comm/compute overlap (typically
1869 via :class:`CommComputeOverlap` for MoE EP A2A). Users register
1870 the callback through :meth:`register_custom_function`.
1871 * ``enable_dxdw_split=True`` (requires ``overlap_b_f``): each
1872 steady-state pair becomes ``(BWD_INPUT_i, F_{i+1})`` and the
1873 ``BWD_WEIGHT_i`` runs as its own step after the pair's P2P gap
1874 (see :func:`split_overlap_dxdw`), so the input-grad send is
1875 issued once dx and the paired forward finish instead of waiting
1876 out the full backward.
1878 The two overlap flags are independent and can be combined.
1880 Example:
1881 >>> # Plain interleaved 1F1B
1882 >>> sched = ScheduleInterleaved1F1B(stages, 8)
1883 >>> # With B/F overlap (dual-pipe-style comm/compute overlap)
1884 >>> sched = ScheduleInterleaved1F1B(stages, 8, overlap_b_f=True)
1885 >>> sched.register_custom_function(MetaStepType.OVERLAP_B_F, callback)
1886 """
1887 def __init__(self,
1888 stages,
1889 micro_batch_num,
1890 args_batch_dim=None,
1891 kwargs_batch_dim=None,
1892 output_concat_dim=None,
1893 overlap_p2p=False,
1894 overlap_b_f=False,
1895 swap=False,
1896 enable_dxdw_split=False,
1897 p2p_transport="auto"):
1898 super().__init__(stages,
1899 micro_batch_num,
1900 args_batch_dim=args_batch_dim,
1901 kwargs_batch_dim=kwargs_batch_dim,
1902 output_concat_dim=output_concat_dim,
1903 overlap_p2p=overlap_p2p,
1904 swap=swap,
1905 p2p_transport=p2p_transport)
1906 # _overlap_b_f selects between plain F/B emission and OVERLAP_B_F
1907 # pairing in the 1F1B steady-state phase. Must be set before
1908 # ``construct_stage_exec_order`` is called below.
1909 self._overlap_b_f = overlap_b_f
1910 # dx/dw split: ``construct_exec_order`` rewrites each steady-state
1911 # OVERLAP_B_F pair to ``(BWD_INPUT, FWD)`` and re-emits the matching
1912 # BWD_WEIGHT after the pair's P2P gap (``split_overlap_dxdw``), so the
1913 # input-grad send leaves at ``max(dx, fwd)`` and dw overlaps with the
1914 # in-flight P2P instead of delaying it.
1915 self._enable_dxdw_split = enable_dxdw_split
1916 if enable_dxdw_split and not overlap_b_f:
1917 raise ValueError(
1918 "enable_dxdw_split=True requires overlap_b_f=True; the split "
1919 "is only applied to BWD sub-steps inside OVERLAP_B_F composite steps."
1920 )
1921 self._init_round_layout()
1922 self.build_exec_order()
1924 def _init_round_layout(self):
1925 """Compute per-round micro-batch counts used by stage-order emission.
1927 Populates ``n_rounds``, ``n_microbatch_per_round`` and its prefix-sum
1928 ``n_microbatch_per_round_accu`` from ``micro_batch_num``,
1929 ``real_stage_num`` and ``n_local_stages``. Factored out of
1930 ``__init__`` so the pure schedule-construction path (used by offline
1931 unit tests) can be exercised without instantiating stages.
1932 """
1933 self.n_rounds = max(1, self.micro_batch_num // self.real_stage_num)
1934 if self.micro_batch_num < self.real_stage_num:
1935 base = self.micro_batch_num - self.real_stage_num
1936 remainder = 0
1937 else:
1938 n_extra_microbatch = self.micro_batch_num % self.real_stage_num
1939 base = n_extra_microbatch // self.n_rounds
1940 remainder = n_extra_microbatch % self.n_rounds
1941 self.n_microbatch_per_round = \
1942 [self.real_stage_num + base + 1 if i < remainder else
1943 self.real_stage_num + base for i in range(self.n_rounds)]
1944 self.n_microbatch_per_round_accu = \
1945 [x * self.n_local_stages for x in itertools.accumulate(self.n_microbatch_per_round)]
1946 self.n_microbatch_per_round_accu.insert(0, 0)
1948 def construct_exec_order(self):
1949 for stage_index in range(self.real_stage_num):
1950 self.exec_order[stage_index] = self.construct_stage_exec_order(stage_index)
1951 self.exec_order = add_send_recv(self.exec_order, self._stage_num, self.real_stage_num, style='loop')
1952 if self.enable_dxdw_split:
1953 self.exec_order = split_overlap_dxdw(self.exec_order)
1955 def _build_stage_to_rank_index(self) -> None:
1956 self._stage_to_rank_index = generate_stage_to_rank_mapping(
1957 self.real_stage_num, self._stage_num, style='loop'
1958 )
1960 def warmup_ops(self, stage_index):
1961 """warmup phase."""
1962 warmup_ops_last_stage = (self.n_local_stages - 1) * self.n_microbatch_per_round[0]
1963 warmup_ops = warmup_ops_last_stage + 2 * (self.real_stage_num - 1 - stage_index)
1964 return min(warmup_ops, self.micro_batch_num * self.n_local_stages)
1966 def forward_stage_index(self, op_index, stage_index):
1967 """obtain forward stage_index based on op_index."""
1968 accu_index = bisect.bisect_right(self.n_microbatch_per_round_accu, op_index) - 1
1969 local_index = (op_index - self.n_microbatch_per_round_accu[accu_index]) // \
1970 self.n_microbatch_per_round[accu_index]
1971 return (local_index * self.real_stage_num) + stage_index
1973 def backward_stage_index(self, op_index, stage_index):
1974 """obtain backward stage_index based on op_index."""
1975 accu_index = bisect.bisect_right(self.n_microbatch_per_round_accu, op_index) - 1
1976 local_index = (op_index - self.n_microbatch_per_round_accu[accu_index]) // \
1977 self.n_microbatch_per_round[accu_index]
1978 local_index = self.n_local_stages - 1 - local_index
1979 return (local_index * self.real_stage_num) + stage_index
1981 def _short_micro(self) -> bool:
1982 """True when ``micro_batch_num < real_stage_num`` (extra-bubble regime)."""
1983 return self.micro_batch_num < self.real_stage_num
1985 def _trailing_bubble(self) -> int:
1986 """Bubble count appended after a BWD with ``micro == micro_batch_num - 1``
1987 in the short-micro regime.
1988 """
1989 return self.real_stage_num - self.micro_batch_num
1991 def _emit_warmup_ops(self, stage_index, warmup_ops, fwd_stage_micro_index):
1992 """Emit pure-FWD warmup ops with optional short-micro bubble padding."""
1993 ops = []
1994 short = self._short_micro()
1995 last_micro = self.micro_batch_num - 1
1996 last_stage = self.real_stage_num - 1
1997 bubble = self._trailing_bubble()
1998 for op_idx in range(warmup_ops):
1999 fwd_stage_idx = self.forward_stage_index(op_idx, stage_index)
2000 fwd_micro_idx = fwd_stage_micro_index[fwd_stage_idx]
2001 ops.append(MetaStep(fwd_micro_idx, MetaStepType.FWD, fwd_stage_idx))
2002 need_pad = (
2003 short
2004 and fwd_micro_idx == last_micro
2005 and (op_idx != warmup_ops - 1 or stage_index == last_stage)
2006 )
2007 if need_pad:
2008 ops.extend([None] * bubble)
2009 fwd_stage_micro_index[fwd_stage_idx] += 1
2010 return ops
2012 def _emit_cooldown_ops(self, stage_index, warmup_ops, fwd_bwd_ops, total_ops,
2013 bwd_stage_micro_index):
2014 """Emit pure-BWD cooldown ops (each preceded by a bubble) with
2015 optional short-micro trailing padding.
2016 """
2017 ops = []
2018 short = self._short_micro()
2019 last_micro = self.micro_batch_num - 1
2020 # Double the bubble at each chunk's last-micro BWD: one ``bubble`` covers
2021 # the missing ``rs - micro`` micros, the second offsets the next chunk
2022 # by 2 slots so the wrap-around grad (rank 0 stage ``rs`` -> rank
2023 # last_stage stage ``rs - 1``) lands AFTER its producer in column-scan
2024 # time. Matches the +2 cooldown-rhythm offset that non-short Interleaved
2025 # 1F1B naturally has from extra 1F1B ops on rank last_stage.
2026 bubble = 2 * self._trailing_bubble()
2027 for op_idx in range(warmup_ops + fwd_bwd_ops, total_ops):
2028 ops.append(None)
2029 bwd_stage_idx = self.backward_stage_index(op_idx - warmup_ops, stage_index)
2030 bwd_micro_idx = bwd_stage_micro_index[bwd_stage_idx]
2031 ops.append(MetaStep(bwd_micro_idx, MetaStepType.BWD, bwd_stage_idx))
2032 if short and bwd_micro_idx == last_micro:
2033 ops.extend([None] * bubble)
2034 bwd_stage_micro_index[bwd_stage_idx] += 1
2035 return ops
2037 def _emit_1f1b_ops(self, stage_index, warmup_ops, fwd_bwd_ops,
2038 fwd_stage_micro_index, bwd_stage_micro_index):
2039 """Emit interleaved (FWD, BWD) pairs for the 1F1B steady-state phase."""
2040 ops = []
2041 short = self._short_micro()
2042 last_micro = self.micro_batch_num - 1
2043 last_stage = self.real_stage_num - 1
2044 # Double the bubble at the 1F1B->cooldown chunk boundary on rank
2045 # last_stage; see :meth:`_emit_cooldown_ops` for the alignment rationale.
2046 bubble = 2 * self._trailing_bubble()
2047 for op_idx in range(warmup_ops, warmup_ops + fwd_bwd_ops):
2048 fwd_stage_idx = self.forward_stage_index(op_idx, stage_index)
2049 fwd_micro_idx = fwd_stage_micro_index[fwd_stage_idx]
2050 ops.append(MetaStep(fwd_micro_idx, MetaStepType.FWD, fwd_stage_idx))
2051 fwd_stage_micro_index[fwd_stage_idx] += 1
2052 bwd_stage_idx = self.backward_stage_index(op_idx - warmup_ops, stage_index)
2053 bwd_micro_idx = bwd_stage_micro_index[bwd_stage_idx]
2054 ops.append(MetaStep(bwd_micro_idx, MetaStepType.BWD, bwd_stage_idx))
2055 need_pad = (
2056 short
2057 and bwd_micro_idx == last_micro
2058 and stage_index == last_stage
2059 )
2060 if need_pad:
2061 ops.extend([None] * bubble)
2062 bwd_stage_micro_index[bwd_stage_idx] += 1
2063 return ops
2065 @staticmethod
2066 def _collect_fwd_bwd_steps(emit_fwd, emit_bwd, fwd_bwd_ops, warmup_ops):
2067 """Walk the 1F1B range collecting parallel ``fwd_steps`` / ``bwd_steps``.
2069 ``emit_fwd(op_idx)`` and ``emit_bwd(op_idx)`` build a single
2070 :class:`MetaStep` and advance their respective per-stage micro
2071 counters as a side effect.
2072 """
2073 fwd_steps = []
2074 bwd_steps = []
2075 for op_idx in range(warmup_ops, warmup_ops + fwd_bwd_ops):
2076 fwd_steps.append(emit_fwd(op_idx))
2077 bwd_steps.append(emit_bwd(op_idx))
2078 return fwd_steps, bwd_steps
2080 @staticmethod
2081 def _pair_into_overlap_b_f(fwd_steps, bwd_steps):
2082 """Build ``F₁, [B_i, F_{i+1}], B_n`` ordering with OVERLAP_B_F pairs.
2084 ``sub_steps`` carry the ``(bwd, fwd)`` tuple — callbacks access
2085 them via ``step.sub_steps`` to recover per-direction stage /
2086 micro info.
2087 """
2088 ops = []
2089 if fwd_steps:
2090 ops.append(fwd_steps[0]) # F₁ runs alone
2091 for i in range(len(bwd_steps) - 1):
2092 ops.append(MetaStep(
2093 None, MetaStepType.OVERLAP_B_F, None,
2094 sub_steps=(bwd_steps[i], fwd_steps[i + 1]),
2095 ))
2096 if bwd_steps:
2097 ops.append(bwd_steps[-1]) # B_n runs alone
2098 return ops
2100 def _emit_1f1b_overlap_ops(self, stage_index, warmup_ops, fwd_bwd_ops,
2101 fwd_stage_micro_index, bwd_stage_micro_index):
2102 """Emit ``F₁, [B_i, F_{i+1}], B_n`` for the 1F1B phase under
2103 ``overlap_b_f=True``. Each ``[B_i, F_{i+1}]`` becomes an
2104 ``OVERLAP_B_F`` composite step; a registered callback drives the
2105 actual concurrent execution. Short-micro extra-bubble padding
2106 on the last rank is appended after ``B_n``.
2107 """
2108 def emit_fwd(op_idx):
2109 fwd_si = self.forward_stage_index(op_idx, stage_index)
2110 fwd_mi = fwd_stage_micro_index[fwd_si]
2111 fwd_stage_micro_index[fwd_si] += 1
2112 return MetaStep(fwd_mi, MetaStepType.FWD, fwd_si)
2114 def emit_bwd(op_idx):
2115 bwd_si = self.backward_stage_index(op_idx - warmup_ops, stage_index)
2116 bwd_mi = bwd_stage_micro_index[bwd_si]
2117 bwd_stage_micro_index[bwd_si] += 1
2118 return MetaStep(bwd_mi, MetaStepType.BWD, bwd_si)
2120 fwd_steps, bwd_steps = self._collect_fwd_bwd_steps(
2121 emit_fwd, emit_bwd, fwd_bwd_ops, warmup_ops,
2122 )
2123 ops = self._pair_into_overlap_b_f(fwd_steps, bwd_steps)
2125 last_stage = self.real_stage_num - 1
2126 if self._short_micro() and stage_index == last_stage and bwd_steps:
2127 if bwd_steps[-1].micro_index == self.micro_batch_num - 1:
2128 # Double the bubble at the 1F1B->cooldown chunk boundary;
2129 # see :meth:`_emit_cooldown_ops` for the alignment rationale.
2130 ops.extend([None] * (2 * self._trailing_bubble()))
2131 return ops
2133 def construct_stage_exec_order(self, stage_index):
2134 """Construct the execution order for ``stage_index``.
2136 Builds: warmup → bubbles → 1F1B steady state → cooldown. The
2137 1F1B segment switches between :meth:`_emit_1f1b_ops` (plain) and
2138 :meth:`_emit_1f1b_overlap_ops` (OVERLAP_B_F pairing) based on
2139 the ``overlap_b_f`` constructor flag.
2140 """
2141 warmup_ops = self.warmup_ops(stage_index)
2142 fwd_bwd_ops = self.n_local_stages * self.micro_batch_num - warmup_ops
2143 total_ops = 2 * warmup_ops + fwd_bwd_ops
2144 order_list = [None for _ in range(stage_index)]
2145 fwd_stage_micro_index = defaultdict(int)
2146 bwd_stage_micro_index = defaultdict(int)
2147 order_list.extend(self._emit_warmup_ops(stage_index, warmup_ops, fwd_stage_micro_index))
2148 bubbles_before_1f1b = max(
2149 0,
2150 2 * (self.real_stage_num - stage_index - 1) - self.micro_batch_num,
2151 )
2152 order_list.extend([None] * bubbles_before_1f1b)
2153 order_list.extend([None] * (self.real_stage_num - 1 - stage_index))
2154 if self._overlap_b_f:
2155 order_list.extend(self._emit_1f1b_overlap_ops(
2156 stage_index, warmup_ops, fwd_bwd_ops,
2157 fwd_stage_micro_index, bwd_stage_micro_index,
2158 ))
2159 else:
2160 order_list.extend(self._emit_1f1b_ops(
2161 stage_index, warmup_ops, fwd_bwd_ops,
2162 fwd_stage_micro_index, bwd_stage_micro_index,
2163 ))
2164 order_list.extend(self._emit_cooldown_ops(
2165 stage_index, warmup_ops, fwd_bwd_ops, total_ops, bwd_stage_micro_index,
2166 ))
2167 return order_list
2170def detect_cycle_in_graph(ranks_map):
2171 """
2172 Detects a cycle in the directed graph constructed from ranks_map.
2174 Args:
2175 ranks_map: A dictionary where keys are rank names and values are lists of nodes.
2177 Returns:
2178 tuple: (cycle_path, cycle_ranks) where cycle_path is a list of nodes forming the cycle and cycle_ranks
2179 is a list of rank transitions corresponding to the cycle path.
2180 """
2181 graph = defaultdict(list)
2182 rank_edges = {}
2184 for rank, nodes in ranks_map.items():
2185 for i in range(len(nodes) - 1):
2186 u, v = nodes[i], nodes[i + 1]
2187 graph[u].append(v)
2188 rank_edges[(u, v)] = rank
2190 visited = set()
2191 path = []
2192 node_indices = {}
2193 cycle_path = []
2194 cycle_ranks = []
2196 stack = []
2197 for node in list(graph.keys()):
2198 if node not in visited:
2199 stack.append((node, False))
2200 while stack:
2201 current_node, is_processed = stack.pop()
2203 if is_processed:
2204 path.pop()
2205 del node_indices[current_node]
2206 continue
2208 if current_node in node_indices:
2209 cycle_start = node_indices[current_node]
2210 cycle_path = path[cycle_start:] + [current_node]
2211 for i in range(cycle_start, len(path)):
2212 u = path[i]
2213 v = path[i + 1] if i + 1 < len(path) else current_node
2214 cycle_ranks.append(f"{rank_edges[(u, v)]} {u} -> {v}")
2215 return cycle_path, cycle_ranks
2217 if current_node in visited:
2218 continue
2220 visited.add(current_node)
2221 node_indices[current_node] = len(path)
2222 path.append(current_node)
2224 stack.append((current_node, True))
2225 for neighbor in reversed(graph[current_node]):
2226 stack.append((neighbor, False))
2228 return None, None
2231def output_cycle_results(cycle_path, cycle_ranks):
2232 """
2233 Helper function to output cycle detection results.
2235 Args:
2236 cycle_path (list): List of nodes forming a cycle, if any.
2237 cycle_ranks (list): List of ranks involved in the cycle.
2239 Returns:
2240 None: Outputs results to the console.
2241 """
2242 if cycle_path:
2243 logger.error("Cycle detected:")
2244 path_str = " -> ".join(str(node) for node in cycle_path)
2245 logger.error("%s -> %s", path_str, cycle_path[0]) # Close the cycle
2246 logger.error("Involving ranks:")
2247 for rank in cycle_ranks:
2248 logger.error(rank)
2249 else:
2250 logger.warning("Cycle Check succeeded. There is no cycle in the graph.")
2253def parse_and_validate(data: dict, all_rank: bool = True):
2254 """
2255 Parse and validate execution orders in a directed graph structure.
2257 This function checks the integrity and consistency of a given dataset, ensuring all required
2258 keys are present and correctly referenced. It also validates the structure of the input data
2259 and parses string values to extract meaningful components.
2261 Args:
2262 data (dict): A dictionary where keys are string identifiers and values are lists of strings.
2263 Each value represents a dependency or reference to other keys.
2264 all_rank (bool): If True, checks that all elements referenced in the data are present as keys
2265 in the dictionary. If False, only checks intersections.
2267 Returns:
2268 None: Log error messages to the console if validation fails, otherwise completes silently.
2270 Raises:
2271 ValueError: Raised indirectly if `parse_elements` encounters malformed input strings.
2272 TypeError: Raised indirectly if data contains unexpected types.
2273 """
2275 def parse_elements(value: str, max_groups: int = 2) -> set:
2276 """Extract unique elements inside the first one or two parentheses from a string."""
2278 groups = re.findall(r'\((\d+)\)', value)
2279 limited_groups = groups[:max_groups] # Limit to the first `max_groups` matches
2281 return {item.strip() for item in limited_groups}
2283 if not isinstance(data, dict):
2284 logger.error("Input must be a dictionary with string keys and lists of strings as values.")
2285 return
2287 key_to_values = {key: set(values) for key, values in data.items() if
2288 isinstance(values, list) and all(isinstance(v, str) for v in values)}
2290 for key, values in data.items():
2291 if not isinstance(values, list) or not all(isinstance(v, str) for v in values):
2292 logger.error("Values for key '%s' must be a list of strings.", key)
2293 continue
2295 for value in values:
2296 try:
2297 elements = parse_elements(value)
2298 except (ValueError, TypeError, AttributeError) as e:
2299 logger.error("Unable to parse elements from value '%s' in key '%s'. Error: %s", value, key, e)
2300 continue
2302 # Check for missing keys if all_rank is True
2303 if all_rank:
2304 missing_keys = elements - key_to_values.keys()
2305 if missing_keys:
2306 logger.error("The following keys are missing for value '%s': %s", value, missing_keys)
2307 continue
2309 # Check if the value is present in the referenced keys
2310 for element in elements & key_to_values.keys() if not all_rank else elements:
2311 if value not in key_to_values[element]:
2312 logger.error("Key '%s' is missing the value '%s'.", element, value)
2315def generate_operations(order_list: dict[int, list[MetaStep]],
2316 chunk_num: int,
2317 com_type: str = 'loop') -> dict[str, list[str]]:
2318 """
2319 Generate formatted operations dictionary from pipeline execution order.
2321 Args:
2322 order_list (dict): Dictionary where keys are rank IDs and values are MetaStep execution sequences
2323 chunk_num (int): Number of chunks (virtual pipeline stages)
2324 com_type (str): Stage-to-rank mapping type ('loop' for cyclic, 'v' for V-shaped)
2326 Returns:
2327 Dictionary where keys are rank IDs (as strings) and values are lists of formatted operation strings
2328 """
2330 def stage_to_rank(stage_index, style, stage_num, real_stage_num):
2331 """Map stage index to rank"""
2332 if style == 'loop':
2333 return stage_index % real_stage_num
2334 if style == 'v':
2335 if stage_index < real_stage_num:
2336 return stage_index
2337 return stage_num - 1 - stage_index
2338 raise ValueError("Invalid style")
2340 def find_send_target(stage_idx, op_type):
2341 """Find target stage for SEND operation"""
2342 if op_type == MetaStepType.FWD_SEND:
2343 return forward_comm.get(stage_idx)
2344 return backward_comm.get(stage_idx)
2346 def find_recv_source(stage_idx, op_type):
2347 """Find source stage for RECV operation"""
2348 if op_type == MetaStepType.FWD_RECV:
2349 # Reverse lookup in forward_comm
2350 for src, dst in forward_comm.items():
2351 if dst == stage_idx:
2352 return src
2353 else:
2354 # Reverse lookup in backward_comm
2355 for src, dst in backward_comm.items():
2356 if dst == stage_idx:
2357 return src
2358 return None
2360 real_stage = len(order_list)
2361 total_stages = real_stage * chunk_num
2363 # Build communication rules
2364 forward_comm = {}
2365 backward_comm = {}
2367 for i in range(total_stages):
2368 if i + 1 < total_stages:
2369 forward_comm[i] = i + 1
2370 if i - 1 >= 0:
2371 backward_comm[i] = i - 1
2373 formatted_operations = defaultdict(list)
2375 for rank, steps in order_list.items():
2376 operation_counter = defaultdict(int)
2378 for step in steps:
2379 if step.type in [MetaStepType.FWD_SEND, MetaStepType.BWD_SEND]:
2380 target_stage = find_send_target(step.stage_index, step.type)
2381 if target_stage is not None:
2382 target_rank = stage_to_rank(target_stage, com_type, total_stages, real_stage)
2383 comm_pair = (rank, target_rank, step.micro_index)
2384 operation_counter[comm_pair] += 1
2385 count = operation_counter[comm_pair]
2386 formatted_op = f"Send_Receive_({rank})->({target_rank})_micro{step.micro_index}_{count}th"
2387 formatted_operations[str(rank)].append(formatted_op)
2389 elif step.type in [MetaStepType.FWD_RECV, MetaStepType.BWD_RECV]:
2390 source_stage = find_recv_source(step.stage_index, step.type)
2391 if source_stage is not None:
2392 source_rank = stage_to_rank(source_stage, com_type, total_stages, real_stage)
2393 comm_pair = (source_rank, rank, step.micro_index)
2394 operation_counter[comm_pair] += 1
2395 count = operation_counter[comm_pair]
2396 formatted_op = f"Send_Receive_({source_rank})->({rank})_micro{step.micro_index}_{count}th"
2397 formatted_operations[str(rank)].append(formatted_op)
2399 # Convert defaultdict to dict
2400 return dict(formatted_operations)
2403def validate_pipeline_execution(order_list: dict[int, list[MetaStep]],
2404 chunk_num: int,
2405 com_type: str = 'loop') -> dict[str, any]:
2406 """
2407 Comprehensive validation function for pipeline parallel execution order.
2409 This function validates the execution order of pipeline parallelism by:
2410 1. Checking SEND/RECV communication pair matching
2411 2. Detecting duplicate operations
2412 3. Detecting cycles in communication graphs
2413 4. Verifying computation-SEND matching
2415 Args:
2416 order_list: Dictionary where keys are rank IDs and values are MetaStep execution sequences
2417 chunk_num: Number of chunks (virtual pipeline stages)
2418 com_type: Stage-to-rank mapping type ('loop' for cyclic, 'v' for V-shaped)
2420 Returns:
2421 Dictionary containing validation results with the following keys:
2422 - validation: Communication pair validation results
2423 - cycle_detection: Cycle detection results
2424 - computation_send_matching: Computation-SEND matching validation results
2425 - has_errors: Boolean indicating if any errors were found
2426 - error_messages: List of all error messages found
2427 - formatted_operations: Generated formatted operations
2428 """
2430 # Generate operations
2431 formatted_operations = generate_operations(order_list, chunk_num, com_type)
2433 parse_and_validate(formatted_operations, True)
2435 # Detect cycles
2436 cycle_path, cycle_ranks = detect_cycle_in_graph(formatted_operations)
2438 # Output results
2439 output_cycle_results(cycle_path, cycle_ranks)
2441 result = {
2442 'formatted_operations': formatted_operations,
2443 'cycle_path': cycle_path,
2444 'cycle_ranks': cycle_ranks,
2445 'has_cycle': bool(cycle_path)
2446 }
2447 return result
2450_COMPUTE_META_STEP_TYPES = frozenset({
2451 MetaStepType.FWD,
2452 MetaStepType.BWD,
2453 MetaStepType.BWD_INPUT,
2454 MetaStepType.BWD_WEIGHT,
2455})
2458def _next_active_stage_indices(actions, start_index, max_active_stages, managed_stage_indices):
2459 """Find the next distinct managed stages that will execute compute work.
2461 Send/recv and previously injected FSDP control steps are skipped so that the
2462 lookahead window only counts real compute, otherwise communication-only
2463 actions would consume the budget and shrink the effective prefetch depth.
2464 """
2465 stage_indices = []
2466 seen = set()
2467 for action in actions[start_index:]:
2468 for leaf_step in iter_leaf_meta_steps(action):
2469 if leaf_step.type not in _COMPUTE_META_STEP_TYPES:
2470 continue
2471 if leaf_step.stage_index not in managed_stage_indices or leaf_step.stage_index in seen:
2472 continue
2473 seen.add(leaf_step.stage_index)
2474 stage_indices.append(leaf_step.stage_index)
2475 if len(stage_indices) == max_active_stages:
2476 return stage_indices
2477 return stage_indices
2480def add_fsdp_unshard_reshard(actions, managed_stage_indices, max_active_stages=3):
2481 """Insert FSDP unshard/reshard actions for locally managed stages."""
2482 if not managed_stage_indices:
2483 return actions
2485 fsdp_actions = []
2486 active_stages = []
2487 for index, action in enumerate(actions):
2488 next_stage_indices = _next_active_stage_indices(
2489 actions, index, max_active_stages, managed_stage_indices
2490 )
2491 evicted_stages = [stage_index for stage_index in active_stages if stage_index not in next_stage_indices]
2492 fetched_stages = [stage_index for stage_index in next_stage_indices if stage_index not in active_stages]
2493 for stage_index in evicted_stages:
2494 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_RESHARD, stage_index))
2495 active_stages.remove(stage_index)
2496 for stage_index in fetched_stages:
2497 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_UNSHARD, stage_index))
2498 active_stages.append(stage_index)
2499 fsdp_actions.append(action)
2501 while active_stages:
2502 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_RESHARD, active_stages.pop(0)))
2503 return fsdp_actions
2506def add_fsdp_reduce_grad(actions, managed_stage_indices, micro_batch_num):
2507 """Insert FSDP reduce-grad actions after the last backward-like action of each stage."""
2508 if not managed_stage_indices:
2509 return actions
2511 fsdp_actions = []
2512 for action in actions:
2513 fsdp_actions.append(action)
2514 reduced_stage_indices = []
2515 for leaf_step in iter_leaf_meta_steps(action):
2516 if leaf_step.stage_index not in managed_stage_indices:
2517 continue
2518 if leaf_step.type not in (MetaStepType.BWD, MetaStepType.BWD_WEIGHT):
2519 continue
2520 if leaf_step.micro_index != micro_batch_num - 1:
2521 continue
2522 if leaf_step.stage_index not in reduced_stage_indices:
2523 reduced_stage_indices.append(leaf_step.stage_index)
2524 for stage_index in reduced_stage_indices:
2525 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_REDUCE_GRAD, stage_index))
2526 return fsdp_actions