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