Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / parallel / pp_pass.py: 93%
715 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 Pass - Pipeline-Parallel Graph Split Pass
18Splits the (already FSDP-transformed) joint fwd+bwd FX graph into this
19rank's pipeline stage, entirely at graph level — no model deepcopy, no
20``torch.distributed.pipelining`` dependency. Where torchtitan splits at the
21*module* level (``_split_module`` deepcopy + layer deletion, then wraps each
22chunk in a ``PipelineStage``), this pass:
241. Attributes every FX node to a stage via ``nn_module_stack`` metadata
25 (parameters are static placeholders, attributed by FQN ancestor walk).
26 Root-level nodes the stack cannot resolve (residual adds written in the
27 parent's ``forward``, loss glue) are attributed by their CONSUMERS, so a
28 residual add sandwiched between two stages rides with the stage that
29 consumes it.
302. Detects the stage boundary as the set of values crossing each cut —
31 activations forward, gradients backward. Real ATen graphs cross with
32 more than the trunk tensor: dynamic-shape ``sym_size`` scalars and
33 parent-level views all cross, so every crossing value is shipped (int
34 scalars are packed as 0-d int64 tensors by the schedule and unwrapped
35 with ``item()`` on arrival).
363. Slices stage ``k`` into a forward subgraph and a backward subgraph. The
37 forward outputs (all crossing activations plus every intermediate the
38 backward references) become backward placeholders, so the GPipe schedule
39 can run all forwards before all backwards.
404. Prunes foreign-stage submodules from the live model (torch's
41 ``named_parameters`` then yields only this stage's entries, keeping the
42 trainer / optimizer FSDP- and PP-agnostic).
435. Installs a self-contained pipeline schedule as a ``call_module`` node in
44 a stub graph — the trainer's ``graph_module(*flat_inputs)`` dispatches to
45 it unchanged, and the stub survives any later ``recompile()`` because the
46 schedule lives in the graph itself, not in an overwritten ``forward``.
47 ``ScheduleGPipe`` / ``Schedule1F1B`` drive one stage per rank;
48 ``ScheduleInterleaved1F1B`` drives ``PassConfig.pp_chunks`` virtual stages
49 per rank (VPP) using the loop mapping ``rank, rank + pp_degree, ...``.
51Run order in ``PassPipeline``: AFTER ``FSDPPass``. The FSDP collectives are
52value-level FX nodes that follow the data flow: this stage's all_gather /
53reduce_scatter nodes land inside its slices untouched, while foreign
54stages' collectives are simply never copied. Splitting first (PP before
55FSDP) would instead hand ``FSDPPass`` two subgraphs and a loss-less output
56list, breaking its grad-index contract. This ordering is also what makes
57the PP+DP hybrid work with no extra machinery: with the 2-D ``(pp, DP)``
58mesh (``GraphCompiler._init_pp_dp_mesh``), ``FSDPPass`` shards parameters
59over the intra-stage DP group and ``PpPass`` pipelines each DP replica with
60the same DP replica of the neighbouring stages over the cross-stage
61``"pp"`` group. FSDP's backward re-gather / remat nodes are tagged
62``autograd_backward`` by ``FSDPPass`` precisely so this pass classifies
63them into the backward slice (lazy re-gather), and forward-phase reshard
64nodes stay in the forward slice.
65"""
67__all__ = ["PpPass", "_auto_stage_split"]
69import logging
70from dataclasses import dataclass
71from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple
73import torch
74import torch.distributed as dist
75from torch import fx, nn
76from torch.distributed.distributed_c10d import _resolve_process_group
78from ...pass_config import PassConfig, normalize_dp_mode
79from ..base import GraphPass
80from ...graph_parallel_plan import GraphParallelPlan
81from .pp_schedule import (
82 PipelineScheduleBase,
83 ScheduleInterleaved1F1B,
84 get_schedule_class,
85)
87_LOG = logging.getLogger(__name__)
89_FWD = "fwd"
90_BWD = "bwd"
93def _sanitize(name: str) -> str:
94 """Turn a dotted FQN into a codegen-safe placeholder name."""
95 return name.replace(".", "_")
98def _iter_node_args(node: fx.Node) -> Iterator[fx.Node]:
99 """Yield every ``fx.Node`` reachable in ``node.args`` / ``node.kwargs``."""
100 stack: List[Any] = list(node.args) + list(node.kwargs.values())
101 while stack:
102 item = stack.pop()
103 if isinstance(item, fx.Node):
104 yield item
105 elif isinstance(item, (tuple, list)):
106 stack.extend(item)
107 elif isinstance(item, dict):
108 stack.extend(item.values())
111def _tensor_arg_stages(node: fx.Node, node_stage: Dict[fx.Node, int]) -> List[int]:
112 """Stages of ``node``'s TENSOR-valued resolved args.
114 Shape-metadata args (dynamic-shape ``sym_size`` scalars) are excluded:
115 they are born wherever shapes are known, carry no data dependency, and
116 would otherwise drag a node's placement against its dataflow (e.g. a
117 boundary-gradient view whose ``sym_size`` args are stage-0 born while
118 the gradient itself is born downstream).
119 """
120 stages = []
121 for arg in _iter_node_args(node):
122 stage = node_stage.get(arg)
123 if stage is None:
124 continue
125 val = arg.meta.get("val")
126 if val is not None and not isinstance(val, torch.Tensor):
127 continue
128 stages.append(stage)
129 return stages
132def _innermost_fqns(node: fx.Node) -> List[str]:
133 """Module FQNs on ``node``'s ``nn_module_stack``, innermost first."""
134 stack = node.meta.get("nn_module_stack") or {}
135 return [fqn for fqn, _ in reversed(list(stack.values()))]
138def _even_split(names: List[str], pp_degree: int, prefix: str = "") -> List[List[str]]:
139 """Distribute ``names`` across ``pp_degree`` stages as evenly as possible.
141 The first ``len(names) % pp_degree`` stages get one extra entry; each
142 entry is prefixed (``prefix``) so bare child names become FQNs.
143 """
144 base, rem = divmod(len(names), pp_degree)
145 stages: List[List[str]] = [[] for _ in range(pp_degree)]
146 cursor = 0
147 for s in range(pp_degree):
148 size = base + (1 if s < rem else 0)
149 stages[s] = [f"{prefix}{name}" for name in names[cursor : cursor + size]]
150 cursor += size
151 return stages
154def _auto_stage_split(model: nn.Module, pp_degree: int) -> List[List[str]]:
155 """Default PP split: even distribution of the model's layer container.
157 Mirrors torchtitan's ``_generate_llm_fqn_per_model_part`` weighting
158 without hardcoding LLM module names:
160 - The first top-level child that is an ``nn.ModuleList`` with at least
161 ``pp_degree`` entries is the *layer container* (typically ``layers``).
162 - Top-level children registered before it go to stage 0 (input side,
163 e.g. embeddings); children after it go to the last stage (output
164 side, e.g. norm / lm_head).
165 - Container elements are distributed as evenly as possible (see
166 ``_even_split``).
168 Without a qualifying container, all top-level children are split evenly
169 in registration order.
171 Returns:
172 ``pp_degree`` lists of exact module FQNs, one per stage.
173 """
174 if pp_degree < 2:
175 raise ValueError(f"pp_degree must be >= 2 for a PP split, got {pp_degree}")
177 children = [name for name, _ in model.named_children()]
179 container: Optional[str] = None
180 for name, child in model.named_children():
181 if isinstance(child, nn.ModuleList) and len(child) >= pp_degree:
182 container = name
183 break
185 if container is None:
186 if len(children) < pp_degree:
187 raise ValueError(
188 f"Cannot split {len(children)} top-level children across "
189 f"{pp_degree} PP stages; provide a manual stage plan via "
190 f"GraphParallelPlan.pp_stage()"
191 )
192 return _even_split(children, pp_degree)
194 container_mod = getattr(model, container)
195 if isinstance(container_mod, nn.ModuleDict):
196 elem_keys = list(container_mod.keys())
197 else:
198 elem_keys = [str(i) for i in range(len(container_mod))]
200 idx = children.index(container)
201 before, after = children[:idx], children[idx + 1 :]
203 stages = _even_split(elem_keys, pp_degree, prefix=f"{container}.")
204 stages[0] = before + stages[0]
205 stages[-1] = stages[-1] + after
206 return stages
209@dataclass
210class _RunContext:
211 """This rank's PP execution context, resolved and guarded for ``run``.
213 ``pp_degree`` is the number of PHYSICAL ranks in the PP group, while
214 ``num_stages`` is the number of VIRTUAL stages (``pp_degree * chunks``).
215 Rank ``stage_idx`` owns the virtual stages in ``stage_indices`` using the
216 interleaved (loop) assignment ``stage_idx, stage_idx + pp_degree, ...``.
217 With the default ``chunks == 1`` the two coincide.
218 """
220 model: nn.Module
221 pp_degree: int
222 group: Any
223 stage_idx: int
224 chunks: int
225 num_stages: int
226 stage_indices: List[int]
229@dataclass
230class _StageGraphs:
231 """One virtual stage's sliced graphs and boundary contract."""
233 stage_idx: int
234 fwd_gm: fx.GraphModule
235 bwd_gm: fx.GraphModule
236 boundaries: "_Boundaries"
237 state_indices: List[int]
238 num_trainable: int
241@dataclass
242class _StageSplit:
243 """Stage plan and per-node classification produced by ``_classify_stage``.
245 Bundles what the boundary-detection and subgraph-construction phases
246 need, so ``run`` orchestrates phases instead of threading a dozen
247 parallel locals through every call.
248 """
250 state_fqns: List[str]
251 state_is_param: List[bool]
252 stage_plan: List[List[str]]
253 node_stage: Dict[fx.Node, int]
254 node_phase: Dict[fx.Node, str]
255 user_input_phs: List[fx.Node]
256 user_input_stages: List[int]
257 user_input_stage_sets: List[Tuple[int, ...]]
258 state_phs: List[fx.Node]
259 stage_state_indices: List[int]
260 stage_state_fqns: List[str]
263@dataclass
264class _GraphClassification:
265 """Per-node stage/phase classification produced by ``_classify_graph``."""
267 stage_state_indices: List[int]
268 node_stage: Dict[fx.Node, int]
269 node_phase: Dict[fx.Node, str]
270 user_input_phs: List[fx.Node]
271 user_input_stages: List[int]
272 user_input_stage_sets: List[Tuple[int, ...]]
273 state_phs: List[fx.Node]
276@dataclass
277class _Boundaries:
278 """Boundary value lists crossing this stage's cuts.
280 All lists are in joint-graph topological order so the producer and
281 consumer of a cut agree on P2P ordering:
283 - ``act_in``: stage-(k-1) values consumed by stage k (tensors,
284 dynamic-shape scalars, parent-level views — everything a real ATen
285 graph pushes across a cut).
286 - ``act_out``: stage-k values consumed by stage k+1.
287 - ``grad_in``: stage-(k+1) backward values consumed by stage k.
288 - ``grad_out``: stage-k backward values consumed by stage k-1.
289 - ``saved``: forward-phase values referenced by this stage's backward
290 nodes — exported as forward outputs and replayed as backward inputs.
291 """
293 act_in: List[fx.Node]
294 act_out: List[fx.Node]
295 grad_in: List[fx.Node]
296 grad_out: List[fx.Node]
297 saved: List[fx.Node]
300@dataclass
301class _SubgraphBuilder:
302 """Accumulator for assembling one stage subgraph.
304 ``env`` maps already-copied same-slice nodes; ``foreign`` maps
305 boundary values to their placeholder stand-ins; ``used`` carries the
306 taken placeholder names for ``_unique``.
307 """
309 g: fx.Graph
310 used: Set[str]
311 env: Dict[fx.Node, fx.Node]
312 foreign: Dict[fx.Node, fx.Node]
315class PpPass(GraphPass):
316 """Pipeline-parallel partitioning pass (graph-level stage split)."""
318 name = "pp_parallel"
320 def __init__(self, parallel_plan: Optional[GraphParallelPlan] = None) -> None:
321 """Initialize PP pass state.
323 Args:
324 parallel_plan: Declarative plan; ``pp_module_fqns_per_stage`` (from
325 ``pp_stage()``) overrides the automatic even-by-layers
326 split.
327 """
328 super().__init__()
329 self._parallel_plan = parallel_plan
331 def run(
332 self,
333 graph_module: fx.GraphModule,
334 pass_config: PassConfig,
335 **kwargs: Any,
336 ) -> fx.GraphModule:
337 """Split the joint graph to this rank's stage and install the schedule.
339 Args:
340 graph_module: Joint fwd+bwd FX graph (post-``FSDPPass``).
341 pass_config: Parallel configuration; ``pp_degree`` (``None``
342 resolves to ``world_size``) and ``pp_microbatch_size`` are
343 read directly.
344 **kwargs: Must include ``model`` (the live ``nn.Module``); may
345 include ``parallel_plan``.
347 Returns:
348 The same graph module, rewritten into a stage stub whose
349 ``call_module`` node invokes the installed pipeline schedule
350 (``ScheduleGPipe`` / ``Schedule1F1B`` /
351 ``ScheduleInterleaved1F1B`` per ``pass_config.pp_schedule``).
352 """
353 run_ctx = self._resolve_run_context(pass_config, kwargs)
354 if run_ctx is None:
355 return graph_module
357 split = self._classify_stage(graph_module, run_ctx)
358 self._build_and_install_stage(graph_module, pass_config, run_ctx, split)
359 return graph_module
361 def _resolve_run_context(
362 self,
363 pass_config: PassConfig,
364 kwargs: Dict[str, Any],
365 ) -> Optional[_RunContext]:
366 """Guard and resolve the model, PP group, and this rank's stage.
368 Returns:
369 The run context, or ``None`` when the pass is skipped
370 (distributed not initialized, ``world_size == 1``, or
371 ``pp_degree < 2``).
372 """
373 if not dist.is_initialized() or dist.get_world_size() == 1:
374 _LOG.info("Skipped: distributed not initialized or world_size=1")
375 return None
377 model = kwargs.get("model")
378 if model is None:
379 raise ValueError(
380 "PpPass requires the live model via kwargs (model=...) so it "
381 "can resolve the stage plan and prune foreign-stage modules"
382 )
383 self._parallel_plan = kwargs.get("parallel_plan", self._parallel_plan)
385 # Skip BEFORE _resolve_group_and_stage so a disabled pass does not
386 # execute the collective dist.new_group.
387 pp_degree = pass_config.pp_degree or dist.get_world_size()
388 if pp_degree < 2:
389 _LOG.info("Skipped: pp_degree=%s < 2", pp_degree)
390 return None
391 pp_degree, group, stage_idx = self._resolve_group_and_stage(pass_config, kwargs)
392 chunks = pass_config.pp_chunks
393 num_stages = pp_degree * chunks
394 stage_indices = list(range(stage_idx, num_stages, pp_degree))
395 _LOG.info(
396 "Running: rank %s/%s owns virtual stages %s/%s, world_size=%s",
397 stage_idx,
398 pp_degree,
399 stage_indices,
400 num_stages,
401 dist.get_world_size(),
402 )
403 return _RunContext(
404 model=model,
405 pp_degree=pp_degree,
406 group=group,
407 stage_idx=stage_idx,
408 chunks=chunks,
409 num_stages=num_stages,
410 stage_indices=stage_indices,
411 )
413 def _classify_stage(
414 self, graph_module: fx.GraphModule, run_ctx: _RunContext
415 ) -> _StageSplit:
416 """Resolve the stage plan and classify every node's stage + phase.
418 Chains the state-list contract checks, the stage-plan resolution,
419 node classification, the gradient-output anchoring with its stage
420 propagation, and the adjacent-dataflow validation.
422 Args:
423 graph_module: Joint fwd+bwd FX graph (post-``FSDPPass``).
424 run_ctx: This rank's PP execution context.
426 Returns:
427 The ``_StageSplit`` bundling the plan and per-node
428 classification for the construction phase.
429 """
430 state_fqns = self._require_state_list(graph_module)
431 state_is_param = self._require_is_param(graph_module, state_fqns)
432 num_state_inputs = getattr(graph_module, "num_state_inputs", len(state_fqns))
434 stage_plan = self._resolve_stage_plan(run_ctx.model, run_ctx.num_stages)
435 stage_of_fqn = {fqn: s for s, fqns in enumerate(stage_plan) for fqn in fqns}
437 graph_cls = self._classify_graph(
438 graph_module,
439 run_ctx,
440 stage_of_fqn,
441 state_fqns,
442 num_state_inputs,
443 )
445 self._anchor_grad_outputs(
446 graph_module,
447 graph_cls.node_stage,
448 stage_of_fqn,
449 state_fqns,
450 run_ctx.model,
451 )
452 self._propagate_anchor_stages(
453 graph_module, graph_cls.node_stage, graph_cls.node_phase
454 )
455 self._validate_adjacent_dataflow(graph_module, graph_cls.node_stage)
457 return _StageSplit(
458 state_fqns=state_fqns,
459 state_is_param=state_is_param,
460 stage_plan=stage_plan,
461 node_stage=graph_cls.node_stage,
462 node_phase=graph_cls.node_phase,
463 user_input_phs=graph_cls.user_input_phs,
464 user_input_stages=graph_cls.user_input_stages,
465 user_input_stage_sets=graph_cls.user_input_stage_sets,
466 state_phs=graph_cls.state_phs,
467 stage_state_indices=graph_cls.stage_state_indices,
468 stage_state_fqns=[state_fqns[i] for i in graph_cls.stage_state_indices],
469 )
471 def _build_and_install_stage(
472 self,
473 graph_module: fx.GraphModule,
474 pass_config: PassConfig,
475 run_ctx: _RunContext,
476 split: _StageSplit,
477 ) -> None:
478 """Build the owned virtual-stage subgraphs, prune, and install the stub.
480 Trainable-state resolution must precede ``_prune_live_model``: the
481 gradient-output nodes map positionally onto the FULL trainable list,
482 so pruning foreign-stage parameters first would shift the positions.
484 Args:
485 graph_module: Joint fwd+bwd FX graph, rewritten in place.
486 pass_config: Parallel configuration (microbatch size read).
487 run_ctx: This rank's PP execution context.
488 split: Stage plan and per-node classification.
489 """
490 trainable_indices = self._trainable_state_indices(
491 split.state_fqns, split.state_is_param, run_ctx.model
492 )
493 state_positions = {
494 idx: pos for pos, idx in enumerate(split.stage_state_indices)
495 }
497 stage_graphs: List[_StageGraphs] = []
498 for stage_idx in run_ctx.stage_indices:
499 boundaries = self._find_boundaries(
500 graph_module,
501 split.node_stage,
502 split.node_phase,
503 run_ctx.num_stages,
504 stage_idx,
505 split.user_input_phs,
506 split.user_input_stage_sets,
507 )
508 stage_state_indices = [
509 idx
510 for idx in split.stage_state_indices
511 if split.node_stage[split.state_phs[idx]] == stage_idx
512 ]
513 stage_state_set = set(stage_state_indices)
514 stage_trainable = [i for i in trainable_indices if i in stage_state_set]
515 fwd_gm, bwd_gm = self._build_stage_graphs(
516 graph_module,
517 split,
518 boundaries,
519 stage_idx,
520 run_ctx.num_stages,
521 stage_state_indices,
522 trainable_indices,
523 )
524 stage_graphs.append(
525 _StageGraphs(
526 stage_idx=stage_idx,
527 fwd_gm=fwd_gm,
528 bwd_gm=bwd_gm,
529 boundaries=boundaries,
530 state_indices=stage_state_indices,
531 num_trainable=len(stage_trainable),
532 )
533 )
535 keep_fqns: List[str] = []
536 for stage_idx in run_ctx.stage_indices:
537 keep_fqns.extend(split.stage_plan[stage_idx])
538 self._prune_live_model(run_ctx.model, keep_fqns)
540 sched = self._build_schedule(
541 stage_graphs,
542 pass_config,
543 run_ctx,
544 split,
545 trainable_indices,
546 state_positions,
547 )
548 self._install_stub(
549 graph_module,
550 sched,
551 split.stage_state_fqns,
552 split.state_is_param,
553 split.stage_state_indices,
554 len(split.user_input_phs),
555 )
557 _LOG.info(
558 "Completed: rank %s owns virtual stages %s, %d state entries, "
559 "%d trainable params",
560 run_ctx.stage_idx,
561 run_ctx.stage_indices,
562 len(split.stage_state_fqns),
563 sum(stage.num_trainable for stage in stage_graphs),
564 )
566 def _build_stage_graphs(
567 self,
568 graph_module: fx.GraphModule,
569 split: _StageSplit,
570 boundaries: _Boundaries,
571 stage_idx: int,
572 num_stages: int,
573 stage_state_indices: List[int],
574 trainable_indices: List[int],
575 ) -> Tuple[fx.GraphModule, fx.GraphModule]:
576 """Build one virtual stage's forward and backward subgraphs."""
577 stage_state_fqns = [split.state_fqns[i] for i in stage_state_indices]
578 fwd_gm, fwd_out_orig = self._build_fwd_graph(
579 graph_module,
580 split,
581 boundaries,
582 stage_idx,
583 num_stages,
584 stage_state_indices,
585 stage_state_fqns,
586 )
587 bwd_gm = self._build_bwd_graph(
588 graph_module,
589 split,
590 boundaries,
591 stage_idx,
592 num_stages,
593 stage_state_indices,
594 stage_state_fqns,
595 fwd_out_orig,
596 trainable_indices,
597 )
598 return fwd_gm, bwd_gm
600 # ------------------------------------------------------------------
601 # Tracer contract helpers
602 # ------------------------------------------------------------------
604 @staticmethod
605 def _require_state_list(graph_module: fx.GraphModule) -> List[str]:
606 """Return the graph's live ``state_fqns`` list (shared with JointGraph).
608 The tracer attaches the SAME list object to the GraphModule and the
609 ``JointGraph``; mutating it in place (done in ``_install_stub``)
610 keeps ``run_traced_graph``'s validation consistent with the pruned
611 model. Copying here would desynchronize the two.
612 """
613 state_fqns = getattr(graph_module, "state_fqns", None)
614 if state_fqns is None:
615 raise ValueError(
616 "PpPass requires the tracer's state_fqns attribute on the "
617 "joint graph (attach it via trace_model_graph)"
618 )
619 return state_fqns
621 @staticmethod
622 def _require_is_param(
623 graph_module: fx.GraphModule, state_fqns: List[str]
624 ) -> List[bool]:
625 """Return the graph's live ``state_is_param`` list, defaulting to
626 all-params (matching the tracer's fallback for older traces)."""
627 state_is_param = getattr(graph_module, "state_is_param", None)
628 if state_is_param is None:
629 state_is_param = [True] * len(state_fqns)
630 graph_module.state_is_param = state_is_param
631 return state_is_param
633 @staticmethod
634 def _loss_node(graph_module: fx.GraphModule) -> Optional[fx.Node]:
635 """The original joint graph's loss node (output index 0)."""
636 output = next((n for n in graph_module.graph.nodes if n.op == "output"), None)
637 if output is None:
638 return None
639 returned = output.args[0]
640 if isinstance(returned, (list, tuple)) and returned:
641 first = returned[0]
642 if isinstance(first, fx.Node):
643 return first
644 return None
646 # ------------------------------------------------------------------
647 # Group / plan resolution
648 # ------------------------------------------------------------------
650 def _resolve_group_and_stage(
651 self, pass_config: PassConfig, kwargs: Dict[str, Any]
652 ) -> Tuple[int, Any, int]:
653 """Resolve PP degree, process group, and this rank's stage index.
655 Two layouts are supported:
657 - **PP+DP hybrid** (``pp_group_name`` kwarg, registered by
658 ``GraphCompiler._init_pp_dp_mesh``): the group spans one rank per
659 stage at this rank's DP position, so the stage index is simply this
660 rank's position within the group and every DP replica of a stage
661 pipelines with the same DP replica of its neighbours. Requires
662 ``pp_degree * dp_degree == world_size`` (validated here).
663 - **Pure PP** (no ``pp_group_name``): the group is the whole world and
664 each rank owns one or more virtual stages. ``fsdp_enabled`` is
665 rejected here — a 1-D fallback FSDP mesh spans the whole world and
666 would shard parameters across stage boundaries; the hybrid must go
667 through the 2-D mesh path.
669 A single PP process group carries both forward activations and
670 backward gradients; the schedule keeps the two ends of every
671 neighbour pair's FIFO complementary (see ``Schedule1F1B`` and
672 ``ScheduleInterleaved1F1B``).
673 """
674 pp_group_name = kwargs.get("pp_group_name")
675 if pp_group_name is not None:
676 return self._resolve_from_pp_group(pass_config, pp_group_name)
678 world_size = dist.get_world_size()
679 rank = dist.get_rank()
680 pp_degree = pass_config.pp_degree or world_size
682 if pp_degree not in (1, world_size):
683 raise ValueError(
684 f"Without a pp process group only pure-PP is supported: "
685 f"pp_degree={pp_degree} must equal world_size={world_size} "
686 f"(or be 1 to disable PP). For the PP+DP hybrid pass "
687 f"pp_group_name (registered by GraphCompiler's 2-D pp mesh)."
688 )
689 if pass_config.fsdp_enabled:
690 raise ValueError(
691 "PP+FSDP hybrid requires a 2-D (pp, DP) mesh expressing both "
692 "dimensions: let GraphCompiler build it (fsdp_enabled=True + "
693 "pp_enabled=True, which registers the cross-stage 'pp' group "
694 "and passes pp_group_name to this pass). The 1-D fallback "
695 "FSDP mesh spans the whole world and cannot host PP — FSDP "
696 "would shard parameters across stage boundaries. Run pure-PP "
697 "(fsdp_enabled=False, pp_degree == world_size) or the mesh "
698 "path."
699 )
701 group = dist.new_group(list(range(world_size)))
702 # Pure-PP (pp_degree == world_size, checked above), so each rank's
703 # physical stage index is its rank. Keep the general formula so a
704 # future DP-replication layout (pp_degree < world_size) maps
705 # rank->stage contiguously rather than silently assuming 1:1.
706 stage_idx = rank * pp_degree // world_size
707 return pp_degree, group, stage_idx
709 def _resolve_from_pp_group(
710 self, pass_config: PassConfig, group_name: str
711 ) -> Tuple[int, Any, int]:
712 """Resolve stage/pp-degree from the registered ``"pp"`` group.
714 Validates the hybrid topology: the configured ``pp_degree`` (when
715 set) must match the group size, and with the DP axis enabled
716 ``pp_degree * dp_degree`` must tile the world. HSDP needs a 3-D mesh
717 and is rejected outright.
718 """
719 group = _resolve_process_group(group_name)
720 pp_degree = dist.get_world_size(group)
721 stage_idx = dist.get_rank(group)
723 configured = pass_config.pp_degree
724 if configured is not None and configured != pp_degree:
725 raise ValueError(
726 f"pass_config.pp_degree={configured} does not match the "
727 f"registered pp group size ({pp_degree})"
728 )
730 if pass_config.fsdp_enabled:
731 dp_mode = normalize_dp_mode(pass_config.dp_mode)
732 if dp_mode == "hsdp":
733 raise ValueError(
734 "PP + HSDP requires a 3-D (pp, dp_replicate, fsdp_shard) "
735 "mesh, which is not supported yet; run PP with "
736 "dp_mode='fsdp' or 'ddp'"
737 )
738 dp_degree = (
739 pass_config.fsdp_degree
740 if dp_mode == "fsdp"
741 else pass_config.dp_replicate_degree
742 )
743 world_size = dist.get_world_size()
744 if dp_degree is None:
745 if world_size % pp_degree != 0:
746 raise ValueError(
747 f"world_size ({world_size}) not divisible by the pp "
748 f"group size ({pp_degree}); cannot derive the DP "
749 f"degree"
750 )
751 dp_degree = world_size // pp_degree
752 if pp_degree * dp_degree != world_size:
753 raise ValueError(
754 f"PP+DP degrees must tile the world: pp_degree "
755 f"({pp_degree}) x dp_degree ({dp_degree}) != world_size "
756 f"({world_size})"
757 )
759 _LOG.info(
760 "Resolved stage %s/%s from pp group '%s'", stage_idx, pp_degree, group_name
761 )
762 return pp_degree, group, stage_idx
764 def _resolve_stage_plan(self, model: nn.Module, num_stages: int) -> List[List[str]]:
765 """Resolve the per-virtual-stage module FQN lists (manual or auto)."""
766 manual = (
767 self._parallel_plan.pp_module_fqns_per_stage
768 if self._parallel_plan is not None
769 else None
770 )
771 if manual is not None:
772 if len(manual) != num_stages:
773 raise ValueError(
774 f"Manual PP plan declares {len(manual)} stages but the "
775 f"pipeline has {num_stages} virtual stages "
776 f"(pp_degree * pp_chunks)"
777 )
778 seen: Set[str] = set()
779 for s, fqns in enumerate(manual):
780 for fqn in fqns:
781 if fqn in seen:
782 raise ValueError(
783 f"Module '{fqn}' is assigned to multiple stages "
784 f"(stage {s} repeats it)"
785 )
786 seen.add(fqn)
787 stage_plan = [list(fqns) for fqns in manual]
788 else:
789 stage_plan = _auto_stage_split(model, num_stages)
791 known = dict(model.named_modules())
792 for s, fqns in enumerate(stage_plan):
793 for fqn in fqns:
794 if fqn not in known:
795 raise ValueError(f"Stage {s} declares unknown module '{fqn}'")
796 self._warn_unassigned_modules(model, stage_plan)
797 return stage_plan
799 @staticmethod
800 def _warn_unassigned_modules(model: nn.Module, stage_plan: List[List[str]]) -> None:
801 """Warn about live modules no stage declares.
803 A module that is neither declared itself, nor lives inside a
804 declared one, nor is an ancestor of a declared one (a container
805 whose elements are declared) has its graph nodes attributed by
806 the dataflow fallbacks (consumer stage, then last stage) — rarely
807 what a manual plan intends. Auto splits declare every top-level
808 child, so they never warn.
809 """
810 declared = [fqn for fqns in stage_plan for fqn in fqns]
811 unassigned: List[str] = []
812 for fqn, _ in model.named_modules():
813 if not fqn:
814 # Root: parent-level glue is attributed by dataflow.
815 continue
816 covered = any(
817 fqn == d or fqn.startswith(d + ".") or d.startswith(fqn + ".")
818 for d in declared
819 )
820 if not covered:
821 unassigned.append(fqn)
822 if not unassigned:
823 return
824 shown = unassigned[:10]
825 if len(unassigned) > len(shown):
826 shown.append(f"... {len(unassigned) - len(shown)} more")
827 _LOG.warning(
828 "No PP stage declares modules %s — their nodes ride with the "
829 "dataflow attribution (consumer stage, then last stage). "
830 "Declare them via GraphParallelPlan.pp_stage() if that is not intended.",
831 shown,
832 )
834 # ------------------------------------------------------------------
835 # Node classification (stage + phase, topo order)
836 # ------------------------------------------------------------------
838 def _classify_graph( # pylint: disable=too-many-locals
839 self,
840 graph_module: fx.GraphModule,
841 run_ctx: _RunContext,
842 stage_of_fqn: Dict[str, int],
843 state_fqns: Sequence[str],
844 num_state_inputs: int,
845 ) -> _GraphClassification:
846 """Resolve every node's stage and fwd/bwd phase in topological order.
848 Phase: explicit ``autograd_backward`` tag, else inherited from args
849 (post-trace inserts such as FSDP's reduce_scatter/wait inherit the
850 bwd phase through their grad arguments).
852 Stage resolution, in priority order:
854 1. ``nn_module_stack`` walk (traced nodes, innermost module first)
855 2. argument inheritance (post-trace inserts — state placeholders
856 win, since backward compute nodes reference their params)
857 3. consumer attribution for stack-less root-level glue (residual
858 adds written in the parent's forward): the stage of their
859 consumers, minimum when several — resolved in reverse topo
860 4. last stage as the final fallback (loss chain, dead nodes)
862 User-input placeholders are NOT anchors: they are routed AFTER the
863 fwd/bwd classification by dataflow (``_route_user_inputs``) — each
864 rides with the single stage that consumes it, so kwargs-general
865 inputs (multiple features on stage 0, several labels on the last
866 stage, an auxiliary input on a middle stage) all work without
867 positional heuristics.
869 ``stage_state_indices`` lists the state entries owned by the virtual
870 stages THIS rank owns (``run_ctx.stage_indices``); every state
871 placeholder still gets a stage so foreign consumers can be
872 attributed.
874 Root-level state (no stage-owning module ancestor) is assigned to
875 the stage of its (single) consumer; multi-stage consumption is the
876 RoPE-cache-shaped error case and fails with an actionable message.
877 """
878 graph = graph_module.graph
879 placeholders = [n for n in graph.nodes if n.op == "placeholder"]
880 state_ph_set = set(placeholders[:num_state_inputs])
881 num_stages = run_ctx.num_stages
883 node_stage: Dict[fx.Node, int] = {}
884 node_phase: Dict[fx.Node, str] = {}
886 (
887 state_phs,
888 deferred_state,
889 stage_state_indices,
890 user_input_phs,
891 ) = self._classify_placeholders(
892 placeholders,
893 node_stage,
894 node_phase,
895 state_fqns,
896 stage_of_fqn,
897 num_state_inputs,
898 run_ctx,
899 )
900 self._classify_forward_nodes(
901 graph, state_ph_set, stage_of_fqn, num_stages, node_stage, node_phase
902 )
903 self._classify_backward_nodes(
904 graph, state_ph_set, num_stages, node_stage, node_phase
905 )
906 user_input_stages, user_input_stage_sets = self._route_user_inputs(
907 user_input_phs, node_stage, num_stages
908 )
909 self._assign_deferred_state(
910 deferred_state,
911 state_phs,
912 state_fqns,
913 node_stage,
914 run_ctx.stage_indices,
915 stage_state_indices,
916 )
918 return _GraphClassification(
919 stage_state_indices=sorted(stage_state_indices),
920 node_stage=node_stage,
921 node_phase=node_phase,
922 user_input_phs=user_input_phs,
923 user_input_stages=user_input_stages,
924 user_input_stage_sets=user_input_stage_sets,
925 state_phs=state_phs,
926 )
928 def _classify_placeholders( # pylint: disable=too-many-locals
929 self,
930 placeholders: List[fx.Node],
931 node_stage: Dict[fx.Node, int],
932 node_phase: Dict[fx.Node, str],
933 state_fqns: Sequence[str],
934 stage_of_fqn: Dict[str, int],
935 num_state_inputs: int,
936 run_ctx: _RunContext,
937 ) -> Tuple[List[fx.Node], List[int], List[int], List[fx.Node]]:
938 """Classify placeholder nodes and collect this stage's state entries.
940 Only STATE placeholders get a stage here (via their FQN's module
941 ancestor); user-input placeholders are routed later by dataflow
942 (``_route_user_inputs``), once their consumers' stages are known.
944 Args:
945 placeholders: All placeholder nodes in graph order.
946 node_stage: Stage assignment accumulator to fill.
947 node_phase: Phase assignment accumulator to fill.
948 state_fqns: FQNs of the traced state entries (params/buffers).
949 stage_of_fqn: Module FQN -> owning stage map.
950 num_state_inputs: Number of leading state placeholders.
951 run_ctx: This rank's PP execution context.
953 Returns:
954 ``(state_phs, deferred_state, stage_state_indices,
955 user_input_phs)``. ``deferred_state`` lists root-level state
956 indices whose stage is resolved later from their consumers.
957 """
958 state_phs: List[fx.Node] = []
959 deferred_state: List[int] = []
960 stage_state_indices: List[int] = []
961 for idx, ph in enumerate(placeholders):
962 node_phase[ph] = _FWD
963 if idx < num_state_inputs:
964 stage = self._state_stage(state_fqns[idx], stage_of_fqn)
965 state_phs.append(ph)
966 if stage is None:
967 deferred_state.append(idx)
968 else:
969 node_stage[ph] = stage
970 if stage in run_ctx.stage_indices:
971 stage_state_indices.append(idx)
972 user_input_phs = list(placeholders[num_state_inputs:])
973 return state_phs, deferred_state, stage_state_indices, user_input_phs
975 @staticmethod
976 def _route_user_inputs(
977 user_input_phs: List[fx.Node],
978 node_stage: Dict[fx.Node, int],
979 num_stages: int,
980 ) -> Tuple[List[int], List[Tuple[int, ...]]]:
981 """Assign each user-input placeholder to the stages consuming it.
983 Runs after the fwd/bwd classification, so every consumer node
984 already carries its final stage (user inputs are deliberately not
985 anchors — see ``_classify_graph``). Every rank receives the full
986 flattened input list, so an input consumed on several stages (an
987 embedding's token ids reused by its backward, a mask/position ids
988 shared by many layers) is replicated directly to each consumer with
989 no P2P. An input consumed by NO node (an ignored kwarg — the
990 hand-built test fixtures' ``y``, an ``epoch`` scalar train_fn never
991 reads) is warned about and routed to the last stage, where unused
992 inputs have always ridden.
994 Args:
995 user_input_phs: User-input placeholders in stub arg order.
996 node_stage: Final per-node stage assignment (fwd+bwd done).
997 num_stages: Total PP stage count.
999 Returns:
1000 ``(primary_stages, stage_sets)`` in stub arg order.
1001 ``primary_stages[i]`` is the first consumer stage (kept for
1002 diagnostics and backward compatibility); ``stage_sets[i]`` is
1003 every stage that consumes the input.
1004 """
1005 primary: List[int] = []
1006 stage_sets: List[Tuple[int, ...]] = []
1007 for ph in user_input_phs:
1008 consumer_stages = {
1009 node_stage[u]
1010 for u in ph.users
1011 if u.op != "output" and node_stage.get(u) is not None
1012 }
1013 if not consumer_stages:
1014 _LOG.warning(
1015 "Model input '%s' is not consumed by the traced graph — "
1016 "routing it to the last stage with the other unused "
1017 "inputs (a misspelled keyword argument would leave its "
1018 "tensor unused too)",
1019 ph.name,
1020 )
1021 consumers = (num_stages - 1,)
1022 else:
1023 consumers = tuple(sorted(consumer_stages))
1024 node_stage[ph] = consumers[0]
1025 primary.append(consumers[0])
1026 stage_sets.append(consumers)
1027 _LOG.info(
1028 "User input '%s' routed to stages %s of %s",
1029 ph.name,
1030 consumers,
1031 num_stages,
1032 )
1033 return primary, stage_sets
1035 def _classify_forward_nodes(
1036 self,
1037 graph: fx.Graph,
1038 state_ph_set: Set[fx.Node],
1039 stage_of_fqn: Dict[str, int],
1040 num_stages: int,
1041 node_stage: Dict[fx.Node, int],
1042 node_phase: Dict[fx.Node, str],
1043 ) -> None:
1044 """Assign stages to forward nodes in topological order.
1046 Anchors: module stack, then state-ph args, then max-args over
1047 TENSOR args (data flows forward; shape-metadata scalars carry no
1048 data dependency). Leftover root-level glue rides with its
1049 consumers (reverse topo), else the last stage.
1050 """
1051 floating: List[fx.Node] = []
1052 for node in graph.nodes:
1053 if node.op in ("placeholder", "output"):
1054 continue
1055 phase = _BWD if node.meta.get("autograd_backward", False) else _FWD
1056 if phase == _FWD:
1057 for arg in _iter_node_args(node):
1058 if node_phase.get(arg) == _BWD:
1059 phase = _BWD
1060 break
1061 node_phase[node] = phase
1062 if phase == _FWD:
1063 stage = self._resolve_fwd_stage(
1064 node, node_stage, state_ph_set, stage_of_fqn
1065 )
1066 if stage is None:
1067 floating.append(node)
1068 else:
1069 node_stage[node] = stage
1071 self._assign_floating_fwd_nodes(
1072 graph, floating, num_stages, node_stage, node_phase
1073 )
1075 def _resolve_fwd_stage(
1076 self,
1077 node: fx.Node,
1078 node_stage: Dict[fx.Node, int],
1079 state_ph_set: Set[fx.Node],
1080 stage_of_fqn: Dict[str, int],
1081 ) -> Optional[int]:
1082 """Resolve a forward node's stage anchor.
1084 Priority: module stack, then state-ph args, then max-args over
1085 TENSOR args (data flows forward; shape-metadata scalars carry no
1086 data dependency). ``None`` means root-level glue to be resolved
1087 from its consumers (see ``_assign_floating_fwd_nodes``).
1088 """
1089 stage = self._stage_from_stack(node, stage_of_fqn)
1090 if stage is None:
1091 stage = self._stage_from_state_arg(node, node_stage, state_ph_set)
1092 if stage is None:
1093 arg_stages = _tensor_arg_stages(node, node_stage)
1094 if arg_stages:
1095 stage = max(arg_stages)
1096 return stage
1098 @staticmethod
1099 def _assign_floating_fwd_nodes(
1100 graph: fx.Graph,
1101 floating: List[fx.Node],
1102 num_stages: int,
1103 node_stage: Dict[fx.Node, int],
1104 node_phase: Dict[fx.Node, str],
1105 ) -> None:
1106 """Leftover fwd glue: consumers first (reverse topo), else last."""
1107 for node in reversed(list(graph.nodes)):
1108 if node not in floating or node_phase[node] != _FWD:
1109 continue
1110 consumer_stages = {
1111 node_stage[u]
1112 for u in node.users
1113 if u.op != "output" and node_stage.get(u) is not None
1114 }
1115 node_stage[node] = (
1116 min(consumer_stages) if consumer_stages else num_stages - 1
1117 )
1119 def _classify_backward_nodes(
1120 self,
1121 graph: fx.Graph,
1122 state_ph_set: Set[fx.Node],
1123 num_stages: int,
1124 node_stage: Dict[fx.Node, int],
1125 node_phase: Dict[fx.Node, str],
1126 ) -> None:
1127 """Assign stages to backward nodes in topological order.
1129 A bwd node runs where its earliest TENSOR input lives: gradients
1130 born downstream are shipped down during the backward sweep, and
1131 upstream activations arrive as replayed forward outputs. Weight-
1132 gradient nodes that consume early activations are corrected by
1133 the positional anchor (the param's stage owns its grad).
1134 """
1135 for node in graph.nodes:
1136 if node.op in ("placeholder", "output"):
1137 continue
1138 if node_phase[node] != _BWD:
1139 continue
1140 stage = self._stage_from_state_arg(node, node_stage, state_ph_set)
1141 if stage is None:
1142 arg_stages = _tensor_arg_stages(node, node_stage)
1143 if arg_stages:
1144 stage = min(arg_stages)
1145 if stage is None:
1146 # Dead-end bwd glue: ride with consumers if any, else last.
1147 consumer_stages = {
1148 node_stage[u]
1149 for u in node.users
1150 if u.op != "output" and node_stage.get(u) is not None
1151 }
1152 stage = min(consumer_stages) if consumer_stages else num_stages - 1
1153 node_stage[node] = stage
1155 def _assign_deferred_state(
1156 self,
1157 deferred_state: List[int],
1158 state_phs: List[fx.Node],
1159 state_fqns: Sequence[str],
1160 node_stage: Dict[fx.Node, int],
1161 owned_stages: Sequence[int],
1162 stage_state_indices: List[int],
1163 ) -> None:
1164 """Assign root-level state to the stage of its (single) consumer.
1166 Multi-stage consumption is the RoPE-cache-shaped error case and
1167 fails with an actionable message.
1168 """
1169 owned = set(owned_stages)
1170 for idx in deferred_state:
1171 ph = state_phs[idx]
1172 consumer_stages = {node_stage[u] for u in ph.users if u.op != "output"}
1173 if len(consumer_stages) > 1:
1174 raise ValueError(
1175 f"Root-level state '{state_fqns[idx]}' is consumed by "
1176 f"stages {sorted(consumer_stages)} — v1 requires state "
1177 f"to live under a single stage's module (e.g. move a "
1178 f"shared RoPE cache under a per-stage module or split "
1179 f"it per stage)"
1180 )
1181 stage = next(iter(consumer_stages), 0)
1182 node_stage[ph] = stage
1183 if stage in owned:
1184 stage_state_indices.append(idx)
1185 _LOG.info(
1186 "Root-level state '%s' assigned to stage %s by its consumer",
1187 state_fqns[idx],
1188 stage,
1189 )
1191 def _state_stage(self, fqn: str, stage_of_fqn: Dict[str, int]) -> Optional[int]:
1192 """Stage of a state entry via longest-ancestor walk over its FQN."""
1193 parts = fqn.split(".")
1194 for i in range(len(parts) - 1, 0, -1):
1195 stage = stage_of_fqn.get(".".join(parts[:i]))
1196 if stage is not None:
1197 return stage
1198 return stage_of_fqn.get(fqn)
1200 def _stage_from_stack(
1201 self, node: fx.Node, stage_of_fqn: Dict[str, int]
1202 ) -> Optional[int]:
1203 """Stage from ``nn_module_stack``, innermost module first."""
1204 for fqn in _innermost_fqns(node):
1205 stage = stage_of_fqn.get(fqn)
1206 if stage is not None:
1207 return stage
1208 return None
1210 def _stage_from_state_arg(
1211 self,
1212 node: fx.Node,
1213 node_stage: Dict[fx.Node, int],
1214 state_ph_set: Set[fx.Node],
1215 ) -> Optional[int]:
1216 """Stage from a state-placeholder argument, if any.
1218 Post-trace inserts (FSDP collectives, weight-grad compute) reference
1219 their owning stage's parameters — a state placeholder arg pins the
1220 node to that stage. Generic argument inheritance is deliberately
1221 NOT done here: it would let an early-stage argument pull a node
1222 against the dataflow (see the max/min-args rules in
1223 ``_classify_graph``).
1224 """
1225 for arg in _iter_node_args(node):
1226 if arg in state_ph_set:
1227 return node_stage.get(arg)
1228 return None
1230 @staticmethod
1231 def _propagate_anchor_stages(
1232 graph_module: fx.GraphModule,
1233 node_stage: Dict[fx.Node, int],
1234 node_phase: Dict[fx.Node, str],
1235 ) -> None:
1236 """Pull bwd args up to anchored consumers (one reverse-topo pass).
1238 A bwd value consumed by a DOWNSTREAM-stage bwd node must be
1239 recomputed there: gradients flow downstream -> upstream, so a
1240 downstream consumer cannot receive an upstream-born bwd value in
1241 time (its own bwd sweep runs first). Fwd-phase args are untouched —
1242 they arrive as replayed forward outputs; downstream-born args are
1243 untouched — they arrive via the grad crossings.
1245 A single REVERSE topological sweep reaches the fixpoint: a node is
1246 visited only after every user of it has been visited and pulled it
1247 to its final stage, so the node propagates its final value to its
1248 own args. The previous iterate-to-fixpoint loop was quadratic in
1249 the worst case on large joint graphs.
1250 """
1251 for node in reversed(list(graph_module.graph.nodes)):
1252 if node.op in ("placeholder", "output"):
1253 continue
1254 if node_phase.get(node) != _BWD:
1255 continue
1256 stage = node_stage.get(node)
1257 if stage is None:
1258 continue
1259 for arg in _iter_node_args(node):
1260 if node_phase.get(arg) != _BWD or node_stage.get(arg) is None:
1261 continue
1262 if node_stage[arg] < stage:
1263 node_stage[arg] = stage
1265 # ------------------------------------------------------------------
1266 # Boundary detection
1267 # ------------------------------------------------------------------
1269 @staticmethod
1270 def _validate_adjacent_dataflow(
1271 graph_module: fx.GraphModule,
1272 node_stage: Dict[fx.Node, int],
1273 ) -> None:
1274 """Reject skip-stage dataflow before it reaches the slice copy.
1276 v1 ships boundary values between NEIGHBOURING stages only: a value
1277 produced on stage ``j`` and consumed on stage ``k`` with
1278 ``abs(j - k) > 1`` has no route (it would have to transit
1279 intermediate stages untouched). Without this check such graphs die
1280 later inside ``_copy_node`` / the fwd-output assembly with a
1281 generic "outside this stage's slice" error that hides the root
1282 cause.
1284 Raises:
1285 ValueError: When any producer-consumer pair spans more than
1286 one stage cut (e.g. a dense skip connection across two
1287 stages, or a stage plan that leaves a module's nodes to
1288 the dataflow fallbacks).
1289 """
1290 for node in graph_module.graph.nodes:
1291 if node.op == "output":
1292 continue
1293 producer = node_stage.get(node)
1294 if producer is None:
1295 continue
1296 for user in node.users:
1297 if user.op == "output":
1298 continue
1299 consumer = node_stage.get(user)
1300 if consumer is None or abs(consumer - producer) <= 1:
1301 continue
1302 raise ValueError(
1303 f"Value '{node.name}' is produced on stage {producer} "
1304 f"but consumed on stage {consumer}: v1 ships boundary "
1305 f"values between neighbouring stages only. A skip "
1306 f"connection spanning 2+ stages or a stage plan that "
1307 f"leaves modules unattributed produces this — adjust "
1308 f"GraphParallelPlan.pp_stage() so the dataflow crosses one cut "
1309 f"at a time."
1310 )
1312 def _anchor_grad_outputs(
1313 self,
1314 graph_module: fx.GraphModule,
1315 node_stage: Dict[fx.Node, int],
1316 stage_of_fqn: Dict[str, int],
1317 state_fqns: Sequence[str],
1318 model: Optional[nn.Module],
1319 ) -> None:
1320 """Pin each gradient-output node to its parameter's stage.
1322 Gradient-output node i differentiates the i-th trainable parameter
1323 (positional mapping, mirroring the tracer); it runs on the stage
1324 owning that parameter even when min-args would place it upstream
1325 (a later stage's weight-grad consumes an earlier stage's
1326 activation, replayed from the forward outputs).
1327 """
1328 output = next((n for n in graph_module.graph.nodes if n.op == "output"), None)
1329 if output is None or model is None:
1330 return
1331 returned = output.args[0]
1332 if not isinstance(returned, (list, tuple)):
1333 return
1334 state_is_param = getattr(
1335 graph_module, "state_is_param", [True] * len(state_fqns)
1336 )
1337 trainable = self._trainable_state_indices(state_fqns, state_is_param, model)
1338 for i, grad_node in enumerate(returned[1:]):
1339 if i >= len(trainable) or not isinstance(grad_node, fx.Node):
1340 continue
1341 pstage = self._state_stage(state_fqns[trainable[i]], stage_of_fqn)
1342 if pstage is not None:
1343 node_stage[grad_node] = pstage
1345 def _find_boundaries( # pylint: disable=too-many-locals
1346 self,
1347 graph_module: fx.GraphModule,
1348 node_stage: Dict[fx.Node, int],
1349 node_phase: Dict[fx.Node, str],
1350 num_stages: int,
1351 stage_idx: int,
1352 user_input_phs: Sequence[fx.Node],
1353 user_input_stage_sets: Sequence[Sequence[int]],
1354 ) -> _Boundaries:
1355 """Locate the value lists crossing this virtual stage's cuts.
1357 Returns:
1358 The ``_Boundaries`` lists, all in joint-graph topological order
1359 so the producer and consumer of a cut agree on P2P ordering.
1360 See ``_Boundaries`` for the per-field semantics.
1361 """
1362 k = stage_idx
1363 is_last = k == num_stages - 1
1364 nodes = [
1365 n for n in graph_module.graph.nodes if n.op not in ("placeholder", "output")
1366 ]
1368 def crossing(
1369 producer: int, consumer: int, phase: Optional[str]
1370 ) -> List[fx.Node]:
1371 """Values produced by one stage and consumed by another.
1373 ``phase=None`` matches any phase (backward crossings may carry
1374 gradients as well as functionalized fwd views).
1375 """
1376 found = []
1377 for v in nodes:
1378 if node_stage[v] != producer:
1379 continue
1380 if phase is not None and node_phase[v] != phase:
1381 continue
1382 if any(node_stage.get(u) == consumer for u in v.users):
1383 found.append(v)
1384 return found
1386 # FWD-phase only: values born in this stage's forward sweep and
1387 # consumed downstream. Bwd-phase values cannot ship forward (they
1388 # do not exist yet during the forward sweep); downstream-born
1389 # gradients flow back via grad crossings below.
1390 act_in_list = crossing(k - 1, k, _FWD) if k > 0 else []
1391 act_out_list = [] if is_last else crossing(k, k + 1, _FWD)
1392 # Backward crossings are phase-AGNOSTIC: besides plain gradients,
1393 # functionalized views (fwd nodes) can be consumed by an earlier
1394 # stage's backward (weight-grad compute) — those values flow back
1395 # during the backward sweep just the same.
1396 grad_in_list = [] if is_last else crossing(k + 1, k, None)
1397 grad_out_list = crossing(k, k - 1, None) if k > 0 else []
1399 # Forward-phase values the backward references. Stage-owned
1400 # placeholders (state / routed user inputs of this stage) may appear
1401 # and round-trip through the forward outputs; foreign placeholders
1402 # cannot be shipped and fail loudly during the slice copy.
1403 stage_owned_phs = {
1404 p
1405 for p in graph_module.graph.nodes
1406 if p.op == "placeholder" and node_stage.get(p) == k
1407 }
1408 stage_owned_phs.update(
1409 ph
1410 for ph, stages in zip(user_input_phs, user_input_stage_sets)
1411 if k in stages
1412 )
1413 saved = self._collect_saved_values(
1414 nodes,
1415 node_stage,
1416 node_phase,
1417 k,
1418 stage_owned_phs,
1419 set(act_in_list),
1420 grad_out_list,
1421 )
1422 return _Boundaries(
1423 act_in=act_in_list,
1424 act_out=act_out_list,
1425 grad_in=grad_in_list,
1426 grad_out=grad_out_list,
1427 saved=saved,
1428 )
1430 @staticmethod
1431 def _collect_saved_values(
1432 nodes: List[fx.Node],
1433 node_stage: Dict[fx.Node, int],
1434 node_phase: Dict[fx.Node, str],
1435 stage_idx: int,
1436 stage_owned_phs: Set[fx.Node],
1437 act_in_set: Set[fx.Node],
1438 grad_out_list: List[fx.Node],
1439 ) -> List[fx.Node]:
1440 """Forward-phase values this stage's backward references.
1442 Stage-owned placeholders (state / routed user inputs of this stage)
1443 may appear and round-trip through the forward outputs; foreign
1444 placeholders cannot be shipped and fail loudly during the slice
1445 copy. Fwd-phase values this stage must ship BACKWARD (consumed by
1446 the previous stage's backward) ride the fwd outputs as
1447 pass-throughs: computed once in the fwd subgraph, replayed into
1448 the bwd subgraph, and appended to its outputs.
1449 """
1450 saved: List[fx.Node] = []
1451 seen: Set[fx.Node] = set()
1452 for node in nodes:
1453 if node_phase[node] != _BWD or node_stage[node] != stage_idx:
1454 continue
1455 for arg in _iter_node_args(node):
1456 if not isinstance(arg, fx.Node) or arg in seen:
1457 continue
1458 # Fwd-phase values replay from the fwd outputs; bwd-phase
1459 # values are only shippable when they arrived as boundary
1460 # inputs (e.g. gradient-shape sym_size scalars).
1461 if node_phase.get(arg) != _FWD and arg not in act_in_set:
1462 continue
1463 if arg.op == "placeholder" and arg not in stage_owned_phs:
1464 continue
1465 seen.add(arg)
1466 saved.append(arg)
1467 for v in grad_out_list:
1468 if node_phase[v] == _FWD and v not in seen:
1469 seen.add(v)
1470 saved.append(v)
1471 return saved
1473 # ------------------------------------------------------------------
1474 # Subgraph construction
1475 # ------------------------------------------------------------------
1477 @staticmethod
1478 def _unique(name: str, used: Set[str]) -> str:
1479 """Return ``name`` (or a suffixed variant) not yet in ``used``."""
1480 candidate = name
1481 suffix = 1
1482 while candidate in used:
1483 candidate = f"{name}_{suffix}"
1484 suffix += 1
1485 used.add(candidate)
1486 return candidate
1488 def _copy_node(
1489 self,
1490 builder: _SubgraphBuilder,
1491 node: fx.Node,
1492 ) -> None:
1493 """Copy ``node`` into ``builder.g``, remapping args through the
1494 builder's env/foreign maps.
1496 ``env`` maps already-copied same-slice nodes; ``foreign`` maps
1497 boundary values to their placeholder stand-ins. The copy is
1498 registered under ``builder.env[node]``; nothing is returned —
1499 callers consume the copy through the builder's maps.
1500 """
1501 if node.op == "get_attr":
1502 raise ValueError(
1503 f"Unexpected get_attr node '{node.name}' — the tracer "
1504 f"contracts parameters/buffers to be static placeholders"
1505 )
1507 def map_fn(a: Any) -> Any:
1508 """Remap a single arg: env first, then foreign placeholders."""
1509 if isinstance(a, fx.Node):
1510 if a in builder.env:
1511 return builder.env[a]
1512 if a in builder.foreign:
1513 return builder.foreign[a]
1514 raise ValueError(
1515 f"Node '{node.name}' references '{a.name}' which lies "
1516 f"outside this stage's slice — boundary detection or "
1517 f"the stage plan is inconsistent"
1518 )
1519 return a
1521 new_args = fx.map_arg(node.args, map_fn)
1522 new_kwargs = fx.map_arg(node.kwargs, map_fn)
1523 new_node = builder.g.create_node(node.op, node.target, new_args, new_kwargs)
1524 new_node.meta = dict(node.meta)
1525 builder.env[node] = new_node
1527 def _build_fwd_graph(
1528 self,
1529 graph_module: fx.GraphModule,
1530 split: _StageSplit,
1531 boundaries: _Boundaries,
1532 stage_idx: int,
1533 num_stages: int,
1534 state_indices: Sequence[int],
1535 state_fqns: Sequence[str],
1536 ) -> Tuple[fx.GraphModule, List[fx.Node]]:
1537 """Build one virtual stage's forward subgraph.
1539 Outputs: ``(*act_out, *saved)`` for non-last stages, ``(loss,
1540 *saved)`` for the last stage — where ``saved`` contains every
1541 forward-phase value the backward references that is not already
1542 crossing. The schedule ships the crossing prefix and replays the
1543 full output tuple into the backward subgraph.
1545 Returns:
1546 ``(fwd_gm, fwd_out_orig)`` — the subgraph module together with
1547 the ORIGINAL nodes at each output position, so the backward
1548 builder can key its placeholder mapping on them.
1549 """
1550 is_last = stage_idx == num_stages - 1
1551 builder = _SubgraphBuilder(g=fx.Graph(), used=set(), env={}, foreign={})
1553 self._create_fwd_placeholders(
1554 builder, split, boundaries, stage_idx, state_indices, state_fqns
1555 )
1557 for node in graph_module.graph.nodes:
1558 if node.op in ("placeholder", "output"):
1559 continue
1560 if split.node_phase[node] != _FWD or split.node_stage[node] != stage_idx:
1561 continue
1562 self._copy_node(builder, node)
1564 loss_node = self._loss_node(graph_module)
1565 out_orig = self._resolve_fwd_outputs(
1566 boundaries.act_out, boundaries.saved, loss_node, is_last
1567 )
1568 out_nodes = [
1569 builder.env[n] if n in builder.env else builder.foreign[n] for n in out_orig
1570 ]
1571 builder.g.output(tuple(out_nodes))
1572 return fx.GraphModule(nn.Module(), builder.g), out_orig
1574 def _create_fwd_placeholders(
1575 self,
1576 builder: _SubgraphBuilder,
1577 split: _StageSplit,
1578 boundaries: _Boundaries,
1579 stage_idx: int,
1580 state_indices: Sequence[int],
1581 state_fqns: Sequence[str],
1582 ) -> None:
1583 """Create the fwd subgraph's placeholders (state / boundary / inputs).
1585 Incoming boundary values (virtual stage > 0) become placeholders;
1586 each stage additionally takes the user inputs routed to it (stage 0
1587 typically the feature batch, the last stage the labels; see
1588 ``_route_user_inputs``).
1589 """
1590 for i, orig_idx in enumerate(state_indices):
1591 ph = builder.g.placeholder(
1592 self._unique(_sanitize(state_fqns[i]), builder.used)
1593 )
1594 ph.meta = dict(split.state_phs[orig_idx].meta)
1595 builder.env[split.state_phs[orig_idx]] = ph
1597 for i, act in enumerate(boundaries.act_in):
1598 ph = builder.g.placeholder(self._unique(f"pp_act_in_{i}", builder.used))
1599 ph.meta = dict(act.meta)
1600 builder.foreign[act] = ph
1602 for i, (ph, stages) in enumerate(
1603 zip(split.user_input_phs, split.user_input_stage_sets)
1604 ):
1605 if stage_idx not in stages:
1606 continue
1607 in_ph = builder.g.placeholder(self._unique(f"pp_input_{i}", builder.used))
1608 in_ph.meta = dict(ph.meta)
1609 builder.env[ph] = in_ph
1611 @staticmethod
1612 def _resolve_fwd_outputs(
1613 act_out_list: List[fx.Node],
1614 saved: List[fx.Node],
1615 loss_node: Optional[fx.Node],
1616 is_last: bool,
1617 ) -> List[fx.Node]:
1618 """Order the fwd subgraph's outputs: boundary values first, then saved."""
1619 if is_last:
1620 if loss_node is None:
1621 raise ValueError("Last PP stage requires a loss node in the output")
1622 out_orig: List[fx.Node] = [loss_node]
1623 else:
1624 if not act_out_list:
1625 raise ValueError(
1626 "Stage ships no boundary values but is not the last "
1627 "stage — the stage plan does not cut the graph"
1628 )
1629 out_orig = list(act_out_list)
1631 crossing_ids = {id(v) for v in act_out_list}
1632 for node in saved:
1633 if id(node) in crossing_ids or (is_last and node is loss_node):
1634 continue
1635 out_orig.append(node)
1636 return out_orig
1638 def _build_bwd_graph(
1639 self,
1640 graph_module: fx.GraphModule,
1641 split: _StageSplit,
1642 boundaries: _Boundaries,
1643 stage_idx: int,
1644 num_stages: int,
1645 state_indices: Sequence[int],
1646 state_fqns: Sequence[str],
1647 fwd_out_orig: List[fx.Node],
1648 trainable_indices: List[int],
1649 ) -> fx.GraphModule:
1650 """Build one virtual stage's backward subgraph.
1652 Placeholders: ``(*state, *grad_in, *fwd_outputs)`` — the schedule
1653 passes the received boundary gradients and then the forward
1654 outputs verbatim (on the last stage a single ``ones_like(loss)``
1655 takes the gradient slots' place). Outputs: ``(*param_grads,)``
1656 plus the boundary gradients for every stage except stage 0.
1657 """
1658 builder = _SubgraphBuilder(g=fx.Graph(), used=set(), env={}, foreign={})
1659 self._create_bwd_placeholders(
1660 builder,
1661 graph_module,
1662 split,
1663 boundaries,
1664 stage_idx,
1665 num_stages,
1666 state_indices,
1667 state_fqns,
1668 fwd_out_orig,
1669 )
1671 stage_bwd_set = self._collect_stage_bwd_nodes(
1672 graph_module.graph, split.node_stage, split.node_phase, stage_idx
1673 )
1675 # Gradient-output nodes of this stage's trainable params (by state
1676 # index identity, not module attribution — post-trace inserts like
1677 # zeros_like carry no module stack). A LIST — output order is the
1678 # positional gradient contract; only membership uses the set.
1679 stage_grad_nodes = self._stage_grad_nodes(
1680 graph_module.graph, trainable_indices, set(state_indices)
1681 )
1682 stage_grad_set = set(stage_grad_nodes)
1684 for node in graph_module.graph.nodes:
1685 if node.op in ("placeholder", "output"):
1686 continue
1687 if node not in stage_bwd_set and node not in stage_grad_set:
1688 continue
1689 self._copy_node(builder, node)
1691 out_nodes = [builder.env[gn] for gn in stage_grad_nodes]
1692 out_nodes.extend(builder.env[gn] for gn in boundaries.grad_out)
1693 builder.g.output(tuple(out_nodes))
1694 return fx.GraphModule(nn.Module(), builder.g)
1696 @staticmethod
1697 def _collect_stage_bwd_nodes(
1698 graph: fx.Graph,
1699 node_stage: Dict[fx.Node, int],
1700 node_phase: Dict[fx.Node, str],
1701 stage_idx: int,
1702 ) -> Set[fx.Node]:
1703 """Bwd-phase nodes owned by ``stage_idx``, as a set.
1705 Membership is probed once per graph node by ``_build_bwd_graph``
1706 and real joint graphs carry tens of thousands of nodes.
1707 """
1708 return {
1709 node
1710 for node in graph.nodes
1711 if node.op not in ("placeholder", "output")
1712 and node_phase[node] == _BWD
1713 and node_stage[node] == stage_idx
1714 }
1716 def _create_bwd_placeholders(
1717 self,
1718 builder: _SubgraphBuilder,
1719 graph_module: fx.GraphModule,
1720 split: _StageSplit,
1721 boundaries: _Boundaries,
1722 stage_idx: int,
1723 num_stages: int,
1724 state_indices: Sequence[int],
1725 state_fqns: Sequence[str],
1726 fwd_out_orig: List[fx.Node],
1727 ) -> None:
1728 """Create the bwd subgraph's placeholders (state / grads / fwd outs).
1730 The schedule passes the received boundary gradients and then the
1731 forward outputs verbatim (on the last stage a single
1732 ``ones_like(loss)`` takes the gradient slots' place). The
1733 ``foreign`` entries double as the remap for backward nodes
1734 referencing forward values.
1735 """
1736 for i, orig_idx in enumerate(state_indices):
1737 ph = builder.g.placeholder(
1738 self._unique(_sanitize(state_fqns[i]), builder.used)
1739 )
1740 ph.meta = dict(split.state_phs[orig_idx].meta)
1741 builder.env[split.state_phs[orig_idx]] = ph
1743 for i, grad in enumerate(boundaries.grad_in):
1744 ph = builder.g.placeholder(self._unique(f"pp_grad_in_{i}", builder.used))
1745 ph.meta = dict(grad.meta)
1746 builder.foreign[grad] = ph
1747 if not boundaries.grad_in and stage_idx == num_stages - 1:
1748 # Last stage: the schedule feeds ones_like(loss) into the
1749 # single gradient slot.
1750 ph = builder.g.placeholder(self._unique("pp_grad_ones", builder.used))
1751 loss = self._loss_node(graph_module)
1752 if loss is not None:
1753 ph.meta = dict(loss.meta)
1755 # One placeholder per forward output, same order — the schedule
1756 # passes the forward outputs verbatim after the gradients.
1757 for orig in fwd_out_orig:
1758 ph = builder.g.placeholder(
1759 self._unique(f"pp_fwd_{orig.name}", builder.used)
1760 )
1761 ph.meta = dict(orig.meta)
1762 builder.foreign[orig] = ph
1764 @staticmethod
1765 def _stage_grad_nodes(
1766 graph: fx.Graph, trainable_indices: List[int], stage_state_set: Set[int]
1767 ) -> List[fx.Node]:
1768 """Gradient-output nodes of the stage's trainable params."""
1769 output = next((n for n in graph.nodes if n.op == "output"), None)
1770 returned = output.args[0] if output is not None else []
1771 stage_grad_nodes: List[fx.Node] = []
1772 for i, grad_node in enumerate(returned[1:]):
1773 if i >= len(trainable_indices):
1774 break
1775 if trainable_indices[i] in stage_state_set:
1776 stage_grad_nodes.append(grad_node)
1777 return stage_grad_nodes
1779 # ------------------------------------------------------------------
1780 # Live model pruning
1781 # ------------------------------------------------------------------
1783 def _prune_live_model(self, model: nn.Module, keep_fqns: List[str]) -> None:
1784 """Remove foreign-stage submodules from the live model in place.
1786 Mirrors torchtitan's ``_split_module`` (``None``-out unkept
1787 children, filter ModuleList/ModuleDict entries) while PRESERVING
1788 original index keys — reindexing a ModuleList would rename the
1789 surviving parameters' FQNs and break ``run_traced_graph``'s state
1790 validation against the sliced ``state_fqns``. Manual plans may cut
1791 inside a container element; such elements are recursed into
1792 (mixed subtree), exactly like plain modules.
1793 """
1794 keep = set(keep_fqns)
1796 def is_kept(fqn: str) -> bool:
1797 """True when ``fqn`` equals or descends from a stage module."""
1798 return any(fqn == kf or fqn.startswith(kf + ".") for kf in keep)
1800 def prune_container_child(child: nn.Module, cfqn: str) -> None:
1801 """Prune a ModuleList/ModuleDict child, preserving index keys.
1803 Reindexing would rename the surviving parameters' FQNs and break
1804 ``run_traced_graph``'s state validation. An element is kept
1805 whole, deleted, or recursed into when a manual plan cuts inside
1806 it (same rule as plain modules).
1807 """
1808 children = child._modules # pylint: disable=protected-access
1809 keys = list(children.keys())
1810 if all(is_kept(f"{cfqn}.{key}") for key in keys):
1811 return
1812 for key in keys:
1813 elem_fqn = f"{cfqn}.{key}"
1814 if is_kept(elem_fqn):
1815 continue
1816 elem = children[key]
1817 descendants = [f for f, _ in elem.named_modules() if f]
1818 if any(is_kept(f"{elem_fqn}.{f}") for f in descendants):
1819 prune(elem, elem_fqn + ".")
1820 else:
1821 del children[key]
1823 def prune_plain_child(
1824 mod: nn.Module, name: str, child: nn.Module, cfqn: str
1825 ) -> None:
1826 """Prune a non-container child: delete it or recurse when mixed."""
1827 descendants = [f for f, _ in child.named_modules() if f]
1828 if not descendants:
1829 # Leaf module (e.g. a bare Linear): keep only if declared.
1830 if not is_kept(cfqn):
1831 setattr(mod, name, None)
1832 return
1833 if not any(is_kept(f"{cfqn}.{f}") for f in descendants):
1834 setattr(mod, name, None)
1835 else:
1836 # Mixed subtree (manual plans may cut inside a child).
1837 prune(child, cfqn + ".")
1839 def prune(mod: nn.Module, prefix: str) -> None:
1840 """Recursively prune non-stage children under ``prefix``."""
1841 for name, child in list(mod.named_children()):
1842 cfqn = f"{prefix}{name}"
1843 if is_kept(cfqn):
1844 # Declared (or ancestor of a declared) stage module —
1845 # keep the whole subtree.
1846 continue
1847 if isinstance(child, (nn.ModuleList, nn.ModuleDict)):
1848 prune_container_child(child, cfqn)
1849 continue
1850 prune_plain_child(mod, name, child, cfqn)
1852 prune(model, "")
1854 # ------------------------------------------------------------------
1855 # Schedule + stub installation
1856 # ------------------------------------------------------------------
1858 def _trainable_state_indices(
1859 self,
1860 state_fqns: Sequence[str],
1861 state_is_param: Sequence[bool],
1862 model: Optional[nn.Module],
1863 ) -> List[int]:
1864 """Trainable-parameter state indices, mirroring the tracer's order."""
1865 param_lookup = (
1866 dict(model.named_parameters(remove_duplicate=False))
1867 if model is not None
1868 else {}
1869 )
1870 trainable: List[int] = []
1871 for idx, fqn in enumerate(state_fqns):
1872 if idx < len(state_is_param) and not state_is_param[idx]:
1873 continue
1874 param = param_lookup.get(fqn)
1875 if param is not None and param.requires_grad:
1876 trainable.append(idx)
1877 return trainable
1879 @staticmethod
1880 def _value_spec(
1881 nodes: Sequence[fx.Node], stage_idx: int, kind: str
1882 ) -> List[Tuple[Any, ...]]:
1883 """Build the schedule's recv descriptor for a boundary value list.
1885 Tensor values yield ``("tensor", shape, dtype, device)``; int/SymInt
1886 scalars (dynamic-shape ``sym_size`` nodes) yield ``("scalar",)``
1887 and are shipped as 0-d int64 tensors, unwrapped via ``item()``.
1888 Scalars carry no device of their own — the schedule anchors them to
1889 the first tensor entry's device, so a non-empty list must contain
1890 at least one tensor.
1891 """
1892 spec: List[Tuple[Any, ...]] = []
1893 for v in nodes:
1894 val = v.meta.get("val")
1895 if isinstance(val, torch.Tensor):
1896 # int(s): trace-time dims are SymInts (symbolic over the
1897 # sample input); the schedule allocates real buffers from
1898 # this spec, and ATen rejects symbolic sizes at runtime.
1899 # Training uses static shapes per compile, so the hinted
1900 # concrete values are stable.
1901 shape = tuple(int(s) for s in val.shape)
1902 spec.append(("tensor", shape, val.dtype, val.device))
1903 elif val is None:
1904 raise ValueError(
1905 f"Boundary {kind} value '{v.name}' on stage {stage_idx} "
1906 f"has no 'val' meta (traced without FakeTensor "
1907 f"metadata) — cannot size the P2P receive buffer"
1908 )
1909 else:
1910 spec.append(("scalar",))
1911 if nodes and not any(entry[0] == "tensor" for entry in spec):
1912 raise ValueError(
1913 f"Boundary {kind} list on stage {stage_idx} carries only "
1914 f"scalar values — no tensor entry to anchor the P2P "
1915 f"receive device on"
1916 )
1917 return spec
1919 def _build_schedule(
1920 self,
1921 stage_graphs: Sequence[_StageGraphs],
1922 pass_config: PassConfig,
1923 run_ctx: _RunContext,
1924 split: _StageSplit,
1925 trainable_indices: Sequence[int],
1926 state_positions: Dict[int, int],
1927 ) -> PipelineScheduleBase:
1928 """Assemble this rank's pipeline schedule (``pass_config.pp_schedule``).
1930 Receive descriptors come from the boundary values' fake-tensor
1931 ``val`` metas; the send side packs by value type at runtime. The
1932 single-stage schedules take one graph pair; the interleaved schedule
1933 takes every virtual stage this rank owns.
1934 """
1935 schedule_cls = get_schedule_class(pass_config.pp_schedule)
1936 if issubclass(schedule_cls, ScheduleInterleaved1F1B):
1937 return self._build_interleaved_schedule(
1938 schedule_cls,
1939 stage_graphs,
1940 pass_config,
1941 run_ctx,
1942 split,
1943 trainable_indices,
1944 state_positions,
1945 )
1946 if run_ctx.chunks != 1:
1947 raise ValueError(
1948 f"pp_chunks={run_ctx.chunks} requires the interleaved "
1949 f"schedule, got {pass_config.pp_schedule!r}"
1950 )
1951 return self._build_single_stage_schedule(
1952 schedule_cls,
1953 stage_graphs[0],
1954 pass_config,
1955 run_ctx,
1956 split,
1957 )
1959 def _build_single_stage_schedule(
1960 self,
1961 schedule_cls: Any,
1962 stage: _StageGraphs,
1963 pass_config: PassConfig,
1964 run_ctx: _RunContext,
1965 split: _StageSplit,
1966 ) -> PipelineScheduleBase:
1967 """Build one of the single-virtual-stage schedules (GPipe / 1F1B)."""
1968 boundaries = stage.boundaries
1969 stage_idx = stage.stage_idx
1970 recv_spec = (
1971 self._value_spec(boundaries.act_in, stage_idx, "activation")
1972 if stage_idx > 0
1973 else []
1974 )
1975 grad_recv_spec = (
1976 self._value_spec(boundaries.grad_in, stage_idx, "gradient")
1977 if stage_idx < run_ctx.num_stages - 1
1978 else []
1979 )
1980 return schedule_cls(
1981 stage.fwd_gm,
1982 stage.bwd_gm,
1983 stage_idx=stage_idx,
1984 pp_degree=run_ctx.num_stages,
1985 num_state=len(split.stage_state_fqns),
1986 num_trainable=stage.num_trainable,
1987 num_send=len(boundaries.act_out),
1988 grad_send_count=len(boundaries.grad_out),
1989 microbatch_size=pass_config.pp_microbatch_size,
1990 pp_group=run_ctx.group,
1991 recv_spec=recv_spec,
1992 grad_recv_spec=grad_recv_spec,
1993 user_input_stages=split.user_input_stages,
1994 user_input_stage_sets=split.user_input_stage_sets,
1995 )
1997 def _build_interleaved_schedule( # pylint: disable=too-many-locals
1998 self,
1999 schedule_cls: Any,
2000 stage_graphs: Sequence[_StageGraphs],
2001 pass_config: PassConfig,
2002 run_ctx: _RunContext,
2003 split: _StageSplit,
2004 trainable_indices: Sequence[int],
2005 state_positions: Dict[int, int],
2006 ) -> PipelineScheduleBase:
2007 """Build the VPP schedule from this rank's virtual-stage graph list."""
2008 owned_pos = {v: i for i, v in enumerate(run_ctx.stage_indices)}
2009 grad_owner = [
2010 owned_pos[split.node_stage[split.state_phs[idx]]]
2011 for idx in trainable_indices
2012 if idx in state_positions
2013 ]
2014 state_positions_per_stage = [
2015 [state_positions[i] for i in stage.state_indices] for stage in stage_graphs
2016 ]
2017 recv_specs = [
2018 self._recv_spec_for(stage, run_ctx.num_stages, gradient=False)
2019 for stage in stage_graphs
2020 ]
2021 grad_recv_specs = [
2022 self._recv_spec_for(stage, run_ctx.num_stages, gradient=True)
2023 for stage in stage_graphs
2024 ]
2025 return schedule_cls(
2026 [stage.fwd_gm for stage in stage_graphs],
2027 [stage.bwd_gm for stage in stage_graphs],
2028 stage_indices=list(run_ctx.stage_indices),
2029 num_stages=run_ctx.num_stages,
2030 pp_degree=run_ctx.pp_degree,
2031 state_positions=state_positions_per_stage,
2032 num_states=[len(stage.state_indices) for stage in stage_graphs],
2033 num_trainables=[stage.num_trainable for stage in stage_graphs],
2034 num_sends=[len(stage.boundaries.act_out) for stage in stage_graphs],
2035 grad_send_counts=[len(stage.boundaries.grad_out) for stage in stage_graphs],
2036 microbatch_size=pass_config.pp_microbatch_size,
2037 pp_group=run_ctx.group,
2038 recv_specs=recv_specs,
2039 grad_recv_specs=grad_recv_specs,
2040 user_input_stages=split.user_input_stages,
2041 user_input_stage_sets=split.user_input_stage_sets,
2042 grad_owner=grad_owner,
2043 )
2045 def _recv_spec_for(
2046 self,
2047 stage: _StageGraphs,
2048 num_stages: int,
2049 gradient: bool,
2050 ) -> List[Tuple[Any, ...]]:
2051 """Build one virtual stage's activation or gradient receive specs."""
2052 if gradient:
2053 if stage.stage_idx >= num_stages - 1:
2054 return []
2055 return self._value_spec(
2056 stage.boundaries.grad_in, stage.stage_idx, "gradient"
2057 )
2058 if stage.stage_idx == 0:
2059 return []
2060 return self._value_spec(stage.boundaries.act_in, stage.stage_idx, "activation")
2062 def _install_stub(
2063 self,
2064 graph_module: fx.GraphModule,
2065 sched: PipelineScheduleBase,
2066 stage_state_fqns: List[str],
2067 state_is_param: Sequence[bool],
2068 stage_state_indices: List[int],
2069 num_user_inputs: int,
2070 ) -> None:
2071 """Rewrite the graph module into a schedule-dispatching stub.
2073 The schedule is added as a submodule and invoked through a
2074 ``call_module`` node, so it survives any later ``recompile()`` and
2075 the trainer's ``graph_module(*flat_inputs)`` dispatches to it
2076 unchanged. The stub takes ``(*state, *user_inputs)`` — every rank
2077 receives the full flattened user-input list; the schedule picks
2078 the inputs routed to this stage. ``state_fqns`` /
2079 ``state_is_param`` are updated IN PLACE (the JointGraph shares the
2080 list objects).
2081 """
2082 g = fx.Graph()
2083 used: Set[str] = set()
2084 state_phs = [
2085 g.placeholder(self._unique(_sanitize(fqn), used))
2086 for fqn in stage_state_fqns
2087 ]
2088 user_phs = [
2089 g.placeholder(self._unique(f"pp_input_{i}", used))
2090 for i in range(num_user_inputs)
2091 ]
2093 graph_module.add_module("pp_schedule", sched)
2094 call = g.call_module("pp_schedule", args=tuple(state_phs) + tuple(user_phs))
2095 g.output(call)
2097 graph_module.graph = g
2098 graph_module.state_fqns[:] = stage_state_fqns
2099 graph_module.state_is_param[:] = [
2100 state_is_param[i] for i in stage_state_indices
2101 ]
2102 graph_module.num_state_inputs = len(stage_state_fqns)
2103 graph_module.recompile()