[graph_trainer] Add CooR precompile support for aot_fx_trace compile mode (#2975)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.14.0)
(oldest at bottom):
* __->__ #2975

Introduce aot_fx_trace as an alternate precompilation path alongside
the existing AOT joint graph export. aot_fx_trace uses make_fx to
trace the full fwd+loss+bwd step as a single FX graph, then serializes
the GraphModule for loading on all ranks during training. Because
cudagraphs is a default pass, we also include cudagraph composability
related changes in this PR to get to bitwise equivalence between
aot_fx_trace with and without precompile.

Changes:
- precompile_main.py: Refactor into _common_setup() + _precompile_aot()
  + _precompile_aot_fx_trace(), route by --compile.mode. Use Python
  float for global_valid_tokens (matching dist_sum's return type) so
  make_fx bakes it as a graph constant, identical to the runtime path.
- precompile.py: PrecompiledFxTraceArtifact dataclass, save/load fns.
  Uses GraphPickler (not plain pickle) to preserve SymInt expressions
  (_runtime_compute_coordinate_on_dim) through serialization — plain
  pickle evaluates SymInts to concrete trace-time values, baking in
  rank-specific embedding vocab offsets.
- compile.py: Mode-aware validation (aot_fx_trace doesn't need
  serializable passes), logging for precompile load vs runtime trace
- trainer.py: _load_precompiled_fx_trace(), branching in
  _make_fx_forward_backward_step()
- make_fx_tracer.py: Use set_multithreading_enabled(False) so backward
  tracing runs on the calling thread, keeping the compile_on_one_rank
  ContextVar visible
- cudagraph.py: Defer _input_indices_to_copy when example_inputs is
  empty (precompile load path), accept opaque values (DeviceMesh) in
  input validation, skip non-tensors in static address checks.
  Squashed the separate cudagraph CooR PR into this one because
  cudagraph activates automatically for aot_fx_trace graphs (all
  inputs are desugared plain tensors + opaques, no HOPs), so bitwise
  equivalence testing required both fixes together.
- passes.py: Guard cudagraph_pass against non-GraphModule inputs

## Test plan
- `pytest
torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py
-x` — 6 passed, 4 skipped (Llama3, DSv3 × {eager self-deterministic,
eager vs aot_fx_trace, precompile vs trace})
- `pytest torchtitan/experiments/graph_trainer/tests/test_precompile.py
-x` — all pass
- `pytest
torchtitan/experiments/graph_trainer/tests/test_trace_module.py -x` —
all pass
- 8-GPU integration tests via CI (ciflow/8gpu)

Author: bobrenjc93

torchtitan/experiments/graph_trainer/compile.py

@@ -246,15 +246,18 @@ def apply_compile(
         parallelism.fsdp_reshard_after_forward, parallel_dims.pp_enabled
     )
 
-    if compile_config.precompile_artifact_dir and mode != "aot":
+    if compile_config.precompile_artifact_dir and mode not in ("aot", "aot_fx_trace"):
         logger.warning(
             "--compile.precompile_artifact_dir is only supported with "
-            f"--compile.mode=aot, but mode is '{mode}'. Ignoring precompile."
+            f"--compile.mode=aot or aot_fx_trace, but mode is '{mode}'. "
+            "Ignoring precompile."
         )
         compile_config = dataclasses.replace(compile_config, precompile_artifact_dir="")
 
-    if compile_config.precompile_artifact_dir and not (
-        _SERIALIZABLE_PASSES & set(compile_config.passes)
+    if (
+        compile_config.precompile_artifact_dir
+        and mode == "aot"
+        and not (_SERIALIZABLE_PASSES & set(compile_config.passes))
     ):
         raise ValueError(
             "--compile.precompile_artifact_dir requires at least one pass that "
@@ -281,10 +284,18 @@ def apply_compile(
         )
     elif mode == "aot_fx_trace":
         # aot_fx_trace traces fwd+loss+bwd together inside forward_backward_step,
-        # so no model-level wrapping is needed here.
-        logger.info(
-            "aot_fx_trace compile mode: graph capture will happen at training time"
-        )
+        # so no model-level wrapping is needed here. If precompile_artifact_dir
+        # is set, the precompiled artifact will be loaded lazily in
+        # GraphTrainer._make_fx_forward_backward_step.
+        if compile_config.precompile_artifact_dir:
+            logger.info(
+                "aot_fx_trace compile mode: precompiled artifact will be loaded "
+                f"from {compile_config.precompile_artifact_dir}"
+            )
+        else:
+            logger.info(
+                "aot_fx_trace compile mode: graph capture will happen at training time"
+            )
         return model
     else:
         raise ValueError(

torchtitan/experiments/graph_trainer/cudagraph.py

@@ -124,6 +124,7 @@ def __init__(
         example_inputs: Sequence[Any],
         static_input_indices: tuple[int] | None = None,
         should_check_address: bool = False,
+        tensor_input_indices: list[int] | None = None,
     ):
         _cg_manager.maybe_initialize()
         _cg_manager.register_wrapper(self)
@@ -132,11 +133,16 @@ def __init__(
         self._static_input_indices = OrderedSet(
             static_input_indices if static_input_indices is not None else []
         )
-        self._input_indices_to_copy = [
-            i
-            for i, inp in enumerate(example_inputs)
-            if isinstance(inp, torch.Tensor) and i not in self._static_input_indices
-        ]
+        if tensor_input_indices is not None:
+            self._input_indices_to_copy = [
+                i for i in tensor_input_indices if i not in self._static_input_indices
+            ]
+        else:
+            self._input_indices_to_copy = [
+                i
+                for i, inp in enumerate(example_inputs)
+                if isinstance(inp, torch.Tensor) and i not in self._static_input_indices
+            ]
         self._cudagraph: torch.cuda.CUDAGraph | None = None
         self._has_warmup = False
 
@@ -157,19 +163,26 @@ def _copy_non_static_inputs(self, *args):
         for i in self._input_indices_to_copy:
             self._args[i].copy_(args[i])
 
-    def _check_input_types(self, inputs) -> None:
+    def _validate_inputs(self, inputs) -> None:
+        """Validate that all inputs are of supported types.
+
+        Opaque inputs (e.g. DeviceMesh from SimpleFSDP/DTensor) are
+        inherently static and already excluded from copying (only
+        tensors appear in ``_input_indices_to_copy``), so no special
+        handling is needed beyond accepting them here.
+        """
         for i, inp in enumerate(inputs):
-            if not (
-                isinstance(inp, (torch.Tensor, int, float, torch._C.Generator))
-                or is_opaque_value(inp)
-            ):
-                raise ValueError(
-                    "args must be tensor, integer (for dynamic shapes), "
-                    "float (for scalar constants), "
-                    "Generator (for random number generator), "
-                    "or opaque object, "
-                    f"but found {type(inp)} with value {inp!r} at index {i}"
-                )
+            if isinstance(inp, (torch.Tensor, int, float, torch._C.Generator)):
+                continue
+            if is_opaque_value(inp):
+                continue
+            raise ValueError(
+                "args must be tensor, integer (for dynamic shapes), "
+                "float (for scalar constants), "
+                "Generator (for random number generator), "
+                "or opaque object, "
+                f"but found {type(inp)} with value {inp!r} at index {i}"
+            )
 
     def _check_static_inputs_address(self) -> None:
         for i in self._static_input_indices:
@@ -194,7 +207,7 @@ def __call__(self, *args):
             return out
 
         if self._cudagraph is None:
-            self._check_input_types(args)
+            self._validate_inputs(args)
             self._args = args
             self._input_addresses = [
                 x.data_ptr() if isinstance(x, torch.Tensor) else None for x in args

torchtitan/experiments/graph_trainer/make_fx_tracer.py

@@ -6,7 +6,7 @@
 
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import torch
@@ -307,6 +307,7 @@ class TracedResult:
     num_flat_outputs: int
     output_subclass_layouts: dict[int, SubclassLayout]
     output_spec: pytree.TreeSpec
+    tensor_input_indices: list[int] = field(default_factory=list)
 
     @property
     def num_static_inputs(self) -> int:
@@ -412,7 +413,23 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
 
         ctx = TracingContext(fake_mode)
         # preserve_node_meta propagates fx.traceback.annotate metadata to traced nodes
-        with fake_mode, tracing(ctx), preserve_node_meta(), _skip_nested_compile():
+
+        # Disable autograd multithreading so that backward tracing
+        # runs on the calling thread.  Without this, the C++ autograd
+        # engine dispatches backward to a worker thread that has a
+        # fresh contextvars.Context, making the compile_on_one_rank
+        # ContextVar invisible and causing _sym_get_coordinate to
+        # bake rank 0's concrete coordinates into the backward graph.
+        # TODO: Move set_multithreading_enabled(False) to global init.
+        # Forcing backward onto the main CPU thread is a good default
+        # for both tracing and runtime, not just the tracing path.
+        with (
+            fake_mode,
+            tracing(ctx),
+            preserve_node_meta(),
+            _skip_nested_compile(),
+            torch.autograd.set_multithreading_enabled(False),
+        ):
             traced = make_fx(
                 fn_with_subclass_handling,
                 record_stack_traces=True,
@@ -435,6 +452,9 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
             num_flat_outputs=num_flat_outputs,
             output_subclass_layouts=output_layouts,
             output_spec=output_spec,
+            tensor_input_indices=[
+                i for i, x in enumerate(fake_args) if isinstance(x, torch.Tensor)
+            ],
         )
 
     return _trace_with_args

torchtitan/experiments/graph_trainer/passes.py

@@ -50,24 +50,20 @@
 from torchtitan.tools.logging import logger
 
 
-def construct_default_graph_passes(
+def compile_time_passes(
     traced_result: "TracedResult",
 ) -> list[Callable]:
-    """Build the default pass list for the aot_fx_trace compile path.
-
-    Per-pass configuration (e.g. ``static_input_indices`` for cudagraph) is
-    bound here via ``functools.partial`` so that ``apply_graph_passes``
-    stays a generic pass runner with no pass-specific parameters.
+    """Cleanup, FlexAttention annotation, and regional_inductor passes.
 
-    Args:
-        traced_result: The traced graph and metadata from ``trace_train_step``.
+    If precompile is enabled, these are applied before serialization so
+    that compiled Triton kernels are baked into the artifact. Otherwise
+    they run at trace time via ``construct_default_graph_passes``.
 
-    Returns:
-        An ordered list of graph passes ready to apply.
+    cudagraph is excluded — it needs real tensors and devices at runtime.
     """
     from torchtitan.models.common.attention import FlexAttention
 
-    passes: list[Callable] = [
+    return [
         functools.partial(tlparse_log_graph_pass, graph_name="make_fx_graph_traced"),
         remove_detach_pass,
         remove_identity_view_pass,
@@ -93,19 +89,37 @@ def construct_default_graph_passes(
         custom_codegen_pass,
     ]
 
-    # cudagraph should be the last pass.
+
+def construct_default_graph_passes(
+    traced_result: "TracedResult",
+    *,
+    precompiled: bool = False,
+) -> list[Callable]:
+    """Build the pass list for the aot_fx_trace path.
+
+    When ``precompiled=False`` (default), returns the full list: cleanup,
+    FlexAttention annotation, regional_inductor, and cudagraph.
+
+    When ``precompiled=True``, the artifact already has cleanup and
+    regional_inductor baked in, so only cudagraph is returned.
+    """
     from torchtitan.experiments.graph_trainer.cudagraph import is_cudagraph_compatible
 
+    passes: list[Callable] = []
+    if not precompiled:
+        passes.extend(compile_time_passes(traced_result))
+
+    # cudagraph should be the last pass.
     if is_cudagraph_compatible(traced_result.gm):
         static_input_indices = list(range(traced_result.num_static_inputs))
         passes.append(
             functools.partial(
                 cudagraph_pass,
                 is_forward=True,
                 static_input_indices=static_input_indices,
+                tensor_input_indices=traced_result.tensor_input_indices,
             )
         )
-
     return passes
 
 
@@ -277,6 +291,7 @@ def cudagraph_pass(
     *,
     is_forward: bool,
     static_input_indices: list[int] | None = None,
+    tensor_input_indices: list[int] | None = None,
 ) -> torch.fx.GraphModule:
     """
     Apply cudagraph.
@@ -295,7 +310,18 @@ def cudagraph_pass(
             when ``static_input_indices`` is not provided.
         static_input_indices: Explicit list of input indices with stable tensor
             addresses. When provided, ``is_forward`` is not used for inference.
+        tensor_input_indices: Indices of graph inputs that are tensors (as
+            opposed to opaque values like DeviceMesh). Used to compute which
+            inputs need copying for cudagraph replay. When not provided, this
+            is inferred from ``example_inputs``.
     """
+    if not isinstance(gm, torch.fx.GraphModule):
+        raise TypeError(
+            f"cudagraph_pass requires a GraphModule but got {type(gm).__name__}. "
+            f"Ensure cudagraph is not combined with passes that replace the "
+            f"GraphModule (e.g. full_inductor_compilation)."
+        )
+
     # Lazy import: cudagraph.py runs init_global_graph_pool() at import time,
     # which must happen after torch.cuda.set_device(local_rank).
     from torchtitan.experiments.graph_trainer.cudagraph import (
@@ -305,7 +331,12 @@ def cudagraph_pass(
 
     if static_input_indices is None:
         static_input_indices = get_static_input_indices(gm, is_forward)
-    gm.forward = CUDAGraphWrapper(gm.forward, example_inputs, static_input_indices)
+    gm.forward = CUDAGraphWrapper(
+        gm.forward,
+        example_inputs,
+        static_input_indices,
+        tensor_input_indices=tensor_input_indices,
+    )
     return gm