Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / parallel / pp_schedule.py: 98%
461 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""
16PP Schedules - Graph-mode pipeline schedule drivers.
18Single module per the house convention (``core/pipeline_parallel/scheduler.py``
19keeps its eager-mode schedules the same way): shared P2P / microbatch
20machinery in :class:`PipelineScheduleBase` plus one class per algorithm —
21:class:`ScheduleGPipe`, :class:`Schedule1F1B`, and the VPP
22:class:`ScheduleInterleaved1F1B` today — dispatched by name from
23:data:`SCHEDULE_REGISTRY` / ``PassConfig.pp_schedule``.
25Both schedules run the per-stage FX subgraphs produced by ``PpPass`` across
26microbatches and share the same P2P / loss / gradient contract:
28- Model inputs are routed by dataflow: each user input is fed directly to
29 every stage that consumes it (resolved by ``PpPass``; a shared input is
30 replicated, not shipped over P2P) — every rank receives the full flattened
31 user-input list, slices its OWNED tensor inputs per microbatch, and ignores
32 the rest. All tensor user inputs must carry the leading batch dim.
33- Boundary values are exchanged as ORDERED LISTS: tensors go as-is, int
34 scalars (dynamic-shape ``sym_size`` nodes crossing the cut) are packed as
35 0-d int64 tensors and unwrapped with ``item()`` on arrival — mirroring how
36 ``torch.distributed.pipelining``'s ``PipelineStage`` ships full argument
37 lists. The P2P exchange uses eager ``dist.isend``/``irecv`` on the PP
38 process group; sends are async so the sweeps overlap receive/compute across
39 stages, receives wait before first use, and each microbatch's boundary sends
40 are retained under that microbatch's key — forward and backward boundaries
41 share the key — until their ``Work`` handles are waited (1F1B: one turn
42 after the backward posts, so a forward send rides along; GPipe: at the end
43 of the step).
44- Microbatch gradients are un-normalized sums while sweeping, then divided by
45 ``num_microbatches`` at the end, so the step gradient matches the non-PP
46 semantics of a full-batch mean loss.
48The schedules differ in their action order — and in the activation lifetime
49that order allows:
51``ScheduleGPipe``
52 1. Forward sweep: run every microbatch's forward (shipping boundary
53 activations downstream), holding all forward outputs.
54 2. Backward sweep (reverse order): seed ``ones_like(loss)`` on the last
55 stage, run each backward against the activations saved by the forward
56 sweep, and pass boundary gradients one stage back.
58``Schedule1F1B``
59 Warmup runs ``min(num_microbatches, pp_degree - stage_idx)`` forwards (a
60 stage nearer the front keeps one more microbatch in flight), then the
61 steady state alternates one backward with one forward per turn, and a
62 cooldown drains the remaining backwards. Interleaving backward earlier
63 shrinks the pipeline bubble relative to GPipe. Requires
64 ``num_microbatches >= pp_degree``.
66Forward activations and backward gradients share one PP process group, whose
67per-pair FIFO must stay complementary — at each position one end sends while
68the other receives — or the group deadlocks. ``Schedule1F1B`` staggers its
69boundary sends to hold that invariant despite the ``pp_degree``-asymmetric
70warmup; the placement mirrors ``core/pipeline_parallel/scheduler.py`` and is
71spelled out in ``Schedule1F1B.forward``.
72"""
74__all__ = [
75 "PipelineScheduleBase",
76 "ScheduleGPipe",
77 "Schedule1F1B",
78 "ScheduleInterleaved1F1B",
79 "SCHEDULE_REGISTRY",
80 "get_schedule_class",
81]
83import bisect
84import itertools
85from collections import defaultdict
86from dataclasses import dataclass
87from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type
89import torch
90import torch.distributed as dist
91from torch import nn
94class PipelineScheduleBase(nn.Module):
95 """Shared P2P / microbatch machinery for the graph-mode PP schedules.
97 Installed by ``PpPass`` as a submodule of the compiled GraphModule and
98 invoked through a ``call_module`` node, so the trainer's
99 ``graph_module(*flat_inputs)`` dispatches to the concrete schedule
100 unchanged. Subclasses implement :meth:`forward` (the action order) from
101 the per-microbatch building blocks defined here.
103 Args:
104 fwd_gm: Stage forward subgraph. Signature per stage:
105 stage 0 -> ``(*state, *owned_inputs)`` returning
106 ``(*act_out, *saved)``;
107 other -> ``(*state, *act_in, *owned_inputs)`` returning
108 ``(*act_out, *saved)`` for non-last stages and
109 ``(loss, *saved)`` for the last one, where
110 ``owned_inputs`` are the user inputs routed to
111 this stage (may be empty).
112 bwd_gm: Stage backward subgraph. Signature:
113 ``(*state, *grad_in, *fwd_outs)`` where ``grad_in`` are the
114 received boundary gradients (``ones_like(loss)`` on the last
115 stage) and ``fwd_outs`` are the forward subgraph's outputs
116 verbatim. Returns ``(*param_grads,)`` plus the boundary
117 gradients for every stage except stage 0.
118 stage_idx: This rank's stage index (0-based).
119 pp_degree: Number of pipeline stages.
120 num_state: Number of leading state tensors (stage params/buffers).
121 num_trainable: Number of trainable parameters on this stage; the
122 backward subgraph returns exactly this many parameter gradients
123 first.
124 num_send: Number of forward boundary values to ship downstream
125 (the forward outputs' prefix); 0 on the last stage.
126 grad_send_count: Number of backward boundary values to ship
127 upstream (the backward outputs' suffix); 0 on stage 0.
128 microbatch_size: Samples per microbatch (leading dim).
129 pp_group: PP process group used for all boundary P2P (activations and
130 gradients) and for stage->global rank mapping.
131 recv_spec: Per-value descriptors for the incoming forward
132 boundary: ``("tensor", shape, dtype, device)`` or
133 ``("scalar",)``. Empty on stage 0.
134 grad_recv_spec: Same, for the incoming backward boundary
135 gradients. Empty on the last stage.
136 user_input_stages: Per flattened user input (stub arg order), its
137 primary consuming stage.
138 user_input_stage_sets: Optional per-input owner sets; when omitted,
139 each input is owned by its single ``user_input_stages`` entry.
140 This stage slices and feeds the inputs mapped to itself.
141 """
143 def __init__(
144 self,
145 fwd_gm: nn.Module,
146 bwd_gm: nn.Module,
147 stage_idx: int,
148 pp_degree: int,
149 num_state: int,
150 num_trainable: int,
151 num_send: int,
152 grad_send_count: int,
153 microbatch_size: int,
154 pp_group: Any = None,
155 recv_spec: Sequence[Tuple[Any, ...]] = (),
156 grad_recv_spec: Sequence[Tuple[Any, ...]] = (),
157 user_input_stages: Sequence[int] = (),
158 user_input_stage_sets: Optional[Sequence[Sequence[int]]] = None,
159 ) -> None:
160 """Store the stage subgraphs, P2P group, and boundary specs."""
161 super().__init__()
162 self.fwd_gm = fwd_gm
163 self.bwd_gm = bwd_gm
164 self.stage_idx = stage_idx
165 self.pp_degree = pp_degree
166 self.num_state = num_state
167 self.num_trainable = num_trainable
168 self.num_send = num_send
169 self.grad_send_count = grad_send_count
170 self.microbatch_size = microbatch_size
171 self.pp_group = pp_group
172 self.recv_spec = list(recv_spec)
173 self.grad_recv_spec = list(grad_recv_spec)
174 self.user_input_stages = list(user_input_stages)
175 self.user_input_stage_sets = (
176 [list(stages) for stages in user_input_stage_sets]
177 if user_input_stage_sets is not None
178 else [[stage] for stage in self.user_input_stages]
179 )
180 self.is_first = stage_idx == 0
181 self.is_last = stage_idx == pp_degree - 1
182 # In-flight isend Work handles keyed by microbatch index — forward and
183 # backward boundaries of one microbatch share the key; the sent
184 # tensors stay referenced so their buffers outlive the send. Schedules
185 # release a microbatch's buffers once its sends provably completed
186 # (1F1B: next turn's gradient receipt) or at the end of the step.
187 self._pending_sends: Dict[int, List[Tuple[Any, torch.Tensor]]] = {}
189 def forward(self, *flat_inputs: Any) -> Tuple[torch.Tensor, ...]:
190 """Run one pipeline step.
192 Args:
193 *flat_inputs: ``(*state, *user_inputs)`` for this stage.
195 Returns:
196 ``(loss, *param_grads)``.
198 Raises:
199 NotImplementedError: Always; concrete schedules override this.
200 """
201 raise NotImplementedError
203 def map_stage_graphs(self, fn: Callable[[nn.Module], nn.Module]) -> None:
204 """Apply ``fn`` to each owned stage subgraph and store the result.
206 Terminal backend compilation (``InductorPass``) runs after ``PpPass``
207 replaced the joint graph with this schedule stub: the stage subgraphs
208 — not the stub — carry the FX compute, so the pass compiles each of
209 them through this hook and writes the returned module back (in-place
210 compilers may return the same object).
211 """
212 self.fwd_gm = fn(self.fwd_gm)
213 self.bwd_gm = fn(self.bwd_gm)
215 # ------------------------------------------------------------------
216 # Input / microbatch resolution
217 # ------------------------------------------------------------------
219 def _split_inputs(
220 self, flat_inputs: Sequence[Any]
221 ) -> Tuple[Sequence[Any], List[Any]]:
222 """Split ``(*state, *user_inputs)`` and validate the input arity."""
223 state = flat_inputs[: self.num_state]
224 user_inputs = list(flat_inputs[self.num_state :])
225 if len(user_inputs) != len(self.user_input_stages):
226 raise RuntimeError(
227 f"PP schedule expects {len(self.user_input_stages)} user "
228 f"inputs (per its routing) but received {len(user_inputs)}"
229 )
230 return state, user_inputs
232 def _resolve_microbatches(
233 self, user_inputs: Sequence[Any]
234 ) -> Tuple[int, List[Any]]:
235 """Derive the microbatch count and this stage's owned user inputs.
237 Every rank receives every user input, so the batch size is derived
238 identically on all stages from the first tensor input.
239 """
240 batch_size = next(
241 (v.shape[0] for v in user_inputs if isinstance(v, torch.Tensor)), None
242 )
243 if batch_size is None:
244 raise ValueError(
245 "PP requires at least one tensor model input to derive the "
246 "microbatch split"
247 )
248 if batch_size % self.microbatch_size != 0:
249 raise ValueError(
250 f"PP microbatch mismatch: batch size {batch_size} is not "
251 f"divisible by pp_microbatch_size {self.microbatch_size}"
252 )
253 num_microbatches = batch_size // self.microbatch_size
254 owned_inputs = [
255 v
256 for v, stages in zip(user_inputs, self.user_input_stage_sets)
257 if self.stage_idx in stages
258 ]
259 return num_microbatches, owned_inputs
261 def _slice_inputs(self, owned_inputs: Sequence[Any], mb_index: int) -> List[Any]:
262 """Slice this stage's owned inputs for microbatch ``mb_index``.
264 Tensors are sliced along the leading batch dim; non-tensors pass
265 through unchanged.
266 """
267 lo = mb_index * self.microbatch_size
268 hi = lo + self.microbatch_size
269 return [v[lo:hi] if isinstance(v, torch.Tensor) else v for v in owned_inputs]
271 # ------------------------------------------------------------------
272 # Per-microbatch building blocks
273 # ------------------------------------------------------------------
274 #
275 # Forward / backward are split into a receive step and a compute step so
276 # a schedule can interleave boundary sends with compute independently;
277 # ``_forward_microbatch`` / ``_backward_microbatch`` compose them for
278 # callers that don't need that control.
280 def _recv_activation(self) -> List[Any]:
281 """Receive this stage's incoming boundary activations (none on stage 0)."""
282 if self.is_first:
283 return []
284 return self._recv_values(self.recv_spec, src=self.stage_idx - 1)
286 def _apply_forward(
287 self,
288 state: Sequence[Any],
289 act_in: Sequence[Any],
290 owned_inputs: Sequence[Any],
291 mb_index: int,
292 ) -> Tuple[Any, ...]:
293 """Run one microbatch's forward subgraph against ``act_in``.
295 ``owned_inputs`` are this stage's user inputs (see
296 ``user_input_stages``): tensors are sliced per microbatch along the
297 leading batch dim, non-tensors pass through. They are appended after
298 the received activations in the fwd subgraph's arg order.
299 """
300 owned_mb = self._slice_inputs(owned_inputs, mb_index)
301 if self.is_first:
302 out = self.fwd_gm(*state, *owned_mb)
303 else:
304 out = self.fwd_gm(*state, *act_in, *owned_mb)
305 return tuple(out)
307 def _forward_microbatch(
308 self, state: Sequence[Any], owned_inputs: Sequence[Any], mb_index: int
309 ) -> Tuple[Any, ...]:
310 """Receive activations then run one microbatch's forward subgraph."""
311 act_in = self._recv_activation()
312 return self._apply_forward(state, act_in, owned_inputs, mb_index)
314 def _send_forward(self, out: Sequence[Any], mb_index: int) -> None:
315 """Ship microbatch ``mb_index``'s boundary activations downstream."""
316 if self.num_send:
317 self._send_values(
318 out[: self.num_send], dst=self.stage_idx + 1, mb_index=mb_index
319 )
321 def _recv_gradient(self, fwd_outs: Sequence[Any]) -> List[Any]:
322 """Receive the next stage's output grads (a seed on the last stage).
324 The traced backward graph seeds its own loss gradient internally, so
325 the ``ones_like`` seed on the last stage only fixes the shape.
326 """
327 if self.is_last:
328 return [torch.ones_like(fwd_outs[0])]
329 return self._recv_values(self.grad_recv_spec, src=self.stage_idx + 1)
331 def _apply_backward(
332 self,
333 state: Sequence[Any],
334 grad_in: Sequence[Any],
335 fwd_outs: Sequence[Any],
336 mb_index: int,
337 ) -> Sequence[torch.Tensor]:
338 """Run microbatch ``mb_index``'s backward subgraph and ship grads.
340 Gradients flow UPSTREAM, so the previous stage is the destination,
341 mirroring the forward direction. The grads returned are un-normalized
342 sums over the microbatches; the shipped boundary grads are keyed by
343 ``mb_index`` so the schedule can release their buffers later.
344 """
345 out = self.bwd_gm(*state, *grad_in, *fwd_outs)
346 param_grads = out[: self.num_trainable]
347 if self.grad_send_count:
348 self._send_values(
349 out[self.num_trainable :], dst=self.stage_idx - 1, mb_index=mb_index
350 )
351 return param_grads
353 def _backward_microbatch(
354 self, state: Sequence[Any], fwd_outs: Sequence[Any], mb_index: int
355 ) -> Sequence[torch.Tensor]:
356 """Receive boundary grads then run one microbatch's backward subgraph."""
357 grad_in = self._recv_gradient(fwd_outs)
358 return self._apply_backward(state, grad_in, fwd_outs, mb_index)
360 @staticmethod
361 def _accumulate_grads(
362 accumulated: Sequence[torch.Tensor], param_grads: Sequence[torch.Tensor]
363 ) -> List[torch.Tensor]:
364 """Sum one microbatch's param grads into the running accumulation."""
365 if not accumulated:
366 return list(param_grads)
367 return [acc + grad for acc, grad in zip(accumulated, param_grads)]
369 def _finalize_loss(
370 self, fwd_outs_per_mb: Sequence[Sequence[Any]], user_inputs: Sequence[Any]
371 ) -> torch.Tensor:
372 """Return the step loss (last stage) or a zero placeholder."""
373 if self.is_last:
374 # On the last stage each forward output tuple starts with the
375 # per-microbatch loss.
376 losses = [outs[0] for outs in fwd_outs_per_mb]
377 return torch.stack(losses).mean()
378 return torch.zeros((), device=self._anchor_device(user_inputs))
380 @staticmethod
381 def _average_grads(
382 grads: Sequence[torch.Tensor], num_microbatches: int
383 ) -> List[torch.Tensor]:
384 """Turn un-normalized sums over microbatches into the mean."""
385 return [grad / num_microbatches for grad in grads]
387 # ------------------------------------------------------------------
388 # P2P helpers
389 # ------------------------------------------------------------------
391 def _flush_pending_sends(self) -> None:
392 """Wait every in-flight isend and drop all buffer references."""
393 for sends in self._pending_sends.values():
394 for work, _ in sends:
395 work.wait()
396 self._pending_sends = {}
398 def _wait_sends(self, mb_index: int) -> None:
399 """Wait microbatch ``mb_index``'s in-flight sends, freeing its buffers.
401 Callers must only invoke this once the sends are provably complete
402 (e.g. the peer consumed the values), so the wait never blocks.
403 """
404 for work, _ in self._pending_sends.pop(mb_index, ()):
405 work.wait()
407 @staticmethod
408 def _anchor_device(user_inputs: Sequence[Any]) -> torch.device:
409 """Device for the zero-loss placeholder: first tensor user input."""
410 for value in user_inputs:
411 if isinstance(value, torch.Tensor):
412 return value.device
413 return torch.device("cpu")
415 def _send_values(self, values: Sequence[Any], dst: int, mb_index: int) -> None:
416 """Async-send boundary values to the neighbouring stage ``dst``.
418 Tensors go as-is (made contiguous, with THAT buffer retained); int
419 scalars (dynamic-shape ``sym_size`` nodes) are packed as 0-d int64
420 tensors on the first tensor value's device — CPU tensors cannot ride
421 an NCCL/NPU-backend group — and unwrapped with ``item()`` on the
422 receiving side. The ``Work`` handles are queued under ``mb_index``
423 (forward and backward boundaries of a microbatch share the key) and
424 the sent tensors kept referenced so their buffers outlive the send;
425 the schedule waits them (per turn or at step end) before the buffers
426 are released.
427 """
428 dst_rank = self._global_rank(dst)
429 anchor_device = next(
430 (v.device for v in values if isinstance(v, torch.Tensor)), None
431 )
432 for value in values:
433 if isinstance(value, torch.Tensor):
434 tensor = value.contiguous()
435 else:
436 tensor = torch.tensor(value, dtype=torch.int64, device=anchor_device)
437 work = dist.isend(tensor, dst=dst_rank, group=self.pp_group)
438 self._pending_sends.setdefault(mb_index, []).append((work, tensor))
440 def _recv_values(self, spec: Sequence[Tuple[Any, ...]], src: int) -> List[Any]:
441 """Receive boundary values from the neighbouring stage ``src``, in order.
443 The buffer for each entry comes from ``spec`` (built by ``PpPass``
444 from the boundary values' fake-tensor ``val`` metas, including the
445 device each value was traced on); scalar entries arrive as 0-d
446 int64 tensors and are unwrapped to Python ints so the graph sees
447 the same types it was traced with. Scalar buffers are anchored to
448 the first tensor entry's device (``PpPass._value_spec`` guarantees
449 a non-empty list carries at least one tensor).
450 """
451 src_rank = self._global_rank(src)
452 anchor_device = next((entry[3] for entry in spec if entry[0] == "tensor"), None)
453 values: List[Any] = []
454 for entry in spec:
455 if entry[0] == "scalar":
456 buffer = torch.empty((), dtype=torch.int64, device=anchor_device)
457 else:
458 _, shape, dtype, device = entry
459 buffer = torch.empty(shape, dtype=dtype, device=device)
460 work = dist.irecv(buffer, src=src_rank, group=self.pp_group)
461 work.wait()
462 values.append(int(buffer.item()) if entry[0] == "scalar" else buffer)
463 return values
465 def _global_rank(self, stage_idx: int) -> int:
466 """Map a stage index to its global rank within the PP group."""
467 if self.pp_group is None:
468 return stage_idx
469 return dist.get_global_rank(self.pp_group, stage_idx)
472class ScheduleGPipe(PipelineScheduleBase):
473 """GPipe driver over two per-stage FX subgraphs (forward / backward).
475 Runs ALL forwards before ANY backward, so every microbatch's forward
476 outputs (boundary activations + saved values) are held until the backward
477 sweep completes — activation memory scales with ``num_microbatches``.
478 Larger ``microbatch_size`` trades pipeline bubble for fewer P2P
479 round-trips and fewer retained activations.
480 """
482 def forward(self, *flat_inputs: Any) -> Tuple[torch.Tensor, ...]:
483 """Run one full GPipe step (fwd sweep then bwd sweep).
485 Args:
486 flat_inputs: ``(*state, *user_inputs)`` — the same surface
487 ``run_traced_graph`` feeds the compiled graph, with
488 ``user_inputs`` the flattened model inputs in trace order.
490 Returns:
491 ``(loss, *param_grads)``. On non-last stages the loss is a zero
492 scalar placeholder (the real loss lives on the last stage);
493 parameter gradients are the microbatch-averaged accumulation.
494 """
495 state, user_inputs = self._split_inputs(flat_inputs)
496 num_microbatches, owned_inputs = self._resolve_microbatches(user_inputs)
498 fwd_outs_per_mb = self._forward_sweep(state, owned_inputs, num_microbatches)
499 grads = self._backward_sweep(state, fwd_outs_per_mb, num_microbatches)
501 self._flush_pending_sends()
503 loss = self._finalize_loss(fwd_outs_per_mb, user_inputs)
504 return (loss, *grads)
506 def _forward_sweep(
507 self,
508 state: Sequence[Any],
509 owned_inputs: Sequence[Any],
510 num_microbatches: int,
511 ) -> List[Tuple[Any, ...]]:
512 """Run forward microbatch 0..N-1, shipping boundary values downstream.
514 Returns per-microbatch forward output tuples (replayed verbatim
515 into the backward subgraph by ``_backward_sweep``).
516 """
517 fwd_outs_per_mb: List[Tuple[Any, ...]] = []
518 for mb_index in range(num_microbatches):
519 out = self._forward_microbatch(state, owned_inputs, mb_index)
520 fwd_outs_per_mb.append(out)
521 self._send_forward(out, mb_index)
522 return fwd_outs_per_mb
524 def _backward_sweep(
525 self,
526 state: Sequence[Any],
527 fwd_outs_per_mb: Sequence[Tuple[Any, ...]],
528 num_microbatches: int,
529 ) -> List[torch.Tensor]:
530 """Run backward microbatches N-1..0, accumulating full-batch gradients.
532 The traced backward graph seeds its own loss gradient internally
533 (an ``ones_like`` over the saved loss), so the grads that come back
534 are un-normalized sums over the microbatches. Dividing the
535 accumulation by ``num_microbatches`` turns that into the mean, which
536 is exactly the full-batch mean-loss gradient the non-PP path
537 produces.
538 """
539 grads: List[torch.Tensor] = []
540 for mb_index in reversed(range(num_microbatches)):
541 grads = self._accumulate_grads(
542 grads,
543 self._backward_microbatch(state, fwd_outs_per_mb[mb_index], mb_index),
544 )
545 return self._average_grads(grads, num_microbatches)
548class Schedule1F1B(PipelineScheduleBase):
549 """1F1B driver over two per-stage FX subgraphs (forward / backward).
551 Warmup runs ``min(num_microbatches, pp_degree - stage_idx)`` forwards
552 (one on the last stage), the steady state alternates one backward then
553 one forward per turn, and the cooldown drains the remaining backwards.
554 Running each backward as soon as its forward output is available shrinks
555 the pipeline bubble relative to GPipe's fill-drain order.
557 Requires ``num_microbatches >= pp_degree`` (checked at runtime), the same
558 constraint ``torch.distributed.pipelining.Schedule1F1B`` enforces.
560 Note:
561 Unlike GPipe, each microbatch's saved activations are freed as soon
562 as its backward consumes them. Its boundary sends share one key for
563 both directions and are released one turn after the backward posts —
564 the next turn's gradient receipt proves the previous microbatch's
565 transfers completed, so the wait cannot deadlock. Peak memory
566 therefore tracks the warmup depth (``pp_degree - stage_idx``) instead
567 of ``num_microbatches``; only the final microbatch's sends are still
568 waited at the end of the step. The last stage has no downstream
569 gradient to mark that progress, so its per-turn wait may briefly block
570 on the upstream peer.
571 """
573 def forward(self, *flat_inputs: Any) -> Tuple[torch.Tensor, ...]:
574 """Run one full 1F1B step (warmup -> steady 1B1F -> cooldown).
576 Args:
577 flat_inputs: ``(*state, *user_inputs)`` — the same surface
578 ``run_traced_graph`` feeds the compiled graph, with
579 ``user_inputs`` the flattened model inputs in trace order.
581 Returns:
582 ``(loss, *param_grads)``. On non-last stages the loss is a zero
583 scalar placeholder (the real loss lives on the last stage);
584 parameter gradients are the microbatch-averaged accumulation.
586 Raises:
587 ValueError: When ``num_microbatches < pp_degree`` (the schedule
588 cannot fill the pipeline).
589 """
590 state, user_inputs = self._split_inputs(flat_inputs)
591 num_microbatches, owned_inputs = self._resolve_microbatches(user_inputs)
592 if num_microbatches < self.pp_degree:
593 raise ValueError(
594 f"1F1B requires at least pp_degree={self.pp_degree} "
595 f"microbatches, got {num_microbatches} "
596 f"(batch {num_microbatches * self.microbatch_size} with "
597 f"pp_microbatch_size={self.microbatch_size}); lower "
598 f"pp_microbatch_size or use the gpipe schedule"
599 )
601 fwd_outs_per_mb: List[Any] = [None] * num_microbatches
602 grads: List[torch.Tensor] = []
603 fwd_index = 0
604 bwd_index = 0
605 # Forward output whose boundary activation has not shipped yet, with
606 # its microbatch index; the send is deferred one slot to keep the
607 # pair's FIFO complementary.
608 pending_index = -1
609 pending_out: Any = None
610 even = self.stage_idx % 2 == 0
612 # Warmup: fill the pipeline with ``pp_degree - stage_idx`` in-flight
613 # microbatches (one on the last stage) before the first backward.
614 # Even stages ship each activation right after its forward, except
615 # the last; odd stages ship the previous one just before the next
616 # forward. The parity stagger keeps neighbours complementary.
617 warmup = min(num_microbatches, self.pp_degree - self.stage_idx)
618 for k in range(warmup):
619 act_in = self._recv_activation()
620 if not even and k > 0:
621 self._send_forward(pending_out, pending_index)
622 pending_out = None
623 out = self._apply_forward(state, act_in, owned_inputs, fwd_index)
624 fwd_outs_per_mb[fwd_index] = out
625 fwd_index += 1
626 if even:
627 if k != warmup - 1:
628 self._send_forward(out, fwd_index - 1)
629 else:
630 pending_index, pending_out = fwd_index - 1, out
631 else:
632 pending_index, pending_out = fwd_index - 1, out
634 # Steady turn: recv_grad (+ release the previous turn's boundary
635 # buffers), send_activation, backward, send_grad, recv_activation,
636 # forward.
637 for _ in range(num_microbatches - warmup):
638 grad_in = self._recv_gradient(fwd_outs_per_mb[bwd_index])
639 if bwd_index > 0:
640 # This gradient's arrival proves the peer consumed microbatch
641 # ``bwd_index - 1``'s boundary values on both directions, so
642 # its in-flight sends completed and their buffers can go.
643 self._wait_sends(bwd_index - 1)
644 if pending_out is not None:
645 self._send_forward(pending_out, pending_index)
646 pending_out = None
647 grads = self._accumulate_grads(
648 grads,
649 self._apply_backward(
650 state, grad_in, fwd_outs_per_mb[bwd_index], bwd_index
651 ),
652 )
653 # Backward was the last consumer of these outputs; keep only the
654 # last stage's loss scalar for _finalize_loss.
655 fwd_outs_per_mb[bwd_index] = (
656 fwd_outs_per_mb[bwd_index][:1] if self.is_last else None
657 )
658 bwd_index += 1
659 act_in = self._recv_activation()
660 out = self._apply_forward(state, act_in, owned_inputs, fwd_index)
661 fwd_outs_per_mb[fwd_index] = out
662 fwd_index += 1
663 pending_index, pending_out = fwd_index - 1, out
664 # Drop the frame-local reference to the final forward's outputs (the
665 # warmup's when the steady state never runs); it would otherwise pin
666 # them through the cooldown.
667 out = None
669 # Cooldown: drain the remaining backwards; the final forward's
670 # activation ships on the first turn.
671 for _ in range(warmup):
672 grad_in = self._recv_gradient(fwd_outs_per_mb[bwd_index])
673 if bwd_index > 0:
674 # Same release invariant as the steady state.
675 self._wait_sends(bwd_index - 1)
676 if pending_out is not None:
677 self._send_forward(pending_out, pending_index)
678 pending_out = None
679 grads = self._accumulate_grads(
680 grads,
681 self._apply_backward(
682 state, grad_in, fwd_outs_per_mb[bwd_index], bwd_index
683 ),
684 )
685 # Backward was the last consumer of these outputs; keep only the
686 # last stage's loss scalar for _finalize_loss.
687 fwd_outs_per_mb[bwd_index] = (
688 fwd_outs_per_mb[bwd_index][:1] if self.is_last else None
689 )
690 bwd_index += 1
692 self._flush_pending_sends()
694 loss = self._finalize_loss(fwd_outs_per_mb, user_inputs)
695 return (loss, *self._average_grads(grads, num_microbatches))
698@dataclass(frozen=True)
699class _Action:
700 """One graph-mode pipeline action."""
702 kind: str
703 stage: int
704 micro: int
707@dataclass
708class _InterleavedStep:
709 """Mutable buffers shared by the VPP step's per-action handlers."""
711 state: Sequence[Any]
712 user_inputs: Sequence[Any]
713 act_cache: Dict[Tuple[int, int], List[Any]]
714 grad_cache: Dict[Tuple[int, int], List[Any]]
715 fwd_outs: Dict[Tuple[int, int], Tuple[Any, ...]]
716 bwd_outs: Dict[Tuple[int, int], Tuple[Any, ...]]
717 grads: List[List[torch.Tensor]]
718 losses: List[Optional[torch.Tensor]]
721_FWD = "fwd"
722_BWD = "bwd"
723_FWD_SEND = "fwd_send"
724_FWD_RECV = "fwd_recv"
725_BWD_SEND = "bwd_send"
726_BWD_RECV = "bwd_recv"
729def _round_layout(
730 num_microbatches: int, pp_degree: int, local_stages: int
731) -> Tuple[List[int], List[int]]:
732 """Per-round microbatch counts and their local-stage prefix sum.
734 Mirrors ``core.pipeline_parallel.scheduler.ScheduleInterleaved1F1B`` so
735 the graph-mode VPP order matches the eager schedule's warmup/cooldown
736 rhythm, including the ``num_microbatches < pp_degree`` case.
737 """
738 rounds = max(1, num_microbatches // pp_degree)
739 if num_microbatches < pp_degree:
740 base = num_microbatches - pp_degree
741 remainder = 0
742 else:
743 extra = num_microbatches % pp_degree
744 base = extra // rounds
745 remainder = extra % rounds
746 per_round = [pp_degree + base + (1 if i < remainder else 0) for i in range(rounds)]
747 accu = [x * local_stages for x in itertools.accumulate(per_round)]
748 accu.insert(0, 0)
749 return per_round, accu
752def _warmup_ops(
753 stage_index: int,
754 num_microbatches: int,
755 pp_degree: int,
756 local_stages: int,
757 per_round: Sequence[int],
758) -> int:
759 """Number of forwards this physical rank runs before its first backward."""
760 last_stage_warmup = (local_stages - 1) * per_round[0]
761 depth = pp_degree - 1 - stage_index
762 return min(last_stage_warmup + 2 * depth, num_microbatches * local_stages)
765def _forward_stage_index(
766 op_index: int,
767 stage_index: int,
768 per_round: Sequence[int],
769 accu: Sequence[int],
770 pp_degree: int,
771) -> int:
772 """Virtual stage a forward op at ``op_index`` belongs to."""
773 round_index = bisect.bisect_right(accu, op_index) - 1
774 local_index = (op_index - accu[round_index]) // per_round[round_index]
775 return local_index * pp_degree + stage_index
778def _backward_stage_index(
779 op_index: int,
780 stage_index: int,
781 per_round: Sequence[int],
782 accu: Sequence[int],
783 pp_degree: int,
784 local_stages: int,
785) -> int:
786 """Virtual stage a backward op at ``op_index`` belongs to."""
787 round_index = bisect.bisect_right(accu, op_index) - 1
788 local_index = (op_index - accu[round_index]) // per_round[round_index]
789 local_index = local_stages - 1 - local_index
790 return local_index * pp_degree + stage_index
793def _emit_warmup_actions(
794 stage_index: int,
795 warmup_ops: int,
796 num_microbatches: int,
797 pp_degree: int,
798 per_round: Sequence[int],
799 accu: Sequence[int],
800 fwd_micro: Dict[int, int],
801) -> List[Optional[_Action]]:
802 """Emit pure-forward warmup actions with the short-pipeline padding."""
803 actions: List[Optional[_Action]] = []
804 short = num_microbatches < pp_degree
805 last_micro = num_microbatches - 1
806 last_stage = pp_degree - 1
807 bubble = pp_degree - num_microbatches
808 for op_index in range(warmup_ops):
809 fwd_stage = _forward_stage_index(
810 op_index, stage_index, per_round, accu, pp_degree
811 )
812 micro = fwd_micro[fwd_stage]
813 actions.append(_Action(_FWD, fwd_stage, micro))
814 pad = (
815 short
816 and micro == last_micro
817 and (op_index != warmup_ops - 1 or stage_index == last_stage)
818 )
819 if pad:
820 actions.extend([None] * bubble)
821 fwd_micro[fwd_stage] += 1
822 return actions
825def _emit_1f1b_actions(
826 stage_index: int,
827 warmup_ops: int,
828 fwd_bwd_ops: int,
829 num_microbatches: int,
830 pp_degree: int,
831 local_stages: int,
832 per_round: Sequence[int],
833 accu: Sequence[int],
834 fwd_micro: Dict[int, int],
835 bwd_micro: Dict[int, int],
836) -> List[Optional[_Action]]:
837 """Emit the interleaved forward/backward steady state."""
838 actions: List[Optional[_Action]] = []
839 short = num_microbatches < pp_degree
840 last_micro = num_microbatches - 1
841 last_stage = pp_degree - 1
842 bubble = 2 * (pp_degree - num_microbatches)
843 for op_index in range(warmup_ops, warmup_ops + fwd_bwd_ops):
844 fwd_stage = _forward_stage_index(
845 op_index, stage_index, per_round, accu, pp_degree
846 )
847 actions.append(_Action(_FWD, fwd_stage, fwd_micro[fwd_stage]))
848 fwd_micro[fwd_stage] += 1
849 bwd_stage = _backward_stage_index(
850 op_index - warmup_ops,
851 stage_index,
852 per_round,
853 accu,
854 pp_degree,
855 local_stages,
856 )
857 actions.append(_Action(_BWD, bwd_stage, bwd_micro[bwd_stage]))
858 bwd_micro[bwd_stage] += 1
859 if (
860 short
861 and bwd_micro[bwd_stage] - 1 == last_micro
862 and stage_index == last_stage
863 ):
864 actions.extend([None] * bubble)
865 return actions
868def _emit_cooldown_actions(
869 stage_index: int,
870 warmup_ops: int,
871 fwd_bwd_ops: int,
872 total_ops: int,
873 num_microbatches: int,
874 pp_degree: int,
875 local_stages: int,
876 per_round: Sequence[int],
877 accu: Sequence[int],
878 bwd_micro: Dict[int, int],
879) -> List[Optional[_Action]]:
880 """Emit pure-backward cooldown actions with the matching padding."""
881 actions: List[Optional[_Action]] = []
882 short = num_microbatches < pp_degree
883 last_micro = num_microbatches - 1
884 bubble = 2 * (pp_degree - num_microbatches)
885 for op_index in range(warmup_ops + fwd_bwd_ops, total_ops):
886 actions.append(None)
887 bwd_stage = _backward_stage_index(
888 op_index - warmup_ops,
889 stage_index,
890 per_round,
891 accu,
892 pp_degree,
893 local_stages,
894 )
895 micro = bwd_micro[bwd_stage]
896 actions.append(_Action(_BWD, bwd_stage, micro))
897 bwd_micro[bwd_stage] += 1
898 if short and micro == last_micro:
899 actions.extend([None] * bubble)
900 return actions
903def _interleaved_stage_order(
904 stage_index: int,
905 num_microbatches: int,
906 pp_degree: int,
907 num_stages: int,
908) -> List[Optional[_Action]]:
909 """Build one physical rank's compute order for VPP (interleaved 1F1B)."""
910 local_stages = num_stages // pp_degree
911 per_round, accu = _round_layout(num_microbatches, pp_degree, local_stages)
912 warmup_ops = _warmup_ops(
913 stage_index, num_microbatches, pp_degree, local_stages, per_round
914 )
915 fwd_bwd_ops = local_stages * num_microbatches - warmup_ops
916 total_ops = 2 * warmup_ops + fwd_bwd_ops
917 fwd_micro: Dict[int, int] = defaultdict(int)
918 bwd_micro: Dict[int, int] = defaultdict(int)
920 actions: List[Optional[_Action]] = [None] * stage_index
921 actions.extend(
922 _emit_warmup_actions(
923 stage_index,
924 warmup_ops,
925 num_microbatches,
926 pp_degree,
927 per_round,
928 accu,
929 fwd_micro,
930 )
931 )
932 depth = pp_degree - 1 - stage_index
933 actions.extend([None] * max(0, 2 * depth - num_microbatches))
934 actions.extend([None] * (pp_degree - 1 - stage_index))
935 actions.extend(
936 _emit_1f1b_actions(
937 stage_index,
938 warmup_ops,
939 fwd_bwd_ops,
940 num_microbatches,
941 pp_degree,
942 local_stages,
943 per_round,
944 accu,
945 fwd_micro,
946 bwd_micro,
947 )
948 )
949 actions.extend(
950 _emit_cooldown_actions(
951 stage_index,
952 warmup_ops,
953 fwd_bwd_ops,
954 total_ops,
955 num_microbatches,
956 pp_degree,
957 local_stages,
958 per_round,
959 accu,
960 bwd_micro,
961 )
962 )
963 return actions
966def _insert_interleaved_comms(
967 schedule: Dict[int, List[Optional[_Action]]],
968 num_stages: int,
969 pp_degree: int,
970) -> Dict[int, List[Optional[_Action]]]:
971 """Insert SEND/RECV actions with the eager schedule's column scan.
973 Neighbouring ranks alternate send-first / recv-first so the per-peer P2P
974 FIFO stays complementary on the shared PP group; a naive per-action
975 insertion would make both ends send-first and deadlock.
976 """
977 ordered = {rank: [] for rank in range(pp_degree)}
979 def _peer_of(stage: int, delta: int) -> Optional[int]:
980 target = stage + delta
981 if target < 0 or target >= num_stages:
982 return None
983 peer = target % pp_degree
984 return peer if peer != stage % pp_degree else None
986 def _insert(
987 action: Optional[_Action], rank: int, out: Dict[int, List[Optional[_Action]]]
988 ) -> None:
989 if action is None:
990 return
991 if action.kind == _FWD:
992 peer = _peer_of(action.stage, 1)
993 if peer is not None:
994 out[rank].append(_Action(_FWD_SEND, action.stage, action.micro))
995 out[peer].append(_Action(_FWD_RECV, action.stage + 1, action.micro))
996 elif action.kind == _BWD:
997 peer = _peer_of(action.stage, -1)
998 if peer is not None:
999 out[rank].append(_Action(_BWD_SEND, action.stage, action.micro))
1000 out[peer].append(_Action(_BWD_RECV, action.stage - 1, action.micro))
1002 max_length = max(len(actions) for actions in schedule.values())
1003 for time_step in range(max_length):
1004 current: Dict[int, Optional[_Action]] = {}
1005 for rank in range(pp_degree):
1006 if time_step < len(schedule[rank]):
1007 action = schedule[rank][time_step]
1008 current[rank] = action
1009 ordered[rank].append(action)
1010 else:
1011 current[rank] = None
1012 for rank in range(0, pp_degree, 2):
1013 _insert(current.get(rank), rank, ordered)
1014 for rank in range(1, pp_degree, 2):
1015 _insert(current.get(rank), rank, ordered)
1016 return ordered
1019def _interleaved_action_order(
1020 num_stages: int, pp_degree: int, num_microbatches: int
1021) -> Dict[int, List[Optional[_Action]]]:
1022 """Return the full per-physical-rank VPP action schedule."""
1023 if num_stages % pp_degree != 0:
1024 raise ValueError(
1025 f"num_stages ({num_stages}) must be divisible by pp_degree ({pp_degree})"
1026 )
1027 schedule = {
1028 rank: _interleaved_stage_order(rank, num_microbatches, pp_degree, num_stages)
1029 for rank in range(pp_degree)
1030 }
1031 return _insert_interleaved_comms(schedule, num_stages, pp_degree)
1034class ScheduleInterleaved1F1B(PipelineScheduleBase):
1035 """Graph-mode VPP (Interleaved 1F1B) schedule.
1037 Each physical rank owns ``num_stages / pp_degree`` virtual stages. The
1038 model is split into ``num_stages`` chunks and rank ``r`` owns virtual
1039 stages ``r, r + pp_degree, ...``; the action order is the classic
1040 interleaved-1F1B warmup / steady / cooldown sequence, and P2P actions are
1041 inserted with the same complementary column scan as the eager schedule
1042 (``core/pipeline_parallel/scheduler.py``).
1044 Args:
1045 fwd_gms: Forward subgraphs, ordered by ``stage_indices``.
1046 bwd_gms: Backward subgraphs, same order.
1047 stage_indices: Virtual stages this physical rank owns, ascending.
1048 num_stages: Total virtual stage count (``pp_degree * chunks``).
1049 pp_degree: Number of physical pipeline ranks.
1050 state_positions: Per-stage positions into the flat state argument.
1051 num_states: Per-stage state-tensor count.
1052 num_trainables: Per-stage trainable-parameter count.
1053 num_sends: Per-stage forward boundary value count.
1054 grad_send_counts: Per-stage backward boundary value count.
1055 microbatch_size: Samples per microbatch.
1056 pp_group: PP process group.
1057 recv_specs: Per-stage forward receive descriptors.
1058 grad_recv_specs: Per-stage gradient receive descriptors.
1059 user_input_stages: Per user input, its consuming virtual stage.
1060 grad_owner: For each emitted gradient in flat state order, the index
1061 in ``stage_indices`` that owns it.
1062 """
1064 def __init__(
1065 self,
1066 fwd_gms: Sequence[nn.Module],
1067 bwd_gms: Sequence[nn.Module],
1068 stage_indices: Sequence[int],
1069 num_stages: int,
1070 pp_degree: int,
1071 state_positions: Sequence[Sequence[int]],
1072 num_states: Sequence[int],
1073 num_trainables: Sequence[int],
1074 num_sends: Sequence[int],
1075 grad_send_counts: Sequence[int],
1076 microbatch_size: int,
1077 pp_group: Any = None,
1078 recv_specs: Sequence[Sequence[Tuple[Any, ...]]] = (),
1079 grad_recv_specs: Sequence[Sequence[Tuple[Any, ...]]] = (),
1080 user_input_stages: Sequence[int] = (),
1081 grad_owner: Sequence[int] = (),
1082 user_input_stage_sets: Optional[Sequence[Sequence[int]]] = None,
1083 ) -> None:
1084 """Store the per-chunk subgraphs and boundary contracts."""
1085 super().__init__(
1086 fwd_gms[0],
1087 bwd_gms[0],
1088 stage_idx=stage_indices[0],
1089 pp_degree=pp_degree,
1090 num_state=sum(num_states),
1091 num_trainable=num_trainables[0],
1092 num_send=num_sends[0],
1093 grad_send_count=grad_send_counts[0],
1094 microbatch_size=microbatch_size,
1095 pp_group=pp_group,
1096 recv_spec=recv_specs[0] if recv_specs else (),
1097 grad_recv_spec=grad_recv_specs[0] if grad_recv_specs else (),
1098 user_input_stages=user_input_stages,
1099 user_input_stage_sets=user_input_stage_sets,
1100 )
1101 self.fwd_gm = None
1102 self.bwd_gm = None
1103 self.fwd_gms = nn.ModuleList(fwd_gms)
1104 self.bwd_gms = nn.ModuleList(bwd_gms)
1105 self.stage_indices = list(stage_indices)
1106 self.num_stages = num_stages
1107 self.pp_degree = pp_degree
1108 self.rank = self.stage_indices[0] % pp_degree if self.stage_indices else 0
1109 self.state_positions = [list(positions) for positions in state_positions]
1110 self.num_states = list(num_states)
1111 self.num_trainables = list(num_trainables)
1112 self.num_sends = list(num_sends)
1113 self.grad_send_counts = list(grad_send_counts)
1114 self.microbatch_size = microbatch_size
1115 self.pp_group = pp_group
1116 self.recv_specs = [list(spec) for spec in recv_specs]
1117 self.grad_recv_specs = [list(spec) for spec in grad_recv_specs]
1118 self.user_input_stages = list(user_input_stages)
1119 self.user_input_stage_sets = (
1120 [list(stages) for stages in user_input_stage_sets]
1121 if user_input_stage_sets is not None
1122 else [[stage] for stage in self.user_input_stages]
1123 )
1124 self.grad_owner = list(grad_owner)
1125 self.num_state = sum(self.num_states)
1126 self._owned_pos = {v: i for i, v in enumerate(self.stage_indices)}
1127 self._pending_sends: Dict[Any, List[Tuple[Any, torch.Tensor]]] = {}
1129 def map_stage_graphs(self, fn: Callable[[nn.Module], nn.Module]) -> None:
1130 """Apply ``fn`` to every virtual stage's fwd and bwd subgraph."""
1131 self.fwd_gms = nn.ModuleList([fn(gm) for gm in self.fwd_gms])
1132 self.bwd_gms = nn.ModuleList([fn(gm) for gm in self.bwd_gms])
1134 def _resolve_num_microbatches(self, user_inputs: Sequence[Any]) -> int:
1135 """Derive the microbatch count from the first tensor model input."""
1136 batch_size = next(
1137 (v.shape[0] for v in user_inputs if isinstance(v, torch.Tensor)), None
1138 )
1139 if batch_size is None:
1140 raise ValueError(
1141 "PP requires at least one tensor model input to derive the "
1142 "microbatch split"
1143 )
1144 if batch_size % self.microbatch_size != 0:
1145 raise ValueError(
1146 f"PP microbatch mismatch: batch size {batch_size} is not "
1147 f"divisible by pp_microbatch_size {self.microbatch_size}"
1148 )
1149 return batch_size // self.microbatch_size
1151 def _apply_stage_forward(
1152 self,
1153 index: int,
1154 stage_idx: int,
1155 state: Sequence[Any],
1156 act_in: Sequence[Any],
1157 owned_inputs: Sequence[Any],
1158 mb_index: int,
1159 ) -> Tuple[Any, ...]:
1160 """Run one virtual stage's forward subgraph for one microbatch."""
1161 owned_mb = self._slice_inputs(owned_inputs, mb_index)
1162 if stage_idx == 0:
1163 out = self.fwd_gms[index](*state, *owned_mb)
1164 else:
1165 out = self.fwd_gms[index](*state, *act_in, *owned_mb)
1166 return tuple(out)
1168 def forward(self, *flat_inputs: Any) -> Tuple[torch.Tensor, ...]:
1169 """Run one full interleaved-1F1B step across this rank's chunks."""
1170 state, user_inputs = self._split_inputs(flat_inputs)
1171 num_microbatches = self._resolve_num_microbatches(user_inputs)
1172 order = _interleaved_action_order(
1173 self.num_stages, self.pp_degree, num_microbatches
1174 )[self.rank]
1175 step = _InterleavedStep(
1176 state=state,
1177 user_inputs=user_inputs,
1178 act_cache={},
1179 grad_cache={},
1180 fwd_outs={},
1181 bwd_outs={},
1182 grads=[[] for _ in self.stage_indices],
1183 losses=[None] * num_microbatches,
1184 )
1186 for action in order:
1187 if action is not None:
1188 self._run_action(action, step)
1190 self._flush_pending_sends()
1191 loss = self._finalize_interleaved_loss(step.losses, user_inputs)
1192 return (loss, *self._ordered_grads(step.grads, num_microbatches))
1194 def _run_action(self, action: _Action, step: _InterleavedStep) -> None:
1195 """Apply one forward or backward action to the step buffers."""
1196 stage_idx = action.stage
1197 mb_index = action.micro
1198 key = (stage_idx, mb_index)
1199 index = self._owned_pos[stage_idx]
1200 if action.kind == _FWD_RECV:
1201 step.act_cache[key] = self._recv_values(
1202 self.recv_specs[index], src=(stage_idx - 1) % self.pp_degree
1203 )
1204 elif action.kind == _FWD:
1205 state_v = [step.state[pos] for pos in self.state_positions[index]]
1206 owned_inputs = [
1207 value
1208 for value, owners in zip(step.user_inputs, self.user_input_stage_sets)
1209 if stage_idx in owners
1210 ]
1211 step.fwd_outs[key] = self._apply_stage_forward(
1212 index,
1213 stage_idx,
1214 state_v,
1215 step.act_cache.pop(key, []),
1216 owned_inputs,
1217 mb_index,
1218 )
1219 elif action.kind == _FWD_SEND:
1220 self._send_values(
1221 step.fwd_outs[key][: self.num_sends[index]],
1222 dst=(stage_idx + 1) % self.pp_degree,
1223 mb_index=(_FWD_SEND, stage_idx, mb_index),
1224 )
1225 elif action.kind == _BWD_RECV:
1226 step.grad_cache[key] = self._recv_values(
1227 self.grad_recv_specs[index],
1228 src=(stage_idx + 1) % self.pp_degree,
1229 )
1230 elif action.kind == _BWD:
1231 state_v = [step.state[pos] for pos in self.state_positions[index]]
1232 if stage_idx == self.num_stages - 1:
1233 grad_in = [torch.ones_like(step.fwd_outs[key][0])]
1234 else:
1235 grad_in = step.grad_cache.pop(key)
1236 out = self.bwd_gms[index](*state_v, *grad_in, *step.fwd_outs[key])
1237 trainable = self.num_trainables[index]
1238 step.grads[index] = self._accumulate_grads(
1239 step.grads[index], out[:trainable]
1240 )
1241 if self.grad_send_counts[index]:
1242 step.bwd_outs[key] = tuple(out[trainable:])
1243 if stage_idx == self.num_stages - 1:
1244 step.losses[mb_index] = step.fwd_outs[key][0]
1245 step.fwd_outs.pop(key, None)
1246 elif action.kind == _BWD_SEND:
1247 self._send_values(
1248 step.bwd_outs.pop(key),
1249 dst=(stage_idx - 1) % self.pp_degree,
1250 mb_index=(_BWD_SEND, stage_idx, mb_index),
1251 )
1253 def _finalize_interleaved_loss(
1254 self,
1255 losses: Sequence[Optional[torch.Tensor]],
1256 user_inputs: Sequence[Any],
1257 ) -> torch.Tensor:
1258 """Return the mean last-stage loss or this rank's zero placeholder."""
1259 if self.num_stages - 1 in self._owned_pos:
1260 return torch.stack(losses).mean()
1261 return torch.zeros((), device=self._anchor_device(user_inputs))
1263 def _ordered_grads(
1264 self,
1265 grads: Sequence[Sequence[torch.Tensor]],
1266 num_microbatches: int,
1267 ) -> List[torch.Tensor]:
1268 """Average per-stage grads and flatten them into state-FQN order."""
1269 averaged = [
1270 self._average_grads(stage_grads, num_microbatches) for stage_grads in grads
1271 ]
1272 pointers = [0] * len(averaged)
1273 ordered: List[torch.Tensor] = []
1274 for owner in self.grad_owner:
1275 ordered.append(averaged[owner][pointers[owner]])
1276 pointers[owner] += 1
1277 return ordered
1280SCHEDULE_REGISTRY: Dict[str, Type[PipelineScheduleBase]] = {
1281 "gpipe": ScheduleGPipe,
1282 "1f1b": Schedule1F1B,
1283 "interleaved_1f1b": ScheduleInterleaved1F1B,
1284 "vpp": ScheduleInterleaved1F1B,
1285}
1288def get_schedule_class(name: str) -> Type[PipelineScheduleBase]:
1289 """Resolve a ``PassConfig.pp_schedule`` name to its schedule class.
1291 Args:
1292 name: Schedule name, case-insensitive (``"gpipe"`` / ``"1f1b"``).
1294 Returns:
1295 The schedule class.
1297 Raises:
1298 ValueError: When ``name`` is not a registered schedule.
1299 """
1300 key = str(name).lower()
1301 schedule_cls = SCHEDULE_REGISTRY.get(key)
1302 if schedule_cls is None:
1303 raise ValueError(
1304 f"Unknown pp_schedule {name!r}; expected one of {sorted(SCHEDULE_REGISTRY)}"
1305 )
1306 return schedule_cls