Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / scheduler.py: 74%
1276 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""pipeline schedule"""
16from abc import ABC, abstractmethod
17from contextlib import nullcontext
18from collections import defaultdict
19import itertools
20import bisect
21import logging
22import re
23from typing import Any, Iterator
25import torch
26import torch.distributed as dist
28import hyper_parallel
29from hyper_parallel.core.fully_shard.api import HSDPModule
30from hyper_parallel.core.pipeline_parallel.pipeline_swap import (
31 PipelineSwapSession,
32 inject_pipeline_swap_steps,
33 swap_launch_load,
34 swap_launch_offload,
35 swap_wait_load,
36 swap_wait_offload,
37 unregister_layer_swap_hooks,
38)
39from hyper_parallel.core.pipeline_parallel.utils import BatchDimSpec, MetaStep, MetaStepType
40from hyper_parallel.core.pipeline_parallel._microbatch import _MicroBatch
41from hyper_parallel.core.pipeline_parallel._p2p import (
42 create_p2p_multi_stream_groups,
43 prepare_batch_p2p_group,
44)
46logger = logging.getLogger(__name__)
49def generate_stage_to_rank_mapping(real_stage_num, stage_num, style='loop'):
50 """Generate stage to rank mapping for loop or V schedules."""
51 if style == 'loop':
52 return {stage_index: stage_index % real_stage_num for stage_index in range(stage_num)}
53 if style == 'v':
54 if stage_num % real_stage_num != 0:
55 raise ValueError(
56 f"stage_num {stage_num} must be evenly divisible by real_stage_num {real_stage_num} for V schedules."
57 )
58 mapping = {}
59 rank_index = 0
60 for stage_index in range(stage_num):
61 mapping[stage_index] = rank_index
62 if (stage_index + 1) % real_stage_num == 0:
63 continue
64 if (stage_index // real_stage_num) % 2 == 0:
65 rank_index += 1
66 else:
67 rank_index -= 1
68 return mapping
69 raise ValueError(f"Unsupported stage rank mapping style: {style}")
72def generate_rank_to_stage_mapping(real_stage_num, stage_num, style='loop'):
73 """Invert the stage to rank mapping."""
74 stage_to_rank = generate_stage_to_rank_mapping(real_stage_num, stage_num, style)
75 rank_to_stages = defaultdict(list)
76 for stage_index, rank in stage_to_rank.items():
77 rank_to_stages[rank].append(stage_index)
79 for stages in rank_to_stages.values():
80 stages.sort()
81 return dict(rank_to_stages)
84def iter_leaf_meta_steps(step):
85 """Yield leaf MetaSteps, recursively expanding OVERLAP containers.
87 Both ``OVERLAP_F_B`` and ``OVERLAP_B_F`` carry their real FWD/BWD work in
88 ``sub_steps``; the FSDP unshard/reshard injection pass relies on this to
89 see the FWD/BWD buried inside an overlap. Missing ``OVERLAP_B_F`` here let
90 an overlapped FWD run against a resharded stage → "expected HSDPModule
91 parameters in unsharded state". Mirror the other expansion sites
92 (run/_expand/add_fsdp_*) which already handle both composite types.
93 """
94 if step is None:
95 return
96 if step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F) and step.sub_steps:
97 for sub_step in step.sub_steps:
99 yield from iter_leaf_meta_steps(sub_step)
100 return
101 yield step
104class PipelineContext:
105 """Per-run state handed to a custom execution function (see
106 :meth:`PipelineScheduleRuntime.register_custom_function`).
108 A plain data carrier for one :meth:`PipelineScheduleRuntime.run_microbatches`
109 call. The P2P helpers (``wait_fwd_recv`` / ``wait_bwd_recv`` / ``send_fwd``
110 / ``send_bwd``) and the ``enable_dxdw_split`` flag live on the schedule, so a
111 callback reaches them through :attr:`schedule`, e.g.
112 ``ctx.schedule.send_bwd(stage, micro_index)``.
114 Attributes:
115 schedule: The owning :class:`PipelineScheduleRuntime`.
116 arg_mbs: Per-micro-batch positional args.
117 kwarg_mbs: Per-micro-batch keyword args.
118 losses: Mutable list collecting per-step losses.
119 """
121 def __init__(self, schedule: "PipelineScheduleRuntime", arg_mbs: list,
122 kwarg_mbs: list, losses: list) -> None:
123 """Bundle the active schedule with one run's micro-batch inputs and losses."""
124 self.schedule = schedule
125 self.arg_mbs = arg_mbs
126 self.kwarg_mbs = kwarg_mbs
127 self.losses = losses
130def _exec_fsdp_unshard(stage):
131 """Unshard every HSDPModule in the stage's submodule tree."""
132 for _, module in stage.submodule.named_modules():
133 if isinstance(module, HSDPModule):
134 module.unshard()
137def _exec_fsdp_reshard(stage):
138 """Reshard every HSDPModule in the stage's submodule tree."""
139 for _, module in stage.submodule.named_modules():
140 if isinstance(module, HSDPModule):
141 module.reshard()
144def _exec_fsdp_reduce_grad(stage):
145 """Run the stage's FSDP post-backward gradient reduction."""
146 stage.execute_reduce_grad()
149# FSDP control MetaStep -> handler(stage). Membership also marks which
150# MetaStepTypes are FSDP control steps, so the runtime loop dispatches with a
151# single table lookup instead of re-switching on the step type.
152_FSDP_STEP_HANDLERS = {
153 MetaStepType.FSDP_UNSHARD: _exec_fsdp_unshard,
154 MetaStepType.FSDP_RESHARD: _exec_fsdp_reshard,
155 MetaStepType.FSDP_REDUCE_GRAD: _exec_fsdp_reduce_grad,
156}
159class PipelineScheduleRuntime(ABC):
160 """
161 Base class for pipeline schedule.
162 Implements the `split_microbatches` and `run_microbatches` method.
163 Derived classes should implement `run_microbatches` method and `run` method.
165 Supports registering **custom execution functions** for any
166 :class:`MetaStepType` via :meth:`register_custom_function`. When
167 ``run_microbatches`` encounters a step whose type has a registered
168 handler, it creates a :class:`PipelineContext` and delegates execution
169 to the handler instead of using the built-in logic.
171 Args:
172 stages (list[PipelineStage], PipelineStage): PipelineStage used to run_microbatches.
173 micro_batch_num (int): The number of micro-batch.
174 args_batch_dim (int | BatchDimSpec | list | tuple, optional): Per
175 positional-arg batch dim, indexed by arg position. Entries may be
176 plain ``int`` (or ``None`` to keep the default); a single-input
177 model may pass a bare ``int``/``BatchDimSpec`` instead of a
178 one-element list (wrapped automatically). Default ``None``.
179 kwargs_batch_dim (dict, optional): Per keyword-arg batch dim, mapping
180 arg name to a plain ``int`` or ``BatchDimSpec``. Default ``None``.
181 swap (bool, optional): Whether to inject pipeline activation swap
182 control steps. Supported by ``ScheduleGPipe``, ``Schedule1F1B``,
183 and ``ScheduleInterleaved1F1B``. Default ``False``.
184 p2p_transport (str, optional): How pipeline send/recv are issued.
185 ``"auto"`` (default) — gap-time duplex batching on overlap_b_f
186 schedules (``coalesce_p2p``: same-peer send+recv as one
187 ``batch_isend_irecv``, TX||RX; hardware-validated and measured a
188 net win on real workloads), plain per-op ``isend``/``irecv``
189 everywhere else. ``"plain"`` — force per-op ``isend``/``irecv``
190 (escape hatch for transports or topologies where batching
191 misbehaves). ``"batch"`` — duplex batching explicitly (what
192 ``"auto"`` picks under overlap_b_f). ``"boundary"`` —
193 EXPERIMENTAL fwd-boundary batching: each overlap's ``F_SEND`` +
194 the next slot's recvs go out mid-overlap, right after the forward,
195 as per-op solo batches; only ``B_SEND`` waits for the backward.
196 Avoids the duplex handle's send-coupling (a2a-friendly) and posts
197 the activation send ~half a slot early, but is not yet
198 hardware-validated — opt in deliberately. Must be set identically
199 on every rank — HCCL cannot match a batched op against a plain one
200 (EI0005). On backends that require full-group participation before
201 the first batched P2P call, batch-backed modes prepare the PP group
202 once at their first run boundary.
203 ``"multi_stream"`` — EXPERIMENTAL: the same gap-time duplex
204 batching as ``"batch"``, but each adjacent PP peer pair uses its
205 own two-rank communication group so the previous and next PP edges
206 can progress on separate communication streams.
207 """
209 _P2P_TRANSPORTS = ("auto", "plain", "batch", "boundary", "multi_stream")
211 def __init__(self,
212 stages,
213 micro_batch_num,
214 args_batch_dim=None,
215 kwargs_batch_dim=None,
216 output_concat_dim=None,
217 overlap_p2p=False,
218 swap=False,
219 p2p_transport="auto"):
220 if p2p_transport not in self._P2P_TRANSPORTS:
221 raise ValueError(
222 f"p2p_transport must be one of {self._P2P_TRANSPORTS}, got "
223 f"{p2p_transport!r}"
224 )
225 self.stages = self._check_stages(stages)
226 self.micro_batch_num = micro_batch_num
227 self._args_batch_dim = self._normalize_args_batch_dim(args_batch_dim)
228 self._kwargs_batch_dim = self._normalize_kwargs_batch_dim(kwargs_batch_dim)
229 # Every key on the wire is one DATA_LOAD actually populated. MPipe
230 # overrides this in ``construct_exec_order``.
231 self._data_keys = tuple(self._kwargs_batch_dim or ())
232 self._output_concat_dim = output_concat_dim
233 self.split_micro_batch = _MicroBatch(
234 self.micro_batch_num, self._args_batch_dim, self._kwargs_batch_dim)
235 self.n_local_stages = len(self.stages)
236 self._stage_dict = self.convert_stages_dict()
237 self.real_stage_num = self.stages[0].stage_num // self.n_local_stages
238 self._stage_num = self.stages[0].stage_num
239 self._stage_to_rank_index = None
240 self._overlap_p2p = overlap_p2p
241 self.exec_order = {}
242 self._init_stages()
243 self._build_stage_to_rank_index()
244 self.fwd_handle_cache = {}
245 self.bwd_handle_cache = {}
246 self._custom_fn_map = {}
247 self._pp_swap_enabled = swap
248 self._swap_keys = frozenset()
249 self._swap_session = None
250 self.data_iterator = None
251 self._pp_fsdp_composed = False
252 self._dp_group_info = None
253 self.last_local_tokens = 0
254 # Outstanding async send handle groups for the in-flight
255 # ``run_microbatches`` call; reset per run and drained at its end.
256 self._send_handles = []
257 # ``p2p_transport`` resolves in ``build_exec_order`` (it needs the
258 # subclass's ``_overlap_b_f``) to one of:
259 #
260 # * ``"batch"`` (the ``"auto"`` default on overlap_b_f schedules) —
261 # gap-time duplex via ``coalesce_p2p``: same-peer send+recv as one
262 # ``batch_isend_irecv`` (TX||RX on the full-duplex link).
263 # Waiting the whole batch before consuming its recv also waits for
264 # the riding send, which can reduce EP all-to-all overlap.
265 # * ``"plain"`` (the ``"auto"`` default elsewhere) — per-op
266 # ``isend``/``irecv``, the upstream-original path.
267 # * ``"boundary"`` (EXPERIMENTAL, explicit opt-in only) — fwd-boundary
268 # batching. ``attach_fwd_boundary_p2p`` hangs each steady gap's
269 # ``F_SEND`` (its data is ready when the overlap's forward finishes
270 # — the backward, ~2x FLOPs, is the long pole) plus the next slot's
271 # recvs on the OVERLAP_B_F step; the stage's after-forward hook
272 # fires ``exec_boundary_p2p`` at the fwd/bwd boundary, while the
273 # backward is still running. The send leaves roughly half a slot
274 # early — the only mode that moves the SENDER's post time — and the
275 # recv handles carry no send (a2a-friendly). Every op is a per-op
276 # solo batch and ``coalesce_p2p`` is NOT run: a hoisted recv cannot
277 # stay duplexed with a send that is only ready later, and
278 # asymmetric shapes (one end duplex, other end split) hang —
279 # all-solo keeps per-pair batch sequences complementary
280 # ([S,R] vs [R,S]), safe under both candidate HCCL pairing
281 # semantics. Promote to the auto default only after it earns both
282 # a hardware accuracy pass and a perf win over "batch".
283 # * ``"multi_stream"`` (EXPERIMENTAL, explicit opt-in only) — keeps the
284 # ``"batch"`` schedule and same-peer duplex fusion, but routes each
285 # physical PP edge through a fixed two-rank group so neighboring
286 # edges can use separate communication streams. Group count is
287 # bounded by physical PP degree (at most two per rank), not by the
288 # number of micro-batches.
289 #
290 # Two transport invariants for any future rewrite: plain per-pair
291 # streams need the gap's recv-first/send-first complementarity (a
292 # send-crossing hoist made both ends recv-first -> rendezvous deadlock,
293 # 2026-06), and batch pairing needs per-pair shape mirroring (a
294 # one-sided split made 2 solos face 1 duplex -> hang, 2026-06).
295 self._p2p_transport = p2p_transport
296 # Effective mode + per-op batch gating; set by ``build_exec_order``.
297 self._p2p_mode = None
298 self._batch_p2p = False
299 self._batch_p2p_group = None
300 self._batch_p2p_group_initialized = False
301 # ``peer_global_rank -> two-rank process group``. Populated only for
302 # the experimental ``p2p_transport="multi_stream"`` mode.
303 self._p2p_multi_stream_groups = {}
304 self._p2p_multi_stream_groups_initialized = False
305 # OVERLAP steps whose boundary_p2p was already issued this run (the
306 # stage after-forward hook and the post-step safety net are both
307 # allowed to call exec_boundary_p2p; reset per run_microbatches call).
308 self._boundary_issued = set()
309 # (fwd stage_index, micro_index) -> armed OVERLAP step, consumed by the
310 # stage after-forward hook to fire the boundary issue mid-overlap.
311 self._pending_boundary = {}
313 @property
314 def kwargs_batch_dim(self) -> dict:
315 """Per-kwarg batch-dim split spec declared by the model's pipeline setup.
317 Read by the trainer's PP step to forward exactly the batch keys the
318 schedule knows how to split into per-micro kwargs (text: attention_mask /
319 position_ids; VL also pixel_values / image_grid_thw / mm_token_type_ids).
320 """
321 return self._kwargs_batch_dim or {}
323 def register_custom_function(self, step_type: MetaStepType, fn) -> None:
324 """Register a custom execution function for the given step type.
326 When :meth:`run_microbatches` encounters a :class:`MetaStep` whose
327 ``type`` matches ``step_type``, it calls ``fn(step, ctx)`` instead
328 of the built-in logic.
330 Args:
331 step_type: The :class:`MetaStepType` to intercept.
332 fn: A callable with signature ``(step: MetaStep, ctx: PipelineContext) -> None``.
334 When pipeline activation swap is enabled, callbacks that
335 execute FWD leaves must call ``ctx.schedule.execute_fwd_leaf``
336 so swap collection runs in the matching chunk context.
338 Example:
339 >>> def my_overlap_callback(step, ctx):
340 ... fwd_step, bwd_step = step.sub_steps
341 ... # custom parallel execution logic
342 >>> schedule.register_custom_function(MetaStepType.OVERLAP_F_B, my_overlap_callback)
343 """
344 self._custom_fn_map[step_type] = fn
346 @staticmethod
347 def _stage_hsdp_modules(stage):
348 """Return the independently configurable HSDP roots in a stage."""
349 if isinstance(stage.submodule, HSDPModule):
350 return [stage.submodule]
351 modules = []
352 seen = set()
353 for _, module in stage.submodule.named_modules():
354 if not isinstance(module, HSDPModule) or id(module) in seen:
355 continue
356 seen.add(id(module))
357 modules.append(module)
358 return modules
360 def _inject_local_fsdp_actions(self):
361 """Annotate the local rank schedule with optional FSDP control actions."""
362 current_rank = self._stage_to_rank_index[self.stages[0].stage_index]
363 managed_stage_indices = {
364 stage.stage_index
365 for stage in self.stages
366 if self._stage_hsdp_modules(stage)
367 }
368 if not managed_stage_indices:
369 return
370 if len(managed_stage_indices) != len(self.stages):
371 raise RuntimeError(
372 "When injecting fsdp_action, expect all stages to be HSDPModule. "
373 "Check whether all separated modules are wrapped with 'fully_shard'."
374 )
375 for stage in self.stages:
376 for module in self._stage_hsdp_modules(stage):
377 module.set_reshard_after_forward(False)
378 rank_actions = add_fsdp_unshard_reshard(self.exec_order[current_rank], managed_stage_indices)
379 self.exec_order[current_rank] = add_fsdp_reduce_grad(
380 rank_actions,
381 managed_stage_indices,
382 self.micro_batch_num,
383 )
385 def _prepare_fsdp_backward(self):
386 """Disable HSDP post-backward actions once for the current pipeline run.
388 Pipeline accumulates gradients across all micro-batches and executes
389 the reduction through ``FSDP_REDUCE_GRAD`` after the last backward.
390 ``execute_reduce_grad`` restores both flags, so they must be disabled
391 once again at the beginning of the next run.
392 """
393 for stage in self.stages:
394 if not stage.has_backward:
395 continue
396 for module in self._stage_hsdp_modules(stage):
397 module.set_reshard_after_backward(False)
398 module.set_requires_gradient_sync(False)
400 def _inject_local_pp_swap_actions(self):
401 """Annotate the local rank schedule with pipeline activation-swap actions."""
402 self._swap_keys = frozenset()
403 if not self._pp_swap_enabled:
404 return
405 current_rank = self._stage_to_rank_index[self.stages[0].stage_index]
406 unregister_layer_swap_hooks(self.stages)
407 self.exec_order[current_rank] = inject_pipeline_swap_steps(self.exec_order[current_rank])
408 self._swap_keys = frozenset(
409 (step.stage_index, step.micro_index)
410 for step in self.exec_order[current_rank]
411 if step is not None and step.type == MetaStepType.SWAP_LAUNCH_OFFLOAD
412 )
414 @abstractmethod
415 def _build_stage_to_rank_index(self) -> None:
416 """
417 Build attribute of _stage_to_rank_index.
418 Each subclass constructs it according to its own schedule style.
419 """
421 @abstractmethod
422 def construct_exec_order(self) -> None:
423 """Build exec order, PP cmopute and PP comms(Send/Recv)"""
425 def build_exec_order(self) -> None:
426 """Build the execution order and inject optional FSDP / PP-swap actions.
428 Meta-Step Ordering Contract
429 --------------------------
430 The per-rank execution schedule is assembled in layers:
432 1. ``construct_exec_order()``
433 Pure compute (FWD/BWD) + P2P communication (SEND/RECV).
435 2. ``_inject_local_fsdp_actions()``
436 FSDP parameter management: UNSHARD before compute, RESHARD after,
437 REDUCE_GRAD after the last backward of each stage. This layer must
438 run *before* swap injection so that swap can see the already-placed
439 FSDP steps and optimise its placement relative to them:
441 * After FWD: ``SWAP_LAUNCH_OFFLOAD`` before ``FSDP_RESHARD``
442 (start D2H early, freeing parameter memory is fast).
443 * Before BWD: if a direct lookahead ``FSDP_UNSHARD`` exists,
444 ``SWAP_LAUNCH_LOAD`` precedes it so H2D overlaps all-gather.
445 If parameters are already unsharded, load is delayed until the
446 BWD container to avoid extending activation residency.
448 3. ``coalesce_p2p()`` or ``attach_fwd_boundary_p2p()``
449 P2P batching — runs before swap injection so P2P containers are
450 finalized first.
452 4. ``_inject_local_pp_swap_actions()``
453 Activation swap: leaf-local collection around ``FWD``, followed by
454 ``LAUNCH_OFFLOAD → WAIT_OFFLOAD → LAUNCH_LOAD → WAIT_LOAD``.
455 ``WAIT_LOAD`` executes on the scheduler thread immediately before
456 the backward consumer container.
457 When ``enable_dxdw_split`` is active, ``BWD_INPUT`` is used as
458 the backward anchor because it is the activation consumer;
459 ``BWD_WEIGHT`` receives no swap control steps.
461 Swap steps are inserted only into before/after slots around
462 top-level containers and never inside ``BATCH_SEND_RECV.sub_steps``.
463 This preserves symmetric, paired fusion-block members and avoids
464 deadlocks caused by splitting a fused P2P block.
466 Also resolves ``p2p_transport``: ``"auto"`` becomes ``"batch"`` (the
467 measured-beneficial duplex) on schedules running with ``overlap_b_f``
468 and ``"plain"`` everywhere else. The matching P2P order-rewrite pass
469 runs after FSDP injection and before swap injection, so swap sees the
470 finalized top-level P2P containers.
471 """
472 mode = self._p2p_transport
473 if mode == "auto":
474 mode = "batch" if getattr(self, "_overlap_b_f", False) else "plain"
475 self._p2p_mode = mode
476 self._batch_p2p = mode != "plain"
477 # Multi-stream tensor P2P uses edge groups, but dynamic-shape metadata still uses the full PP group.
478 self._batch_p2p_group = (
479 self._resolve_batch_p2p_group() if self._batch_p2p else None
480 )
481 if mode == "multi_stream":
482 self._init_p2p_multi_stream_groups()
483 self.construct_exec_order()
484 # FSDP must inject before PP-swap so that swap can optimise
485 # placement relative to FSDP steps (LAUNCH_OFFLOAD before
486 # RESHARD after FWD; LAUNCH_LOAD before UNSHARD for BWD).
487 self._inject_local_fsdp_actions()
488 if mode == "boundary":
489 # fwd-boundary mode: hang the forward's F_SEND + next slot's recvs
490 # on the OVERLAP step (issued mid-overlap, right after the
491 # forward). Everything stays per-op solo batches — deliberately
492 # NO coalesce_p2p (see __init__).
493 self.exec_order = attach_fwd_boundary_p2p(self.exec_order)
494 elif mode in ("batch", "multi_stream"):
495 # Coalesce contiguous P2P runs into BATCH_SEND_RECV so the runtime
496 # issues same-peer send+recv as one duplex batch. NOTE: couples
497 # the riding send into the compute-gating recv wait — see
498 # __init__.
499 self.exec_order = coalesce_p2p(self.exec_order)
500 self._inject_local_pp_swap_actions()
502 def _init_p2p_multi_stream_groups(self) -> None:
503 """Create adjacent two-rank groups for multi-stream P2P transport."""
504 first_stage = self.stages[0]
505 mesh = getattr(first_stage, "mesh", None)
506 pp_rank_list = (
507 list(mesh.rank_list)
508 if mesh is not None
509 else dist.get_process_group_ranks(first_stage.pp_group)
510 )
511 self._p2p_multi_stream_groups = create_p2p_multi_stream_groups(
512 pp_rank_list,
513 include_wrap=self.n_local_stages > 1,
514 )
516 def _p2p_op(self, op_type, tensor, peer):
517 """Build a P2P descriptor, selecting the multi-stream group when enabled."""
518 group = self._batch_p2p_group
519 if self._p2p_mode == "multi_stream":
520 group = self._p2p_multi_stream_groups.get(peer)
521 if group is None:
522 raise RuntimeError(
523 f"No P2P multi-stream group was created for peer global rank {peer}. "
524 f"Available peers are {sorted(self._p2p_multi_stream_groups)}."
525 )
526 if op_type not in ("isend", "irecv"):
527 raise ValueError(f"Unsupported pipeline P2P operation: {op_type!r}.")
528 op = dist.isend if op_type == "isend" else dist.irecv
529 return dist.P2POp(op, tensor, peer, group=group)
531 def convert_stages_dict(self):
532 """convert stages to dict."""
533 stage_dict = {}
534 for stage in self.stages:
535 stage_dict[stage.stage_index] = stage
536 return stage_dict
538 def split_microbatches(self, args, kwargs):
539 """split_microbatches."""
540 if args or kwargs:
541 args_split, kwargs_split = self.split_micro_batch(args, kwargs)
542 return args_split, kwargs_split
543 return [[] for _ in range(self.micro_batch_num)], [{} for _ in range(self.micro_batch_num)]
545 @staticmethod
546 def _to_spec(elem):
547 """Normalize one batch-dim entry: ``int`` -> ``BatchDimSpec``.
549 ``None`` and ``BatchDimSpec`` pass through unchanged. ``bool`` is
550 rejected even though it is an ``int`` subclass, so ``True``/``False``
551 are not silently read as dims 1/0.
552 """
553 if elem is None or isinstance(elem, BatchDimSpec):
554 return elem
555 if isinstance(elem, int) and not isinstance(elem, bool):
556 return BatchDimSpec(elem)
557 raise TypeError(
558 f"batch-dim entry must be int, BatchDimSpec or None, but got {type(elem)}.")
560 @staticmethod
561 def _normalize_args_batch_dim(args_batch_dim):
562 """Accept a plain ``int``/``BatchDimSpec`` or a ``list``/``tuple`` of them.
564 ``args_batch_dim`` is a per-arg spec indexed by positional-arg position
565 (see ``_MicroBatch``). A single-input model can pass a bare ``int`` /
566 ``BatchDimSpec`` instead of the awkward one-element
567 ``BatchDimSpec.from_tuple((0,))``; elements may be plain ``int`` (or
568 ``None`` to keep the default). Always returns ``None`` or a
569 ``tuple[BatchDimSpec | None]`` so downstream per-arg indexing is
570 unchanged.
571 """
572 if args_batch_dim is None:
573 return None
574 if isinstance(args_batch_dim, BatchDimSpec) or \
575 (isinstance(args_batch_dim, int) and not isinstance(args_batch_dim, bool)):
576 args_batch_dim = (args_batch_dim,)
577 if isinstance(args_batch_dim, (list, tuple)):
578 return tuple(PipelineScheduleRuntime._to_spec(e) for e in args_batch_dim)
579 raise TypeError(
580 f"args_batch_dim must be int, BatchDimSpec or a list/tuple of them, "
581 f"but got {type(args_batch_dim)}.")
583 @staticmethod
584 def _normalize_kwargs_batch_dim(kwargs_batch_dim):
585 """Accept plain ``int`` dict values: ``{\"x\": 0}`` -> ``{\"x\": BatchDimSpec(0)}``.
587 ``kwargs_batch_dim`` maps each keyword-arg name to its batch dim.
588 Returns ``None`` or a ``dict[str, BatchDimSpec | None]`` so downstream
589 per-key indexing is unchanged.
590 """
591 if kwargs_batch_dim is None:
592 return None
593 if not isinstance(kwargs_batch_dim, dict):
594 raise TypeError(
595 f"kwargs_batch_dim must be a dict[str, int | BatchDimSpec], "
596 f"but got {type(kwargs_batch_dim)}.")
597 return {k: PipelineScheduleRuntime._to_spec(v) for k, v in kwargs_batch_dim.items()}
599 def _check_stages(self, stages):
600 """check stages type."""
601 if isinstance(stages, hyper_parallel.PipelineStage):
602 return [stages]
603 if isinstance(stages, (list, tuple)):
604 for stage in stages:
605 if not isinstance(stage, hyper_parallel.PipelineStage):
606 raise TypeError(f"Argument 'stages' must be type of PipelineStage, \
607 list or tuple of PipelineStage, but got list or tuple of {type(stage)}.")
608 return stages
609 raise TypeError(f"Argument 'stages' must be type of PipelineStage, \
610 list or tuple of PipelineStage, but got type of {type(stages)}.")
612 def _init_stages(self):
613 """init stages."""
614 pp_group = None
615 for stage in self.stages:
616 # Virtual chunks on one rank must share the automatically created
617 # communicator; native new_group does not cache equivalent groups.
618 if stage.pp_group is None and stage.mesh is None:
619 stage.pp_group = pp_group
620 stage.init(self.n_local_stages)
621 pp_group = stage.pp_group
622 # After-forward hook: lets the schedule issue fwd-boundary P2P the
623 # moment a forward chunk completes (no-op unless an OVERLAP step
624 # with boundary_p2p was armed for that (stage, micro)).
625 stage._after_forward_chunk = self._on_forward_chunk_done # pylint: disable=W0212
627 def _on_forward_chunk_done(self, stage_index, micro_index):
628 """Stage after-forward hook: fire the armed boundary P2P, if any.
630 Runs on the thread executing the forward (the overlap callback's
631 ``fwd_fn`` / the main thread), at the fwd/bwd boundary — the paired
632 backward is still running, so the boundary ops overlap it. Keyed by
633 the overlap's forward ``(stage_index, micro_index)``; unrelated
634 forwards (warm-up steps, recompute re-runs of past micros) miss the
635 key and no-op.
636 """
637 step = self._pending_boundary.pop((stage_index, micro_index), None)
638 if step is not None:
639 self.exec_boundary_p2p(step)
641 def run(self, *args: Any, **kwargs: Any) -> list:
642 """schedule run."""
643 losses = []
644 try:
645 if self._pp_swap_enabled:
646 self._swap_session = PipelineSwapSession(self._swap_keys)
647 split_args, split_kwargs = self.split_microbatches(args, kwargs)
648 self.run_microbatches(split_args, split_kwargs, losses)
649 finally:
650 self._drain_inflight_p2p()
651 if self._swap_session is not None:
652 self._swap_session.close()
653 self._swap_session = None
654 return losses
656 def run_with_dataiterator(self, data_iterator: Iterator, pp_fsdp_composed: bool,
657 dp_group_info: object = None, **kwargs: object) -> list:
658 """Run the schedule, pulling each micro-batch from ``data_iterator``.
660 Args:
661 data_iterator (Iterator): Yields one micro-batch per DATA_LOAD.
662 pp_fsdp_composed (bool): Whether PP composes with (F/H)SDP.
663 dp_group_info (object): DP group used to all-reduce the token count.
664 **kwargs: Accepted for signature parity with :meth:`run`; unused.
666 Returns:
667 list: The last stage's per-micro-batch losses (empty elsewhere).
668 """
669 del kwargs
670 self.data_iterator = data_iterator
671 self._pp_fsdp_composed = pp_fsdp_composed
672 # DATA_LOAD all-reduces n_valid over this group under PP+FSDP; the
673 # trainer owns the handle, so it has to be threaded in.
674 self._dp_group_info = dp_group_info
675 self.last_local_tokens = 0
676 rargs = [[] for _ in range(self.micro_batch_num)]
677 rkwargs = [{} for _ in range(self.micro_batch_num)]
678 losses = []
679 try:
680 self.run_microbatches(rargs, rkwargs, losses)
681 finally:
682 # An exception unwinds past the end-of-iteration drain, leaving
683 # handles un-waited; wait them so the contract holds on that path too.
684 self._drain_inflight_p2p()
685 if self._swap_session is not None:
686 self._swap_session.close()
687 self._swap_session = None
688 return losses
690 def sync_shared_parameters_grad(self):
691 """sync_shared_parameters_grad."""
692 for stage in self.stages:
693 stage.sync_shared_parameters_grad()
695 def update_losses(self, stage, loss, losses):
696 """update_losses."""
697 if stage.is_last_stage:
698 losses.append(loss)
700 @property
701 def enable_dxdw_split(self) -> bool:
702 """Whether this schedule splits ``OVERLAP_B_F`` backward into dx/dw."""
703 return getattr(self, "_enable_dxdw_split", False)
705 def _wait_p2p(self, handles):
706 for handle in handles:
707 if handle is not None:
708 handle.wait()
710 def _resolve_batch_p2p_group(self):
711 """Return the common PP group required by every local stage."""
712 first_stage = self.stages[0]
713 group = first_stage.pp_group
714 for stage in self.stages[1:]:
715 if stage.pp_group is group or stage.pp_group == group:
716 continue
717 raise ValueError(
718 "Batch P2P requires all local stages to use the same PP process group, "
719 f"but stages {first_stage.stage_index} and {stage.stage_index} use different groups."
720 )
721 return group
723 def _ensure_batch_p2p_group_initialized(self) -> None:
724 """Run backend-specific PP-group preparation once before batched P2P."""
725 if (not self._batch_p2p or self._batch_p2p_group_initialized
726 or self.real_stage_num <= 1):
727 return
728 prepare_batch_p2p_group(self._batch_p2p_group)
729 self._batch_p2p_group_initialized = True
731 def _ensure_p2p_multi_stream_groups_initialized(self) -> None:
732 """Prepare multi-stream edge groups in their stable creation order."""
733 if (getattr(self, "_p2p_mode", None) != "multi_stream"
734 or getattr(self, "_p2p_multi_stream_groups_initialized", False)):
735 return
736 for group in self._p2p_multi_stream_groups.values():
737 prepare_batch_p2p_group(group)
738 self._p2p_multi_stream_groups_initialized = True
740 def _drain_inflight_p2p(self):
741 """Wait every P2P handle still in flight — error-path cleanup.
743 run_microbatches waits its deferred sends only in the end-of-iteration
744 drain; an exception mid-iteration unwinds past that drain, leaving issued
745 isend/irecv handles un-waited in ``_send_handles`` and the recv caches
746 (the ``CommHandle destroyed without calling wait()`` warning). run()'s
747 finally calls this so every handle is still ``wait()``-ed — honoring the
748 comm contract — on the error path too, and pops them so a later run()
749 does not re-wait stale handles (the recv caches are never reset per run).
750 No-op on the normal path: the drain already emptied ``_send_handles`` and
751 every cached recv was consumed.
752 """
753 while self._send_handles:
754 self._wait_p2p(self._send_handles.pop())
755 while self.fwd_handle_cache:
756 self._wait_p2p(self.fwd_handle_cache.popitem()[1])
757 while self.bwd_handle_cache:
758 self._wait_p2p(self.bwd_handle_cache.popitem()[1])
760 def _batched_issue(self, specs):
761 """Launch same-peer P2P ``specs`` as one ``batch_isend_irecv`` group.
763 ``specs`` are ``(op_type, tensor, peer_global_rank)`` from the stage's
764 ``*_specs`` builders (which carry the meta/bookkeeping side effects).
765 Returns the native list of work handles or ``[]`` — shaped like
766 the per-op ``exec_*_ops`` return so the cache / drain paths are
767 unchanged. Only the launch is coalesced; matching stays per-peer FIFO.
768 """
769 if not specs:
770 return []
771 ops = [self._p2p_op(op_type, tensor, peer) for op_type, tensor, peer in specs]
772 return dist.batch_isend_irecv(ops)
774 # --- P2P step primitives ------------------------------------------------
775 # One method per cross-rank comm action, used both by the runtime loop
776 # (``_exec_step``) and by OVERLAP callbacks (via ``ctx.schedule``). With
777 # ``overlap_p2p=True`` comm is decoupled from its compute: a recv caches its
778 # handles for the consuming step to ``wait_*`` later, and a send defers its
779 # handles to the end-of-iteration drain. With ``overlap_p2p=False`` every
780 # op waits inline.
782 def recv_fwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
783 """Post the FWD recv for ``micro_index``; cache it (overlap_p2p) or wait now."""
784 handles = (self._batched_issue(stage.fwd_recv_specs(micro_index))
785 if self._batch_p2p else stage.exec_fwd_recv_ops(micro_index))
786 if self._overlap_p2p:
787 self.fwd_handle_cache[(stage.stage_index, micro_index)] = handles
788 else:
789 self._wait_p2p(handles)
791 def recv_bwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
792 """Post the BWD recv for ``micro_index``; cache it (overlap_p2p) or wait now."""
793 handles = (self._batched_issue(stage.bwd_recv_specs(micro_index))
794 if self._batch_p2p else stage.exec_bwd_recv_ops(micro_index))
795 if self._overlap_p2p:
796 self.bwd_handle_cache[(stage.stage_index, micro_index)] = handles
797 else:
798 self._wait_p2p(handles)
800 def wait_fwd_recv(self, stage_index: int, micro_index: int) -> None:
801 """Wait the FWD recv cached by :meth:`recv_fwd`; no-op if nothing is cached."""
802 handles = self.fwd_handle_cache.pop((stage_index, micro_index), None)
803 if handles:
804 self._wait_p2p(handles)
806 def wait_bwd_recv(self, stage_index: int, micro_index: int) -> None:
807 """Wait the BWD recv cached by :meth:`recv_bwd`; no-op if nothing is cached."""
808 handles = self.bwd_handle_cache.pop((stage_index, micro_index), None)
809 if handles:
810 self._wait_p2p(handles)
812 def send_fwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
813 """Send this stage's forward output for ``micro_index`` to the next stage."""
814 handles = (self._batched_issue(stage.fwd_send_specs(micro_index))
815 if self._batch_p2p else stage.exec_fwd_send_ops(micro_index)) or []
816 if self._overlap_p2p:
817 # Append the whole handle group: run_microbatches drains _send_handles
818 # group by group, so a bare handle would be wrongly iterated as a list.
819 self._send_handles.append(handles)
820 else:
821 self._wait_p2p(handles)
823 def send_bwd(self, stage: "hyper_parallel.PipelineStage", micro_index: int) -> None:
824 """Send this stage's input-gradient for ``micro_index`` to the previous stage.
826 Driven by the scheduler's ``BWD_SEND`` step. It pops the input grad that
827 the backward (unified ``backward_one_chunk`` or, under
828 ``enable_dxdw_split=True``, ``backward_input_one_chunk``) wrote to the
829 stage's ``bwd_cache``. Calling it manually in addition to the scheduled
830 ``BWD_SEND`` would double-send the gradient.
831 """
832 handles = (self._batched_issue(stage.bwd_send_specs(micro_index))
833 if self._batch_p2p else stage.exec_bwd_send_ops(micro_index)) or []
834 if self._overlap_p2p:
835 self._send_handles.append(handles)
836 else:
837 self._wait_p2p(handles)
839 def _arm_boundary(self, step):
840 """Register ``step`` for the stage after-forward hook; return its key.
842 No-op (returns ``None``) unless ``step`` carries ``boundary_p2p``. The
843 key is the overlap's forward ``(stage_index, micro_index)`` — exactly
844 what the hook receives when that forward chunk completes.
845 """
846 if not getattr(step, "boundary_p2p", None) or not step.sub_steps:
847 return None
848 fwd_sub = next((s for s in step.sub_steps
849 if s.type == MetaStepType.FWD), None)
850 if fwd_sub is None:
851 return None
852 key = (fwd_sub.stage_index, fwd_sub.micro_index)
853 self._pending_boundary[key] = step
854 return key
856 def _finish_boundary(self, step, armed_key) -> None:
857 """Post-step safety net: issue any boundary P2P the hook missed."""
858 self.exec_boundary_p2p(step)
859 if armed_key is not None:
860 self._pending_boundary.pop(armed_key, None)
862 def exec_boundary_p2p(self, step) -> None:
863 """Issue ``step.boundary_p2p`` (fwd-boundary P2P) once per run.
865 Fired by the stage after-forward hook (``_on_forward_chunk_done``) the
866 moment the overlap's forward chunk completes — the backward is still
867 running on its own thread, so the F_SEND leaves ~half a slot early and
868 the next slot's recvs are already posted when the peers' sends arrive.
869 Idempotent per ``run_microbatches`` call: the post-step safety net
870 (``_finish_boundary``) also invokes it, so an overlap whose forward
871 never went through ``forward_one_chunk`` degrades to gap-time issue
872 order instead of dropping the ops.
874 Dispatches through the existing per-op helpers (``send_fwd`` /
875 ``recv_fwd`` / ``recv_bwd``), so batching, handle caching for
876 ``wait_*_recv`` and deferred-send bookkeeping behave exactly like the
877 scheduled steps they replace. No-op for steps without
878 ``boundary_p2p``.
879 """
880 ops = getattr(step, "boundary_p2p", None)
881 if not ops or id(step) in self._boundary_issued:
882 return
883 self._boundary_issued.add(id(step))
884 for sub in ops:
885 stage = self._stage_dict[sub.stage_index]
886 if sub.type == MetaStepType.FWD_SEND:
887 self.send_fwd(stage, sub.micro_index)
888 elif sub.type == MetaStepType.FWD_RECV:
889 self.recv_fwd(stage, sub.micro_index)
890 elif sub.type == MetaStepType.BWD_RECV:
891 self.recv_bwd(stage, sub.micro_index)
892 elif sub.type == MetaStepType.BWD_SEND:
893 # attach_fwd_boundary_p2p never hoists BWD_SEND (its grad is
894 # produced by the backward still in flight); defensive only.
895 self.send_bwd(stage, sub.micro_index)
897 # ``op_type -> (specs builder name, route kind)`` for a coalesced sub-step.
898 # ``route`` is the recv-cache kind (so wait_*_recv finds the handle), or
899 # ``None`` for a send (no local consumer).
900 _BATCH_SUB_DISPATCH = {
901 MetaStepType.FWD_RECV: ("fwd_recv_specs", "fwd"),
902 MetaStepType.BWD_RECV: ("bwd_recv_specs", "bwd"),
903 MetaStepType.FWD_SEND: ("fwd_send_specs", None),
904 MetaStepType.BWD_SEND: ("bwd_send_specs", None),
905 }
907 def _exec_batch_send_recv(self, step) -> None:
908 """Execute a coalesced P2P run: one ``batch_isend_irecv`` per peer.
910 Builds each sub-step's specs (same meta/bookkeeping side effects as the
911 per-step ``recv_fwd`` / ``send_fwd`` / ...), groups every op by peer
912 global rank, and issues one batch per peer so a same-peer send+recv runs
913 duplex. Handle routing mirrors the per-step path: under
914 ``overlap_p2p`` a batch carrying a recv is cached for ``wait_*_recv``
915 (its send rides along), a send-only batch defers to ``_send_handles``;
916 without ``overlap_p2p`` every batch waits inline.
917 """
918 # (op_type, tensor, peer, route) per op; route = (kind, stage, micro)
919 # for a recv, else None.
920 tagged = []
921 for sub in step.sub_steps:
922 builder_name, kind = self._BATCH_SUB_DISPATCH[sub.type]
923 stage = self._stage_dict[sub.stage_index]
924 specs = getattr(stage, builder_name)(sub.micro_index)
925 route = (kind, sub.stage_index, sub.micro_index) if kind is not None else None
926 for op_type, tensor, peer in specs:
927 tagged.append((op_type, tensor, peer, route))
929 by_peer = {}
930 for item in tagged:
931 by_peer.setdefault(item[2], []).append(item)
933 for items in by_peer.values():
934 ops = [self._p2p_op(op_type, tensor, peer) for op_type, tensor, peer, _ in items]
935 handles = dist.batch_isend_irecv(ops)
936 if not handles:
937 continue
938 if not self._overlap_p2p:
939 self._wait_p2p(handles)
940 continue
941 recv_routes = [route for *_, route in items if route is not None]
942 if recv_routes:
943 for kind, si, mi in recv_routes:
944 cache = self.fwd_handle_cache if kind == "fwd" else self.bwd_handle_cache
945 cache[(si, mi)] = handles
946 else:
947 self._send_handles.append(handles)
949 def _assert_in_unshard_if_needed(self, stage, check_step):
950 if not isinstance(stage.submodule, HSDPModule):
951 return
952 submodule_hsdp_scheduler = stage.submodule.hsdp_scheduler
953 scheduler_state = submodule_hsdp_scheduler.hsdp_state
954 if scheduler_state.is_shard:
955 raise RuntimeError(
956 f"Executing MetaStep: {check_step}, expected HSDPModule parameters in unsharded "
957 f"state, but got sharded parameters."
958 )
960 def _exec_step(self, cur_step, arg_mbs, kwarg_mbs, losses):
961 """Execute one built-in step (non-custom, non-composite).
963 Each comm step dispatches to a single P2P primitive; each compute step
964 first waits its cached recv (a no-op under ``overlap_p2p=False``) and
965 then runs.
966 """
967 # MPipe prestages use stage_index=-1 which isn't in self._stage_dict;
968 # their branches below never touch ``stage``.
969 stage_index = cur_step.stage_index
970 stage = self._stage_dict[stage_index] if stage_index >= 0 else None
971 micro_index = cur_step.micro_index
972 step_type = cur_step.type
974 if step_type in (
975 MetaStepType.DATA_LOAD,
976 MetaStepType.DATA_SEND,
977 MetaStepType.DATA_RECV,
978 ):
979 if self.data_iterator is not None:
980 self._exec_pipeline_data_steps(step_type, stage_index, micro_index, arg_mbs, kwarg_mbs)
982 elif step_type in (
983 MetaStepType.SWAP_LAUNCH_OFFLOAD,
984 MetaStepType.SWAP_WAIT_OFFLOAD,
985 MetaStepType.SWAP_LAUNCH_LOAD,
986 MetaStepType.SWAP_WAIT_LOAD,
987 ):
988 self._exec_pipeline_swap_step(cur_step, arg_mbs, kwarg_mbs)
990 elif step_type == MetaStepType.FWD_RECV:
991 self.recv_fwd(stage, micro_index)
993 elif step_type == MetaStepType.FWD:
994 self.execute_fwd_leaf(cur_step, arg_mbs, kwarg_mbs, losses)
996 elif step_type == MetaStepType.FWD_SEND:
997 self.send_fwd(stage, micro_index)
999 elif step_type == MetaStepType.BWD_RECV:
1000 self.recv_bwd(stage, micro_index)
1002 elif step_type == MetaStepType.BWD_INPUT:
1003 self._assert_in_unshard_if_needed(stage, cur_step)
1004 self.wait_bwd_recv(stage_index, micro_index)
1005 stage.backward_input_one_chunk(micro_index)
1007 elif step_type == MetaStepType.BWD_WEIGHT:
1008 self._assert_in_unshard_if_needed(stage, cur_step)
1009 self.wait_bwd_recv(stage_index, micro_index)
1010 stage.backward_weight_one_chunk(micro_index)
1012 elif step_type == MetaStepType.BWD:
1013 self._assert_in_unshard_if_needed(stage, cur_step)
1014 self.wait_bwd_recv(stage_index, micro_index)
1015 stage.backward_one_chunk(micro_index)
1017 elif step_type == MetaStepType.BWD_SEND:
1018 self.send_bwd(stage, micro_index)
1020 else:
1021 # FSDP control steps dispatch via the handler table; any other type
1022 # is a no-op here (composite/custom types are handled upstream).
1023 fsdp_handler = _FSDP_STEP_HANDLERS.get(step_type)
1024 if fsdp_handler is not None:
1025 fsdp_handler(stage)
1027 def execute_fwd_leaf(self, step: MetaStep, arg_mbs: list, kwarg_mbs: list, losses: list) -> None:
1028 """Execute a forward leaf within its run-scoped swap group context."""
1029 stage = self._stage_dict[step.stage_index]
1030 micro_index = step.micro_index
1031 self._assert_in_unshard_if_needed(stage, step)
1032 self.wait_fwd_recv(stage.stage_index, micro_index)
1033 swap_managed = self._swap_session is not None and self._swap_session.manages(step)
1034 group_context = self._swap_session.group_context(step) if swap_managed else nullcontext()
1035 with group_context:
1036 out = stage.forward_one_chunk(micro_index, arg_mbs[micro_index], kwarg_mbs[micro_index])
1037 if swap_managed:
1038 # Boundary transport may pop the stage output cache from the
1039 # after-forward hook before this leaf returns. Protect the
1040 # local output while it is still directly available.
1041 self._swap_session.protect_aliases(step, out)
1042 self.update_losses(stage, out, losses)
1044 def _exec_pipeline_data_steps(self,
1045 step_type: MetaStepType,
1046 stage_index: int,
1047 micro_index: int,
1048 arg_mbs: list,
1049 kwarg_mbs: list):
1050 """Execute a pipeline data management step.
1052 active only under ``run_with_dataiterator``; the legacy ``run()``
1053 skip these step
1054 ie. None self.data_iterator will skip these steps
1055 """
1056 # Only active under ``run_with_dataiterator``; the legacy ``run()``
1057 # path pre-stages every per-micro kwarg on every rank already.
1058 if self.data_iterator is None:
1059 return
1060 if step_type == MetaStepType.DATA_LOAD:
1061 if stage_index <= 0:
1062 self._exec_data_load(micro_index, arg_mbs, kwarg_mbs)
1063 elif step_type == MetaStepType.DATA_SEND:
1064 self._exec_data_send(stage_index, micro_index, arg_mbs, kwarg_mbs)
1065 elif step_type == MetaStepType.DATA_RECV:
1066 self._exec_data_recv(stage_index, micro_index, arg_mbs, kwarg_mbs)
1068 def _exec_data_load(self, micro_index: int, arg_mbs: list, kwarg_mbs: list) -> None:
1069 """Load the next micro-batch onto stage 0's device and stage it as ``micro_index``'s input."""
1070 device = self.stages[0].device
1071 micro_batch = next(self.data_iterator)
1072 if isinstance(micro_batch, list):
1073 micro_batch = micro_batch[micro_index]
1074 micro_batch = {
1075 key: (value.to(device, non_blocking=True) if hasattr(value, "to") else value)
1076 for key, value in micro_batch.items()
1077 }
1078 input_ids = micro_batch["input_ids"]
1079 labels = micro_batch["labels"]
1080 # Ignore the final token after shifting next-token targets.
1081 pad_col = torch.full_like(labels[..., :1], -100)
1082 targets = torch.cat([labels[..., 1:], pad_col], dim=-1).contiguous()
1083 n_valid = max(int((targets != -100).sum().item()), 1)
1084 if getattr(self, "_pp_fsdp_composed", False):
1085 nt = torch.full((1,), n_valid).to(device)
1086 dist.all_reduce(nt, group=self._dp_group_info.group)
1087 n_valid = max(int(nt.item()), 1)
1088 self.last_local_tokens += n_valid
1090 micro_batch["targets"] = targets
1091 for key in self.kwargs_batch_dim:
1092 if key in kwarg_mbs[micro_index].keys():
1093 kwarg_mbs[micro_index][key] = torch.cat(
1094 [kwarg_mbs[micro_index][key], micro_batch[key]], dim=0)
1095 else:
1096 kwarg_mbs[micro_index][key] = micro_batch[key]
1097 arg_mbs[micro_index] = [input_ids]
1099 def _exec_data_send(self, stage_index: int, micro_index: int, arg_mbs: list, kwarg_mbs: list) -> None:
1100 """Send ``micro_index``'s data keys to the stage that consumes them."""
1101 if getattr(self, "_data_dst", False):
1102 dst_stage_idx = self._data_dst[stage_index][micro_index]
1103 else:
1104 dst_stage_idx = self.stages[0].dst_stage
1105 dst = self.stages[0]._global_rank(dst_stage_idx) # pylint: disable=protected-access
1106 # One meta round-trip per DATA_SEND instead of K; the matching
1107 # DATA_RECV unpacks the list in the same order.
1108 metas = []
1109 tensors = []
1110 for key in self._data_keys:
1111 tensor = (arg_mbs[micro_index][0] if key == "input_ids"
1112 else kwarg_mbs[micro_index][key])
1113 metas.append([tuple(tensor.shape), tensor.dtype])
1114 tensors.append(tensor)
1115 dist.send_object_list(metas, dst)
1116 handles = [dist.isend(t, dst) for t in tensors]
1117 self._wait_p2p(handles)
1119 def _exec_data_recv(self, stage_index: int, micro_index: int, arg_mbs: list, kwarg_mbs: list) -> None:
1120 """Receive ``micro_index``'s data keys and stage them as its input."""
1121 if getattr(self, "_data_src", False):
1122 src_stage_idx = self._data_src[stage_index][micro_index]
1123 else:
1124 src_stage_idx = self.stages[0].src_stage
1125 src = self.stages[0]._global_rank(src_stage_idx) # pylint: disable=protected-access
1126 device = self.stages[0].device
1127 # One recv_object_list unpacks every key's (shape, dtype) at
1128 # once, matching the packed DATA_SEND above.
1129 metas: list = [None] * len(self._data_keys)
1130 dist.recv_object_list(metas, src)
1131 handles = []
1132 for key, meta in zip(self._data_keys, metas):
1133 shape, dtype = meta
1134 buffer = torch.empty(shape, dtype=dtype, device=device)
1135 handles.append(dist.irecv(buffer, src))
1136 if key == "input_ids":
1137 arg_mbs[micro_index] = [buffer]
1138 else:
1139 kwarg_mbs[micro_index][key] = buffer
1140 self._wait_p2p(handles)
1142 def _exec_pipeline_swap_step(self, cur_step, arg_mbs, kwarg_mbs):
1143 """Execute a pipeline activation-swap control step."""
1144 if self._swap_session is None:
1145 raise RuntimeError("Pipeline swap step executed without an active run session.")
1146 if cur_step.type == MetaStepType.SWAP_LAUNCH_OFFLOAD:
1147 swap_launch_offload(cur_step, self, arg_mbs, kwarg_mbs, self._swap_session)
1148 elif cur_step.type == MetaStepType.SWAP_WAIT_OFFLOAD:
1149 swap_wait_offload(cur_step, self._swap_session)
1150 elif cur_step.type == MetaStepType.SWAP_LAUNCH_LOAD:
1151 swap_launch_load(cur_step, self._swap_session)
1152 elif cur_step.type == MetaStepType.SWAP_WAIT_LOAD:
1153 swap_wait_load(cur_step, self._swap_session)
1155 def run_microbatches(self, arg_mbs: list, kwarg_mbs: list, losses: list) -> None:
1156 """Execute the schedule step by step.
1158 Steps whose :attr:`MetaStep.type` has a registered custom function
1159 are delegated to that function with a :class:`PipelineContext`.
1160 Composite ``OVERLAP_F_B`` / ``OVERLAP_B_F`` steps without a
1161 registered handler fall back to executing their ``sub_steps``
1162 sequentially via :meth:`_exec_step` — correct but without
1163 comm/compute overlap. All other steps are executed by
1164 :meth:`_exec_step`.
1166 Logs one ``DEBUG`` line per non-bubble step showing the rank's
1167 progress: ``rank=<r> step=<i>/<n> <MetaStep>``. Enable with
1168 ``logging.getLogger('hyper_parallel.core.pipeline_parallel.scheduler')
1169 .setLevel(logging.DEBUG)`` to trace per-rank schedule advancement
1170 (handy when diagnosing deadlocks or callback ordering issues).
1171 """
1172 self._ensure_batch_p2p_group_initialized()
1173 self._ensure_p2p_multi_stream_groups_initialized()
1174 real_stage_index = self.stages[0].stage_index % self.real_stage_num
1175 self._prepare_fsdp_backward()
1176 self._send_handles = []
1177 self._boundary_issued = set()
1178 self._pending_boundary = {}
1179 ctx = None # lazily created
1181 ordered = self.exec_order[real_stage_index]
1182 total_steps = len(ordered)
1183 logger.debug(
1184 "run_microbatches start: rank=%d total_steps=%d micro_batch_num=%d",
1185 real_stage_index, total_steps, self.micro_batch_num,
1186 )
1188 for step_idx, cur_step in enumerate(ordered):
1189 if cur_step is None:
1190 continue
1192 logger.debug(
1193 "rank=%d step=%d/%d %s",
1194 real_stage_index, step_idx, total_steps, cur_step,
1195 )
1197 # Arm the fwd-boundary hook: when this step carries boundary_p2p,
1198 # the stage's after-forward hook fires exec_boundary_p2p the moment
1199 # the overlap's forward chunk completes (works for any callback —
1200 # no callback cooperation needed).
1201 armed_key = self._arm_boundary(cur_step)
1203 # Check for registered custom function
1204 custom_fn = self._custom_fn_map.get(cur_step.type)
1205 if custom_fn is not None:
1206 if ctx is None:
1207 ctx = PipelineContext(self, arg_mbs, kwarg_mbs, losses)
1208 custom_fn(cur_step, ctx)
1209 # Safety net: if the forward hook never fired (custom fwd path),
1210 # issue the boundary ops now (gap-time order) instead of
1211 # dropping them. Idempotent.
1212 self._finish_boundary(cur_step, armed_key)
1213 continue
1215 # Coalesced P2P block: group sub-steps by peer, issue one
1216 # batch_isend_irecv per peer (same-peer send+recv -> duplex).
1217 if cur_step.type == MetaStepType.BATCH_SEND_RECV:
1218 self._exec_batch_send_recv(cur_step)
1219 continue
1221 # Default for composite OVERLAP steps: run sub_steps sequentially.
1222 # P2P send/recv around these steps are already laid out in two
1223 # virtual slots by ``add_send_recv``, so sequential execution is
1224 # semantically equivalent to non-overlapped 1F1B.
1225 if (cur_step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F)
1226 and cur_step.sub_steps):
1227 for sub in cur_step.sub_steps:
1228 self._exec_step(sub, arg_mbs, kwarg_mbs, losses)
1229 self._finish_boundary(cur_step, armed_key)
1230 continue
1232 self._exec_step(cur_step, arg_mbs, kwarg_mbs, losses)
1234 logger.debug(
1235 "run_microbatches end: rank=%d pending_send_handles=%d",
1236 real_stage_index, len(self._send_handles),
1237 )
1238 self.sync_shared_parameters_grad()
1239 while self._send_handles:
1240 self._wait_p2p(self._send_handles.pop())
1243class _OverlapPhantom:
1244 """Internal marker used by :func:`add_send_recv` to expand an
1245 ``OVERLAP_F_B`` or ``OVERLAP_B_F`` step into two virtual time slots.
1247 An overlap step composes two sub-steps (``B + F`` or ``F + B``) that
1248 execute concurrently on the GPU but occupy **two** logical time slots
1249 in the column-scan sender timeline — the sender can only finish
1250 emitting the second sub-step's output after the first sub-step has
1251 completed. Treating an overlap step as a single slot places the RECV
1252 triggered by the second sub-step too early on the receiver.
1254 Each overlap step is expanded into two phantoms:
1255 * ``is_first_half=True`` — represents the first sub-step's emission
1256 slot; the original overlap step is emitted into the output
1257 schedule here (only once).
1258 * ``is_first_half=False`` — represents the second sub-step's emission
1259 slot; only its send/recv comms are inserted.
1260 """
1262 __slots__ = ('obf_step', 'sub_step', 'is_first_half')
1264 def __init__(self, obf_step, sub_step, is_first_half: bool):
1265 self.obf_step = obf_step
1266 self.sub_step = sub_step
1267 self.is_first_half = is_first_half
1270def _expand_overlap_slots(scheduler, real_stage_num):
1271 """Expand OVERLAP steps in a per-rank schedule into 2 virtual time slots.
1273 Returns a new ``{rank: [MetaStep | _OverlapPhantom | None, ...]}`` dict
1274 where each OVERLAP step is replaced by a pair of phantoms. Non-OVERLAP
1275 entries pass through unchanged.
1276 """
1277 expanded = {}
1278 for rank in range(real_stage_num):
1279 order = scheduler[rank]
1280 exp = []
1281 for op in order:
1282 if (op is not None
1283 and op.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F)
1284 and op.sub_steps):
1285 exp.append(_OverlapPhantom(op, op.sub_steps[0], is_first_half=True))
1286 exp.append(_OverlapPhantom(op, op.sub_steps[1], is_first_half=False))
1287 else:
1288 exp.append(op)
1289 expanded[rank] = exp
1290 return expanded
1293def _process_rank_items(real_stage_num, current_items, insert_step_comms, new_schedule):
1294 """Run ``insert_step_comms`` for each rank's current item, even ranks first.
1296 Even-before-odd ordering avoids P2P deadlocks between adjacent ranks.
1297 """
1298 for rank in range(0, real_stage_num, 2):
1299 item = current_items.get(rank)
1300 if item is not None:
1301 sub = item.sub_step if isinstance(item, _OverlapPhantom) else item
1302 insert_step_comms(sub, rank, new_schedule)
1303 for rank in range(1, real_stage_num, 2):
1304 item = current_items.get(rank)
1305 if item is not None:
1306 sub = item.sub_step if isinstance(item, _OverlapPhantom) else item
1307 insert_step_comms(sub, rank, new_schedule)
1310def _insert_dataload_before_fwd(schedule):
1311 """Splice a zero-width ``DATA_LOAD`` immediately before every FWD.
1313 Runs AFTER the column scan so the loads never occupy a schedule column:
1314 a per-FWD DATA_LOAD slot would widen each rank's columns by its
1315 (rank-dependent) warmup depth and desync the cross-rank scan — the
1316 last rank's BWD_RECV then lands after the BWD that consumes it.
1317 DATA_LOAD is local (no comm), so inserting it here cannot affect P2P
1318 placement; the handler no-ops when ``run_with_dataiterator`` isn't
1319 active. Composite OVERLAP steps get one load per FWD sub-step, placed
1320 before the composite.
1321 """
1322 def _loads_for(step):
1323 """Zero-width DATA_LOADs to splice before ``step`` (may be empty)."""
1324 if step is None:
1325 return ()
1326 if step.type == MetaStepType.FWD:
1327 return (MetaStep(step.micro_index, MetaStepType.DATA_LOAD, step.stage_index),)
1328 if step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F):
1329 return tuple(MetaStep(sub.micro_index, MetaStepType.DATA_LOAD, sub.stage_index)
1330 for sub in (step.sub_steps or ())
1331 if sub is not None and sub.type == MetaStepType.FWD)
1332 return ()
1334 for rank, steps in schedule.items():
1335 out = []
1336 for step in steps:
1337 out.extend(_loads_for(step))
1338 out.append(step)
1339 schedule[rank] = out
1340 return schedule
1343def _column_scan_insert_comms(expanded, real_stage_num, insert_step_comms):
1344 """Column-scan over an OVERLAP-expanded schedule to insert SEND/RECV.
1346 Processes ``expanded`` one time slot at a time. Emits the original
1347 overlap step into ``new_schedule`` only once (at the first-half
1348 phantom). Delegates comm insertion to ``insert_step_comms`` for each
1349 plain step or phantom's underlying sub-step.
1351 Even ranks are processed before odd ranks at each time step to avoid
1352 P2P deadlocks between adjacent ranks.
1354 Args:
1355 expanded: Result of :func:`_expand_overlap_slots`.
1356 real_stage_num: Number of physical ranks.
1357 insert_step_comms: Callable ``(step, rank, new_schedule) -> None``
1358 that inserts SEND/RECV for a single FWD/BWD step.
1360 Returns:
1361 ``{rank: [MetaStep, ...]}`` final schedule.
1362 """
1363 max_length = max(len(order) for order in expanded.values())
1364 new_schedule = {rank: [] for rank in range(real_stage_num)}
1366 for time_step in range(max_length):
1367 current_items = {}
1368 for rank in range(real_stage_num):
1369 if time_step < len(expanded[rank]):
1370 item = expanded[rank][time_step]
1371 current_items[rank] = item
1372 if item is None:
1373 # Preserve bubble slots to keep per-rank time-step
1374 # indexing aligned with the column scan. The runtime
1375 # loop skips ``None`` entries, so this is execution-
1376 # semantics-neutral.
1377 new_schedule[rank].append(None)
1378 continue
1379 if isinstance(item, _OverlapPhantom):
1380 # Emit the overlap step only once, at the first-half slot.
1381 if item.is_first_half:
1382 new_schedule[rank].append(item.obf_step)
1383 else:
1384 new_schedule[rank].append(item)
1385 else:
1386 current_items[rank] = None
1388 _process_rank_items(
1389 real_stage_num, current_items, insert_step_comms, new_schedule,
1390 )
1392 return new_schedule
1395_P2P_STEP_TYPES = frozenset({
1396 MetaStepType.FWD_SEND, MetaStepType.FWD_RECV,
1397 MetaStepType.BWD_SEND, MetaStepType.BWD_RECV,
1398})
1401def coalesce_p2p(exec_order):
1402 """Coalesce maximal contiguous runs of >=2 P2P steps into BATCH_SEND_RECV.
1404 A *run* is a maximal sequence of consecutive ``FWD_SEND`` / ``FWD_RECV`` /
1405 ``BWD_SEND`` / ``BWD_RECV`` steps with no compute / overlap / bubble (``None``)
1406 step between them — so no recv in the run is consumed before the batch is
1407 issued, and all sends' data is already produced. Each such run is replaced
1408 by a single :class:`MetaStep` of type ``BATCH_SEND_RECV`` carrying the run as
1409 ``sub_steps`` (order preserved, so per-direction FIFO is kept); the runtime
1410 groups those sub-steps by peer and issues one ``batch_isend_irecv`` per peer
1411 (same-peer send+recv -> duplex). Runs of length 1 are left untouched (the
1412 per-op batched path still batches them, so every transfer is still
1413 batch-vs-batch). Pure ``exec_order -> exec_order`` transform.
1415 Args:
1416 exec_order: ``{rank: [MetaStep | None, ...]}``.
1418 Returns:
1419 A new ``{rank: [...]}`` with contiguous P2P runs coalesced.
1420 """
1421 def _flush(run, new):
1422 if len(run) >= 2:
1423 new.append(MetaStep(None, MetaStepType.BATCH_SEND_RECV, None, sub_steps=tuple(run)))
1424 else:
1425 new.extend(run)
1427 out = {}
1428 for rank, order in exec_order.items():
1429 new = []
1430 run = []
1431 for step in order:
1432 if step is not None and step.type in _P2P_STEP_TYPES:
1433 run.append(step)
1434 continue
1435 _flush(run, new)
1436 run = []
1437 new.append(step)
1438 _flush(run, new)
1439 out[rank] = new
1440 return out
1443_RECV_STEP_TYPES = frozenset({MetaStepType.FWD_RECV, MetaStepType.BWD_RECV})
1446def attach_fwd_boundary_p2p(exec_order):
1447 """Hang each overlap gap's boundary-safe P2P on the OVERLAP_B_F step.
1449 For every ``OVERLAP_B_F`` step, the contiguous P2P run right after it is
1450 split by data readiness at the overlap's fwd/bwd boundary (forward is the
1451 short side; backward, ~2x FLOPs, is the long pole):
1453 * the forward's own ``FWD_SEND`` — its payload exists the moment the
1454 forward sub-step finishes, no need to wait out the backward;
1455 * every ``FWD_RECV`` / ``BWD_RECV`` — no data dependency at all;
1457 are removed from the gap and attached to the OVERLAP step as
1458 ``boundary_p2p`` (order: ``F_SEND`` first, then the recvs in original
1459 order), to be issued by :meth:`PipelineScheduleRuntime.exec_boundary_p2p`
1460 at the boundary, while the backward is still running. ``BWD_SEND`` (its
1461 grad is produced by that backward) and any send not produced by this
1462 overlap's forward stay in the gap.
1464 Pairing shape (the reason this composition is safe where naive
1465 hoist+coalesce hung): with every op issued as a per-op solo batch, each
1466 pair's per-pair batch sequence per slot is ``[F_SEND, B_RECV]`` on the
1467 prev end and ``[F_RECV, B_SEND]`` on the next end — complementary at every
1468 position and equal in count, so it matches under both per-direction FIFO
1469 and per-pair shape-mirroring semantics. Per-direction FIFO data order is
1470 preserved (each direction's ops keep their relative order; they all shift
1471 by the same amount). Pure ``exec_order -> exec_order`` transform.
1473 Args:
1474 exec_order: ``{rank: [MetaStep | None, ...]}``.
1476 Returns:
1477 A new ``{rank: [...]}`` with boundary P2P attached to OVERLAP steps.
1478 """
1479 out = {}
1480 for rank, order in exec_order.items():
1481 new = []
1482 i = 0
1483 while i < len(order):
1484 step = order[i]
1485 if (step is None or step.type != MetaStepType.OVERLAP_B_F
1486 or not step.sub_steps):
1487 new.append(step)
1488 i += 1
1489 continue
1490 run, j = _p2p_run_after(order, i + 1)
1491 boundary, leftover = _split_boundary_run(step, run)
1492 if not boundary:
1493 new.append(step)
1494 i += 1
1495 continue
1496 new.append(MetaStep(step.micro_index, step.type, step.stage_index,
1497 sub_steps=step.sub_steps, boundary_p2p=boundary))
1498 new.extend(leftover)
1499 i = j
1500 out[rank] = new
1501 return out
1504def _p2p_run_after(order, start):
1505 """Collect the contiguous P2P run starting at ``start``.
1507 Returns ``(run, end)`` where ``end`` is the index of the first step past
1508 the run (a compute step, ``None`` bubble, or end of order).
1509 """
1510 run = []
1511 j = start
1512 while j < len(order) and order[j] is not None and order[j].type in _P2P_STEP_TYPES:
1513 run.append(order[j])
1514 j += 1
1515 return run, j
1518def _split_boundary_run(step, run):
1519 """Split an overlap's trailing P2P run by fwd/bwd-boundary data readiness.
1521 Returns ``(boundary, leftover)``: ``boundary`` holds the overlap's own
1522 forward ``FWD_SEND`` (payload ready at the boundary) first, then every
1523 recv (no data dependency), keeping original order; ``leftover`` keeps the
1524 sends produced by the still-running backward, in place.
1525 """
1526 fwd_sub = next((s for s in step.sub_steps
1527 if s.type == MetaStepType.FWD), None)
1529 def _is_own_fwd_send(s):
1530 return (fwd_sub is not None
1531 and s.type == MetaStepType.FWD_SEND
1532 and s.stage_index == fwd_sub.stage_index
1533 and s.micro_index == fwd_sub.micro_index)
1535 boundary = ([s for s in run if _is_own_fwd_send(s)]
1536 + [s for s in run if s.type in _RECV_STEP_TYPES])
1537 taken = {id(s) for s in boundary}
1538 leftover = [s for s in run if id(s) not in taken]
1539 return tuple(boundary), leftover
1542def split_overlap_dxdw(exec_order: dict) -> dict:
1543 """Split each OVERLAP_B_F backward into dx (in the pair) + dw (after the gap).
1545 Rewrites ``(BWD, FWD)`` sub_steps to ``(BWD_INPUT, FWD)`` and inserts the
1546 matching ``BWD_WEIGHT`` after the contiguous P2P run that follows the
1547 overlap. The overlap then joins at ``max(dx, fwd)`` instead of
1548 ``max(dx + dw, fwd)``, so the gap's ``BWD_SEND`` (dx already wrote its
1549 grad to ``bwd_cache``) and the next slot's recvs are issued a dw earlier;
1550 under ``overlap_p2p`` they are async and dw computes while they fly.
1552 Comm placement is untouched: the pass runs after ``add_send_recv`` and
1553 only moves local compute, so the cross-rank matching order is identical,
1554 and the P2P run stays contiguous (dw lands after it, not inside) so
1555 ``coalesce_p2p`` / ``attach_fwd_boundary_p2p`` see the same gap shape.
1557 First-stage (``stage_index == 0``) backwards stay unified: their dx is a
1558 no-op (no input grad to compute or send), so splitting would only move
1559 the whole backward out of the overlap and lose its fwd overlap.
1561 Args:
1562 exec_order: ``{rank: [MetaStep | None, ...]}``.
1564 Returns:
1565 A new ``{rank: [...]}`` with overlap backwards split into dx/dw.
1566 """
1567 out = {}
1568 for rank, order in exec_order.items():
1569 new = []
1570 i = 0
1571 n = len(order)
1572 while i < n:
1573 step = order[i]
1574 i += 1
1575 if (step is None or step.type != MetaStepType.OVERLAP_B_F
1576 or not step.sub_steps):
1577 new.append(step)
1578 continue
1579 bwd_sub, fwd_sub = step.sub_steps
1580 if bwd_sub.type != MetaStepType.BWD or bwd_sub.stage_index == 0:
1581 new.append(step)
1582 continue
1583 dx = MetaStep(bwd_sub.micro_index, MetaStepType.BWD_INPUT, bwd_sub.stage_index)
1584 new.append(MetaStep(step.micro_index, step.type, step.stage_index,
1585 sub_steps=(dx, fwd_sub)))
1586 run, i = _p2p_run_after(order, i)
1587 new.extend(run)
1588 new.append(MetaStep(bwd_sub.micro_index, MetaStepType.BWD_WEIGHT,
1589 bwd_sub.stage_index))
1590 out[rank] = new
1591 return out
1594def add_send_recv(scheduler, stage_num, real_stage_num, style='loop'):
1595 """Insert P2P send/recv operations into a per-rank compute schedule.
1597 For each FWD or BWD step that requires cross-rank communication, a
1598 ``FWD_SEND`` / ``BWD_SEND`` is appended to the sender's schedule and a
1599 ``FWD_RECV`` / ``BWD_RECV`` is appended to the receiver's schedule.
1601 ``OVERLAP_F_B`` / ``OVERLAP_B_F`` composite steps are expanded into
1602 **two** virtual time slots during the column scan so that the RECV
1603 triggered by the **second** sub-step lands in the receiver's schedule
1604 one slot later — matching the fact that the sender can only finish
1605 emitting the second sub-step's output after the first completes.
1607 Even ranks are processed before odd ranks at each time step to avoid
1608 P2P deadlocks between adjacent ranks.
1610 The resulting per-gap op order (steady state:
1611 ``[B_RECV, B_SEND, F_RECV, F_SEND]``) is LOAD-BEARING, not cosmetic.
1612 Each adjacent rank pair shares one comm (HCCL split-by-group) on which
1613 plain send/recv execute in queue order, and this layout makes one end of
1614 every pair recv-first while the other is send-first, so the two queue
1615 heads always match (recv<->send), then the tails match (send<->recv).
1616 Any later pass that reorders P2P ops relative to EACH OTHER breaks this:
1617 a (since removed) hoist variant that moved recvs across sends made both
1618 ends recv-first and deadlocked on hardware (2026-06).
1619 ``attach_fwd_boundary_p2p`` (the ``"boundary"`` transport) is safe: it keeps
1620 every per-direction FIFO and runs on the batch transport with per-op solo
1621 batches, whose per-pair sequences stay complementary.
1623 Args:
1624 scheduler: ``{rank: [MetaStep | None, ...]}`` — compute schedule
1625 with ``None`` for bubble slots.
1626 stage_num: Total number of virtual pipeline stages.
1627 real_stage_num: Number of physical ranks.
1628 style: Topology mapping — ``'loop'`` or ``'v'``.
1630 Returns:
1631 ``{rank: [MetaStep, ...]}`` — schedule with communication ops inserted.
1632 """
1634 def stage_to_rank(stage_index: int) -> int:
1635 """Map a virtual stage index to its physical rank."""
1636 if style == 'loop':
1637 return stage_index % real_stage_num
1638 if style == 'v':
1639 if stage_index < real_stage_num:
1640 return stage_index
1641 return stage_num - 1 - stage_index
1642 raise ValueError(f"Argument 'style' must be 'loop' or 'v', but got {style!r}.")
1644 def _fwd_peer(stage_index: int):
1645 """Return the rank that receives this stage's forward output, or None."""
1646 if stage_index >= stage_num - 1:
1647 return None
1648 peer = stage_to_rank(stage_index + 1)
1649 return peer if peer != stage_to_rank(stage_index) else None
1651 def _bwd_peer(stage_index: int):
1652 """Return the rank that receives this stage's backward gradient, or None."""
1653 if stage_index <= 0:
1654 return None
1655 peer = stage_to_rank(stage_index - 1)
1656 return peer if peer != stage_to_rank(stage_index) else None
1658 def _insert_comms_for_step(step, rank, new_schedule):
1659 """Insert send/recv for a single FWD, BWD, or composite OVERLAP step.
1661 Mirrors ``Schedule1F1B``: every forward that crosses a rank boundary
1662 gets a ``DATA_SEND`` / ``DATA_RECV`` pair immediately followed by
1663 ``FWD_SEND`` / ``FWD_RECV``. DATA_SEND/RECV are no-ops when
1664 ``run_with_dataiterator`` isn't active (the handler short-circuits
1665 on ``self.data_iterator is None``), so emitting them unconditionally
1666 keeps Interleaved 1F1B correct under ``data.load: single`` without
1667 penalising the legacy preload path.
1668 """
1669 if step is None:
1670 return
1672 if step.type == MetaStepType.FWD:
1673 peer = _fwd_peer(step.stage_index)
1674 if peer is not None:
1675 new_schedule[rank].append(
1676 MetaStep(step.micro_index, MetaStepType.DATA_SEND, step.stage_index))
1677 new_schedule[rank].append(
1678 MetaStep(step.micro_index, MetaStepType.FWD_SEND, step.stage_index))
1679 new_schedule[peer].append(
1680 MetaStep(step.micro_index, MetaStepType.DATA_RECV, step.stage_index + 1))
1681 new_schedule[peer].append(
1682 MetaStep(step.micro_index, MetaStepType.FWD_RECV, step.stage_index + 1))
1684 elif step.type == MetaStepType.BWD:
1685 peer = _bwd_peer(step.stage_index)
1686 if peer is not None:
1687 new_schedule[rank].append(
1688 MetaStep(step.micro_index, MetaStepType.BWD_SEND, step.stage_index))
1689 new_schedule[peer].append(
1690 MetaStep(step.micro_index, MetaStepType.BWD_RECV, step.stage_index - 1))
1692 elif step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F) and step.sub_steps:
1693 for sub in step.sub_steps:
1694 _insert_comms_for_step(sub, rank, new_schedule)
1696 # --- Main logic: expand OVERLAP steps into 2 virtual slots, then scan ---
1697 expanded = _expand_overlap_slots(scheduler, real_stage_num)
1698 return _insert_dataload_before_fwd(
1699 _column_scan_insert_comms(expanded, real_stage_num, _insert_comms_for_step))
1702_ALIGN_PAD = object()
1703"""Sentinel marking a forced 1F1B-boundary bubble produced during alignment."""
1706def _step_dep_ready(step, rank, t, done, stage_num, stage_to_rank):
1707 """Cross-rank data dependency check used by the alignment simulator.
1709 A FWD step at stage ``s`` depends on FWD at stage ``s-1`` (on a
1710 different rank); BWD at stage ``s`` depends on BWD at stage ``s+1``.
1711 Steps at boundaries or whose producer lives on the same rank are
1712 always ready.
1713 """
1714 si, mi = step.stage_index, step.micro_index
1715 if step.type == MetaStepType.FWD:
1716 if si == 0 or stage_to_rank(si - 1) == rank:
1717 return True
1718 key = (MetaStepType.FWD, si - 1, mi)
1719 return key in done and done[key] < t
1720 if step.type == MetaStepType.BWD:
1721 if si == stage_num - 1 or stage_to_rank(si + 1) == rank:
1722 return True
1723 key = (MetaStepType.BWD, si + 1, mi)
1724 return key in done and done[key] < t
1725 return True
1728def _simulate_aligned_schedule(padded, stage_num, real_stage_num, stage_to_rank):
1729 """Simulate execution time-step by time-step, inserting bubbles where
1730 a step is not yet ready (cross-rank dep) or where the cooldown
1731 rhythm requires it.
1733 Args:
1734 padded: ``{rank: [step | _ALIGN_PAD | None, ...]}`` after
1735 1F1B-boundary padding.
1736 stage_num: Total number of virtual pipeline stages.
1737 real_stage_num: Number of physical ranks.
1738 stage_to_rank: Topology mapping from stage to rank.
1740 Returns:
1741 ``{rank: [step | None, ...]}`` ready for the column-scan SEND/RECV
1742 insertion phase.
1743 """
1744 remaining_fwd = {
1745 rank: sum(
1746 1 for s in padded[rank]
1747 if s is not _ALIGN_PAD and s is not None and s.type == MetaStepType.FWD
1748 )
1749 for rank in range(real_stage_num)
1750 }
1751 cursors = {r: 0 for r in range(real_stage_num)}
1752 aligned = {r: [] for r in range(real_stage_num)}
1753 done = {}
1754 last_was_cooldown_bwd = {r: False for r in range(real_stage_num)}
1755 max_t = sum(len(v) for v in padded.values()) + real_stage_num * 20
1757 def _emit_bubble(rank):
1758 aligned[rank].append(None)
1759 last_was_cooldown_bwd[rank] = False
1761 def _emit_step(rank, step, t, in_cooldown):
1762 aligned[rank].append(step)
1763 done[(step.type, step.stage_index, step.micro_index)] = t
1764 cursors[rank] += 1
1765 if step.type == MetaStepType.FWD:
1766 remaining_fwd[rank] -= 1
1767 last_was_cooldown_bwd[rank] = in_cooldown and step.type == MetaStepType.BWD
1769 def _step_rank_at(t, rank):
1770 if cursors[rank] >= len(padded[rank]):
1771 return
1772 item = padded[rank][cursors[rank]]
1773 if item is _ALIGN_PAD:
1774 _emit_bubble(rank)
1775 cursors[rank] += 1
1776 return
1777 in_cooldown = remaining_fwd[rank] == 0
1778 # Cooldown rhythm: alternate None / BWD in pure-BWD phase.
1779 cooldown_skip = (
1780 in_cooldown
1781 and item.type == MetaStepType.BWD
1782 and last_was_cooldown_bwd[rank]
1783 )
1784 if cooldown_skip:
1785 _emit_bubble(rank)
1786 return
1787 if not _step_dep_ready(item, rank, t, done, stage_num, stage_to_rank):
1788 _emit_bubble(rank)
1789 return
1790 _emit_step(rank, item, t, in_cooldown)
1792 for t in range(max_t):
1793 if all(cursors[r] >= len(padded[r]) for r in range(real_stage_num)):
1794 break
1795 for rank in range(real_stage_num):
1796 _step_rank_at(t, rank)
1797 return aligned
1800def auto_align_and_add_send_recv(scheduler, stage_num, real_stage_num, style='loop'):
1801 """Auto-insert bubble alignment and P2P send/recv into a pure-compute schedule.
1803 Unlike :func:`add_send_recv` which requires the caller to pre-insert
1804 ``None`` bubble slots for time-step alignment, this function accepts a
1805 **pure compute order** (``FWD`` / ``BWD`` only, no ``None`` needed) and
1806 automatically determines bubble placement via execution simulation.
1808 Three constraints are enforced:
1810 1. **Data dependency** — a ``FWD(stage_k)`` cannot execute until
1811 ``FWD(stage_{k-1})`` on its source rank has completed (and
1812 analogously for ``BWD``).
1813 2. **1F1B transition alignment** — ``real_stage_num - 1 - rank`` padding
1814 slots are inserted at the warmup → 1F1B boundary (detected as the
1815 first ``FWD`` immediately followed by a ``BWD`` in the compute order)
1816 so that all ranks enter the 1F1B steady state in lockstep.
1817 3. **Cooldown rhythm** — once a rank exhausts its ``FWD`` ops and enters
1818 pure-``BWD`` cooldown, consecutive ``BWD`` steps are separated by a
1819 ``None`` slot, maintaining the column-phase-sync property (no rank
1820 does ``BWD`` while another does ``FWD`` at the same time step).
1822 After alignment, a column-scan pass inserts ``FWD_SEND`` / ``FWD_RECV``
1823 and ``BWD_SEND`` / ``BWD_RECV`` with the same prefetch semantics as
1824 :func:`add_send_recv`.
1826 Args:
1827 scheduler: ``{rank: [MetaStep, ...]}`` — pure compute schedule.
1828 ``None`` entries are silently stripped before processing.
1829 stage_num: Total number of virtual pipeline stages.
1830 real_stage_num: Number of physical ranks.
1831 style: Topology mapping — ``'loop'`` or ``'v'``.
1833 Returns:
1834 ``{rank: [MetaStep, ...]}`` — fully aligned schedule with bubbles
1835 and communication ops inserted.
1836 """
1838 # ---- topology helpers (shared with column-scan phase) ----
1840 def stage_to_rank(stage_index: int) -> int:
1841 if style == 'loop':
1842 return stage_index % real_stage_num
1843 if style == 'v':
1844 if stage_index < real_stage_num:
1845 return stage_index
1846 return stage_num - 1 - stage_index
1847 raise ValueError(f"Argument 'style' must be 'loop' or 'v', but got {style!r}.")
1849 def _fwd_peer(stage_index: int):
1850 if stage_index >= stage_num - 1:
1851 return None
1852 peer = stage_to_rank(stage_index + 1)
1853 return peer if peer != stage_to_rank(stage_index) else None
1855 def _bwd_peer(stage_index: int):
1856 if stage_index <= 0:
1857 return None
1858 peer = stage_to_rank(stage_index - 1)
1859 return peer if peer != stage_to_rank(stage_index) else None
1861 # ---- Phase 1: strip None, detect 1F1B boundary, insert transition padding ----
1863 def _find_1f1b_boundary(order):
1864 """Index of the first FWD followed by BWD; ``len(order)`` if absent."""
1865 for i in range(len(order) - 1):
1866 if (order[i].type == MetaStepType.FWD
1867 and order[i + 1].type == MetaStepType.BWD):
1868 return i
1869 return len(order)
1871 padded = {}
1872 for rank in range(real_stage_num):
1873 order = [s for s in scheduler[rank] if s is not None]
1874 boundary = _find_1f1b_boundary(order)
1875 pad_count = real_stage_num - 1 - rank
1876 padded[rank] = order[:boundary] + [_ALIGN_PAD] * pad_count + order[boundary:]
1878 # ---- Phase 2: simulate execution with data deps + cooldown rhythm ----
1880 aligned = _simulate_aligned_schedule(padded, stage_num, real_stage_num, stage_to_rank)
1882 # ---- Phase 3: column-scan SEND/RECV insertion (same as add_send_recv) ----
1884 def _insert_comms_for_step(step, rank, new_schedule):
1885 # The DATA pair precedes the FWD pair so ``data.load: single`` ships raw
1886 # inputs stage-to-stage; it no-ops without run_with_dataiterator.
1887 if step is None:
1888 return
1889 if step.type == MetaStepType.FWD:
1890 peer = _fwd_peer(step.stage_index)
1891 if peer is not None:
1892 new_schedule[rank].append(
1893 MetaStep(step.micro_index, MetaStepType.DATA_SEND, step.stage_index))
1894 new_schedule[rank].append(
1895 MetaStep(step.micro_index, MetaStepType.FWD_SEND, step.stage_index))
1896 new_schedule[peer].append(
1897 MetaStep(step.micro_index, MetaStepType.DATA_RECV, step.stage_index + 1))
1898 new_schedule[peer].append(
1899 MetaStep(step.micro_index, MetaStepType.FWD_RECV, step.stage_index + 1))
1900 elif step.type == MetaStepType.BWD:
1901 peer = _bwd_peer(step.stage_index)
1902 if peer is not None:
1903 new_schedule[rank].append(
1904 MetaStep(step.micro_index, MetaStepType.BWD_SEND, step.stage_index))
1905 new_schedule[peer].append(
1906 MetaStep(step.micro_index, MetaStepType.BWD_RECV, step.stage_index - 1))
1907 elif step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F) and step.sub_steps:
1908 for sub in step.sub_steps:
1909 _insert_comms_for_step(sub, rank, new_schedule)
1911 # Expand OVERLAP steps into 2 virtual slots before the column scan so
1912 # the RECV triggered by an overlap's second sub-step lands one slot
1913 # later on the receiver — matching the fact that the sender can only
1914 # finish emitting the second sub-step after the first completes.
1915 expanded = _expand_overlap_slots(aligned, real_stage_num)
1916 return _insert_dataload_before_fwd(
1917 _column_scan_insert_comms(expanded, real_stage_num, _insert_comms_for_step))
1920class ScheduleGPipe(PipelineScheduleRuntime):
1921 """
1922 The Gpipe schedule.
1923 It first executes all forward micro batches and then execute all backward micro batches.
1924 """
1925 def __init__(self,
1926 stages,
1927 micro_batch_num,
1928 args_batch_dim=None,
1929 kwargs_batch_dim=None,
1930 output_concat_dim=None,
1931 swap=False):
1932 super().__init__(stages,
1933 micro_batch_num,
1934 args_batch_dim=args_batch_dim,
1935 kwargs_batch_dim=kwargs_batch_dim,
1936 output_concat_dim=output_concat_dim,
1937 swap=swap)
1938 self.build_exec_order()
1940 def _build_stage_to_rank_index(self) -> None:
1941 self._stage_to_rank_index = generate_stage_to_rank_mapping(
1942 self.real_stage_num, self._stage_num, style='loop'
1943 )
1945 def construct_exec_order(self):
1946 """construct_exec_order of Gpipe."""
1947 for stage_index in range(self.real_stage_num):
1948 order_list = []
1949 for mb_index in range(self.micro_batch_num):
1950 order_list.append(MetaStep(mb_index, MetaStepType.DATA_LOAD, stage_index))
1951 if stage_index != 0:
1952 order_list.append(MetaStep(mb_index, MetaStepType.DATA_RECV, stage_index))
1953 order_list.append(MetaStep(mb_index, MetaStepType.FWD_RECV, stage_index))
1954 order_list.append(MetaStep(mb_index, MetaStepType.FWD, stage_index))
1955 if stage_index != self.real_stage_num - 1:
1956 order_list.append(MetaStep(mb_index, MetaStepType.DATA_SEND, stage_index))
1957 order_list.append(MetaStep(mb_index, MetaStepType.FWD_SEND, stage_index))
1958 for mb_index in range(self.micro_batch_num):
1959 if stage_index != self.real_stage_num - 1:
1960 order_list.append(MetaStep(mb_index, MetaStepType.BWD_RECV, stage_index))
1961 order_list.append(MetaStep(mb_index, MetaStepType.BWD, stage_index))
1962 if stage_index != 0:
1963 order_list.append(MetaStep(mb_index, MetaStepType.BWD_SEND, stage_index))
1964 self.exec_order[stage_index] = order_list
1967class Schedule1F1B(PipelineScheduleRuntime):
1968 """
1969 The 1F1B schedule.
1970 It will perform one forward and one backward on the micro batches in steady state.
1971 """
1972 def __init__(self,
1973 stages,
1974 micro_batch_num,
1975 args_batch_dim=None,
1976 kwargs_batch_dim=None,
1977 output_concat_dim=None,
1978 swap=False):
1979 super().__init__(stages,
1980 micro_batch_num,
1981 args_batch_dim=args_batch_dim,
1982 kwargs_batch_dim=kwargs_batch_dim,
1983 output_concat_dim=output_concat_dim,
1984 swap=swap)
1985 self.build_exec_order()
1987 def _build_stage_to_rank_index(self) -> None:
1988 self._stage_to_rank_index = generate_stage_to_rank_mapping(
1989 self.real_stage_num, self._stage_num, style='loop'
1990 )
1992 def construct_exec_order(self):
1993 """construct_exec_order of 1F1B."""
1994 for stage_index in range(self.real_stage_num):
1995 order_list = []
1996 fwd_index = 0
1997 bwd_index = 0
1998 # warmup phase
1999 warmup_micro_batches = min(self.real_stage_num - stage_index, self.micro_batch_num)
2000 for _ in range(warmup_micro_batches):
2001 order_list.append(MetaStep(fwd_index, MetaStepType.DATA_LOAD, stage_index))
2002 if stage_index != 0:
2003 order_list.append(MetaStep(fwd_index, MetaStepType.DATA_RECV, stage_index))
2004 order_list.append(MetaStep(fwd_index, MetaStepType.FWD_RECV, stage_index))
2005 if stage_index % 2 == 0:
2006 order_list.append(MetaStep(fwd_index, MetaStepType.FWD, stage_index))
2007 if fwd_index != warmup_micro_batches - 1:
2008 order_list.append(MetaStep(fwd_index, MetaStepType.DATA_SEND, stage_index))
2009 order_list.append(MetaStep(fwd_index, MetaStepType.FWD_SEND, stage_index))
2010 else:
2011 if fwd_index > 0:
2012 order_list.append(MetaStep(fwd_index - 1, MetaStepType.DATA_SEND, stage_index))
2013 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
2014 order_list.append(MetaStep(fwd_index, MetaStepType.FWD, stage_index))
2015 fwd_index += 1
2017 # if warmup phase cannot filled up, then we need to execute fwd send in advance
2018 if self.real_stage_num - stage_index > self.micro_batch_num:
2019 order_list.append(MetaStep(fwd_index - 1, MetaStepType.DATA_SEND, stage_index))
2020 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
2021 fwd_index += 1
2022 # steady phase
2023 steady_micro_batches = self.micro_batch_num - warmup_micro_batches
2024 for _ in range(steady_micro_batches):
2025 if stage_index != self.real_stage_num - 1:
2026 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_RECV, stage_index))
2027 order_list.append(MetaStep(fwd_index - 1, MetaStepType.DATA_SEND, stage_index))
2028 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
2029 order_list.append(MetaStep(bwd_index, MetaStepType.BWD, stage_index))
2031 order_list.append(MetaStep(fwd_index, MetaStepType.DATA_LOAD, stage_index))
2032 if stage_index != 0:
2033 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_SEND, stage_index))
2034 order_list.append(MetaStep(fwd_index, MetaStepType.DATA_RECV, stage_index))
2035 order_list.append(MetaStep(fwd_index, MetaStepType.FWD_RECV, stage_index))
2036 order_list.append(MetaStep(fwd_index, MetaStepType.FWD, stage_index))
2037 fwd_index += 1
2038 bwd_index += 1
2040 # cooldown phase
2041 cooldown_micro_batches = warmup_micro_batches
2042 for _ in range(cooldown_micro_batches):
2043 if stage_index != self.real_stage_num - 1:
2044 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_RECV, stage_index))
2045 if bwd_index == self.micro_batch_num - warmup_micro_batches and fwd_index <= self.micro_batch_num:
2046 order_list.append(MetaStep(fwd_index - 1, MetaStepType.DATA_SEND, stage_index))
2047 order_list.append(MetaStep(fwd_index - 1, MetaStepType.FWD_SEND, stage_index))
2048 order_list.append(MetaStep(bwd_index, MetaStepType.BWD, stage_index))
2050 if stage_index != 0:
2051 order_list.append(MetaStep(bwd_index, MetaStepType.BWD_SEND, stage_index))
2052 bwd_index += 1
2053 self.exec_order[stage_index] = order_list
2056class ScheduleInterleaved1F1B(PipelineScheduleRuntime):
2057 """The Interleaved 1F1B schedule.
2059 Supports multiple stages per rank. In steady state, performs one
2060 forward followed by one backward on each micro-batch. Handles the
2061 cases where ``micro_batch_num`` is less than, equal to, or greater
2062 than the stage count, including non-evenly-divisible micro counts.
2064 Two orthogonal overlap modes can be enabled via constructor flags:
2066 * ``overlap_p2p=True``: defer P2P recv ``handle.wait()`` until the
2067 consuming FWD/BWD step (or the OVERLAP_B_F callback when
2068 ``overlap_b_f=True``), letting recv overlap with prior compute.
2069 * ``overlap_b_f=True``: in the 1F1B steady state, pair consecutive
2070 ``(B_i, F_{i+1})`` steps into ``OVERLAP_B_F`` composite steps so
2071 a registered callback can drive comm/compute overlap (typically
2072 via :class:`CommComputeOverlap` for MoE EP A2A). Users register
2073 the callback through :meth:`register_custom_function`.
2074 * ``enable_dxdw_split=True`` (requires ``overlap_b_f``): each
2075 steady-state pair becomes ``(BWD_INPUT_i, F_{i+1})`` and the
2076 ``BWD_WEIGHT_i`` runs as its own step after the pair's P2P gap
2077 (see :func:`split_overlap_dxdw`), so the input-grad send is
2078 issued once dx and the paired forward finish instead of waiting
2079 out the full backward.
2081 The two overlap flags are independent and can be combined.
2083 ``less_memory=True`` selects the shallow-warmup variant: each stage's
2084 warmup depth becomes ``(V-1)*R + (P-1-i)`` instead of the classic
2085 ``(V-1)*R + 2*(P-1-i)``, so at most ``V*P - i`` micro-batch
2086 activations are in flight per stage instead of ``2*P - 2*i - 1`` extra
2087 on the early stages. At ``V == 1`` this reproduces the plain 1F1B
2088 depth exactly. The steady-state 1F1B rhythm and the cooldown are
2089 unchanged; only the warmup/cooldown lengths shift.
2091 Example:
2092 >>> # Plain interleaved 1F1B
2093 >>> sched = ScheduleInterleaved1F1B(stages, 8)
2094 >>> # With B/F overlap (dual-pipe-style comm/compute overlap)
2095 >>> sched = ScheduleInterleaved1F1B(stages, 8, overlap_b_f=True)
2096 >>> sched.register_custom_function(MetaStepType.OVERLAP_B_F, callback)
2097 """
2099 # Class-level default so schedules built through the pure-construction
2100 # path (``object.__new__`` in offline tests) keep the classic warmup.
2101 _less_memory = False
2103 def __init__(self,
2104 stages,
2105 micro_batch_num,
2106 args_batch_dim=None,
2107 kwargs_batch_dim=None,
2108 output_concat_dim=None,
2109 overlap_p2p=False,
2110 overlap_b_f=False,
2111 swap=False,
2112 enable_dxdw_split=False,
2113 p2p_transport="auto",
2114 less_memory=False):
2115 super().__init__(stages,
2116 micro_batch_num,
2117 args_batch_dim=args_batch_dim,
2118 kwargs_batch_dim=kwargs_batch_dim,
2119 output_concat_dim=output_concat_dim,
2120 overlap_p2p=overlap_p2p,
2121 swap=swap,
2122 p2p_transport=p2p_transport)
2123 # Selects plain F/B vs OVERLAP_B_F pairing in the steady state; must
2124 # be set before ``construct_stage_exec_order`` below.
2125 self._overlap_b_f = overlap_b_f
2126 # Shallow-warmup variant; must be set before build_exec_order below.
2127 self._less_memory = less_memory
2128 # dx/dw split: ``construct_exec_order`` rewrites each steady-state
2129 # OVERLAP_B_F pair to ``(BWD_INPUT, FWD)`` and re-emits the matching
2130 # BWD_WEIGHT after the pair's P2P gap (``split_overlap_dxdw``), so the
2131 # input-grad send leaves at ``max(dx, fwd)`` and dw overlaps with the
2132 # in-flight P2P instead of delaying it.
2133 self._enable_dxdw_split = enable_dxdw_split
2134 if enable_dxdw_split and not overlap_b_f:
2135 raise ValueError(
2136 "enable_dxdw_split=True requires overlap_b_f=True; the split "
2137 "is only applied to BWD sub-steps inside OVERLAP_B_F composite steps."
2138 )
2139 self._init_round_layout()
2140 self.build_exec_order()
2142 def _init_round_layout(self):
2143 """Compute per-round micro-batch counts used by stage-order emission.
2145 Populates ``n_rounds``, ``n_microbatch_per_round`` and its prefix-sum
2146 ``n_microbatch_per_round_accu`` from ``micro_batch_num``,
2147 ``real_stage_num`` and ``n_local_stages``. Factored out of
2148 ``__init__`` so the pure schedule-construction path (used by offline
2149 unit tests) can be exercised without instantiating stages.
2150 """
2151 self.n_rounds = max(1, self.micro_batch_num // self.real_stage_num)
2152 if self.micro_batch_num < self.real_stage_num:
2153 base = self.micro_batch_num - self.real_stage_num
2154 remainder = 0
2155 else:
2156 n_extra_microbatch = self.micro_batch_num % self.real_stage_num
2157 base = n_extra_microbatch // self.n_rounds
2158 remainder = n_extra_microbatch % self.n_rounds
2159 self.n_microbatch_per_round = \
2160 [self.real_stage_num + base + 1 if i < remainder else
2161 self.real_stage_num + base for i in range(self.n_rounds)]
2162 self.n_microbatch_per_round_accu = \
2163 [x * self.n_local_stages for x in itertools.accumulate(self.n_microbatch_per_round)]
2164 self.n_microbatch_per_round_accu.insert(0, 0)
2166 def construct_exec_order(self):
2167 # One line for the whole warmup vector, pp-rank-0 hosts only
2168 # (per-stage lines are too chatty at high rank counts).
2169 own = getattr(self, "stages", None)
2170 if not own or getattr(own[0], "stage_index", 0) == 0:
2171 logger.info("interleaved warmup: warmup_ops=%s less_memory=%s",
2172 [self.warmup_ops(s) for s in range(self.real_stage_num)],
2173 self._less_memory)
2174 for stage_index in range(self.real_stage_num):
2175 self.exec_order[stage_index] = self.construct_stage_exec_order(stage_index)
2176 self.exec_order = add_send_recv(self.exec_order, self._stage_num, self.real_stage_num, style='loop')
2177 if self.enable_dxdw_split:
2178 self.exec_order = split_overlap_dxdw(self.exec_order)
2180 def _build_stage_to_rank_index(self) -> None:
2181 self._stage_to_rank_index = generate_stage_to_rank_mapping(
2182 self.real_stage_num, self._stage_num, style='loop'
2183 )
2185 def warmup_ops(self, stage_index):
2186 """warmup phase."""
2187 warmup_ops_last_stage = (self.n_local_stages - 1) * self.n_microbatch_per_round[0]
2188 depth = self.real_stage_num - 1 - stage_index
2189 # less_memory: plain-1F1B stagger (depth) instead of the classic
2190 # doubled stagger (2*depth) — fewer activations in flight per stage.
2191 warmup_ops = warmup_ops_last_stage + (depth if self._less_memory else 2 * depth)
2192 return min(warmup_ops, self.micro_batch_num * self.n_local_stages)
2194 def forward_stage_index(self, op_index, stage_index):
2195 """obtain forward stage_index based on op_index."""
2196 accu_index = bisect.bisect_right(self.n_microbatch_per_round_accu, op_index) - 1
2197 local_index = (op_index - self.n_microbatch_per_round_accu[accu_index]) // \
2198 self.n_microbatch_per_round[accu_index]
2199 return (local_index * self.real_stage_num) + stage_index
2201 def backward_stage_index(self, op_index, stage_index):
2202 """obtain backward stage_index based on op_index."""
2203 accu_index = bisect.bisect_right(self.n_microbatch_per_round_accu, op_index) - 1
2204 local_index = (op_index - self.n_microbatch_per_round_accu[accu_index]) // \
2205 self.n_microbatch_per_round[accu_index]
2206 local_index = self.n_local_stages - 1 - local_index
2207 return (local_index * self.real_stage_num) + stage_index
2209 def _short_micro(self) -> bool:
2210 """True when ``micro_batch_num < real_stage_num`` (extra-bubble regime)."""
2211 return self.micro_batch_num < self.real_stage_num
2213 def _trailing_bubble(self) -> int:
2214 """Bubble count appended after a BWD with ``micro == micro_batch_num - 1``
2215 in the short-micro regime.
2216 """
2217 return self.real_stage_num - self.micro_batch_num
2219 def _emit_warmup_ops(self, stage_index, warmup_ops, fwd_stage_micro_index):
2220 """Emit pure-FWD warmup ops with optional short-micro bubble padding."""
2221 ops = []
2222 short = self._short_micro()
2223 last_micro = self.micro_batch_num - 1
2224 last_stage = self.real_stage_num - 1
2225 bubble = self._trailing_bubble()
2226 for op_idx in range(warmup_ops):
2227 fwd_stage_idx = self.forward_stage_index(op_idx, stage_index)
2228 fwd_micro_idx = fwd_stage_micro_index[fwd_stage_idx]
2229 # No DATA_LOAD slot here: it is spliced in after comm insertion, else this
2230 # rank's extra columns desync the cross-rank column scan in add_send_recv.
2231 ops.append(MetaStep(fwd_micro_idx, MetaStepType.FWD, fwd_stage_idx))
2232 need_pad = (
2233 short
2234 and fwd_micro_idx == last_micro
2235 and (op_idx != warmup_ops - 1 or stage_index == last_stage)
2236 )
2237 if need_pad:
2238 ops.extend([None] * bubble)
2239 fwd_stage_micro_index[fwd_stage_idx] += 1
2240 return ops
2242 def _emit_cooldown_ops(self, stage_index, warmup_ops, fwd_bwd_ops, total_ops,
2243 bwd_stage_micro_index):
2244 """Emit pure-BWD cooldown ops (each preceded by a bubble) with
2245 optional short-micro trailing padding.
2246 """
2247 ops = []
2248 short = self._short_micro()
2249 last_micro = self.micro_batch_num - 1
2250 # Double the bubble at each chunk's last-micro BWD: one ``bubble`` covers
2251 # the missing ``rs - micro`` micros, the second offsets the next chunk
2252 # by 2 slots so the wrap-around grad (rank 0 stage ``rs`` -> rank
2253 # last_stage stage ``rs - 1``) lands AFTER its producer in column-scan
2254 # time. Matches the +2 cooldown-rhythm offset that non-short Interleaved
2255 # 1F1B naturally has from extra 1F1B ops on rank last_stage.
2256 bubble = 2 * self._trailing_bubble()
2257 for op_idx in range(warmup_ops + fwd_bwd_ops, total_ops):
2258 ops.append(None)
2259 bwd_stage_idx = self.backward_stage_index(op_idx - warmup_ops, stage_index)
2260 bwd_micro_idx = bwd_stage_micro_index[bwd_stage_idx]
2261 ops.append(MetaStep(bwd_micro_idx, MetaStepType.BWD, bwd_stage_idx))
2262 if short and bwd_micro_idx == last_micro:
2263 ops.extend([None] * bubble)
2264 bwd_stage_micro_index[bwd_stage_idx] += 1
2265 return ops
2267 def _emit_1f1b_ops(self, stage_index, warmup_ops, fwd_bwd_ops,
2268 fwd_stage_micro_index, bwd_stage_micro_index):
2269 """Emit interleaved (FWD, BWD) pairs for the 1F1B steady-state phase."""
2270 ops = []
2271 short = self._short_micro()
2272 last_micro = self.micro_batch_num - 1
2273 last_stage = self.real_stage_num - 1
2274 # Double the bubble at the 1F1B->cooldown chunk boundary on rank
2275 # last_stage; see :meth:`_emit_cooldown_ops` for the alignment rationale.
2276 bubble = 2 * self._trailing_bubble()
2277 for op_idx in range(warmup_ops, warmup_ops + fwd_bwd_ops):
2278 fwd_stage_idx = self.forward_stage_index(op_idx, stage_index)
2279 fwd_micro_idx = fwd_stage_micro_index[fwd_stage_idx]
2280 # DATA_LOAD spliced in after comm insertion — see _emit_warmup_ops.
2281 ops.append(MetaStep(fwd_micro_idx, MetaStepType.FWD, fwd_stage_idx))
2282 fwd_stage_micro_index[fwd_stage_idx] += 1
2283 bwd_stage_idx = self.backward_stage_index(op_idx - warmup_ops, stage_index)
2284 bwd_micro_idx = bwd_stage_micro_index[bwd_stage_idx]
2285 ops.append(MetaStep(bwd_micro_idx, MetaStepType.BWD, bwd_stage_idx))
2286 need_pad = (
2287 short
2288 and bwd_micro_idx == last_micro
2289 and stage_index == last_stage
2290 )
2291 if need_pad:
2292 ops.extend([None] * bubble)
2293 bwd_stage_micro_index[bwd_stage_idx] += 1
2294 return ops
2296 @staticmethod
2297 def _collect_fwd_bwd_steps(emit_fwd, emit_bwd, fwd_bwd_ops, warmup_ops):
2298 """Walk the 1F1B range collecting parallel ``fwd_steps`` / ``bwd_steps``.
2300 ``emit_fwd(op_idx)`` and ``emit_bwd(op_idx)`` build a single
2301 :class:`MetaStep` and advance their respective per-stage micro
2302 counters as a side effect.
2303 """
2304 fwd_steps = []
2305 bwd_steps = []
2306 for op_idx in range(warmup_ops, warmup_ops + fwd_bwd_ops):
2307 fwd_steps.append(emit_fwd(op_idx))
2308 bwd_steps.append(emit_bwd(op_idx))
2309 return fwd_steps, bwd_steps
2311 @staticmethod
2312 def _pair_into_overlap_b_f(fwd_steps, bwd_steps):
2313 """Build ``F₁, [B_i, F_{i+1}], B_n`` ordering with OVERLAP_B_F pairs.
2315 ``sub_steps`` carry the ``(bwd, fwd)`` tuple — callbacks access
2316 them via ``step.sub_steps`` to recover per-direction stage /
2317 micro info.
2318 """
2319 ops = []
2320 if fwd_steps:
2321 ops.append(fwd_steps[0]) # F₁ runs alone
2322 for i in range(len(bwd_steps) - 1):
2323 ops.append(MetaStep(
2324 None, MetaStepType.OVERLAP_B_F, None,
2325 sub_steps=(bwd_steps[i], fwd_steps[i + 1]),
2326 ))
2327 if bwd_steps:
2328 ops.append(bwd_steps[-1]) # B_n runs alone
2329 return ops
2331 def _emit_1f1b_overlap_ops(self, stage_index, warmup_ops, fwd_bwd_ops,
2332 fwd_stage_micro_index, bwd_stage_micro_index):
2333 """Emit ``F₁, [B_i, F_{i+1}], B_n`` for the 1F1B phase under
2334 ``overlap_b_f=True``. Each ``[B_i, F_{i+1}]`` becomes an
2335 ``OVERLAP_B_F`` composite step; a registered callback drives the
2336 actual concurrent execution. Short-micro extra-bubble padding
2337 on the last rank is appended after ``B_n``.
2338 """
2339 def emit_fwd(op_idx):
2340 fwd_si = self.forward_stage_index(op_idx, stage_index)
2341 fwd_mi = fwd_stage_micro_index[fwd_si]
2342 fwd_stage_micro_index[fwd_si] += 1
2343 return MetaStep(fwd_mi, MetaStepType.FWD, fwd_si)
2345 def emit_bwd(op_idx):
2346 bwd_si = self.backward_stage_index(op_idx - warmup_ops, stage_index)
2347 bwd_mi = bwd_stage_micro_index[bwd_si]
2348 bwd_stage_micro_index[bwd_si] += 1
2349 return MetaStep(bwd_mi, MetaStepType.BWD, bwd_si)
2351 fwd_steps, bwd_steps = self._collect_fwd_bwd_steps(
2352 emit_fwd, emit_bwd, fwd_bwd_ops, warmup_ops,
2353 )
2354 ops = self._pair_into_overlap_b_f(fwd_steps, bwd_steps)
2356 last_stage = self.real_stage_num - 1
2357 if self._short_micro() and stage_index == last_stage and bwd_steps:
2358 if bwd_steps[-1].micro_index == self.micro_batch_num - 1:
2359 # Double the bubble at the 1F1B->cooldown chunk boundary;
2360 # see :meth:`_emit_cooldown_ops` for the alignment rationale.
2361 ops.extend([None] * (2 * self._trailing_bubble()))
2362 return ops
2364 def construct_stage_exec_order(self, stage_index):
2365 """Construct the execution order for ``stage_index``.
2367 Builds: warmup → bubbles → 1F1B steady state → cooldown. The
2368 1F1B segment switches between :meth:`_emit_1f1b_ops` (plain) and
2369 :meth:`_emit_1f1b_overlap_ops` (OVERLAP_B_F pairing) based on
2370 the ``overlap_b_f`` constructor flag.
2371 """
2372 warmup_ops = self.warmup_ops(stage_index)
2373 fwd_bwd_ops = self.n_local_stages * self.micro_batch_num - warmup_ops
2374 total_ops = 2 * warmup_ops + fwd_bwd_ops
2375 order_list = [None for _ in range(stage_index)]
2376 fwd_stage_micro_index = defaultdict(int)
2377 bwd_stage_micro_index = defaultdict(int)
2378 order_list.extend(self._emit_warmup_ops(stage_index, warmup_ops, fwd_stage_micro_index))
2379 # One bubble per warmup op lost to the short-micro cap: the uncapped
2380 # stagger is depth (less_memory) or 2*depth (classic) beyond V*M.
2381 depth = self.real_stage_num - 1 - stage_index
2382 bubbles_before_1f1b = max(
2383 0,
2384 (depth if self._less_memory else 2 * depth) - self.micro_batch_num,
2385 )
2386 order_list.extend([None] * bubbles_before_1f1b)
2387 order_list.extend([None] * (self.real_stage_num - 1 - stage_index))
2388 if self._overlap_b_f:
2389 order_list.extend(self._emit_1f1b_overlap_ops(
2390 stage_index, warmup_ops, fwd_bwd_ops,
2391 fwd_stage_micro_index, bwd_stage_micro_index,
2392 ))
2393 else:
2394 order_list.extend(self._emit_1f1b_ops(
2395 stage_index, warmup_ops, fwd_bwd_ops,
2396 fwd_stage_micro_index, bwd_stage_micro_index,
2397 ))
2398 order_list.extend(self._emit_cooldown_ops(
2399 stage_index, warmup_ops, fwd_bwd_ops, total_ops, bwd_stage_micro_index,
2400 ))
2401 return order_list
2404def detect_cycle_in_graph(ranks_map):
2405 """
2406 Detects a cycle in the directed graph constructed from ranks_map.
2408 Args:
2409 ranks_map: A dictionary where keys are rank names and values are lists of nodes.
2411 Returns:
2412 tuple: (cycle_path, cycle_ranks) where cycle_path is a list of nodes forming the cycle and cycle_ranks
2413 is a list of rank transitions corresponding to the cycle path.
2414 """
2415 graph = defaultdict(list)
2416 rank_edges = {}
2418 for rank, nodes in ranks_map.items():
2419 for i in range(len(nodes) - 1):
2420 u, v = nodes[i], nodes[i + 1]
2421 graph[u].append(v)
2422 rank_edges[(u, v)] = rank
2424 visited = set()
2425 path = []
2426 node_indices = {}
2427 cycle_path = []
2428 cycle_ranks = []
2430 stack = []
2431 for node in list(graph.keys()):
2432 if node not in visited:
2433 stack.append((node, False))
2434 while stack:
2435 current_node, is_processed = stack.pop()
2437 if is_processed:
2438 path.pop()
2439 del node_indices[current_node]
2440 continue
2442 if current_node in node_indices:
2443 cycle_start = node_indices[current_node]
2444 cycle_path = path[cycle_start:] + [current_node]
2445 for i in range(cycle_start, len(path)):
2446 u = path[i]
2447 v = path[i + 1] if i + 1 < len(path) else current_node
2448 cycle_ranks.append(f"{rank_edges[(u, v)]} {u} -> {v}")
2449 return cycle_path, cycle_ranks
2451 if current_node in visited:
2452 continue
2454 visited.add(current_node)
2455 node_indices[current_node] = len(path)
2456 path.append(current_node)
2458 stack.append((current_node, True))
2459 for neighbor in reversed(graph[current_node]):
2460 stack.append((neighbor, False))
2462 return None, None
2465def output_cycle_results(cycle_path, cycle_ranks):
2466 """
2467 Helper function to output cycle detection results.
2469 Args:
2470 cycle_path (list): List of nodes forming a cycle, if any.
2471 cycle_ranks (list): List of ranks involved in the cycle.
2473 Returns:
2474 None: Outputs results to the console.
2475 """
2476 if cycle_path:
2477 logger.error("Cycle detected:")
2478 path_str = " -> ".join(str(node) for node in cycle_path)
2479 logger.error("%s -> %s", path_str, cycle_path[0]) # Close the cycle
2480 logger.error("Involving ranks:")
2481 for rank in cycle_ranks:
2482 logger.error(rank)
2483 else:
2484 logger.warning("Cycle Check succeeded. There is no cycle in the graph.")
2487def parse_and_validate(data: dict, all_rank: bool = True):
2488 """
2489 Parse and validate execution orders in a directed graph structure.
2491 This function checks the integrity and consistency of a given dataset, ensuring all required
2492 keys are present and correctly referenced. It also validates the structure of the input data
2493 and parses string values to extract meaningful components.
2495 Args:
2496 data (dict): A dictionary where keys are string identifiers and values are lists of strings.
2497 Each value represents a dependency or reference to other keys.
2498 all_rank (bool): If True, checks that all elements referenced in the data are present as keys
2499 in the dictionary. If False, only checks intersections.
2501 Returns:
2502 None: Log error messages to the console if validation fails, otherwise completes silently.
2504 Raises:
2505 ValueError: Raised indirectly if `parse_elements` encounters malformed input strings.
2506 TypeError: Raised indirectly if data contains unexpected types.
2507 """
2509 def parse_elements(value: str, max_groups: int = 2) -> set:
2510 """Extract unique elements inside the first one or two parentheses from a string."""
2512 groups = re.findall(r'\((\d+)\)', value)
2513 limited_groups = groups[:max_groups] # Limit to the first `max_groups` matches
2515 return {item.strip() for item in limited_groups}
2517 if not isinstance(data, dict):
2518 logger.error("Input must be a dictionary with string keys and lists of strings as values.")
2519 return
2521 key_to_values = {key: set(values) for key, values in data.items() if
2522 isinstance(values, list) and all(isinstance(v, str) for v in values)}
2524 for key, values in data.items():
2525 if not isinstance(values, list) or not all(isinstance(v, str) for v in values):
2526 logger.error("Values for key '%s' must be a list of strings.", key)
2527 continue
2529 for value in values:
2530 try:
2531 elements = parse_elements(value)
2532 except (ValueError, TypeError, AttributeError) as e:
2533 logger.error("Unable to parse elements from value '%s' in key '%s'. Error: %s", value, key, e)
2534 continue
2536 # Check for missing keys if all_rank is True
2537 if all_rank:
2538 missing_keys = elements - key_to_values.keys()
2539 if missing_keys:
2540 logger.error("The following keys are missing for value '%s': %s", value, missing_keys)
2541 continue
2543 # Check if the value is present in the referenced keys
2544 for element in elements & key_to_values.keys() if not all_rank else elements:
2545 if value not in key_to_values[element]:
2546 logger.error("Key '%s' is missing the value '%s'.", element, value)
2549def generate_operations(order_list: dict[int, list[MetaStep]],
2550 chunk_num: int,
2551 com_type: str = 'loop') -> dict[str, list[str]]:
2552 """
2553 Generate formatted operations dictionary from pipeline execution order.
2555 Args:
2556 order_list (dict): Dictionary where keys are rank IDs and values are MetaStep execution sequences
2557 chunk_num (int): Number of chunks (virtual pipeline stages)
2558 com_type (str): Stage-to-rank mapping type ('loop' for cyclic, 'v' for V-shaped)
2560 Returns:
2561 Dictionary where keys are rank IDs (as strings) and values are lists of formatted operation strings
2562 """
2564 def stage_to_rank(stage_index, style, stage_num, real_stage_num):
2565 """Map stage index to rank"""
2566 if style == 'loop':
2567 return stage_index % real_stage_num
2568 if style == 'v':
2569 if stage_index < real_stage_num:
2570 return stage_index
2571 return stage_num - 1 - stage_index
2572 raise ValueError("Invalid style")
2574 def find_send_target(stage_idx, op_type):
2575 """Find target stage for SEND operation"""
2576 if op_type == MetaStepType.FWD_SEND:
2577 return forward_comm.get(stage_idx)
2578 return backward_comm.get(stage_idx)
2580 def find_recv_source(stage_idx, op_type):
2581 """Find source stage for RECV operation"""
2582 if op_type == MetaStepType.FWD_RECV:
2583 # Reverse lookup in forward_comm
2584 for src, dst in forward_comm.items():
2585 if dst == stage_idx:
2586 return src
2587 else:
2588 # Reverse lookup in backward_comm
2589 for src, dst in backward_comm.items():
2590 if dst == stage_idx:
2591 return src
2592 return None
2594 real_stage = len(order_list)
2595 total_stages = real_stage * chunk_num
2597 # Build communication rules
2598 forward_comm = {}
2599 backward_comm = {}
2601 for i in range(total_stages):
2602 if i + 1 < total_stages:
2603 forward_comm[i] = i + 1
2604 if i - 1 >= 0:
2605 backward_comm[i] = i - 1
2607 formatted_operations = defaultdict(list)
2609 for rank, steps in order_list.items():
2610 operation_counter = defaultdict(int)
2612 for step in steps:
2613 if step.type in [MetaStepType.FWD_SEND, MetaStepType.BWD_SEND]:
2614 target_stage = find_send_target(step.stage_index, step.type)
2615 if target_stage is not None:
2616 target_rank = stage_to_rank(target_stage, com_type, total_stages, real_stage)
2617 comm_pair = (rank, target_rank, step.micro_index)
2618 operation_counter[comm_pair] += 1
2619 count = operation_counter[comm_pair]
2620 formatted_op = f"Send_Receive_({rank})->({target_rank})_micro{step.micro_index}_{count}th"
2621 formatted_operations[str(rank)].append(formatted_op)
2623 elif step.type in [MetaStepType.FWD_RECV, MetaStepType.BWD_RECV]:
2624 source_stage = find_recv_source(step.stage_index, step.type)
2625 if source_stage is not None:
2626 source_rank = stage_to_rank(source_stage, com_type, total_stages, real_stage)
2627 comm_pair = (source_rank, rank, step.micro_index)
2628 operation_counter[comm_pair] += 1
2629 count = operation_counter[comm_pair]
2630 formatted_op = f"Send_Receive_({source_rank})->({rank})_micro{step.micro_index}_{count}th"
2631 formatted_operations[str(rank)].append(formatted_op)
2633 # Convert defaultdict to dict
2634 return dict(formatted_operations)
2637def validate_pipeline_execution(order_list: dict[int, list[MetaStep]],
2638 chunk_num: int,
2639 com_type: str = 'loop') -> dict[str, any]:
2640 """
2641 Comprehensive validation function for pipeline parallel execution order.
2643 This function validates the execution order of pipeline parallelism by:
2644 1. Checking SEND/RECV communication pair matching
2645 2. Detecting duplicate operations
2646 3. Detecting cycles in communication graphs
2647 4. Verifying computation-SEND matching
2649 Args:
2650 order_list: Dictionary where keys are rank IDs and values are MetaStep execution sequences
2651 chunk_num: Number of chunks (virtual pipeline stages)
2652 com_type: Stage-to-rank mapping type ('loop' for cyclic, 'v' for V-shaped)
2654 Returns:
2655 Dictionary containing validation results with the following keys:
2656 - validation: Communication pair validation results
2657 - cycle_detection: Cycle detection results
2658 - computation_send_matching: Computation-SEND matching validation results
2659 - has_errors: Boolean indicating if any errors were found
2660 - error_messages: List of all error messages found
2661 - formatted_operations: Generated formatted operations
2662 """
2664 # Generate operations
2665 formatted_operations = generate_operations(order_list, chunk_num, com_type)
2667 parse_and_validate(formatted_operations, True)
2669 # Detect cycles
2670 cycle_path, cycle_ranks = detect_cycle_in_graph(formatted_operations)
2672 # Output results
2673 output_cycle_results(cycle_path, cycle_ranks)
2675 result = {
2676 'formatted_operations': formatted_operations,
2677 'cycle_path': cycle_path,
2678 'cycle_ranks': cycle_ranks,
2679 'has_cycle': bool(cycle_path)
2680 }
2681 return result
2684_COMPUTE_META_STEP_TYPES = frozenset({
2685 MetaStepType.FWD,
2686 MetaStepType.BWD,
2687 MetaStepType.BWD_INPUT,
2688 MetaStepType.BWD_WEIGHT,
2689})
2692def _next_active_stage_indices(actions, start_index, max_active_stages, managed_stage_indices):
2693 """Find the next distinct managed stages that will execute compute work.
2695 Send/recv and previously injected FSDP control steps are skipped so that the
2696 lookahead window only counts real compute, otherwise communication-only
2697 actions would consume the budget and shrink the effective prefetch depth.
2698 """
2699 stage_indices = []
2700 seen = set()
2701 for action in actions[start_index:]:
2702 for leaf_step in iter_leaf_meta_steps(action):
2703 if leaf_step.type not in _COMPUTE_META_STEP_TYPES:
2704 continue
2705 if leaf_step.stage_index not in managed_stage_indices or leaf_step.stage_index in seen:
2706 continue
2707 seen.add(leaf_step.stage_index)
2708 stage_indices.append(leaf_step.stage_index)
2709 if len(stage_indices) == max_active_stages:
2710 return stage_indices
2711 return stage_indices
2714def add_fsdp_unshard_reshard(actions, managed_stage_indices, max_active_stages=3):
2715 """Insert FSDP unshard/reshard actions for locally managed stages."""
2716 if not managed_stage_indices:
2717 return actions
2719 fsdp_actions = []
2720 active_stages = []
2721 for index, action in enumerate(actions):
2722 next_stage_indices = _next_active_stage_indices(
2723 actions, index, max_active_stages, managed_stage_indices
2724 )
2725 evicted_stages = [stage_index for stage_index in active_stages if stage_index not in next_stage_indices]
2726 fetched_stages = [stage_index for stage_index in next_stage_indices if stage_index not in active_stages]
2727 for stage_index in evicted_stages:
2728 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_RESHARD, stage_index))
2729 active_stages.remove(stage_index)
2730 for stage_index in fetched_stages:
2731 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_UNSHARD, stage_index))
2732 active_stages.append(stage_index)
2733 fsdp_actions.append(action)
2735 while active_stages:
2736 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_RESHARD, active_stages.pop(0)))
2737 return fsdp_actions
2740def add_fsdp_reduce_grad(actions, managed_stage_indices, micro_batch_num):
2741 """Insert FSDP reduce-grad actions after the last backward-like action of each stage."""
2742 if not managed_stage_indices:
2743 return actions
2745 fsdp_actions = []
2746 for action in actions:
2747 fsdp_actions.append(action)
2748 reduced_stage_indices = []
2749 for leaf_step in iter_leaf_meta_steps(action):
2750 if leaf_step.stage_index not in managed_stage_indices:
2751 continue
2752 if leaf_step.type not in (MetaStepType.BWD, MetaStepType.BWD_WEIGHT):
2753 continue
2754 if leaf_step.micro_index != micro_batch_num - 1:
2755 continue
2756 if leaf_step.stage_index not in reduced_stage_indices:
2757 reduced_stage_indices.append(leaf_step.stage_index)
2758 for stage_index in reduced_stage_indices:
2759 fsdp_actions.append(MetaStep(None, MetaStepType.FSDP_REDUCE_GRAD, stage_index))
2760 return fsdp_actions