[graph_trainer][aot_fx_trace][graph_pp] Add multiplexed graph coverage and DualPipeV validation

Add GraphPP forward/backward multiplex graph construction and wire it into OVERLAP_F_B runtime execution. The multiplexed graph preserves the flat GraphPP calling convention, prebuilds and compiles before schedule execution, and keeps FakeTensor shape metadata on the destination ShapeEnv while preserving raw symbolic metadata relationships.

Add multiplex unit coverage for metadata transfer, output ordering, prebuild/compile behavior, full-Inductor overlap construction, and unsupported BACKWARD_INPUT overlap sub-actions.

Extend DeepSeek V3 GraphPP numerics and integration coverage for DualPipeV and document the GraphPP execution contract. The numerics comparison uses eager Interleaved1F1B as the baseline because torch.compile is incompatible with the eager ZBV and DualPipeV schedules required by FlexAttention. It disables local-rank gradient clipping so schedule-dependent clip coefficients do not mask graph execution equivalence, and compares rank-local TensorBoard losses from the ranks that actually own loss so core metrics remains unchanged.

stack-info: PR: #3090, branch: sanketpurandare/stack/9

Author: sanketpurandare

torchtitan/experiments/graph_trainer/README.md

@@ -60,6 +60,48 @@ NGPU=8 MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_8b ./run_train.sh
 NGPU=8 MODULE=graph_trainer.deepseek_v3 CONFIG=graph_trainer_deepseek_v3_16b ./run_train.sh --parallelism.data_parallel_shard_degree=4 --parallelism.tensor_parallel_degree=2 --parallelism.expert_parallel_degree=2
 ```
 
+### GraphPP Pipeline Parallelism
+
+GraphPP is the `aot_fx_trace` pipeline-parallel path for GraphTrainer models.
+It reuses TorchTitan's eager PP module splitting and PyTorch PP schedules, then
+traces one representative microbatch per local stage with GraphTrainer's
+`minimal_fx_tracer`. The resulting per-stage graph bundles are reused for later
+microbatches; `GraphPPRunner` only executes the prebuilt callable for each PP
+schedule action.
+
+```bash
+NGPU=8 MODULE=graph_trainer.deepseek_v3 CONFIG=graph_trainer_deepseek_v3_debugmodel ./run_train.sh \
+  --compile.mode aot_fx_trace \
+  --parallelism.pipeline_parallel_degree 2 \
+  --parallelism.pipeline_parallel_schedule Interleaved1F1B \
+  --parallelism.data_parallel_shard_degree 4 \
+  --parallelism.expert_parallel_degree 2
+```
+
+Supported runtime schedules include `Interleaved1F1B`, `ZBVZeroBubble`, and
+`DualPipeV`. GraphPP builds stage-local forward/backward graphs, optional FSDP
+`UNSHARD` / `REDUCE_GRAD` graphs, optional dI/dW graphs for split backward
+schedules, and multiplexed graphs for `OVERLAP_F_B` actions. Regional and full
+Inductor compilation reuse the existing GraphTrainer compilation passes on the
+extracted GraphPP callables; GraphPP-specific handling stays in the GraphPP
+stack before those passes are invoked. Multiplexed graphs keep the forward graph
+as the destination module and ShapeEnv, insert backward placeholders/compute
+before it, and transfer backward metadata into that ShapeEnv. This preserves
+forward collective-size provenance for full Inductor without changing the shared
+compiler passes.
+
+GraphPP follows the same subclass boundary as the non-PP GraphTrainer tracer.
+Extracted graphs run on flat plain tensor leaves. Values exposed to the PP
+runtime are rewrapped from tracer metadata: stage forward outputs, input
+gradients sent to previous stages, and parameter gradients before assignment to
+live `param.grad`. Internal values remain flat because they never leave GraphPP
+graph execution: saved-for-backward tensors, unsharded FSDP params, raw grad
+leaves, reduce-grad inputs, and multiplexed intermediate outputs.
+
+Current limitations: GraphPP does not load precompile artifacts yet, CUDA graph
+capture should target the `GraphPPRunner` steady-state path in a future change,
+and EP-overlap annotations will be composed with GraphPP in a later PR.
+
 ### Compiler Optimizations
 
 The `aot_fx_trace` mode has a built-in pass pipeline controlled by dedicated flags.

torchtitan/experiments/graph_trainer/graph_pp/__init__.py

@@ -4,11 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torchtitan.experiments.graph_trainer.graph_pp.split_fsdp_collectives import (
-    GraphPPFSDPBackwardSplit,
-    GraphPPFSDPForwardSplit,
-    split_backward_fsdp_collectives,
-    split_forward_fsdp_collectives,
+from torchtitan.experiments.graph_trainer.graph_pp.graph_multiplex import (
+    multiplex_fw_bw_graph,
 )
 from torchtitan.experiments.graph_trainer.graph_pp.partition import (
     GraphMeta,
@@ -18,12 +15,19 @@
     GraphPPDiDwSplit,
     split_di_dw_graph,
 )
+from torchtitan.experiments.graph_trainer.graph_pp.split_fsdp_collectives import (
+    GraphPPFSDPBackwardSplit,
+    GraphPPFSDPForwardSplit,
+    split_backward_fsdp_collectives,
+    split_forward_fsdp_collectives,
+)
 
 __all__ = [
     "GraphPPDiDwSplit",
     "GraphPPFSDPBackwardSplit",
     "GraphPPFSDPForwardSplit",
     "GraphMeta",
+    "multiplex_fw_bw_graph",
     "partition_joint_graph",
     "split_backward_fsdp_collectives",
     "split_di_dw_graph",

torchtitan/experiments/graph_trainer/graph_pp/graph_builder.py

@@ -14,7 +14,7 @@
    before assigning to live ``param.grad``.
 3. Internal graph values stay flat because they never escape GraphPP graph
    execution: saved-for-backward values, unsharded FSDP params, raw grad
-   leaves, and reduce-grad inputs.
+   leaves, reduce-grad inputs, and multiplexed intermediate outputs.
 4. DTensor and other traceable tensor subclasses use the existing tracer layout
    metadata. GraphPP must not add a separate DTensor-specific wrapping path.
 """
@@ -43,6 +43,9 @@
     split_backward_fsdp_collectives,
     split_forward_fsdp_collectives,
 )
+from torchtitan.experiments.graph_trainer.graph_pp.graph_multiplex import (
+    multiplex_fw_bw_graph,
+)
 from torchtitan.experiments.graph_trainer.graph_pp.partition import (
     GraphMeta as PartitionGraphMeta,
     partition_joint_graph,
@@ -1083,14 +1086,43 @@ def _build_graph_pp_overlap_graphs(
     *,
     compile_config: GraphTrainerCompileConfig,
 ) -> dict[tuple[int, int], GraphPPOverlapGraphs]:
-    del compile_config
-    required_pairs = _required_multiplex_pairs(schedule)
-    if required_pairs:
-        raise NotImplementedError(
-            "GraphPP OVERLAP_F_B requires multiplexed graph support, which is "
-            "introduced in the follow-up DualPipeV commit."
+    """Build multiplexed graphs required by ``OVERLAP_F_B`` schedule actions."""
+
+    stage_index_to_stage = {
+        stage.stage_index: cast(GraphPipelineStage, stage) for stage in schedule._stages
+    }
+    overlap_graphs: dict[tuple[int, int], GraphPPOverlapGraphs] = {}
+    for fw_stage_idx, bw_stage_idx in _required_multiplex_pairs(schedule):
+        pair = (fw_stage_idx, bw_stage_idx)
+        fw_stage = stage_index_to_stage[fw_stage_idx]
+        bw_stage = stage_index_to_stage[bw_stage_idx]
+        fw_graphs = _require_graph_trainer_stage_graphs(fw_stage, "OVERLAP_F_B")
+        bw_graphs = _require_graph_trainer_stage_graphs(bw_stage, "OVERLAP_F_B")
+        if fw_graphs.compiled or bw_graphs.compiled:
+            raise ValueError(
+                "GraphPP overlap graphs must be built before stage graphs are compiled."
+            )
+        multiplexed_graph = multiplex_fw_bw_graph(
+            fw_graphs.modules.fw,
+            bw_graphs.modules.full_bw,
+        )
+        _annotate_graph_pp_graph(
+            multiplexed_graph,
+            stage_index=fw_stage_idx,
+            callable_name="multiplex",
+            action_name="OVERLAP_F_B",
+        )
+        compiled_graph = _compile_graph_pp_module(
+            multiplexed_graph,
+            compile_config=compile_config,
+            graph_name=f"stage_{fw_stage_idx}_fw_stage_{bw_stage_idx}_bw_multiplex",
+        )
+        overlap_graphs[pair] = GraphTrainerOverlapGraphs(
+            fw_graphs=fw_graphs,
+            bw_graphs=bw_graphs,
+            multiplexed_graph=compiled_graph,
         )
-    return {}
+    return overlap_graphs
 
 
 def _example_args_from_stage_metadata(stage: GraphPipelineStage) -> tuple[Any, ...]:

torchtitan/experiments/graph_trainer/graph_pp/graph_multiplex.py

@@ -0,0 +1,162 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import copy
+from itertools import dropwhile
+
+import torch.fx as fx
+from torch.fx.graph_module import _assign_attr, _get_attr
+
+from torchtitan.experiments.graph_trainer.graph_pp.utils import (
+    _find_fake_mode,
+    _MetaShapeEnvTransfer,
+    trace_graph_pp_graph,
+)
+
+
+def _copy_prefixed_get_attrs(
+    dst: fx.GraphModule,
+    src: fx.GraphModule,
+    *,
+    prefix: str,
+) -> dict[str, str]:
+    remap: dict[str, str] = {}
+    for node in src.graph.find_nodes(op="get_attr"):
+        attr_name = str(node.target)
+        if attr_name in remap:
+            continue
+        base_name = f"{prefix}{attr_name.replace('.', '_')}"
+        new_attr_name = base_name
+        suffix = 0
+        # GraphModule attributes share one namespace; probe it directly to
+        # avoid colliding with existing get_attr targets.
+        while hasattr(dst, new_attr_name):
+            suffix += 1
+            new_attr_name = f"{base_name}_{suffix}"
+        _assign_attr(copy.deepcopy(_get_attr(src, attr_name)), dst, new_attr_name)
+        remap[attr_name] = new_attr_name
+    return remap
+
+
+def multiplex_fw_bw_graph(
+    fw_gm: fx.GraphModule,
+    bw_gm: fx.GraphModule,
+) -> fx.GraphModule:
+    """Concatenate backward and forward graphs into one boxed GraphPP callable.
+
+    Contract:
+      ``OVERLAP_F_B`` schedule actions need one callable that performs a
+      backward action for one stage and a forward action for another stage. The
+      returned graph preserves both flat calling conventions by ordering
+      placeholders as ``bw_inputs + fw_inputs`` and outputs as
+      ``bw_outputs + fw_outputs``.
+
+    Pseudocode:
+      deep-copy the forward graph as the destination
+      transfer backward metadata into the forward FakeTensorMode/ShapeEnv
+      insert copied backward placeholders before the forward placeholders
+      copy backward get_attr targets with a prefix to avoid attr collisions
+      copy backward compute nodes before the forward compute nodes
+      replace the output tuple with backward outputs followed by forward outputs
+
+    The forward graph remains the destination module because its ShapeEnv owns
+    the dynamic collective-size constraints needed by full Inductor for MoE
+    all-to-all outputs.  The backward graph is inserted in topological order
+    before the existing forward compute, with disjoint placeholders and
+    prefixed attributes, so this helper does not need dependency analysis or a
+    scheduling policy.  EP-overlap annotations are intentionally not applied
+    inside this graph; pass ordering keeps EP-overlap on the standalone no-FSDP
+    forward/backward graphs.
+
+    Args:
+        fw_gm (fx.GraphModule): Forward graph whose ShapeEnv and module
+            namespace become the destination for the multiplexed graph.
+        bw_gm (fx.GraphModule): Backward graph copied before the forward graph
+            inside the multiplexed callable.
+
+    Returns:
+        fx.GraphModule: One graph module with placeholders ordered as
+        ``bw_inputs + fw_inputs`` and outputs ordered as
+        ``bw_outputs + fw_outputs``.
+
+    Raises:
+        ValueError: If the forward graph does not contain the placeholders or
+            output node needed to build the multiplexed callable.
+    """
+    old_to_new: dict[fx.Node, fx.Node] = {}
+    # Preserve the forward ShapeEnv exactly as traced.  Reconstructing it from
+    # copied metadata can lose collective-size hints that full Inductor needs.
+    multiplexed_gm = copy.deepcopy(fw_gm)
+    dst_fake_mode = _find_fake_mode(multiplexed_gm)
+    meta_transfer = _MetaShapeEnvTransfer(dst_fake_mode)
+    for node in bw_gm.graph.nodes:
+        meta_transfer.collect(node.meta)
+    meta_transfer.seed()
+    bw_get_attr_remap = _copy_prefixed_get_attrs(
+        multiplexed_gm,
+        bw_gm,
+        prefix="bw",
+    )
+
+    fw_placeholders = multiplexed_gm.graph.find_nodes(op="placeholder")
+    if not fw_placeholders:
+        raise ValueError("GraphPP forward graph has no placeholders to multiplex")
+    first_fw_placeholder = fw_placeholders[0]
+    insert_point: fx.Node | None = None
+    for node in bw_gm.graph.find_nodes(op="placeholder"):
+        if insert_point is None:
+            with multiplexed_gm.graph.inserting_before(first_fw_placeholder):
+                new_placeholder = multiplexed_gm.graph.placeholder(f"bw_{node.name}")
+        else:
+            with multiplexed_gm.graph.inserting_after(insert_point):
+                new_placeholder = multiplexed_gm.graph.placeholder(f"bw_{node.name}")
+        new_placeholder.meta = meta_transfer.copy_meta(node.meta)
+        old_to_new[node] = new_placeholder
+        insert_point = new_placeholder
+
+    first_fw_compute = next(
+        (node for node in multiplexed_gm.graph.nodes if node.op != "placeholder"),
+        None,
+    )
+    if first_fw_compute is None:
+        raise ValueError("GraphPP forward graph has no output node to multiplex")
+
+    bw_nodes = iter(bw_gm.graph.nodes)
+    bw_nodes = dropwhile(lambda node: node.op == "placeholder", bw_nodes)
+    insert_point = None
+    for node in bw_nodes:
+        if node.op == "output":
+            break
+        if insert_point is None:
+            with multiplexed_gm.graph.inserting_before(first_fw_compute):
+                new_node = multiplexed_gm.graph.node_copy(
+                    node, lambda arg: old_to_new[arg]
+                )
+        else:
+            with multiplexed_gm.graph.inserting_after(insert_point):
+                new_node = multiplexed_gm.graph.node_copy(
+                    node, lambda arg: old_to_new[arg]
+                )
+        new_node.meta = meta_transfer.copy_meta(node.meta)
+        if new_node.op == "get_attr" and new_node.target in bw_get_attr_remap:
+            new_node.target = bw_get_attr_remap[str(new_node.target)]
+        old_to_new[node] = new_node
+        insert_point = new_node
+
+    multiplexed_output = multiplexed_gm.graph.find_nodes(op="output")[0]
+    bw_output_node = bw_gm.graph.find_nodes(op="output")[0]
+    bw_outputs = [
+        old_to_new[value] if isinstance(value, fx.Node) else value
+        for value in bw_output_node.args[0]
+    ]
+    fw_outputs = list(multiplexed_output.args[0])
+    multiplexed_output.args = (tuple(bw_outputs + fw_outputs),)
+
+    multiplexed_gm.graph.eliminate_dead_code()
+    multiplexed_gm.graph.lint()
+    multiplexed_gm.recompile()
+    trace_graph_pp_graph("graph_pp_multiplexed_graph", multiplexed_gm)
+    return multiplexed_gm