[graph_trainer] Overlap HybridEP and paging across combined PP graphs

.github/scripts/validate_release_gpu.sh

@@ -299,7 +299,7 @@ run_graph_trainer_h100_tests() {
     "${ARTIFACTS}/graph_trainer_h100" \
     --ngpu 8
   python -m pytest "${tests_root}/test_numerics.py::TestGraphTrainerNumerics" \
-    -v -k "moe or hybridep_eager"
+    -v -k "moe or hybridep_eager or hybridep_combined_pp or hybridep_dual_pipe_v"
   python -m pytest "${tests_root}/test_numerics.py::TestHybridEPContracts" -v
   python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests \
     "${ARTIFACTS}/graph_trainer_h100_precompile" \

.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml

@@ -111,6 +111,9 @@ jobs:
         CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_eager_overlap -v
         CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_eager_paging -v
         CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_eager_no_remat_paging -v
+        CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_combined_pp_overlap_and_paging -v
+        CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_dual_pipe_v_overlap_and_paging -v
+        CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics::test_hybridep_combined_pp_no_remat_paging -v
 
         # Run precompile integration tests (DSv3 with EP; Llama3 runs in default workflow)
         NCCL_NVLS_ENABLE=0 python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests $RUNNER_TEMP/artifacts-to-be-uploaded/precompile --ngpu 8 --test_name aot_fx_trace_deepseek_v3_precompile_fsdp_tp_ep

torchtitan/experiments/graph_trainer/README.md

@@ -105,8 +105,7 @@ 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, and
-EP-overlap annotations will be composed with GraphPP in a later PR.
+Current limitation: GraphPP does not load precompile artifacts yet.
 
 ### Compiler Optimizations
 
@@ -269,6 +268,58 @@ Current limitations:
 - The token count must be even.
 - Additional overlap or sharding modes require their own loss and trace checks.
 
+### Pipeline Communication Overlap
+
+Pipeline overlap uses `Interleaved1F1BCombined` or `DualPipeV` and graph
+multiplexing, not EP token chunking. `Interleaved1F1BCombined` adds one warmup
+forward, then pairs an independent
+forward and backward as an explicit `OVERLAP_F_B` action. Both actions' receives
+and parameter unshards precede the pair; normal send, gradient accumulation and
+reduction ownership remain in the existing pipeline schedule.
+
+`DualPipeV` uses the upstream V-shaped mapping with exactly two stages per
+rank. Its combined actions use full backward; separate input-gradient and
+weight-gradient actions fill the bubbles. Their saved activations remain live
+until weight backward completes. The focused HybridEP recipes use native
+varlen attention for both schedules, including actual dI/dW graph splitting.
+
+```python
+config.parallelism.pipeline_parallel_schedule = "Interleaved1F1BCombined"
+config.compile.enable_graph_pp_comm_compute_overlap = True
+config.compile.ep_overlap.enabled = False
+config.compile.inductor_compilation = "regional"
+config.training.disable_cuda_graphs = False
+```
+
+Use non-blocking `HybridEPTokenDispatcher` capacity. Enabling the GraphPP pass
+selects explicit asynchronous native launches/waits without splitting tokens
+or dividing each microbatch's capacity. Configure `sac_and_paged_stash` to
+include page operations. Source graphs retain explicit forward/backward
+ownership when multiplexed. The pass validates native request/input lifetimes,
+preserves tensor/mutation and paging queue dependencies, launches ready
+communication early and postpones waits behind independent compute.
+Each EP or paging launch gives ready work in the opposite source graph priority,
+including its ready completion waits, instead of draining one source first.
+Full CUDA capture covers the entire pipeline forward/backward step, including every
+microbatch's distinct native routing state.
+
+For long multistream captures, PyTorch's native allocator can reclaim storage
+after graph dependencies prove its asynchronous uses have joined:
+
+```bash
+PYTORCH_ALLOC_CONF=graph_capture_record_stream_reuse:True
+```
+
+This upstream option is experimental. Without it, cross-stream frees can wait
+until capture ends, growing the graph pool with the microbatch count. Keep the
+same allocator setting in matched comparisons; it does not replace request
+ownership, stream joins, or the explicit overlap annotations.
+
+The pass can be disabled with `compile.disable_passes =
+["graph_pp_comm_compute_overlap_pass"]` for a same-schedule comparison. Without
+combined actions, the requested overlap emits a warning. Full-graph Inductor
+and DeepEP are outside this feature's scope.
+
 ### Experimental AutoParallel Sharding
 
 GraphTrainer can use AutoParallel to solve SPMD placement for supported models,

torchtitan/experiments/graph_trainer/configs.py

@@ -283,6 +283,11 @@ class GraphTrainerCompileConfig:
     ep_overlap: EpOverlapConfig = field(default_factory=EpOverlapConfig)
     """Configuration for EP-overlap chunking and scheduling."""
 
+    enable_graph_pp_comm_compute_overlap: bool = False
+    """Overlap HybridEP and paged-stash operations across multiplexed PP F/B
+    graphs. Requires a combined-action schedule and non-blocking HybridEP.
+    This does not chunk tokens or divide per-microbatch capacity budgets."""
+
     precompile_artifact_dir: str = ""
     """
     Directory for precompiled artifacts. Setting this enables precompile:

torchtitan/experiments/graph_trainer/ep_overlap_pass.py

@@ -214,12 +214,14 @@ def _prioritized_order(
     launches: set[fx.Node],
     waits: set[fx.Node],
     regions: list[tuple[set[fx.Node], set[fx.Node]]],
+    yield_to: dict[fx.Node, set[fx.Node]] | None = None,
 ) -> list[fx.Node]:
     """Topologically schedule ordered scopes with early launches and late waits.
 
     Each region contains its eligible scope and the owned body whose completion
-    advances to the next region.
-
+    advances to the next region. A multiplexed pair uses one whole-graph scope.
+    Optional launch boundaries prefer complementary work among equally ranked
+    ready nodes, preventing one source from draining all its completion waits.
     Prefer compute that exposes another launch over unrelated compute. This
     lets input gradients issue their exchange before independent weight grads,
     while waits still follow every available computation. For eager paging,
@@ -284,6 +286,7 @@ def _prioritized_order(
     emitted = set()
     result = []
     active_region = 0
+    preferred = set()
     while ready:
         while active_region < len(regions) and regions[active_region][1] <= emitted:
             active_region += 1
@@ -306,6 +309,7 @@ def _prioritized_order(
                         if candidate in launch_ancestors
                         else 2
                     ),
+                    candidate not in preferred,
                     order[candidate],
                 ),
             )
