Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / parallel / pp_pass.py: 92%
642 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 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 (``ScheduleGPipe`` by
44 default; ``Schedule1F1B`` when ``PassConfig.pp_schedule == "1f1b"``) as a
45 ``call_module`` node in a stub graph — the trainer's
46 ``graph_module(*flat_inputs)`` dispatches to it unchanged, and the stub
47 survives any later ``recompile()`` because the schedule lives in the
48 graph itself, not in an overwritten ``forward``.
50Run order in ``PassPipeline``: AFTER ``FSDPPass``. The FSDP collectives are
51value-level FX nodes that follow the data flow: this stage's all_gather /
52reduce_scatter nodes land inside its slices untouched, while foreign
53stages' collectives are simply never copied. Splitting first (PP before
54FSDP) would instead hand ``FSDPPass`` two subgraphs and a loss-less output
55list, breaking its grad-index contract.
56"""
58__all__ = ["PpPass", "_auto_stage_split"]
60import logging
61from dataclasses import dataclass
62from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple
64import torch
65import torch.distributed as dist
66from torch import fx, nn
68from ...pass_config import PassConfig
69from ..base import GraphPass
70from ...graph_parallel_plan import GraphParallelPlan
71from .pp_schedule import PipelineScheduleBase, get_schedule_class
73_LOG = logging.getLogger(__name__)
75_FWD = "fwd"
76_BWD = "bwd"
79def _sanitize(name: str) -> str:
80 """Turn a dotted FQN into a codegen-safe placeholder name."""
81 return name.replace(".", "_")
84def _iter_node_args(node: fx.Node) -> Iterator[fx.Node]:
85 """Yield every ``fx.Node`` reachable in ``node.args`` / ``node.kwargs``."""
86 stack: List[Any] = list(node.args) + list(node.kwargs.values())
87 while stack:
88 item = stack.pop()
89 if isinstance(item, fx.Node):
90 yield item
91 elif isinstance(item, (tuple, list)):
92 stack.extend(item)
93 elif isinstance(item, dict):
94 stack.extend(item.values())
97def _tensor_arg_stages(node: fx.Node, node_stage: Dict[fx.Node, int]) -> List[int]:
98 """Stages of ``node``'s TENSOR-valued resolved args.
100 Shape-metadata args (dynamic-shape ``sym_size`` scalars) are excluded:
101 they are born wherever shapes are known, carry no data dependency, and
102 would otherwise drag a node's placement against its dataflow (e.g. a
103 boundary-gradient view whose ``sym_size`` args are stage-0 born while
104 the gradient itself is born downstream).
105 """
106 stages = []
107 for arg in _iter_node_args(node):
108 stage = node_stage.get(arg)
109 if stage is None:
110 continue
111 val = arg.meta.get("val")
112 if val is not None and not isinstance(val, torch.Tensor):
113 continue
114 stages.append(stage)
115 return stages
118def _innermost_fqns(node: fx.Node) -> List[str]:
119 """Module FQNs on ``node``'s ``nn_module_stack``, innermost first."""
120 stack = node.meta.get("nn_module_stack") or {}
121 return [fqn for fqn, _ in reversed(list(stack.values()))]
124def _even_split(names: List[str], pp_degree: int, prefix: str = "") -> List[List[str]]:
125 """Distribute ``names`` across ``pp_degree`` stages as evenly as possible.
127 The first ``len(names) % pp_degree`` stages get one extra entry; each
128 entry is prefixed (``prefix``) so bare child names become FQNs.
129 """
130 base, rem = divmod(len(names), pp_degree)
131 stages: List[List[str]] = [[] for _ in range(pp_degree)]
132 cursor = 0
133 for s in range(pp_degree):
134 size = base + (1 if s < rem else 0)
135 stages[s] = [f"{prefix}{name}" for name in names[cursor : cursor + size]]
136 cursor += size
137 return stages
140def _auto_stage_split(model: nn.Module, pp_degree: int) -> List[List[str]]:
141 """Default PP split: even distribution of the model's layer container.
143 Mirrors torchtitan's ``_generate_llm_fqn_per_model_part`` weighting
144 without hardcoding LLM module names:
146 - The first top-level child that is an ``nn.ModuleList`` with at least
147 ``pp_degree`` entries is the *layer container* (typically ``layers``).
148 - Top-level children registered before it go to stage 0 (input side,
149 e.g. embeddings); children after it go to the last stage (output
150 side, e.g. norm / lm_head).
151 - Container elements are distributed as evenly as possible (see
152 ``_even_split``).
154 Without a qualifying container, all top-level children are split evenly
155 in registration order.
157 Returns:
158 ``pp_degree`` lists of exact module FQNs, one per stage.
159 """
160 if pp_degree < 2:
161 raise ValueError(f"pp_degree must be >= 2 for a PP split, got {pp_degree}")
163 children = [name for name, _ in model.named_children()]
165 container: Optional[str] = None
166 for name, child in model.named_children():
167 if isinstance(child, nn.ModuleList) and len(child) >= pp_degree:
168 container = name
169 break
171 if container is None:
172 if len(children) < pp_degree:
173 raise ValueError(
174 f"Cannot split {len(children)} top-level children across "
175 f"{pp_degree} PP stages; provide a manual stage plan via "
176 f"GraphParallelPlan.pp_stage()"
177 )
178 return _even_split(children, pp_degree)
180 container_mod = getattr(model, container)
181 if isinstance(container_mod, nn.ModuleDict):
182 elem_keys = list(container_mod.keys())
183 else:
184 elem_keys = [str(i) for i in range(len(container_mod))]
186 idx = children.index(container)
187 before, after = children[:idx], children[idx + 1 :]
189 stages = _even_split(elem_keys, pp_degree, prefix=f"{container}.")
190 stages[0] = before + stages[0]
191 stages[-1] = stages[-1] + after
192 return stages
195@dataclass
196class _RunContext:
197 """This rank's PP execution context, resolved and guarded for ``run``."""
199 model: nn.Module
200 pp_degree: int
201 group: Any
202 stage_idx: int
205@dataclass
206class _StageSplit:
207 """Stage plan and per-node classification produced by ``_classify_stage``.
209 Bundles what the boundary-detection and subgraph-construction phases
210 need, so ``run`` orchestrates phases instead of threading a dozen
211 parallel locals through every call.
212 """
214 state_fqns: List[str]
215 state_is_param: List[bool]
216 stage_plan: List[List[str]]
217 node_stage: Dict[fx.Node, int]
218 node_phase: Dict[fx.Node, str]
219 user_input_phs: List[fx.Node]
220 user_input_stages: List[int]
221 state_phs: List[fx.Node]
222 stage_state_indices: List[int]
223 stage_state_fqns: List[str]
226@dataclass
227class _GraphClassification:
228 """Per-node stage/phase classification produced by ``_classify_graph``."""
230 stage_state_indices: List[int]
231 node_stage: Dict[fx.Node, int]
232 node_phase: Dict[fx.Node, str]
233 user_input_phs: List[fx.Node]
234 user_input_stages: List[int]
235 state_phs: List[fx.Node]
238@dataclass
239class _Boundaries:
240 """Boundary value lists crossing this stage's cuts.
242 All lists are in joint-graph topological order so the producer and
243 consumer of a cut agree on P2P ordering:
245 - ``act_in``: stage-(k-1) values consumed by stage k (tensors,
246 dynamic-shape scalars, parent-level views — everything a real ATen
247 graph pushes across a cut).
248 - ``act_out``: stage-k values consumed by stage k+1.
249 - ``grad_in``: stage-(k+1) backward values consumed by stage k.
250 - ``grad_out``: stage-k backward values consumed by stage k-1.
251 - ``saved``: forward-phase values referenced by this stage's backward
252 nodes — exported as forward outputs and replayed as backward inputs.
253 """
255 act_in: List[fx.Node]
256 act_out: List[fx.Node]
257 grad_in: List[fx.Node]
258 grad_out: List[fx.Node]
259 saved: List[fx.Node]
262@dataclass
263class _SubgraphBuilder:
264 """Accumulator for assembling one stage subgraph.
266 ``env`` maps already-copied same-slice nodes; ``foreign`` maps
267 boundary values to their placeholder stand-ins; ``used`` carries the
268 taken placeholder names for ``_unique``.
269 """
271 g: fx.Graph
272 used: Set[str]
273 env: Dict[fx.Node, fx.Node]
274 foreign: Dict[fx.Node, fx.Node]
277class PpPass(GraphPass):
278 """Pipeline-parallel partitioning pass (graph-level stage split)."""
280 name = "pp_parallel"
282 def __init__(self, parallel_plan: Optional[GraphParallelPlan] = None) -> None:
283 """Initialize PP pass state.
285 Args:
286 parallel_plan: Declarative plan; ``pp_module_fqns_per_stage`` (from
287 ``pp_stage()``) overrides the automatic even-by-layers
288 split.
289 """
290 super().__init__()
291 self._parallel_plan = parallel_plan
293 def run(
294 self,
295 graph_module: fx.GraphModule,
296 pass_config: PassConfig,
297 **kwargs: Any,
298 ) -> fx.GraphModule:
299 """Split the joint graph to this rank's stage and install the schedule.
301 Args:
302 graph_module: Joint fwd+bwd FX graph (post-``FSDPPass``).
303 pass_config: Parallel configuration; ``pp_degree`` (``None``
304 resolves to ``world_size``) and ``pp_microbatch_size`` are
305 read directly.
306 **kwargs: Must include ``model`` (the live ``nn.Module``); may
307 include ``parallel_plan``.
309 Returns:
310 The same graph module, rewritten into a stage stub whose
311 ``call_module`` node invokes the installed pipeline schedule
312 (``ScheduleGPipe`` / ``Schedule1F1B`` per
313 ``pass_config.pp_schedule``).
314 """
315 run_ctx = self._resolve_run_context(pass_config, kwargs)
316 if run_ctx is None:
317 return graph_module
319 split = self._classify_stage(graph_module, run_ctx)
320 boundaries = self._find_boundaries(
321 graph_module,
322 split.node_stage,
323 split.node_phase,
324 run_ctx.pp_degree,
325 run_ctx.stage_idx,
326 )
327 self._build_and_install_stage(
328 graph_module, pass_config, run_ctx, split, boundaries
329 )
330 return graph_module
332 def _resolve_run_context(
333 self,
334 pass_config: PassConfig,
335 kwargs: Dict[str, Any],
336 ) -> Optional[_RunContext]:
337 """Guard and resolve the model, PP group, and this rank's stage.
339 Returns:
340 The run context, or ``None`` when the pass is skipped
341 (distributed not initialized, ``world_size == 1``, or
342 ``pp_degree < 2``).
343 """
344 if not dist.is_initialized() or dist.get_world_size() == 1:
345 _LOG.info("Skipped: distributed not initialized or world_size=1")
346 return None
348 model = kwargs.get("model")
349 if model is None:
350 raise ValueError(
351 "PpPass requires the live model via kwargs (model=...) so it "
352 "can resolve the stage plan and prune foreign-stage modules"
353 )
354 self._parallel_plan = kwargs.get("parallel_plan", self._parallel_plan)
356 # Skip BEFORE _resolve_group_and_stage so a disabled pass does not
357 # execute the collective dist.new_group.
358 pp_degree = pass_config.pp_degree or dist.get_world_size()
359 if pp_degree < 2:
360 _LOG.info("Skipped: pp_degree=%s < 2", pp_degree)
361 return None
362 pp_degree, group, stage_idx = self._resolve_group_and_stage(pass_config)
363 _LOG.info(
364 "Running: stage %s/%s, world_size=%s",
365 stage_idx,
366 pp_degree,
367 dist.get_world_size(),
368 )
369 return _RunContext(
370 model=model, pp_degree=pp_degree, group=group, stage_idx=stage_idx
371 )
373 def _classify_stage(
374 self, graph_module: fx.GraphModule, run_ctx: _RunContext
375 ) -> _StageSplit:
376 """Resolve the stage plan and classify every node's stage + phase.
378 Chains the state-list contract checks, the stage-plan resolution,
379 node classification, the gradient-output anchoring with its stage
380 propagation, and the adjacent-dataflow validation.
382 Args:
383 graph_module: Joint fwd+bwd FX graph (post-``FSDPPass``).
384 run_ctx: This rank's PP execution context.
386 Returns:
387 The ``_StageSplit`` bundling the plan and per-node
388 classification for the construction phase.
389 """
390 state_fqns = self._require_state_list(graph_module)
391 state_is_param = self._require_is_param(graph_module, state_fqns)
392 num_state_inputs = getattr(graph_module, "num_state_inputs", len(state_fqns))
394 stage_plan = self._resolve_stage_plan(run_ctx.model, run_ctx.pp_degree)
395 stage_of_fqn = {fqn: s for s, fqns in enumerate(stage_plan) for fqn in fqns}
397 graph_cls = self._classify_graph(
398 graph_module,
399 run_ctx,
400 stage_of_fqn,
401 state_fqns,
402 num_state_inputs,
403 )
405 self._anchor_grad_outputs(
406 graph_module,
407 graph_cls.node_stage,
408 stage_of_fqn,
409 state_fqns,
410 run_ctx.model,
411 )
412 self._propagate_anchor_stages(
413 graph_module, graph_cls.node_stage, graph_cls.node_phase
414 )
415 self._validate_adjacent_dataflow(graph_module, graph_cls.node_stage)
417 return _StageSplit(
418 state_fqns=state_fqns,
419 state_is_param=state_is_param,
420 stage_plan=stage_plan,
421 node_stage=graph_cls.node_stage,
422 node_phase=graph_cls.node_phase,
423 user_input_phs=graph_cls.user_input_phs,
424 user_input_stages=graph_cls.user_input_stages,
425 state_phs=graph_cls.state_phs,
426 stage_state_indices=graph_cls.stage_state_indices,
427 stage_state_fqns=[state_fqns[i] for i in graph_cls.stage_state_indices],
428 )
430 def _build_and_install_stage(
431 self,
432 graph_module: fx.GraphModule,
433 pass_config: PassConfig,
434 run_ctx: _RunContext,
435 split: _StageSplit,
436 boundaries: _Boundaries,
437 ) -> None:
438 """Build the stage subgraphs, prune the live model, install the stub.
440 Trainable-state resolution must precede ``_prune_live_model``: the
441 gradient-output nodes map positionally onto the FULL trainable list,
442 so pruning foreign-stage parameters first would shift the positions.
444 Args:
445 graph_module: Joint fwd+bwd FX graph, rewritten in place.
446 pass_config: Parallel configuration (microbatch size read).
447 run_ctx: This rank's PP execution context.
448 split: Stage plan and per-node classification.
449 boundaries: Boundary value lists from ``_find_boundaries``.
450 """
451 trainable_indices = self._trainable_state_indices(
452 split.state_fqns, split.state_is_param, run_ctx.model
453 )
454 stage_state_set = set(split.stage_state_indices)
455 stage_trainable = [i for i in trainable_indices if i in stage_state_set]
457 fwd_gm, bwd_gm = self._build_stage_graphs(
458 graph_module, run_ctx, split, boundaries, trainable_indices
459 )
461 self._prune_live_model(run_ctx.model, split.stage_plan[run_ctx.stage_idx])
463 sched = self._build_schedule(
464 fwd_gm,
465 bwd_gm,
466 pass_config,
467 run_ctx,
468 split,
469 len(stage_trainable),
470 boundaries,
471 )
472 self._install_stub(
473 graph_module,
474 sched,
475 split.stage_state_fqns,
476 split.state_is_param,
477 split.stage_state_indices,
478 len(split.user_input_phs),
479 )
481 _LOG.info(
482 "Completed: stage %s owns %d state entries, %d trainable params, "
483 "%d fwd / %d bwd boundary values",
484 run_ctx.stage_idx,
485 len(split.stage_state_fqns),
486 len(stage_trainable),
487 len(boundaries.act_out),
488 len(boundaries.grad_out),
489 )
491 def _build_stage_graphs(
492 self,
493 graph_module: fx.GraphModule,
494 run_ctx: _RunContext,
495 split: _StageSplit,
496 boundaries: _Boundaries,
497 trainable_indices: List[int],
498 ) -> Tuple[fx.GraphModule, fx.GraphModule]:
499 """Build this stage's forward and backward subgraphs."""
500 fwd_gm, fwd_out_orig = self._build_fwd_graph(
501 graph_module, run_ctx, split, boundaries
502 )
503 bwd_gm = self._build_bwd_graph(
504 graph_module, run_ctx, split, boundaries, fwd_out_orig, trainable_indices
505 )
506 return fwd_gm, bwd_gm
508 # ------------------------------------------------------------------
509 # Tracer contract helpers
510 # ------------------------------------------------------------------
512 @staticmethod
513 def _require_state_list(graph_module: fx.GraphModule) -> List[str]:
514 """Return the graph's live ``state_fqns`` list (shared with JointGraph).
516 The tracer attaches the SAME list object to the GraphModule and the
517 ``JointGraph``; mutating it in place (done in ``_install_stub``)
518 keeps ``run_traced_graph``'s validation consistent with the pruned
519 model. Copying here would desynchronize the two.
520 """
521 state_fqns = getattr(graph_module, "state_fqns", None)
522 if state_fqns is None:
523 raise ValueError(
524 "PpPass requires the tracer's state_fqns attribute on the "
525 "joint graph (attach it via trace_model_graph)"
526 )
527 return state_fqns
529 @staticmethod
530 def _require_is_param(
531 graph_module: fx.GraphModule, state_fqns: List[str]
532 ) -> List[bool]:
533 """Return the graph's live ``state_is_param`` list, defaulting to
534 all-params (matching the tracer's fallback for older traces)."""
535 state_is_param = getattr(graph_module, "state_is_param", None)
536 if state_is_param is None:
537 state_is_param = [True] * len(state_fqns)
538 graph_module.state_is_param = state_is_param
539 return state_is_param
541 @staticmethod
542 def _loss_node(graph_module: fx.GraphModule) -> Optional[fx.Node]:
543 """The original joint graph's loss node (output index 0)."""
544 output = next((n for n in graph_module.graph.nodes if n.op == "output"), None)
545 if output is None:
546 return None
547 returned = output.args[0]
548 if isinstance(returned, (list, tuple)) and returned:
549 first = returned[0]
550 if isinstance(first, fx.Node):
551 return first
552 return None
554 # ------------------------------------------------------------------
555 # Group / plan resolution
556 # ------------------------------------------------------------------
558 def _resolve_group_and_stage(self, pass_config: PassConfig) -> Tuple[int, Any, int]:
559 """Resolve PP degree, process group, and this rank's stage index.
561 A single PP process group carries both forward activations and
562 backward gradients; the schedule keeps the two ends of every
563 neighbour pair's FIFO complementary (see ``Schedule1F1B``).
565 v1 supports the pure-PP layout (``pp_degree == world_size``, one
566 stage per rank); ``pp_degree == 1`` disables the pass. Any other
567 layout is rejected loudly: DP replication over one stage would
568 funnel every replica's P2P into a single peer rank (the group is
569 the whole world), and the PP+FSDP hybrid requires a mesh
570 expressing both dimensions (``mesh_context`` path) — the trainer's
571 1-D fallback mesh shards FSDP across the whole world and cannot
572 host PP, so that combination would silently mis-shard.
573 """
574 world_size = dist.get_world_size()
575 rank = dist.get_rank()
576 pp_degree = pass_config.pp_degree or world_size
578 if pp_degree not in (1, world_size):
579 raise ValueError(
580 f"v1 supports pure-PP only: pp_degree={pp_degree} must "
581 f"equal world_size={world_size} (or be 1 to disable PP). "
582 f"DP-replication and PP+FSDP hybrid layouts require a "
583 f"mesh_context with an explicit pp dim."
584 )
585 if pass_config.fsdp_enabled:
586 raise ValueError(
587 "PP+FSDP hybrid requires a mesh expressing both the PP and "
588 "FSDP dimensions (mesh_context path). The trainer's 1-D "
589 "fallback FSDP mesh spans the whole world and cannot host "
590 "PP — FSDP would shard parameters across stage boundaries. "
591 "Run pure-PP (fsdp_enabled=False, pp_degree == world_size) "
592 "or provide a mesh_context with a pp dim."
593 )
595 group = dist.new_group(list(range(world_size)))
596 # v1 is pure-PP (pp_degree == world_size, checked above), so each
597 # rank owns exactly one stage. Keep the general formula so a future
598 # DP-replication layout (pp_degree < world_size) maps rank->stage
599 # contiguously rather than silently assuming 1:1.
600 stage_idx = rank * pp_degree // world_size
601 return pp_degree, group, stage_idx
603 def _resolve_stage_plan(self, model: nn.Module, pp_degree: int) -> List[List[str]]:
604 """Resolve the per-stage module FQN lists (manual plan or auto split)."""
605 manual = (
606 self._parallel_plan.pp_module_fqns_per_stage
607 if self._parallel_plan is not None
608 else None
609 )
610 if manual is not None:
611 if len(manual) != pp_degree:
612 raise ValueError(
613 f"Manual PP plan declares {len(manual)} stages but "
614 f"pp_degree={pp_degree}"
615 )
616 seen: Set[str] = set()
617 for s, fqns in enumerate(manual):
618 for fqn in fqns:
619 if fqn in seen:
620 raise ValueError(
621 f"Module '{fqn}' is assigned to multiple stages "
622 f"(stage {s} repeats it)"
623 )
624 seen.add(fqn)
625 stage_plan = [list(fqns) for fqns in manual]
626 else:
627 stage_plan = _auto_stage_split(model, pp_degree)
629 known = dict(model.named_modules())
630 for s, fqns in enumerate(stage_plan):
631 for fqn in fqns:
632 if fqn not in known:
633 raise ValueError(f"Stage {s} declares unknown module '{fqn}'")
634 self._warn_unassigned_modules(model, stage_plan)
635 return stage_plan
637 @staticmethod
638 def _warn_unassigned_modules(model: nn.Module, stage_plan: List[List[str]]) -> None:
639 """Warn about live modules no stage declares.
641 A module that is neither declared itself, nor lives inside a
642 declared one, nor is an ancestor of a declared one (a container
643 whose elements are declared) has its graph nodes attributed by
644 the dataflow fallbacks (consumer stage, then last stage) — rarely
645 what a manual plan intends. Auto splits declare every top-level
646 child, so they never warn.
647 """
648 declared = [fqn for fqns in stage_plan for fqn in fqns]
649 unassigned: List[str] = []
650 for fqn, _ in model.named_modules():
651 if not fqn:
652 # Root: parent-level glue is attributed by dataflow.
653 continue
654 covered = any(
655 fqn == d or fqn.startswith(d + ".") or d.startswith(fqn + ".")
656 for d in declared
657 )
658 if not covered:
659 unassigned.append(fqn)
660 if not unassigned:
661 return
662 shown = unassigned[:10]
663 if len(unassigned) > len(shown):
664 shown.append(f"... {len(unassigned) - len(shown)} more")
665 _LOG.warning(
666 "No PP stage declares modules %s — their nodes ride with the "
667 "dataflow attribution (consumer stage, then last stage). "
668 "Declare them via GraphParallelPlan.pp_stage() if that is not intended.",
669 shown,
670 )
672 # ------------------------------------------------------------------
673 # Node classification (stage + phase, topo order)
674 # ------------------------------------------------------------------
676 def _classify_graph( # pylint: disable=too-many-locals
677 self,
678 graph_module: fx.GraphModule,
679 run_ctx: _RunContext,
680 stage_of_fqn: Dict[str, int],
681 state_fqns: Sequence[str],
682 num_state_inputs: int,
683 ) -> _GraphClassification:
684 """Resolve every node's stage and fwd/bwd phase in topological order.
686 Phase: explicit ``autograd_backward`` tag, else inherited from args
687 (post-trace inserts such as FSDP's reduce_scatter/wait inherit the
688 bwd phase through their grad arguments).
690 Stage resolution, in priority order:
692 1. ``nn_module_stack`` walk (traced nodes, innermost module first)
693 2. argument inheritance (post-trace inserts — state placeholders
694 win, since backward compute nodes reference their params)
695 3. consumer attribution for stack-less root-level glue (residual
696 adds written in the parent's forward): the stage of their
697 consumers, minimum when several — resolved in reverse topo
698 4. last stage as the final fallback (loss chain, dead nodes)
700 User-input placeholders are NOT anchors: they are routed AFTER the
701 fwd/bwd classification by dataflow (``_route_user_inputs``) — each
702 rides with the single stage that consumes it, so kwargs-general
703 inputs (multiple features on stage 0, several labels on the last
704 stage, an auxiliary input on a middle stage) all work without
705 positional heuristics.
707 ``stage_state_indices`` lists the state entries owned by THIS
708 stage (``stage_idx``); every state placeholder still gets a stage
709 so foreign consumers can be attributed.
711 Root-level state (no stage-owning module ancestor) is assigned to
712 the stage of its (single) consumer; multi-stage consumption is the
713 RoPE-cache-shaped error case and fails with an actionable message.
714 """
715 graph = graph_module.graph
716 placeholders = [n for n in graph.nodes if n.op == "placeholder"]
717 state_ph_set = set(placeholders[:num_state_inputs])
718 num_stages = run_ctx.pp_degree
720 node_stage: Dict[fx.Node, int] = {}
721 node_phase: Dict[fx.Node, str] = {}
723 (
724 state_phs,
725 deferred_state,
726 stage_state_indices,
727 user_input_phs,
728 ) = self._classify_placeholders(
729 placeholders,
730 node_stage,
731 node_phase,
732 state_fqns,
733 stage_of_fqn,
734 num_state_inputs,
735 run_ctx,
736 )
737 self._classify_forward_nodes(
738 graph, state_ph_set, stage_of_fqn, num_stages, node_stage, node_phase
739 )
740 self._classify_backward_nodes(
741 graph, state_ph_set, num_stages, node_stage, node_phase
742 )
743 user_input_stages = self._route_user_inputs(
744 user_input_phs, node_stage, num_stages
745 )
746 self._assign_deferred_state(
747 deferred_state,
748 state_phs,
749 state_fqns,
750 node_stage,
751 run_ctx.stage_idx,
752 stage_state_indices,
753 )
755 return _GraphClassification(
756 stage_state_indices=sorted(stage_state_indices),
757 node_stage=node_stage,
758 node_phase=node_phase,
759 user_input_phs=user_input_phs,
760 user_input_stages=user_input_stages,
761 state_phs=state_phs,
762 )
764 def _classify_placeholders( # pylint: disable=too-many-locals
765 self,
766 placeholders: List[fx.Node],
767 node_stage: Dict[fx.Node, int],
768 node_phase: Dict[fx.Node, str],
769 state_fqns: Sequence[str],
770 stage_of_fqn: Dict[str, int],
771 num_state_inputs: int,
772 run_ctx: _RunContext,
773 ) -> Tuple[List[fx.Node], List[int], List[int], List[fx.Node]]:
774 """Classify placeholder nodes and collect this stage's state entries.
776 Only STATE placeholders get a stage here (via their FQN's module
777 ancestor); user-input placeholders are routed later by dataflow
778 (``_route_user_inputs``), once their consumers' stages are known.
780 Args:
781 placeholders: All placeholder nodes in graph order.
782 node_stage: Stage assignment accumulator to fill.
783 node_phase: Phase assignment accumulator to fill.
784 state_fqns: FQNs of the traced state entries (params/buffers).
785 stage_of_fqn: Module FQN -> owning stage map.
786 num_state_inputs: Number of leading state placeholders.
787 run_ctx: This rank's PP execution context.
789 Returns:
790 ``(state_phs, deferred_state, stage_state_indices,
791 user_input_phs)``. ``deferred_state`` lists root-level state
792 indices whose stage is resolved later from their consumers.
793 """
794 state_phs: List[fx.Node] = []
795 deferred_state: List[int] = []
796 stage_state_indices: List[int] = []
797 for idx, ph in enumerate(placeholders):
798 node_phase[ph] = _FWD
799 if idx < num_state_inputs:
800 stage = self._state_stage(state_fqns[idx], stage_of_fqn)
801 state_phs.append(ph)
802 if stage is None:
803 deferred_state.append(idx)
804 else:
805 node_stage[ph] = stage
806 if stage == run_ctx.stage_idx:
807 stage_state_indices.append(idx)
808 user_input_phs = list(placeholders[num_state_inputs:])
809 return state_phs, deferred_state, stage_state_indices, user_input_phs
811 @staticmethod
812 def _route_user_inputs(
813 user_input_phs: List[fx.Node],
814 node_stage: Dict[fx.Node, int],
815 num_stages: int,
816 ) -> List[int]:
817 """Assign each user-input placeholder to the stage consuming it.
819 Runs after the fwd/bwd classification, so every consumer node
820 already carries its final stage (user inputs are deliberately not
821 anchors — see ``_classify_graph``). Each consumed input must feed
822 exactly one stage: the schedule ships user inputs directly to their
823 owner (no P2P), so an input consumed on several stages has no
824 route. An input consumed by NO node (an ignored kwarg — the
825 hand-built test fixtures' ``y``, an ``epoch`` scalar train_fn never
826 reads) is warned about and routed to the last stage, where unused
827 inputs have always ridden.
829 Args:
830 user_input_phs: User-input placeholders in stub arg order.
831 node_stage: Final per-node stage assignment (fwd+bwd done).
832 num_stages: Total PP stage count.
834 Returns:
835 The stage index per user input, in stub arg order.
837 Raises:
838 ValueError: When an input is consumed by nodes on more than
839 one stage.
840 """
841 stages: List[int] = []
842 for ph in user_input_phs:
843 consumer_stages = {
844 node_stage[u]
845 for u in ph.users
846 if u.op != "output" and node_stage.get(u) is not None
847 }
848 if not consumer_stages:
849 _LOG.warning(
850 "Model input '%s' is not consumed by the traced graph — "
851 "routing it to the last stage with the other unused "
852 "inputs (a misspelled keyword argument would leave its "
853 "tensor unused too)",
854 ph.name,
855 )
856 stage = num_stages - 1
857 elif len(consumer_stages) > 1:
858 raise ValueError(
859 f"Model input '{ph.name}' is consumed on stages "
860 f"{sorted(consumer_stages)} — PP ships each user input "
861 f"directly to the single stage that consumes it; split "
862 f"the input per stage or restructure the graph"
863 )
864 else:
865 stage = next(iter(consumer_stages))
866 node_stage[ph] = stage
867 stages.append(stage)
868 _LOG.info(
869 "User input '%s' routed to stage %s of %s", ph.name, stage, num_stages
870 )
871 return stages
873 def _classify_forward_nodes(
874 self,
875 graph: fx.Graph,
876 state_ph_set: Set[fx.Node],
877 stage_of_fqn: Dict[str, int],
878 num_stages: int,
879 node_stage: Dict[fx.Node, int],
880 node_phase: Dict[fx.Node, str],
881 ) -> None:
882 """Assign stages to forward nodes in topological order.
884 Anchors: module stack, then state-ph args, then max-args over
885 TENSOR args (data flows forward; shape-metadata scalars carry no
886 data dependency). Leftover root-level glue rides with its
887 consumers (reverse topo), else the last stage.
888 """
889 floating: List[fx.Node] = []
890 for node in graph.nodes:
891 if node.op in ("placeholder", "output"):
892 continue
893 phase = _BWD if node.meta.get("autograd_backward", False) else _FWD
894 if phase == _FWD:
895 for arg in _iter_node_args(node):
896 if node_phase.get(arg) == _BWD:
897 phase = _BWD
898 break
899 node_phase[node] = phase
900 if phase == _FWD:
901 stage = self._resolve_fwd_stage(
902 node, node_stage, state_ph_set, stage_of_fqn
903 )
904 if stage is None:
905 floating.append(node)
906 else:
907 node_stage[node] = stage
909 self._assign_floating_fwd_nodes(
910 graph, floating, num_stages, node_stage, node_phase
911 )
913 def _resolve_fwd_stage(
914 self,
915 node: fx.Node,
916 node_stage: Dict[fx.Node, int],
917 state_ph_set: Set[fx.Node],
918 stage_of_fqn: Dict[str, int],
919 ) -> Optional[int]:
920 """Resolve a forward node's stage anchor.
922 Priority: module stack, then state-ph args, then max-args over
923 TENSOR args (data flows forward; shape-metadata scalars carry no
924 data dependency). ``None`` means root-level glue to be resolved
925 from its consumers (see ``_assign_floating_fwd_nodes``).
926 """
927 stage = self._stage_from_stack(node, stage_of_fqn)
928 if stage is None:
929 stage = self._stage_from_state_arg(node, node_stage, state_ph_set)
930 if stage is None:
931 arg_stages = _tensor_arg_stages(node, node_stage)
932 if arg_stages:
933 stage = max(arg_stages)
934 return stage
936 @staticmethod
937 def _assign_floating_fwd_nodes(
938 graph: fx.Graph,
939 floating: List[fx.Node],
940 num_stages: int,
941 node_stage: Dict[fx.Node, int],
942 node_phase: Dict[fx.Node, str],
943 ) -> None:
944 """Leftover fwd glue: consumers first (reverse topo), else last."""
945 for node in reversed(list(graph.nodes)):
946 if node not in floating or node_phase[node] != _FWD:
947 continue
948 consumer_stages = {
949 node_stage[u]
950 for u in node.users
951 if u.op != "output" and node_stage.get(u) is not None
952 }
953 node_stage[node] = (
954 min(consumer_stages) if consumer_stages else num_stages - 1
955 )
957 def _classify_backward_nodes(
958 self,
959 graph: fx.Graph,
960 state_ph_set: Set[fx.Node],
961 num_stages: int,
962 node_stage: Dict[fx.Node, int],
963 node_phase: Dict[fx.Node, str],
964 ) -> None:
965 """Assign stages to backward nodes in topological order.
967 A bwd node runs where its earliest TENSOR input lives: gradients
968 born downstream are shipped down during the backward sweep, and
969 upstream activations arrive as replayed forward outputs. Weight-
970 gradient nodes that consume early activations are corrected by
971 the positional anchor (the param's stage owns its grad).
972 """
973 for node in graph.nodes:
974 if node.op in ("placeholder", "output"):
975 continue
976 if node_phase[node] != _BWD:
977 continue
978 stage = self._stage_from_state_arg(node, node_stage, state_ph_set)
979 if stage is None:
980 arg_stages = _tensor_arg_stages(node, node_stage)
981 if arg_stages:
982 stage = min(arg_stages)
983 if stage is None:
984 # Dead-end bwd glue: ride with consumers if any, else last.
985 consumer_stages = {
986 node_stage[u]
987 for u in node.users
988 if u.op != "output" and node_stage.get(u) is not None
989 }
990 stage = min(consumer_stages) if consumer_stages else num_stages - 1
991 node_stage[node] = stage
993 def _assign_deferred_state(
994 self,
995 deferred_state: List[int],
996 state_phs: List[fx.Node],
997 state_fqns: Sequence[str],
998 node_stage: Dict[fx.Node, int],
999 stage_idx: int,
1000 stage_state_indices: List[int],
1001 ) -> None:
1002 """Assign root-level state to the stage of its (single) consumer.
1004 Multi-stage consumption is the RoPE-cache-shaped error case and
1005 fails with an actionable message.
1006 """
1007 for idx in deferred_state:
1008 ph = state_phs[idx]
1009 consumer_stages = {node_stage[u] for u in ph.users if u.op != "output"}
1010 if len(consumer_stages) > 1:
1011 raise ValueError(
1012 f"Root-level state '{state_fqns[idx]}' is consumed by "
1013 f"stages {sorted(consumer_stages)} — v1 requires state "
1014 f"to live under a single stage's module (e.g. move a "
1015 f"shared RoPE cache under a per-stage module or split "
1016 f"it per stage)"
1017 )
1018 stage = next(iter(consumer_stages), 0)
1019 node_stage[ph] = stage
1020 if stage == stage_idx:
1021 stage_state_indices.append(idx)
1022 _LOG.info(
1023 "Root-level state '%s' assigned to stage %s by its consumer",
1024 state_fqns[idx],
1025 stage,
1026 )
1028 def _state_stage(self, fqn: str, stage_of_fqn: Dict[str, int]) -> Optional[int]:
1029 """Stage of a state entry via longest-ancestor walk over its FQN."""
1030 parts = fqn.split(".")
1031 for i in range(len(parts) - 1, 0, -1):
1032 stage = stage_of_fqn.get(".".join(parts[:i]))
1033 if stage is not None:
1034 return stage
1035 return stage_of_fqn.get(fqn)
1037 def _stage_from_stack(
1038 self, node: fx.Node, stage_of_fqn: Dict[str, int]
1039 ) -> Optional[int]:
1040 """Stage from ``nn_module_stack``, innermost module first."""
1041 for fqn in _innermost_fqns(node):
1042 stage = stage_of_fqn.get(fqn)
1043 if stage is not None:
1044 return stage
1045 return None
1047 def _stage_from_state_arg(
1048 self,
1049 node: fx.Node,
1050 node_stage: Dict[fx.Node, int],
1051 state_ph_set: Set[fx.Node],
1052 ) -> Optional[int]:
1053 """Stage from a state-placeholder argument, if any.
1055 Post-trace inserts (FSDP collectives, weight-grad compute) reference
1056 their owning stage's parameters — a state placeholder arg pins the
1057 node to that stage. Generic argument inheritance is deliberately
1058 NOT done here: it would let an early-stage argument pull a node
1059 against the dataflow (see the max/min-args rules in
1060 ``_classify_graph``).
1061 """
1062 for arg in _iter_node_args(node):
1063 if arg in state_ph_set:
1064 return node_stage.get(arg)
1065 return None
1067 @staticmethod
1068 def _propagate_anchor_stages(
1069 graph_module: fx.GraphModule,
1070 node_stage: Dict[fx.Node, int],
1071 node_phase: Dict[fx.Node, str],
1072 ) -> None:
1073 """Pull bwd args up to anchored consumers (one reverse-topo pass).
1075 A bwd value consumed by a DOWNSTREAM-stage bwd node must be
1076 recomputed there: gradients flow downstream -> upstream, so a
1077 downstream consumer cannot receive an upstream-born bwd value in
1078 time (its own bwd sweep runs first). Fwd-phase args are untouched —
1079 they arrive as replayed forward outputs; downstream-born args are
1080 untouched — they arrive via the grad crossings.
1082 A single REVERSE topological sweep reaches the fixpoint: a node is
1083 visited only after every user of it has been visited and pulled it
1084 to its final stage, so the node propagates its final value to its
1085 own args. The previous iterate-to-fixpoint loop was quadratic in
1086 the worst case on large joint graphs.
1087 """
1088 for node in reversed(list(graph_module.graph.nodes)):
1089 if node.op in ("placeholder", "output"):
1090 continue
1091 if node_phase.get(node) != _BWD:
1092 continue
1093 stage = node_stage.get(node)
1094 if stage is None:
1095 continue
1096 for arg in _iter_node_args(node):
1097 if node_phase.get(arg) != _BWD or node_stage.get(arg) is None:
1098 continue
1099 if node_stage[arg] < stage:
1100 node_stage[arg] = stage
1102 # ------------------------------------------------------------------
1103 # Boundary detection
1104 # ------------------------------------------------------------------
1106 @staticmethod
1107 def _validate_adjacent_dataflow(
1108 graph_module: fx.GraphModule,
1109 node_stage: Dict[fx.Node, int],
1110 ) -> None:
1111 """Reject skip-stage dataflow before it reaches the slice copy.
1113 v1 ships boundary values between NEIGHBOURING stages only: a value
1114 produced on stage ``j`` and consumed on stage ``k`` with
1115 ``abs(j - k) > 1`` has no route (it would have to transit
1116 intermediate stages untouched). Without this check such graphs die
1117 later inside ``_copy_node`` / the fwd-output assembly with a
1118 generic "outside this stage's slice" error that hides the root
1119 cause.
1121 Raises:
1122 ValueError: When any producer-consumer pair spans more than
1123 one stage cut (e.g. a dense skip connection across two
1124 stages, or a stage plan that leaves a module's nodes to
1125 the dataflow fallbacks).
1126 """
1127 for node in graph_module.graph.nodes:
1128 if node.op == "output":
1129 continue
1130 producer = node_stage.get(node)
1131 if producer is None:
1132 continue
1133 for user in node.users:
1134 if user.op == "output":
1135 continue
1136 consumer = node_stage.get(user)
1137 if consumer is None or abs(consumer - producer) <= 1:
1138 continue
1139 raise ValueError(
1140 f"Value '{node.name}' is produced on stage {producer} "
1141 f"but consumed on stage {consumer}: v1 ships boundary "
1142 f"values between neighbouring stages only. A skip "
1143 f"connection spanning 2+ stages or a stage plan that "
1144 f"leaves modules unattributed produces this — adjust "
1145 f"GraphParallelPlan.pp_stage() so the dataflow crosses one cut "
1146 f"at a time."
1147 )
1149 def _anchor_grad_outputs(
1150 self,
1151 graph_module: fx.GraphModule,
1152 node_stage: Dict[fx.Node, int],
1153 stage_of_fqn: Dict[str, int],
1154 state_fqns: Sequence[str],
1155 model: Optional[nn.Module],
1156 ) -> None:
1157 """Pin each gradient-output node to its parameter's stage.
1159 Gradient-output node i differentiates the i-th trainable parameter
1160 (positional mapping, mirroring the tracer); it runs on the stage
1161 owning that parameter even when min-args would place it upstream
1162 (a later stage's weight-grad consumes an earlier stage's
1163 activation, replayed from the forward outputs).
1164 """
1165 output = next((n for n in graph_module.graph.nodes if n.op == "output"), None)
1166 if output is None or model is None:
1167 return
1168 returned = output.args[0]
1169 if not isinstance(returned, (list, tuple)):
1170 return
1171 state_is_param = getattr(
1172 graph_module, "state_is_param", [True] * len(state_fqns)
1173 )
1174 trainable = self._trainable_state_indices(state_fqns, state_is_param, model)
1175 for i, grad_node in enumerate(returned[1:]):
1176 if i >= len(trainable) or not isinstance(grad_node, fx.Node):
1177 continue
1178 pstage = self._state_stage(state_fqns[trainable[i]], stage_of_fqn)
1179 if pstage is not None:
1180 node_stage[grad_node] = pstage
1182 def _find_boundaries( # pylint: disable=too-many-locals
1183 self,
1184 graph_module: fx.GraphModule,
1185 node_stage: Dict[fx.Node, int],
1186 node_phase: Dict[fx.Node, str],
1187 pp_degree: int,
1188 stage_idx: int,
1189 ) -> _Boundaries:
1190 """Locate the value lists crossing this stage's cuts.
1192 Returns:
1193 The ``_Boundaries`` lists, all in joint-graph topological order
1194 so the producer and consumer of a cut agree on P2P ordering.
1195 See ``_Boundaries`` for the per-field semantics.
1196 """
1197 k = stage_idx
1198 is_last = k == pp_degree - 1
1199 nodes = [
1200 n for n in graph_module.graph.nodes if n.op not in ("placeholder", "output")
1201 ]
1203 def crossing(
1204 producer: int, consumer: int, phase: Optional[str]
1205 ) -> List[fx.Node]:
1206 """Values produced by one stage and consumed by another.
1208 ``phase=None`` matches any phase (backward crossings may carry
1209 gradients as well as functionalized fwd views).
1210 """
1211 found = []
1212 for v in nodes:
1213 if node_stage[v] != producer:
1214 continue
1215 if phase is not None and node_phase[v] != phase:
1216 continue
1217 if any(node_stage.get(u) == consumer for u in v.users):
1218 found.append(v)
1219 return found
1221 # FWD-phase only: values born in this stage's forward sweep and
1222 # consumed downstream. Bwd-phase values cannot ship forward (they
1223 # do not exist yet during the forward sweep); downstream-born
1224 # gradients flow back via grad crossings below.
1225 act_in_list = crossing(k - 1, k, _FWD) if k > 0 else []
1226 act_out_list = [] if is_last else crossing(k, k + 1, _FWD)
1227 # Backward crossings are phase-AGNOSTIC: besides plain gradients,
1228 # functionalized views (fwd nodes) can be consumed by an earlier
1229 # stage's backward (weight-grad compute) — those values flow back
1230 # during the backward sweep just the same.
1231 grad_in_list = [] if is_last else crossing(k + 1, k, None)
1232 grad_out_list = crossing(k, k - 1, None) if k > 0 else []
1234 # Forward-phase values the backward references. Stage-owned
1235 # placeholders (state / routed user inputs of this stage) may appear
1236 # and round-trip through the forward outputs; foreign placeholders
1237 # cannot be shipped and fail loudly during the slice copy.
1238 stage_owned_phs = {
1239 p
1240 for p in graph_module.graph.nodes
1241 if p.op == "placeholder" and node_stage.get(p) == k
1242 }
1243 saved = self._collect_saved_values(
1244 nodes,
1245 node_stage,
1246 node_phase,
1247 k,
1248 stage_owned_phs,
1249 set(act_in_list),
1250 grad_out_list,
1251 )
1252 return _Boundaries(
1253 act_in=act_in_list,
1254 act_out=act_out_list,
1255 grad_in=grad_in_list,
1256 grad_out=grad_out_list,
1257 saved=saved,
1258 )
1260 @staticmethod
1261 def _collect_saved_values(
1262 nodes: List[fx.Node],
1263 node_stage: Dict[fx.Node, int],
1264 node_phase: Dict[fx.Node, str],
1265 stage_idx: int,
1266 stage_owned_phs: Set[fx.Node],
1267 act_in_set: Set[fx.Node],
1268 grad_out_list: List[fx.Node],
1269 ) -> List[fx.Node]:
1270 """Forward-phase values this stage's backward references.
1272 Stage-owned placeholders (state / routed user inputs of this stage)
1273 may appear and round-trip through the forward outputs; foreign
1274 placeholders cannot be shipped and fail loudly during the slice
1275 copy. Fwd-phase values this stage must ship BACKWARD (consumed by
1276 the previous stage's backward) ride the fwd outputs as
1277 pass-throughs: computed once in the fwd subgraph, replayed into
1278 the bwd subgraph, and appended to its outputs.
1279 """
1280 saved: List[fx.Node] = []
1281 seen: Set[fx.Node] = set()
1282 for node in nodes:
1283 if node_phase[node] != _BWD or node_stage[node] != stage_idx:
1284 continue
1285 for arg in _iter_node_args(node):
1286 if not isinstance(arg, fx.Node) or arg in seen:
1287 continue
1288 # Fwd-phase values replay from the fwd outputs; bwd-phase
1289 # values are only shippable when they arrived as boundary
1290 # inputs (e.g. gradient-shape sym_size scalars).
1291 if node_phase.get(arg) != _FWD and arg not in act_in_set:
1292 continue
1293 if arg.op == "placeholder" and arg not in stage_owned_phs:
1294 continue
1295 seen.add(arg)
1296 saved.append(arg)
1297 for v in grad_out_list:
1298 if node_phase[v] == _FWD and v not in seen:
1299 seen.add(v)
1300 saved.append(v)
1301 return saved
1303 # ------------------------------------------------------------------
1304 # Subgraph construction
1305 # ------------------------------------------------------------------
1307 @staticmethod
1308 def _unique(name: str, used: Set[str]) -> str:
1309 """Return ``name`` (or a suffixed variant) not yet in ``used``."""
1310 candidate = name
1311 suffix = 1
1312 while candidate in used:
1313 candidate = f"{name}_{suffix}"
1314 suffix += 1
1315 used.add(candidate)
1316 return candidate
1318 def _copy_node(
1319 self,
1320 builder: _SubgraphBuilder,
1321 node: fx.Node,
1322 ) -> None:
1323 """Copy ``node`` into ``builder.g``, remapping args through the
1324 builder's env/foreign maps.
1326 ``env`` maps already-copied same-slice nodes; ``foreign`` maps
1327 boundary values to their placeholder stand-ins. The copy is
1328 registered under ``builder.env[node]``; nothing is returned —
1329 callers consume the copy through the builder's maps.
1330 """
1331 if node.op == "get_attr":
1332 raise ValueError(
1333 f"Unexpected get_attr node '{node.name}' — the tracer "
1334 f"contracts parameters/buffers to be static placeholders"
1335 )
1337 def map_fn(a: Any) -> Any:
1338 """Remap a single arg: env first, then foreign placeholders."""
1339 if isinstance(a, fx.Node):
1340 if a in builder.env:
1341 return builder.env[a]
1342 if a in builder.foreign:
1343 return builder.foreign[a]
1344 raise ValueError(
1345 f"Node '{node.name}' references '{a.name}' which lies "
1346 f"outside this stage's slice — boundary detection or "
1347 f"the stage plan is inconsistent"
1348 )
1349 return a
1351 new_args = fx.map_arg(node.args, map_fn)
1352 new_kwargs = fx.map_arg(node.kwargs, map_fn)
1353 new_node = builder.g.create_node(node.op, node.target, new_args, new_kwargs)
1354 new_node.meta = dict(node.meta)
1355 builder.env[node] = new_node
1357 def _build_fwd_graph(
1358 self,
1359 graph_module: fx.GraphModule,
1360 run_ctx: _RunContext,
1361 split: _StageSplit,
1362 boundaries: _Boundaries,
1363 ) -> Tuple[fx.GraphModule, List[fx.Node]]:
1364 """Build this stage's forward subgraph.
1366 Outputs: ``(*act_out, *saved)`` for non-last stages, ``(loss,
1367 *saved)`` for the last stage — where ``saved`` contains every
1368 forward-phase value the backward references that is not already
1369 crossing. The schedule ships the crossing prefix and replays the
1370 full output tuple into the backward subgraph.
1372 Returns:
1373 ``(fwd_gm, fwd_out_orig)`` — the subgraph module together with
1374 the ORIGINAL nodes at each output position, so the backward
1375 builder can key its placeholder mapping on them.
1376 """
1377 k = run_ctx.stage_idx
1378 is_last = k == run_ctx.pp_degree - 1
1379 builder = _SubgraphBuilder(g=fx.Graph(), used=set(), env={}, foreign={})
1381 self._create_fwd_placeholders(builder, run_ctx, split, boundaries)
1383 for node in graph_module.graph.nodes:
1384 if node.op in ("placeholder", "output"):
1385 continue
1386 if split.node_phase[node] != _FWD or split.node_stage[node] != k:
1387 continue
1388 self._copy_node(builder, node)
1390 loss_node = self._loss_node(graph_module)
1391 out_orig = self._resolve_fwd_outputs(
1392 boundaries.act_out, boundaries.saved, loss_node, is_last
1393 )
1394 out_nodes = [
1395 builder.env[n] if n in builder.env else builder.foreign[n] for n in out_orig
1396 ]
1397 builder.g.output(tuple(out_nodes))
1398 return fx.GraphModule(nn.Module(), builder.g), out_orig
1400 def _create_fwd_placeholders(
1401 self,
1402 builder: _SubgraphBuilder,
1403 run_ctx: _RunContext,
1404 split: _StageSplit,
1405 boundaries: _Boundaries,
1406 ) -> None:
1407 """Create the fwd subgraph's placeholders (state / boundary / inputs).
1409 Incoming boundary values (stage k > 0) become placeholders; each
1410 stage additionally takes the user inputs routed to it (stage 0
1411 typically the feature batch, the last stage the labels; see
1412 ``_route_user_inputs``).
1413 """
1414 for i, orig_idx in enumerate(split.stage_state_indices):
1415 ph = builder.g.placeholder(
1416 self._unique(_sanitize(split.stage_state_fqns[i]), builder.used)
1417 )
1418 ph.meta = dict(split.state_phs[orig_idx].meta)
1419 builder.env[split.state_phs[orig_idx]] = ph
1421 for i, act in enumerate(boundaries.act_in):
1422 ph = builder.g.placeholder(self._unique(f"pp_act_in_{i}", builder.used))
1423 ph.meta = dict(act.meta)
1424 builder.foreign[act] = ph
1426 for i, (ph, stage) in enumerate(
1427 zip(split.user_input_phs, split.user_input_stages)
1428 ):
1429 if stage != run_ctx.stage_idx:
1430 continue
1431 in_ph = builder.g.placeholder(self._unique(f"pp_input_{i}", builder.used))
1432 in_ph.meta = dict(ph.meta)
1433 builder.env[ph] = in_ph
1435 @staticmethod
1436 def _resolve_fwd_outputs(
1437 act_out_list: List[fx.Node],
1438 saved: List[fx.Node],
1439 loss_node: Optional[fx.Node],
1440 is_last: bool,
1441 ) -> List[fx.Node]:
1442 """Order the fwd subgraph's outputs: boundary values first, then saved."""
1443 if is_last:
1444 if loss_node is None:
1445 raise ValueError("Last PP stage requires a loss node in the output")
1446 out_orig: List[fx.Node] = [loss_node]
1447 else:
1448 if not act_out_list:
1449 raise ValueError(
1450 "Stage ships no boundary values but is not the last "
1451 "stage — the stage plan does not cut the graph"
1452 )
1453 out_orig = list(act_out_list)
1455 crossing_ids = {id(v) for v in act_out_list}
1456 for node in saved:
1457 if id(node) in crossing_ids or (is_last and node is loss_node):
1458 continue
1459 out_orig.append(node)
1460 return out_orig
1462 def _build_bwd_graph(
1463 self,
1464 graph_module: fx.GraphModule,
1465 run_ctx: _RunContext,
1466 split: _StageSplit,
1467 boundaries: _Boundaries,
1468 fwd_out_orig: List[fx.Node],
1469 trainable_indices: List[int],
1470 ) -> fx.GraphModule:
1471 """Build this stage's backward subgraph.
1473 Placeholders: ``(*state, *grad_in, *fwd_outputs)`` — the schedule
1474 passes the received boundary gradients and then the forward
1475 outputs verbatim (on the last stage a single ``ones_like(loss)``
1476 takes the gradient slots' place). Outputs: ``(*param_grads,)``
1477 plus the boundary gradients for every stage except stage 0.
1478 """
1479 k = run_ctx.stage_idx
1480 builder = _SubgraphBuilder(g=fx.Graph(), used=set(), env={}, foreign={})
1481 self._create_bwd_placeholders(
1482 builder, graph_module, run_ctx, split, boundaries, fwd_out_orig
1483 )
1485 stage_bwd_set = self._collect_stage_bwd_nodes(
1486 graph_module.graph, split.node_stage, split.node_phase, k
1487 )
1489 # Gradient-output nodes of this stage's trainable params (by state
1490 # index identity, not module attribution — post-trace inserts like
1491 # zeros_like carry no module stack). A LIST — output order is the
1492 # positional gradient contract; only membership uses the set.
1493 stage_grad_nodes = self._stage_grad_nodes(
1494 graph_module.graph, trainable_indices, set(split.stage_state_indices)
1495 )
1496 stage_grad_set = set(stage_grad_nodes)
1498 for node in graph_module.graph.nodes:
1499 if node.op in ("placeholder", "output"):
1500 continue
1501 if node not in stage_bwd_set and node not in stage_grad_set:
1502 continue
1503 self._copy_node(builder, node)
1505 out_nodes = [builder.env[gn] for gn in stage_grad_nodes]
1506 out_nodes.extend(builder.env[gn] for gn in boundaries.grad_out)
1507 builder.g.output(tuple(out_nodes))
1508 return fx.GraphModule(nn.Module(), builder.g)
1510 @staticmethod
1511 def _collect_stage_bwd_nodes(
1512 graph: fx.Graph,
1513 node_stage: Dict[fx.Node, int],
1514 node_phase: Dict[fx.Node, str],
1515 stage_idx: int,
1516 ) -> Set[fx.Node]:
1517 """Bwd-phase nodes owned by ``stage_idx``, as a set.
1519 Membership is probed once per graph node by ``_build_bwd_graph``
1520 and real joint graphs carry tens of thousands of nodes.
1521 """
1522 return {
1523 node
1524 for node in graph.nodes
1525 if node.op not in ("placeholder", "output")
1526 and node_phase[node] == _BWD
1527 and node_stage[node] == stage_idx
1528 }
1530 def _create_bwd_placeholders(
1531 self,
1532 builder: _SubgraphBuilder,
1533 graph_module: fx.GraphModule,
1534 run_ctx: _RunContext,
1535 split: _StageSplit,
1536 boundaries: _Boundaries,
1537 fwd_out_orig: List[fx.Node],
1538 ) -> None:
1539 """Create the bwd subgraph's placeholders (state / grads / fwd outs).
1541 The schedule passes the received boundary gradients and then the
1542 forward outputs verbatim (on the last stage a single
1543 ``ones_like(loss)`` takes the gradient slots' place). The
1544 ``foreign`` entries double as the remap for backward nodes
1545 referencing forward values.
1546 """
1547 for i, orig_idx in enumerate(split.stage_state_indices):
1548 ph = builder.g.placeholder(
1549 self._unique(_sanitize(split.stage_state_fqns[i]), builder.used)
1550 )
1551 ph.meta = dict(split.state_phs[orig_idx].meta)
1552 builder.env[split.state_phs[orig_idx]] = ph
1554 for i, grad in enumerate(boundaries.grad_in):
1555 ph = builder.g.placeholder(self._unique(f"pp_grad_in_{i}", builder.used))
1556 ph.meta = dict(grad.meta)
1557 builder.foreign[grad] = ph
1558 if not boundaries.grad_in and run_ctx.stage_idx == run_ctx.pp_degree - 1:
1559 # Last stage: the schedule feeds ones_like(loss) into the
1560 # single gradient slot.
1561 ph = builder.g.placeholder(self._unique("pp_grad_ones", builder.used))
1562 loss = self._loss_node(graph_module)
1563 if loss is not None:
1564 ph.meta = dict(loss.meta)
1566 # One placeholder per forward output, same order — the schedule
1567 # passes the forward outputs verbatim after the gradients.
1568 for orig in fwd_out_orig:
1569 ph = builder.g.placeholder(
1570 self._unique(f"pp_fwd_{orig.name}", builder.used)
1571 )
1572 ph.meta = dict(orig.meta)
1573 builder.foreign[orig] = ph
1575 @staticmethod
1576 def _stage_grad_nodes(
1577 graph: fx.Graph, trainable_indices: List[int], stage_state_set: Set[int]
1578 ) -> List[fx.Node]:
1579 """Gradient-output nodes of the stage's trainable params."""
1580 output = next((n for n in graph.nodes if n.op == "output"), None)
1581 returned = output.args[0] if output is not None else []
1582 stage_grad_nodes: List[fx.Node] = []
1583 for i, grad_node in enumerate(returned[1:]):
1584 if i >= len(trainable_indices):
1585 break
1586 if trainable_indices[i] in stage_state_set:
1587 stage_grad_nodes.append(grad_node)
1588 return stage_grad_nodes
1590 # ------------------------------------------------------------------
1591 # Live model pruning
1592 # ------------------------------------------------------------------
1594 def _prune_live_model(self, model: nn.Module, keep_fqns: List[str]) -> None:
1595 """Remove foreign-stage submodules from the live model in place.
1597 Mirrors torchtitan's ``_split_module`` (``None``-out unkept
1598 children, filter ModuleList/ModuleDict entries) while PRESERVING
1599 original index keys — reindexing a ModuleList would rename the
1600 surviving parameters' FQNs and break ``run_traced_graph``'s state
1601 validation against the sliced ``state_fqns``. Manual plans may cut
1602 inside a container element; such elements are recursed into
1603 (mixed subtree), exactly like plain modules.
1604 """
1605 keep = set(keep_fqns)
1607 def is_kept(fqn: str) -> bool:
1608 """True when ``fqn`` equals or descends from a stage module."""
1609 return any(fqn == kf or fqn.startswith(kf + ".") for kf in keep)
1611 def prune_container_child(child: nn.Module, cfqn: str) -> None:
1612 """Prune a ModuleList/ModuleDict child, preserving index keys.
1614 Reindexing would rename the surviving parameters' FQNs and break
1615 ``run_traced_graph``'s state validation. An element is kept
1616 whole, deleted, or recursed into when a manual plan cuts inside
1617 it (same rule as plain modules).
1618 """
1619 children = child._modules # pylint: disable=protected-access
1620 keys = list(children.keys())
1621 if all(is_kept(f"{cfqn}.{key}") for key in keys):
1622 return
1623 for key in keys:
1624 elem_fqn = f"{cfqn}.{key}"
1625 if is_kept(elem_fqn):
1626 continue
1627 elem = children[key]
1628 descendants = [f for f, _ in elem.named_modules() if f]
1629 if any(is_kept(f"{elem_fqn}.{f}") for f in descendants):
1630 prune(elem, elem_fqn + ".")
1631 else:
1632 del children[key]
1634 def prune_plain_child(
1635 mod: nn.Module, name: str, child: nn.Module, cfqn: str
1636 ) -> None:
1637 """Prune a non-container child: delete it or recurse when mixed."""
1638 descendants = [f for f, _ in child.named_modules() if f]
1639 if not descendants:
1640 # Leaf module (e.g. a bare Linear): keep only if declared.
1641 if not is_kept(cfqn):
1642 setattr(mod, name, None)
1643 return
1644 if not any(is_kept(f"{cfqn}.{f}") for f in descendants):
1645 setattr(mod, name, None)
1646 else:
1647 # Mixed subtree (manual plans may cut inside a child).
1648 prune(child, cfqn + ".")
1650 def prune(mod: nn.Module, prefix: str) -> None:
1651 """Recursively prune non-stage children under ``prefix``."""
1652 for name, child in list(mod.named_children()):
1653 cfqn = f"{prefix}{name}"
1654 if is_kept(cfqn):
1655 # Declared (or ancestor of a declared) stage module —
1656 # keep the whole subtree.
1657 continue
1658 if isinstance(child, (nn.ModuleList, nn.ModuleDict)):
1659 prune_container_child(child, cfqn)
1660 continue
1661 prune_plain_child(mod, name, child, cfqn)
1663 prune(model, "")
1665 # ------------------------------------------------------------------
1666 # Schedule + stub installation
1667 # ------------------------------------------------------------------
1669 def _trainable_state_indices(
1670 self,
1671 state_fqns: Sequence[str],
1672 state_is_param: Sequence[bool],
1673 model: Optional[nn.Module],
1674 ) -> List[int]:
1675 """Trainable-parameter state indices, mirroring the tracer's order."""
1676 param_lookup = (
1677 dict(model.named_parameters(remove_duplicate=False))
1678 if model is not None
1679 else {}
1680 )
1681 trainable: List[int] = []
1682 for idx, fqn in enumerate(state_fqns):
1683 if idx < len(state_is_param) and not state_is_param[idx]:
1684 continue
1685 param = param_lookup.get(fqn)
1686 if param is not None and param.requires_grad:
1687 trainable.append(idx)
1688 return trainable
1690 @staticmethod
1691 def _value_spec(
1692 nodes: Sequence[fx.Node], stage_idx: int, kind: str
1693 ) -> List[Tuple[Any, ...]]:
1694 """Build the schedule's recv descriptor for a boundary value list.
1696 Tensor values yield ``("tensor", shape, dtype, device)``; int/SymInt
1697 scalars (dynamic-shape ``sym_size`` nodes) yield ``("scalar",)``
1698 and are shipped as 0-d int64 tensors, unwrapped via ``item()``.
1699 Scalars carry no device of their own — the schedule anchors them to
1700 the first tensor entry's device, so a non-empty list must contain
1701 at least one tensor.
1702 """
1703 spec: List[Tuple[Any, ...]] = []
1704 for v in nodes:
1705 val = v.meta.get("val")
1706 if isinstance(val, torch.Tensor):
1707 # int(s): trace-time dims are SymInts (symbolic over the
1708 # sample input); the schedule allocates real buffers from
1709 # this spec, and ATen rejects symbolic sizes at runtime.
1710 # Training uses static shapes per compile, so the hinted
1711 # concrete values are stable.
1712 shape = tuple(int(s) for s in val.shape)
1713 spec.append(("tensor", shape, val.dtype, val.device))
1714 elif val is None:
1715 raise ValueError(
1716 f"Boundary {kind} value '{v.name}' on stage {stage_idx} "
1717 f"has no 'val' meta (traced without FakeTensor "
1718 f"metadata) — cannot size the P2P receive buffer"
1719 )
1720 else:
1721 spec.append(("scalar",))
1722 if nodes and not any(entry[0] == "tensor" for entry in spec):
1723 raise ValueError(
1724 f"Boundary {kind} list on stage {stage_idx} carries only "
1725 f"scalar values — no tensor entry to anchor the P2P "
1726 f"receive device on"
1727 )
1728 return spec
1730 def _build_schedule(
1731 self,
1732 fwd_gm: fx.GraphModule,
1733 bwd_gm: fx.GraphModule,
1734 pass_config: PassConfig,
1735 run_ctx: _RunContext,
1736 split: _StageSplit,
1737 num_trainable: int,
1738 boundaries: _Boundaries,
1739 ) -> PipelineScheduleBase:
1740 """Assemble this stage's pipeline schedule (``pass_config.pp_schedule``).
1742 Receive descriptors come from the boundary values' fake-tensor
1743 ``val`` metas; the send side packs by value type at runtime.
1744 """
1745 recv_spec = (
1746 self._value_spec(boundaries.act_in, run_ctx.stage_idx, "activation")
1747 if run_ctx.stage_idx > 0
1748 else []
1749 )
1750 grad_recv_spec = (
1751 self._value_spec(boundaries.grad_in, run_ctx.stage_idx, "gradient")
1752 if run_ctx.stage_idx < run_ctx.pp_degree - 1
1753 else []
1754 )
1755 schedule_cls = get_schedule_class(pass_config.pp_schedule)
1756 return schedule_cls(
1757 fwd_gm,
1758 bwd_gm,
1759 stage_idx=run_ctx.stage_idx,
1760 pp_degree=run_ctx.pp_degree,
1761 num_state=len(split.stage_state_fqns),
1762 num_trainable=num_trainable,
1763 num_send=len(boundaries.act_out),
1764 grad_send_count=len(boundaries.grad_out),
1765 microbatch_size=pass_config.pp_microbatch_size,
1766 pp_group=run_ctx.group,
1767 recv_spec=recv_spec,
1768 grad_recv_spec=grad_recv_spec,
1769 user_input_stages=split.user_input_stages,
1770 )
1772 def _install_stub(
1773 self,
1774 graph_module: fx.GraphModule,
1775 sched: PipelineScheduleBase,
1776 stage_state_fqns: List[str],
1777 state_is_param: Sequence[bool],
1778 stage_state_indices: List[int],
1779 num_user_inputs: int,
1780 ) -> None:
1781 """Rewrite the graph module into a schedule-dispatching stub.
1783 The schedule is added as a submodule and invoked through a
1784 ``call_module`` node, so it survives any later ``recompile()`` and
1785 the trainer's ``graph_module(*flat_inputs)`` dispatches to it
1786 unchanged. The stub takes ``(*state, *user_inputs)`` — every rank
1787 receives the full flattened user-input list; the schedule picks
1788 the inputs routed to this stage. ``state_fqns`` /
1789 ``state_is_param`` are updated IN PLACE (the JointGraph shares the
1790 list objects).
1791 """
1792 g = fx.Graph()
1793 used: Set[str] = set()
1794 state_phs = [
1795 g.placeholder(self._unique(_sanitize(fqn), used))
1796 for fqn in stage_state_fqns
1797 ]
1798 user_phs = [
1799 g.placeholder(self._unique(f"pp_input_{i}", used))
1800 for i in range(num_user_inputs)
1801 ]
1803 graph_module.add_module("pp_schedule", sched)
1804 call = g.call_module("pp_schedule", args=tuple(state_phs) + tuple(user_phs))
1805 g.output(call)
1807 graph_module.graph = g
1808 graph_module.state_fqns[:] = stage_state_fqns
1809 graph_module.state_is_param[:] = [
1810 state_is_param[i] for i in stage_state_indices
1811 ]
1812 graph_module.num_state_inputs = len(stage_state_fqns)
1813 graph_module.recompile()