[graph_trainer][aot_fx_trace][graph_pp] Add GraphPP runner integration

Add the GraphPP pipeline entry point, GraphPipelineStage wrapper, and GraphPPRunner runtime that executes prebuilt stage-local graph bundles for PyTorch runtime PP schedule actions. Graph construction uses eager PP metadata inference, GraphTrainer tracing, the metadata-preserving compile-time pass pipeline without terminal Inductor compilation, and the GraphPP partition/FSDP/dI-dW extraction passes.

Reuse GraphTrainer tracer subclass wrapping for values crossing PP boundaries, keep internal GraphPP calling-convention values flat, preserve trace-time buffer state, and share explicit parameter-gradient accumulation helpers with the non-PP path. OVERLAP_F_B is registered but intentionally unsupported until the multiplex/DualPipeV follow-up commit.

Add runner unit coverage for graph construction, forward/backward execution, FSDP unshard/reduce actions, metadata contracts, compile hooks, and supported runtime schedule validation. Add DeepSeek V3 GraphPP numerics and integration coverage for Interleaved1F1B and ZBVZeroBubble.

stack-info: PR: #3089, branch: sanketpurandare/stack/8

Author: sanketpurandare

torchtitan/experiments/graph_trainer/common_utils.py

@@ -6,8 +6,9 @@
 
 import fnmatch
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from contextlib import contextmanager
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -97,6 +98,41 @@ def _maybe_materialize_grad_for_param_layout(
 _NOT_IN_LAYERS = -1
 
 
+def compute_annotated_loss(
+    loss_fn: Callable[..., torch.Tensor],
+    pred: Any,
+    labels: Any,
+    loss_kwargs: dict[str, Any] | None = None,
+) -> torch.Tensor:
+    """Compute the loss tensor with the same FX metadata convention as GraphTrainer."""
+    annotated_loss_fn = annotate_fn({_MODULE_FQN: "loss"})(loss_fn)
+    result = annotated_loss_fn(pred, labels, **(loss_kwargs or {}))
+    if isinstance(result, tuple):
+        if len(result) != 2:
+            raise ValueError(
+                "GraphTrainer loss functions must return a loss tensor or "
+                "(loss tensor, metrics)."
+            )
+        loss, _metrics = result
+        return loss
+    return result
+
+
+def accumulate_param_grads_(
+    params: Iterable[torch.Tensor],
+    grads: Iterable[torch.Tensor | None],
+) -> None:
+    """Accumulate explicit graph-produced gradients into live parameters."""
+    for param, grad in zip(params, grads, strict=True):
+        if grad is None:
+            continue
+        grad = _maybe_materialize_grad_for_param_layout(param, grad)
+        if param.grad is None:
+            param.grad = grad
+        else:
+            param.grad += grad
+
+
 def _is_backward_node(node: torch.fx.Node) -> bool:
     return node.meta.get("autograd_backward", False)
 

torchtitan/experiments/graph_trainer/configs.py

@@ -251,6 +251,7 @@ def to_graph_trainer_config(
     d["model_spec"] = replace(
         base_config.model_spec,
         parallelize_fn=graph_spec.parallelize_fn,
+        pipelining_fn=graph_spec.pipelining_fn,
         model=graph_model,
     )
     d.pop("compile")

torchtitan/experiments/graph_trainer/deepseek_v3/__init__.py

@@ -7,7 +7,7 @@
 from dataclasses import fields
 
 from torchtitan.components.optimizer import register_moe_load_balancing_hook
-from torchtitan.distributed.pipeline_parallel import pipeline_llm
+from torchtitan.experiments.graph_trainer.graph_pp.pipeline import graph_pp_llm
 from torchtitan.models.deepseek_v3 import deepseekv3_configs
 from torchtitan.models.deepseek_v3.state_dict_adapter import DeepSeekV3StateDictAdapter
 from torchtitan.protocols.model_spec import ModelSpec
@@ -47,7 +47,7 @@ def model_registry(
         flavor=flavor,
         model=config,
         parallelize_fn=_parallelize_fn,
-        pipelining_fn=pipeline_llm,
+        pipelining_fn=graph_pp_llm,
         post_optimizer_build_fn=register_moe_load_balancing_hook,
         state_dict_adapter=DeepSeekV3StateDictAdapter,
     )

torchtitan/experiments/graph_trainer/deepseek_v3/config_registry.py

@@ -4,6 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from dataclasses import replace
+
+from torchtitan.distributed.pipeline_parallel import pipeline_llm
 from torchtitan.experiments.graph_trainer.configs import (
     GraphTrainerCompileConfig,
     to_graph_trainer_config,
@@ -26,6 +29,12 @@ def graph_trainer_deepseek_v3_debugmodel() -> GraphTrainer.Config:
     return config
 
 
+def graph_trainer_deepseek_v3_debugmodel_ep() -> GraphTrainer.Config:
+    config = to_graph_trainer_config(deepseek_v3_debugmodel(), model_registry)
+    config.compile = GraphTrainerCompileConfig(enable=True)
+    return config
+
+
 def graph_trainer_deepseek_v3_debugmodel_hybridep() -> GraphTrainer.Config:
     config = to_graph_trainer_config(deepseek_v3_debugmodel(), model_registry)
     config.compile = GraphTrainerCompileConfig(enable=True)
@@ -46,6 +55,26 @@ def graph_trainer_deepseek_v3_debugmodel_minimal_async_ep() -> GraphTrainer.Conf
     return config
 
 
+def graph_trainer_deepseek_v3_debugmodel_flex_attn() -> GraphTrainer.Config:
+    return graph_trainer_deepseek_v3_debugmodel()
+
+
+def graph_trainer_deepseek_v3_debugmodel_flex_attn_ep() -> GraphTrainer.Config:
+    return graph_trainer_deepseek_v3_debugmodel_ep()
+
+
+def graph_trainer_deepseek_v3_debugmodel_eager_pp() -> GraphTrainer.Config:
+    """Test-only FlexAttention baseline that runs through eager pipeline parallelism."""
+    config = graph_trainer_deepseek_v3_debugmodel()
+    config.compile = GraphTrainerCompileConfig(
+        enable=True,
+        components=["loss"],
+        mode=None,
+    )
+    config.model_spec = replace(config.model_spec, pipelining_fn=pipeline_llm)
+    return config
+
+
 def graph_trainer_deepseek_v3_16b() -> GraphTrainer.Config:
     config = to_graph_trainer_config(deepseek_v3_16b(), model_registry)
     config.compile = GraphTrainerCompileConfig(enable=True)