@@ -314,6 +318,8 @@ def _prioritized_order(
         ready.remove(node)
         emitted.add(node)
         result.append(node)
+        if yield_to is not None and node in yield_to:
+            preferred = yield_to[node]
         for successor in successors[node]:
             remaining[successor] -= 1
             if remaining[successor] == 0:

torchtitan/experiments/graph_trainer/graph_builder.py

@@ -299,6 +299,11 @@ def prepare_graphs(
 
         required_overlap_pairs = stage_builder._required_multiplex_pairs(schedule)
         if not required_overlap_pairs:
+            if self.config.compile.enable_graph_pp_comm_compute_overlap:
+                logger.warning(
+                    "GraphPP communication overlap requested, but the schedule "
+                    "contains no combined forward/backward actions."
+                )
             self._overlap_graphs = {}
             overlap_graphs: dict[tuple[int, int], OverlapStageGraphs] = {}
         elif self._overlap_graphs is None:

torchtitan/experiments/graph_trainer/graph_pp/graph_multiplex.py

@@ -8,21 +8,97 @@
 from itertools import dropwhile
 
 import torch.fx as fx
+from torch._dynamo.graph_deduplication import _stable_topological_sort
 from torch.fx.graph_module import _assign_attr, _get_attr
 
+from torchtitan.experiments.graph_trainer.common_utils import _GRAPH_PP_DIRECTION
 from torchtitan.experiments.graph_trainer.debug_utils import tlparse_log_graph_pass
 from torchtitan.experiments.graph_trainer.graph_pp.utils import (
     _find_fake_mode,
     _MetaShapeEnvTransfer,
 )
 
 
+def graph_pp_comm_compute_overlap_pass(
+    gm: fx.GraphModule, example_inputs: tuple | None = None
+) -> fx.GraphModule:
+    """Overlap explicitly owned F/B communications in a multiplexed graph.
+
+    Native exchanges retain their exact inputs through completion. Tensor and
+    mutation edges, plus the paging queue order, define the legal schedule.
+    Both source graphs form one scope: launches first, ready compute next,
+    completion waits last. Each communication launch yields to ready work in
+    the opposite source graph, so neither source drains all its waits first.
+    No tokens or computation are split or duplicated.
+    """
+    del example_inputs
+    from torchtitan.experiments.graph_trainer.ep_overlap_pass import (
+        _collect_native_exchanges,
+        _graph_dependencies,
+        _prioritized_order,
+    )
+    from torchtitan.experiments.graph_trainer.paged_stash_memory_policy import (
+        _paging_dependencies,
+    )
+
+    exchanges = _collect_native_exchanges(gm.graph.nodes, context="GraphPP F/B pair")
+    deps = _graph_dependencies(gm)
+    launches, waits = _paging_dependencies(gm, deps)
+    source_nodes = {
+        direction: {
+            node
+            for node in gm.graph.nodes
+            if node.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION) == direction
+        }
+        for direction in ("forward", "backward")
+    }
+    yield_to = {}
+    for exchange in exchanges:
+        direction = exchange.launch.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION)
+        if direction not in ("forward", "backward") or direction != (
+            exchange.wait.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION)
+        ):
+            raise ValueError(
+                f"GraphPP exchange {exchange.launch.name} lost its source F/B ownership"
+            )
+        launches.add(exchange.launch)
+        waits.add(exchange.wait)
+    for launch in launches:
+        direction = launch.meta.get("custom", {}).get(_GRAPH_PP_DIRECTION)
+        if direction not in source_nodes:
+            raise ValueError(
+                f"GraphPP launch {launch.name} lost its source F/B ownership"
+            )
+        yield_to[launch] = source_nodes[
+            "backward" if direction == "forward" else "forward"
+        ]
+    if not launches and not waits:
+        return gm
+    nodes = set(gm.graph.nodes)
+    planned = _prioritized_order(
+        deps,
+        order={node: index for index, node in enumerate(gm.graph.nodes)},
+        launches=launches,
+        waits=waits,
+        regions=[(nodes, nodes)],
+        yield_to=yield_to,
+    )
+    sequencing = {node: {previous} for previous, node in zip(planned, planned[1:])}
+    _stable_topological_sort(gm.graph, sequencing)
+    gm.graph.lint()
+    gm.recompile()
+    if list(gm.graph.nodes) != planned:
+        raise ValueError("GraphPP communication schedule did not materialize its order")
+    return gm
+
+
 def _copy_prefixed_get_attrs(
     dst: fx.GraphModule,
     src: fx.GraphModule,
     *,
     prefix: str,
 ) -> dict[str, str]:
