Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / parallel / pp_schedule.py: 98%
209 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""
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` and :class:`Schedule1F1B` today — dispatched by name
22from :data:`SCHEDULE_REGISTRY` / ``PassConfig.pp_schedule``.
24Both schedules run the per-stage FX subgraphs produced by ``PpPass`` across
25microbatches and share the same P2P / loss / gradient contract:
27- Model inputs are routed by dataflow: each user input is owned by the single
28 stage that consumes it (resolved by ``PpPass``) and is fed to that stage's
29 forward subgraph directly — every rank receives the full flattened
30 user-input list, slices its OWNED tensor inputs per microbatch, and ignores
31 the rest. All tensor user inputs must carry the leading batch dim.
32- Boundary values are exchanged as ORDERED LISTS: tensors go as-is, int
33 scalars (dynamic-shape ``sym_size`` nodes crossing the cut) are packed as
34 0-d int64 tensors and unwrapped with ``item()`` on arrival — mirroring how
35 ``torch.distributed.pipelining``'s ``PipelineStage`` ships full argument
36 lists. The P2P exchange uses eager ``dist.isend``/``irecv`` on the PP
37 process group; sends are async so the sweeps overlap receive/compute across
38 stages, receives wait before first use, and each microbatch's boundary sends
39 are retained under that microbatch's key — forward and backward boundaries
40 share the key — until their ``Work`` handles are waited (1F1B: one turn
41 after the backward posts, so a forward send rides along; GPipe: at the end
42 of the step).
43- Microbatch gradients are un-normalized sums while sweeping, then divided by
44 ``num_microbatches`` at the end, so the step gradient matches the non-PP
45 semantics of a full-batch mean loss.
47The schedules differ in their action order — and in the activation lifetime
48that order allows:
50``ScheduleGPipe``
51 1. Forward sweep: run every microbatch's forward (shipping boundary
52 activations downstream), holding all forward outputs.
53 2. Backward sweep (reverse order): seed ``ones_like(loss)`` on the last
54 stage, run each backward against the activations saved by the forward
55 sweep, and pass boundary gradients one stage back.
57``Schedule1F1B``
58 Warmup runs ``min(num_microbatches, pp_degree - stage_idx)`` forwards (a
59 stage nearer the front keeps one more microbatch in flight), then the
60 steady state alternates one backward with one forward per turn, and a
61 cooldown drains the remaining backwards. Interleaving backward earlier
62 shrinks the pipeline bubble relative to GPipe. Requires
63 ``num_microbatches >= pp_degree``.
65Forward activations and backward gradients share one PP process group, whose
66per-pair FIFO must stay complementary — at each position one end sends while
67the other receives — or the group deadlocks. ``Schedule1F1B`` staggers its
68boundary sends to hold that invariant despite the ``pp_degree``-asymmetric
69warmup; the placement mirrors ``core/pipeline_parallel/scheduler.py`` and is
70spelled out in ``Schedule1F1B.forward``.
71"""
73__all__ = [
74 "PipelineScheduleBase",
75 "ScheduleGPipe",
76 "Schedule1F1B",
77 "SCHEDULE_REGISTRY",
78 "get_schedule_class",
79]
81from typing import Any, Dict, List, Sequence, Tuple, Type
83import torch
84import torch.distributed as dist
85from torch import nn
88class PipelineScheduleBase(nn.Module):
89 """Shared P2P / microbatch machinery for the graph-mode PP schedules.
91 Installed by ``PpPass`` as a submodule of the compiled GraphModule and
92 invoked through a ``call_module`` node, so the trainer's
93 ``graph_module(*flat_inputs)`` dispatches to the concrete schedule
94 unchanged. Subclasses implement :meth:`forward` (the action order) from
95 the per-microbatch building blocks defined here.
97 Args:
98 fwd_gm: Stage forward subgraph. Signature per stage:
99 stage 0 -> ``(*state, *owned_inputs)`` returning
100 ``(*act_out, *saved)``;
101 other -> ``(*state, *act_in, *owned_inputs)`` returning
102 ``(*act_out, *saved)`` for non-last stages and
103 ``(loss, *saved)`` for the last one, where
104 ``owned_inputs`` are the user inputs routed to
105 this stage (may be empty).
106 bwd_gm: Stage backward subgraph. Signature:
107 ``(*state, *grad_in, *fwd_outs)`` where ``grad_in`` are the
108 received boundary gradients (``ones_like(loss)`` on the last
109 stage) and ``fwd_outs`` are the forward subgraph's outputs
110 verbatim. Returns ``(*param_grads,)`` plus the boundary
111 gradients for every stage except stage 0.
112 stage_idx: This rank's stage index (0-based).
113 pp_degree: Number of pipeline stages.
114 num_state: Number of leading state tensors (stage params/buffers).
115 num_trainable: Number of trainable parameters on this stage; the
116 backward subgraph returns exactly this many parameter gradients
117 first.
118 num_send: Number of forward boundary values to ship downstream
119 (the forward outputs' prefix); 0 on the last stage.
120 grad_send_count: Number of backward boundary values to ship
121 upstream (the backward outputs' suffix); 0 on stage 0.
122 microbatch_size: Samples per microbatch (leading dim).
123 pp_group: PP process group used for all boundary P2P (activations and
124 gradients) and for stage->global rank mapping.
125 recv_spec: Per-value descriptors for the incoming forward
126 boundary: ``("tensor", shape, dtype, device)`` or
127 ``("scalar",)``. Empty on stage 0.
128 grad_recv_spec: Same, for the incoming backward boundary
129 gradients. Empty on the last stage.
130 user_input_stages: Per flattened user input (stub arg order), the
131 stage that consumes it. This stage slices and feeds the inputs
132 mapped to itself and ignores the rest.
133 """
135 def __init__(
136 self,
137 fwd_gm: nn.Module,
138 bwd_gm: nn.Module,
139 stage_idx: int,
140 pp_degree: int,
141 num_state: int,
142 num_trainable: int,
143 num_send: int,
144 grad_send_count: int,
145 microbatch_size: int,
146 pp_group: Any = None,
147 recv_spec: Sequence[Tuple[Any, ...]] = (),
148 grad_recv_spec: Sequence[Tuple[Any, ...]] = (),
149 user_input_stages: Sequence[int] = (),
150 ) -> None:
151 """Store the stage subgraphs, P2P group, and boundary specs."""
152 super().__init__()
153 self.fwd_gm = fwd_gm
154 self.bwd_gm = bwd_gm
155 self.stage_idx = stage_idx
156 self.pp_degree = pp_degree
157 self.num_state = num_state
158 self.num_trainable = num_trainable
159 self.num_send = num_send
160 self.grad_send_count = grad_send_count
161 self.microbatch_size = microbatch_size
162 self.pp_group = pp_group
163 self.recv_spec = list(recv_spec)
164 self.grad_recv_spec = list(grad_recv_spec)
165 self.user_input_stages = list(user_input_stages)
166 self.is_first = stage_idx == 0
167 self.is_last = stage_idx == pp_degree - 1
168 # In-flight isend Work handles keyed by microbatch index — forward and
169 # backward boundaries of one microbatch share the key; the sent
170 # tensors stay referenced so their buffers outlive the send. Schedules
171 # release a microbatch's buffers once its sends provably completed
172 # (1F1B: next turn's gradient receipt) or at the end of the step.
173 self._pending_sends: Dict[int, List[Tuple[Any, torch.Tensor]]] = {}
175 # ------------------------------------------------------------------
176 # Input / microbatch resolution
177 # ------------------------------------------------------------------
179 def _split_inputs(
180 self, flat_inputs: Sequence[Any]
181 ) -> Tuple[Sequence[Any], List[Any]]:
182 """Split ``(*state, *user_inputs)`` and validate the input arity."""
183 state = flat_inputs[: self.num_state]
184 user_inputs = list(flat_inputs[self.num_state :])
185 if len(user_inputs) != len(self.user_input_stages):
186 raise RuntimeError(
187 f"PP schedule expects {len(self.user_input_stages)} user "
188 f"inputs (per its routing) but received {len(user_inputs)}"
189 )
190 return state, user_inputs
192 def _resolve_microbatches(
193 self, user_inputs: Sequence[Any]
194 ) -> Tuple[int, List[Any]]:
195 """Derive the microbatch count and this stage's owned user inputs.
197 Every rank receives every user input, so the batch size is derived
198 identically on all stages from the first tensor input.
199 """
200 batch_size = next(
201 (v.shape[0] for v in user_inputs if isinstance(v, torch.Tensor)), None
202 )
203 if batch_size is None:
204 raise ValueError(
205 "PP requires at least one tensor model input to derive the "
206 "microbatch split"
207 )
208 if batch_size % self.microbatch_size != 0:
209 raise ValueError(
210 f"PP microbatch mismatch: batch size {batch_size} is not "
211 f"divisible by pp_microbatch_size {self.microbatch_size}"
212 )
213 num_microbatches = batch_size // self.microbatch_size
214 owned_inputs = [
215 v
216 for v, s in zip(user_inputs, self.user_input_stages)
217 if s == self.stage_idx
218 ]
219 return num_microbatches, owned_inputs
221 def _slice_inputs(self, owned_inputs: Sequence[Any], mb_index: int) -> List[Any]:
222 """Slice this stage's owned inputs for microbatch ``mb_index``.
224 Tensors are sliced along the leading batch dim; non-tensors pass
225 through unchanged.
226 """
227 lo = mb_index * self.microbatch_size
228 hi = lo + self.microbatch_size
229 return [v[lo:hi] if isinstance(v, torch.Tensor) else v for v in owned_inputs]
231 # ------------------------------------------------------------------
232 # Per-microbatch building blocks
233 # ------------------------------------------------------------------
234 #
235 # Forward / backward are split into a receive step and a compute step so
236 # a schedule can interleave boundary sends with compute independently;
237 # ``_forward_microbatch`` / ``_backward_microbatch`` compose them for
238 # callers that don't need that control.
240 def _recv_activation(self) -> List[Any]:
241 """Receive this stage's incoming boundary activations (none on stage 0)."""
242 if self.is_first:
243 return []
244 return self._recv_values(self.recv_spec, src=self.stage_idx - 1)
246 def _apply_forward(
247 self,
248 state: Sequence[Any],
249 act_in: Sequence[Any],
250 owned_inputs: Sequence[Any],
251 mb_index: int,
252 ) -> Tuple[Any, ...]:
253 """Run one microbatch's forward subgraph against ``act_in``.
255 ``owned_inputs`` are this stage's user inputs (see
256 ``user_input_stages``): tensors are sliced per microbatch along the
257 leading batch dim, non-tensors pass through. They are appended after
258 the received activations in the fwd subgraph's arg order.
259 """
260 owned_mb = self._slice_inputs(owned_inputs, mb_index)
261 if self.is_first:
262 out = self.fwd_gm(*state, *owned_mb)
263 else:
264 out = self.fwd_gm(*state, *act_in, *owned_mb)
265 return tuple(out)
267 def _forward_microbatch(
268 self, state: Sequence[Any], owned_inputs: Sequence[Any], mb_index: int
269 ) -> Tuple[Any, ...]:
270 """Receive activations then run one microbatch's forward subgraph."""
271 act_in = self._recv_activation()
272 return self._apply_forward(state, act_in, owned_inputs, mb_index)
274 def _send_forward(self, out: Sequence[Any], mb_index: int) -> None:
275 """Ship microbatch ``mb_index``'s boundary activations downstream."""
276 if self.num_send:
277 self._send_values(
278 out[: self.num_send], dst=self.stage_idx + 1, mb_index=mb_index
279 )
281 def _recv_gradient(self, fwd_outs: Sequence[Any]) -> List[Any]:
282 """Receive the next stage's output grads (a seed on the last stage).
284 The traced backward graph seeds its own loss gradient internally, so
285 the ``ones_like`` seed on the last stage only fixes the shape.
286 """
287 if self.is_last:
288 return [torch.ones_like(fwd_outs[0])]
289 return self._recv_values(self.grad_recv_spec, src=self.stage_idx + 1)
291 def _apply_backward(
292 self,
293 state: Sequence[Any],
294 grad_in: Sequence[Any],
295 fwd_outs: Sequence[Any],
296 mb_index: int,
297 ) -> Sequence[torch.Tensor]:
298 """Run microbatch ``mb_index``'s backward subgraph and ship grads.
300 Gradients flow UPSTREAM, so the previous stage is the destination,
301 mirroring the forward direction. The grads returned are un-normalized
302 sums over the microbatches; the shipped boundary grads are keyed by
303 ``mb_index`` so the schedule can release their buffers later.
304 """
305 out = self.bwd_gm(*state, *grad_in, *fwd_outs)
306 param_grads = out[: self.num_trainable]
307 if self.grad_send_count:
308 self._send_values(
309 out[self.num_trainable :], dst=self.stage_idx - 1, mb_index=mb_index
310 )
311 return param_grads
313 def _backward_microbatch(
314 self, state: Sequence[Any], fwd_outs: Sequence[Any], mb_index: int
315 ) -> Sequence[torch.Tensor]:
316 """Receive boundary grads then run one microbatch's backward subgraph."""
317 grad_in = self._recv_gradient(fwd_outs)
318 return self._apply_backward(state, grad_in, fwd_outs, mb_index)
320 @staticmethod
321 def _accumulate_grads(
322 accumulated: Sequence[torch.Tensor], param_grads: Sequence[torch.Tensor]
323 ) -> List[torch.Tensor]:
324 """Sum one microbatch's param grads into the running accumulation."""
325 if not accumulated:
326 return list(param_grads)
327 return [acc + grad for acc, grad in zip(accumulated, param_grads)]
329 def _finalize_loss(
330 self, fwd_outs_per_mb: Sequence[Sequence[Any]], user_inputs: Sequence[Any]
331 ) -> torch.Tensor:
332 """Return the step loss (last stage) or a zero placeholder."""
333 if self.is_last:
334 # On the last stage each forward output tuple starts with the
335 # per-microbatch loss.
336 losses = [outs[0] for outs in fwd_outs_per_mb]
337 return torch.stack(losses).mean()
338 return torch.zeros((), device=self._anchor_device(user_inputs))
340 @staticmethod
341 def _average_grads(
342 grads: Sequence[torch.Tensor], num_microbatches: int
343 ) -> List[torch.Tensor]:
344 """Turn un-normalized sums over microbatches into the mean."""
345 return [grad / num_microbatches for grad in grads]
347 # ------------------------------------------------------------------
348 # P2P helpers
349 # ------------------------------------------------------------------
351 def _flush_pending_sends(self) -> None:
352 """Wait every in-flight isend and drop all buffer references."""
353 for sends in self._pending_sends.values():
354 for work, _ in sends:
355 work.wait()
356 self._pending_sends = {}
358 def _wait_sends(self, mb_index: int) -> None:
359 """Wait microbatch ``mb_index``'s in-flight sends, freeing its buffers.
361 Callers must only invoke this once the sends are provably complete
362 (e.g. the peer consumed the values), so the wait never blocks.
363 """
364 for work, _ in self._pending_sends.pop(mb_index, ()):
365 work.wait()
367 @staticmethod
368 def _anchor_device(user_inputs: Sequence[Any]) -> torch.device:
369 """Device for the zero-loss placeholder: first tensor user input."""
370 for value in user_inputs:
371 if isinstance(value, torch.Tensor):
372 return value.device
373 return torch.device("cpu")
375 def _send_values(self, values: Sequence[Any], dst: int, mb_index: int) -> None:
376 """Async-send boundary values to the neighbouring stage ``dst``.
378 Tensors go as-is (made contiguous, with THAT buffer retained); int
379 scalars (dynamic-shape ``sym_size`` nodes) are packed as 0-d int64
380 tensors on the first tensor value's device — CPU tensors cannot ride
381 an NCCL/NPU-backend group — and unwrapped with ``item()`` on the
382 receiving side. The ``Work`` handles are queued under ``mb_index``
383 (forward and backward boundaries of a microbatch share the key) and
384 the sent tensors kept referenced so their buffers outlive the send;
385 the schedule waits them (per turn or at step end) before the buffers
386 are released.
387 """
388 dst_rank = self._global_rank(dst)
389 anchor_device = next(
390 (v.device for v in values if isinstance(v, torch.Tensor)), None
391 )
392 for value in values:
393 if isinstance(value, torch.Tensor):
394 tensor = value.contiguous()
395 else:
396 tensor = torch.tensor(value, dtype=torch.int64, device=anchor_device)
397 work = dist.isend(tensor, dst=dst_rank, group=self.pp_group)
398 self._pending_sends.setdefault(mb_index, []).append((work, tensor))
400 def _recv_values(self, spec: Sequence[Tuple[Any, ...]], src: int) -> List[Any]:
401 """Receive boundary values from the neighbouring stage ``src``, in order.
403 The buffer for each entry comes from ``spec`` (built by ``PpPass``
404 from the boundary values' fake-tensor ``val`` metas, including the
405 device each value was traced on); scalar entries arrive as 0-d
406 int64 tensors and are unwrapped to Python ints so the graph sees
407 the same types it was traced with. Scalar buffers are anchored to
408 the first tensor entry's device (``PpPass._value_spec`` guarantees
409 a non-empty list carries at least one tensor).
410 """
411 src_rank = self._global_rank(src)
412 anchor_device = next((entry[3] for entry in spec if entry[0] == "tensor"), None)
413 values: List[Any] = []
414 for entry in spec:
415 if entry[0] == "scalar":
416 buffer = torch.empty((), dtype=torch.int64, device=anchor_device)
417 else:
418 _, shape, dtype, device = entry
419 buffer = torch.empty(shape, dtype=dtype, device=device)
420 work = dist.irecv(buffer, src=src_rank, group=self.pp_group)
421 work.wait()
422 values.append(int(buffer.item()) if entry[0] == "scalar" else buffer)
423 return values
425 def _global_rank(self, stage_idx: int) -> int:
426 """Map a stage index to its global rank within the PP group."""
427 if self.pp_group is None:
428 return stage_idx
429 return dist.get_global_rank(self.pp_group, stage_idx)
432class ScheduleGPipe(PipelineScheduleBase):
433 """GPipe driver over two per-stage FX subgraphs (forward / backward).
435 Runs ALL forwards before ANY backward, so every microbatch's forward
436 outputs (boundary activations + saved values) are held until the backward
437 sweep completes — activation memory scales with ``num_microbatches``.
438 Larger ``microbatch_size`` trades pipeline bubble for fewer P2P
439 round-trips and fewer retained activations.
440 """
442 def forward(self, *flat_inputs: Any) -> Tuple[torch.Tensor, ...]:
443 """Run one full GPipe step (fwd sweep then bwd sweep).
445 Args:
446 flat_inputs: ``(*state, *user_inputs)`` — the same surface
447 ``run_traced_graph`` feeds the compiled graph, with
448 ``user_inputs`` the flattened model inputs in trace order.
450 Returns:
451 ``(loss, *param_grads)``. On non-last stages the loss is a zero
452 scalar placeholder (the real loss lives on the last stage);
453 parameter gradients are the microbatch-averaged accumulation.
454 """
455 state, user_inputs = self._split_inputs(flat_inputs)
456 num_microbatches, owned_inputs = self._resolve_microbatches(user_inputs)
458 fwd_outs_per_mb = self._forward_sweep(state, owned_inputs, num_microbatches)
459 grads = self._backward_sweep(state, fwd_outs_per_mb, num_microbatches)
461 self._flush_pending_sends()
463 loss = self._finalize_loss(fwd_outs_per_mb, user_inputs)
464 return (loss, *grads)
466 def _forward_sweep(
467 self,
468 state: Sequence[Any],
469 owned_inputs: Sequence[Any],
470 num_microbatches: int,
471 ) -> List[Tuple[Any, ...]]:
472 """Run forward microbatch 0..N-1, shipping boundary values downstream.
474 Returns per-microbatch forward output tuples (replayed verbatim
475 into the backward subgraph by ``_backward_sweep``).
476 """
477 fwd_outs_per_mb: List[Tuple[Any, ...]] = []
478 for mb_index in range(num_microbatches):
479 out = self._forward_microbatch(state, owned_inputs, mb_index)
480 fwd_outs_per_mb.append(out)
481 self._send_forward(out, mb_index)
482 return fwd_outs_per_mb
484 def _backward_sweep(
485 self,
486 state: Sequence[Any],
487 fwd_outs_per_mb: Sequence[Tuple[Any, ...]],
488 num_microbatches: int,
489 ) -> List[torch.Tensor]:
490 """Run backward microbatches N-1..0, accumulating full-batch gradients.
492 The traced backward graph seeds its own loss gradient internally
493 (an ``ones_like`` over the saved loss), so the grads that come back
494 are un-normalized sums over the microbatches. Dividing the
495 accumulation by ``num_microbatches`` turns that into the mean, which
496 is exactly the full-batch mean-loss gradient the non-PP path
497 produces.
498 """
499 grads: List[torch.Tensor] = []
500 for mb_index in reversed(range(num_microbatches)):
501 grads = self._accumulate_grads(
502 grads,
503 self._backward_microbatch(state, fwd_outs_per_mb[mb_index], mb_index),
504 )
505 return self._average_grads(grads, num_microbatches)
508class Schedule1F1B(PipelineScheduleBase):
509 """1F1B driver over two per-stage FX subgraphs (forward / backward).
511 Warmup runs ``min(num_microbatches, pp_degree - stage_idx)`` forwards
512 (one on the last stage), the steady state alternates one backward then
513 one forward per turn, and the cooldown drains the remaining backwards.
514 Running each backward as soon as its forward output is available shrinks
515 the pipeline bubble relative to GPipe's fill-drain order.
517 Requires ``num_microbatches >= pp_degree`` (checked at runtime), the same
518 constraint ``torch.distributed.pipelining.Schedule1F1B`` enforces.
520 Note:
521 Unlike GPipe, each microbatch's saved activations are freed as soon
522 as its backward consumes them. Its boundary sends share one key for
523 both directions and are released one turn after the backward posts —
524 the next turn's gradient receipt proves the previous microbatch's
525 transfers completed, so the wait cannot deadlock. Peak memory
526 therefore tracks the warmup depth (``pp_degree - stage_idx``) instead
527 of ``num_microbatches``; only the final microbatch's sends are still
528 waited at the end of the step. The last stage has no downstream
529 gradient to mark that progress, so its per-turn wait may briefly block
530 on the upstream peer.
531 """
533 def forward(self, *flat_inputs: Any) -> Tuple[torch.Tensor, ...]:
534 """Run one full 1F1B step (warmup -> steady 1B1F -> cooldown).
536 Args:
537 flat_inputs: ``(*state, *user_inputs)`` — the same surface
538 ``run_traced_graph`` feeds the compiled graph, with
539 ``user_inputs`` the flattened model inputs in trace order.
541 Returns:
542 ``(loss, *param_grads)``. On non-last stages the loss is a zero
543 scalar placeholder (the real loss lives on the last stage);
544 parameter gradients are the microbatch-averaged accumulation.
546 Raises:
547 ValueError: When ``num_microbatches < pp_degree`` (the schedule
548 cannot fill the pipeline).
549 """
550 state, user_inputs = self._split_inputs(flat_inputs)
551 num_microbatches, owned_inputs = self._resolve_microbatches(user_inputs)
552 if num_microbatches < self.pp_degree:
553 raise ValueError(
554 f"1F1B requires at least pp_degree={self.pp_degree} "
555 f"microbatches, got {num_microbatches} "
556 f"(batch {num_microbatches * self.microbatch_size} with "
557 f"pp_microbatch_size={self.microbatch_size}); lower "
558 f"pp_microbatch_size or use the gpipe schedule"
559 )
561 fwd_outs_per_mb: List[Any] = [None] * num_microbatches
562 grads: List[torch.Tensor] = []
563 fwd_index = 0
564 bwd_index = 0
565 # Forward output whose boundary activation has not shipped yet, with
566 # its microbatch index; the send is deferred one slot to keep the
567 # pair's FIFO complementary.
568 pending_index = -1
569 pending_out: Any = None
570 even = self.stage_idx % 2 == 0
572 # Warmup: fill the pipeline with ``pp_degree - stage_idx`` in-flight
573 # microbatches (one on the last stage) before the first backward.
574 # Even stages ship each activation right after its forward, except
575 # the last; odd stages ship the previous one just before the next
576 # forward. The parity stagger keeps neighbours complementary.
577 warmup = min(num_microbatches, self.pp_degree - self.stage_idx)
578 for k in range(warmup):
579 act_in = self._recv_activation()
580 if not even and k > 0:
581 self._send_forward(pending_out, pending_index)
582 pending_out = None
583 out = self._apply_forward(state, act_in, owned_inputs, fwd_index)
584 fwd_outs_per_mb[fwd_index] = out
585 fwd_index += 1
586 if even:
587 if k != warmup - 1:
588 self._send_forward(out, fwd_index - 1)
589 else:
590 pending_index, pending_out = fwd_index - 1, out
591 else:
592 pending_index, pending_out = fwd_index - 1, out
594 # Steady turn: recv_grad (+ release the previous turn's boundary
595 # buffers), send_activation, backward, send_grad, recv_activation,
596 # forward.
597 for _ in range(num_microbatches - warmup):
598 grad_in = self._recv_gradient(fwd_outs_per_mb[bwd_index])
599 if bwd_index > 0:
600 # This gradient's arrival proves the peer consumed microbatch
601 # ``bwd_index - 1``'s boundary values on both directions, so
602 # its in-flight sends completed and their buffers can go.
603 self._wait_sends(bwd_index - 1)
604 if pending_out is not None:
605 self._send_forward(pending_out, pending_index)
606 pending_out = None
607 grads = self._accumulate_grads(
608 grads,
609 self._apply_backward(
610 state, grad_in, fwd_outs_per_mb[bwd_index], bwd_index
611 ),
612 )
613 # Backward was the last consumer of these outputs; keep only the
614 # last stage's loss scalar for _finalize_loss.
615 fwd_outs_per_mb[bwd_index] = (
616 fwd_outs_per_mb[bwd_index][:1] if self.is_last else None
617 )
618 bwd_index += 1
619 act_in = self._recv_activation()
620 out = self._apply_forward(state, act_in, owned_inputs, fwd_index)
621 fwd_outs_per_mb[fwd_index] = out
622 fwd_index += 1
623 pending_index, pending_out = fwd_index - 1, out
624 # Drop the frame-local reference to the final forward's outputs (the
625 # warmup's when the steady state never runs); it would otherwise pin
626 # them through the cooldown.
627 out = None
629 # Cooldown: drain the remaining backwards; the final forward's
630 # activation ships on the first turn.
631 for _ in range(warmup):
632 grad_in = self._recv_gradient(fwd_outs_per_mb[bwd_index])
633 if bwd_index > 0:
634 # Same release invariant as the steady state.
635 self._wait_sends(bwd_index - 1)
636 if pending_out is not None:
637 self._send_forward(pending_out, pending_index)
638 pending_out = None
639 grads = self._accumulate_grads(
640 grads,
641 self._apply_backward(
642 state, grad_in, fwd_outs_per_mb[bwd_index], bwd_index
643 ),
644 )
645 # Backward was the last consumer of these outputs; keep only the
646 # last stage's loss scalar for _finalize_loss.
647 fwd_outs_per_mb[bwd_index] = (
648 fwd_outs_per_mb[bwd_index][:1] if self.is_last else None
649 )
650 bwd_index += 1
652 self._flush_pending_sends()
654 loss = self._finalize_loss(fwd_outs_per_mb, user_inputs)
655 return (loss, *self._average_grads(grads, num_microbatches))
658SCHEDULE_REGISTRY: Dict[str, Type[PipelineScheduleBase]] = {
659 "gpipe": ScheduleGPipe,
660 "1f1b": Schedule1F1B,
661}
664def get_schedule_class(name: str) -> Type[PipelineScheduleBase]:
665 """Resolve a ``PassConfig.pp_schedule`` name to its schedule class.
667 Args:
668 name: Schedule name, case-insensitive (``"gpipe"`` / ``"1f1b"``).
670 Returns:
671 The schedule class.
673 Raises:
674 ValueError: When ``name`` is not a registered schedule.
675 """
676 key = str(name).lower()
677 schedule_cls = SCHEDULE_REGISTRY.get(key)
678 if schedule_cls is None:
679 raise ValueError(
680 f"Unknown pp_schedule {name!r}; expected one of {sorted(SCHEDULE_REGISTRY)}"
681 )
682 return schedule_cls