+    """Copy referenced attributes under unique names and return the remapping."""
     remap: dict[str, str] = {}
     for node in src.graph.find_nodes(op="get_attr"):
         attr_name = str(node.target)
@@ -63,13 +139,10 @@ def multiplex_fw_bw_graph(
       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.
+    the traced collective-size constraints. The backward graph is inserted
+    before forward compute, with disjoint placeholders and prefixed attributes.
+    Explicit source ownership allows the optional communication/compute pass
+    to reorder the pair; multiplexing itself preserves each source's order.
 
     Args:
         fw_gm (fx.GraphModule): Forward graph whose ShapeEnv and module
@@ -90,6 +163,11 @@ def multiplex_fw_bw_graph(
     # 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)
+    for node in multiplexed_gm.graph.nodes:
+        node.meta["custom"] = {
+            **node.meta.get("custom", {}),
+            _GRAPH_PP_DIRECTION: "forward",
+        }
     dst_fake_mode = _find_fake_mode(multiplexed_gm)
     meta_transfer = _MetaShapeEnvTransfer(dst_fake_mode)
     for node in bw_gm.graph.nodes:
@@ -141,6 +219,10 @@ def multiplex_fw_bw_graph(
                     node, lambda arg: old_to_new[arg]
                 )
         new_node.meta = meta_transfer.copy_meta(node.meta)
+        new_node.meta["custom"] = {
+            **new_node.meta.get("custom", {}),
+            _GRAPH_PP_DIRECTION: "backward",
+        }
         if new_node.op == "get_attr":
             attr_name = str(new_node.target)
             if attr_name not in bw_get_attr_remap:

torchtitan/experiments/graph_trainer/graph_pp/pipeline.py

@@ -56,6 +56,8 @@
 from torchtitan.experiments.graph_trainer.graph_pp.stage import GraphPipelineStage
 from torchtitan.experiments.graph_trainer.registry import PASS_PIPELINE_REGISTRY
 from torchtitan.models.base.model import BaseModel
+from torchtitan.models.common.moe import MoE
+from torchtitan.models.common.moe.token_dispatcher import HybridEPTokenDispatcher
 
 
 if TYPE_CHECKING:
@@ -196,6 +198,19 @@ def _validate_graph_pp_config(
     compile_config: GraphTrainerCompileConfig,
     parallelism: ParallelismConfig,
 ) -> None:
+    if compile_config.enable_graph_pp_comm_compute_overlap:
+        if compile_config.ep_overlap.enabled:
+            raise ValueError(
+                "GraphPP communication overlap does not use EP token chunking"
+            )
+        if compile_config.inductor_compilation != "regional":
+            raise ValueError(
+                "GraphPP communication overlap requires regional compilation"
+            )
+        if parallelism.expert_parallel_degree <= 1:
+            raise ValueError(
+                "GraphPP HybridEP overlap requires expert_parallel_degree > 1"
+            )
     if compile_config.precompile_artifact_dir:
         raise ValueError(
             "GraphPP does not support compile.precompile_artifact_dir yet. "
@@ -261,6 +276,8 @@ def resolve_graph_execution_plan(
         pp_enabled: Whether pipeline parallelism is enabled.
         fsdp_enabled: Whether FSDP is enabled.
     """
+    if compile_config.enable_graph_pp_comm_compute_overlap and not pp_enabled:
+        raise ValueError("GraphPP communication overlap requires pipeline parallelism")
     if pp_enabled:
         _validate_graph_pp_config(
             compile_config=compile_config,
@@ -850,6 +867,18 @@ def graph_pipeline_llm(
             activation_remat=activation_remat,
             dump_folder=dump_folder,
         )
+        if compile_config.enable_graph_pp_comm_compute_overlap:
+            for module in model_part.modules():
+                if not isinstance(module, MoE):
+                    continue
+                dispatcher = module.routed_experts.token_dispatcher
+                if not isinstance(dispatcher, HybridEPTokenDispatcher):
+                    raise ValueError("GraphPP communication overlap requires HybridEP")
+                if dispatcher.non_blocking_capacity_factor is None:
+                    raise ValueError(
+                        "GraphPP overlap requires non-blocking HybridEP capacity"
+                    )
+                dispatcher.async_op = True
         logger.info(
             "PP rank %s is building GraphPP stage_idx %s with modules %s",
             pp_mesh.get_local_rank(),

torchtitan/experiments/graph_trainer/graph_pp/stage_builder.py

@@ -19,6 +19,7 @@
 
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
 from torchtitan.experiments.graph_trainer.graph_pp.graph_multiplex import (
+    graph_pp_comm_compute_overlap_pass,
     multiplex_fw_bw_graph,
 )
 from torchtitan.experiments.graph_trainer.graph_pp.runner import (
@@ -35,8 +36,10 @@
     SplitStageGraphs,
 )
 from torchtitan.experiments.graph_trainer.graph_pp.utils import (
+    example_inputs_from_placeholders,
     overlap_fw_bw_sub_actions,
 )
+from torchtitan.experiments.graph_trainer.passes import apply_graph_passes
 
 
 class _StageGraphModules(Protocol):
@@ -246,6 +249,13 @@ def _build_graph_pp_overlap_graphs(
             fw_graphs.modules.fw,
             bw_graphs.modules.full_bw,
         )
+        if compile_config.enable_graph_pp_comm_compute_overlap:
+            multiplexed_graph = apply_graph_passes(
+                multiplexed_graph,
+                example_inputs_from_placeholders(multiplexed_graph),
+                [graph_pp_comm_compute_overlap_pass],
+                compile_config=compile_config,
+            )
         compiled_graph = compile_graph_module(
             multiplexed_graph,
             compile_config=compile_config,