[graph_trainer] Add torch.no_grad() and graph-based SAC to traced execution

Execute traced graph under torch.no_grad() since it already contains
explicit backward ops (from torch.autograd.grad traced by make_fx).
Without this, PyTorch builds a redundant autograd graph keeping all
forward intermediates alive via grad_fn references.

Replace monkey-patched _CachingTorchDispatchMode AC approach with clean
graph-based SAC: annotate_ac_regions before tracing, apply_sac_pass
(which now skips backward-tagged nodes) for post-hoc tagging, then
remat_using_tags_for_fwd_loss_bwd_graph for the remat transform.

apply_ac_remat_pass now takes GraphModule and returns GraphModule,
following the standard pass signature convention.

Results on Llama 1B (H100): traced graph SAC uses 18.23 GB vs eager SAC
19.60 GB (0.93x ratio), with bitwise identical losses and gradients.

ghstack-source-id: 157213f
Pull Request resolved: #2766

Author: tugsbayasgalan

.github/workflows/integration_test_8gpu_graph_trainer.yaml

@@ -78,7 +78,8 @@ jobs:
         # Run precompile unit tests
         pytest torchtitan/experiments/graph_trainer/tests/test_precompile.py -v
 
-        # Run bitwise deterministic guardrail test
+        # Run bitwise deterministic and SAC peak-memory guardrail tests
         pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v
+        pytest torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py -v
 
         rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/*/checkpoint

.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml

@@ -68,7 +68,8 @@ jobs:
         # Run the MoE numerics tests
         NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics -v -k "moe"
 
-        # Run bitwise deterministic guardrail test (includes H100-only hardcoded-hash tests)
+        # Run bitwise deterministic and SAC peak-memory guardrail tests
         pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v
+        pytest torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py -v
 
         rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/*/checkpoint

tests/integration_tests/run_tests.py

@@ -91,9 +91,13 @@ def run_single_test(
             )
 
 
-def run_tests(args, test_list: list[OverrideDefinitions], module=None, config=None):
+def run_tests(
+    args,
+    test_list: list[OverrideDefinitions],
+    module=None,
+    config=None,
+):
     """Run all integration tests to test the core features of TorchTitan"""
-
     exclude_set = set()
     if hasattr(args, "exclude") and args.exclude:
         exclude_set = {name.strip() for name in args.exclude.split(",")}
@@ -123,7 +127,12 @@ def run_tests(args, test_list: list[OverrideDefinitions], module=None, config=No
             )
         else:
             try:
-                run_single_test(test_flavor, args.output_dir, module, config)
+                run_single_test(
+                    test_flavor,
+                    args.output_dir,
+                    module,
+                    config,
+                )
             except Exception as e:
                 logger.error(str(e))
                 failed_tests.append((test_flavor.test_name, str(e)))

torchtitan/experiments/graph_trainer/make_fx_tracer.py

@@ -12,9 +12,6 @@
 import torch
 import torch.nn as nn
 import torch.utils._pytree as pytree
-from torch._functorch._aot_autograd.logging_utils import (
-    setup_stacktrace_preservation_hooks,
-)
 from torch._guards import tracing, TracingContext
 from torch._subclasses import FakeTensorMode
 from torch.fx.experimental.proxy_tensor import make_fx
@@ -200,52 +197,6 @@ def _remove_cpu_shadow_chains(gm: torch.fx.GraphModule) -> None:
     gm.recompile()
 
 
-@contextmanager
-def _patch_engine_run_backward() -> Generator[None, None, None]:
-    """Patch _engine_run_backward to install stacktrace preservation hooks.
-
-    Why this is needed:
-    When make_fx traces a function that calls loss.backward(), the backward
-    pass is decomposed into primitive ATen ops. Normally (in eager autograd),
-    ``setup_stacktrace_preservation_hooks`` is called by the autograd engine
-    to propagate ``seq_nr`` from forward ops to their corresponding backward
-    ops. Under make_fx tracing, this hook setup doesn't happen automatically
-    because the engine path differs, so backward FX nodes end up without
-    ``seq_nr`` metadata. Without ``seq_nr``, we can't correlate backward
-    nodes back to their forward counterparts (needed by
-    ``_copy_fwd_metadata_to_bw_nodes``).
-
-    This context manager patches ``_engine_run_backward`` to call
-    ``setup_stacktrace_preservation_hooks`` before the autograd engine runs,
-    restoring ``seq_nr`` propagation during tracing.
-
-    We must patch the name in both modules since ``torch.autograd.__init__``
-    imports it via ``from .graph import``.
-    """
-    import torch.autograd
-    import torch.autograd.graph
-
-    _orig_fn = torch.autograd.graph._engine_run_backward
-
-    def _patched(t_outputs, *args, **kwargs):  # type: ignore[no-untyped-def]
-        roots = [
-            t.grad_fn
-            for t in t_outputs
-            if isinstance(t, torch.Tensor) and t.grad_fn is not None
-        ]
-        if roots:
-            setup_stacktrace_preservation_hooks(roots)
-        return _orig_fn(t_outputs, *args, **kwargs)
-
-    torch.autograd.graph._engine_run_backward = _patched  # type: ignore[assignment]
-    torch.autograd._engine_run_backward = _patched  # type: ignore[assignment]
-    try:
-        yield
-    finally:
-        torch.autograd.graph._engine_run_backward = _orig_fn  # type: ignore[assignment]
-        torch.autograd._engine_run_backward = _orig_fn  # type: ignore[assignment]
-
-
 def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
     """Copy forward node metadata (custom) to later nodes sharing the same seq_nr.
 
@@ -401,18 +352,27 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
 
             state_for_fn = dict(zip(state_fqns, state_wrapped, strict=True))
             user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
-
-            with _patch_engine_run_backward():
+            with torch.compiler._patch_autograd_grad():
                 result = fn(state_for_fn, *user_list)
-
             flat_outs, output_spec = pytree.tree_flatten(result)
             num_flat_outputs = len(flat_outs)
             unwrapped_outs, output_layouts = _unwrap_subclasses(flat_outs)
             return unwrapped_outs
 
         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():
+        # _skip_nested_compile lets this tracer run when an outer dynamo trace
+        # reaches torch.compile'd FlexAttention kernels.
+        # _non_strict_tracing_context is required by _patch_autograd_grad() and
+        # marks this make_fx pass as the non-strict tracing flow, distinct from
+        # other make_fx-based entry points such as non-strict export.
+        with (
+            fake_mode,
+            tracing(ctx),
+            preserve_node_meta(),
+            _skip_nested_compile(),
+            torch.compiler._non_strict_tracing_context(),
+        ):
             traced = make_fx(
                 fn_with_subclass_handling,
                 record_stack_traces=True,
@@ -450,6 +410,10 @@ def run_traced(
     This is a reference implementation of traced-graph execution. It keeps the
     state handling, subclass unwrapping, and output reconstruction logic
     explicit instead of baking those semantics into ``TracedResult`` itself.
+    Runs under ``torch.no_grad()`` because the graph already contains explicit
+    backward ops (from ``torch.autograd.grad`` traced by make_fx). Without
+    this, PyTorch would build a redundant autograd graph on top, keeping all
+    forward intermediates alive via ``grad_fn`` references.
     """
     state_flat = list(state.values())
     user_args_flat, _ = pytree.tree_flatten(list(args))
@@ -461,7 +425,8 @@ def run_traced(
     all_args = list(state_flat) + list(user_args_flat)
     flat_inputs, _ = _unwrap_subclasses(all_args)
 
-    flat_outputs = traced_result.gm(*flat_inputs)
+    with torch.no_grad():
+        flat_outputs = traced_result.gm(*flat_inputs)
     wrapped = _wrap_subclasses(
         flat_outputs,
         traced_result.num_flat_outputs,

torchtitan/experiments/graph_trainer/passes.py

@@ -49,6 +49,10 @@
 from torchtitan.tools.logging import logger
 
 
+def _is_backward_node(node: torch.fx.Node) -> bool:
+    return node.meta.get("autograd_backward", False)
+
+
 def construct_default_graph_passes(
     traced_result: "TracedResult",
 ) -> list[Callable]:
@@ -68,6 +72,7 @@ def construct_default_graph_passes(
 
     passes: list[Callable] = [
         functools.partial(tlparse_log_graph_pass, graph_name="make_fx_graph_traced"),
+        apply_ac_on_fwd_bwd_graph,
         remove_detach_pass,
         remove_identity_view_pass,
         remove_identity_slice_pass,
@@ -409,6 +414,11 @@ def apply_sac_pass(
         if node.op != "call_function":
             continue
 
+        # Skip backward nodes — they must not carry recompute tags,
+        # otherwise the remat pass would try to duplicate backward ops.
+        if _is_backward_node(node):
+            continue
+
         if node.target in (
             operator.getitem,
             torch.ops._c10d_functional.wait_tensor.default,
@@ -426,7 +436,6 @@ def apply_sac_pass(
                 node.meta["recompute"] = parent.meta["recompute"]
                 node.meta["ac_graph_id"] = parent.meta.get("ac_graph_id", 0)
             continue
-
         custom_meta = node.meta.get("custom", {})
         ac_region_id = custom_meta.get(_AC_REGION_ID, 0)
         node.meta["ac_graph_id"] = ac_region_id
@@ -461,6 +470,27 @@ def apply_sac_pass(
     return gm
 
 
+def apply_ac_on_fwd_bwd_graph(
+    gm: torch.fx.GraphModule, example_inputs: tuple | None = None
+) -> torch.fx.GraphModule:
+    """Apply graph-based SAC to a traced fwd+loss+bwd graph.
+
+    Tags forward nodes with recompute policy via apply_sac_pass (backward
+    nodes are skipped automatically via ``node.meta["autograd_backward"]``), then
+    applies remat_using_tags_for_fwd_loss_bwd_graph to duplicate
+    PREFER_RECOMPUTE forward ops before backward and DCE originals.
+
+    The model must have been annotated with annotate_ac_regions before
+    tracing so that nodes have custom["ac_region_id"] metadata.
+    """
+    from torch._functorch._activation_checkpointing.remat_using_tags_for_fwd_loss_bwd_graph_pass import (
+        remat_using_tags_for_fwd_loss_bwd_graph,
+    )
+
+    apply_sac_pass(gm)
+    return remat_using_tags_for_fwd_loss_bwd_graph(gm)
+
+
 # Apply activation checkpointing on joint graph before partitioner
 def fsdp_reshard_after_fwd_pass(
     gm: torch.fx.GraphModule,
@@ -561,7 +591,7 @@ def inductor_decomposition_pass(
             f"Placeholder count mismatch: {len(orig_placeholders)} vs {len(decomp_placeholders)}"
         )
 
-    for orig, decomp in zip(orig_placeholders, decomp_placeholders):
+    for orig, decomp in zip(orig_placeholders, decomp_placeholders, strict=True):
         # Copy all metadata from original to decomposed
         for key, value in orig.meta.items():
             if key not in decomp.meta:

torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py

@@ -0,0 +1,59 @@
+# 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.
+
+from types import SimpleNamespace
+
+import torch
+import torch.nn as nn
+
+from torchtitan.components.loss import cross_entropy_loss
+from torchtitan.config import ActivationCheckpointConfig
+from torchtitan.distributed.utils import get_train_context
+from torchtitan.experiments.graph_trainer.trainer import GraphTrainer
+from torchtitan.trainer import Trainer
+
+
+def build_minimal_trainer(
+    model: nn.Module,
+    model_config,
+    trainer_cls: type[Trainer],
+    *,
+    activation_checkpoint_mode: str = "none",
+    compile_enable_passes: bool = True,
+    compile_passes: list[str] | None = None,
+    compile_joint_passes: list[str] | None = None,
+    tokenizer=None,
+) -> Trainer:
+    """Build the minimal Trainer/GraphTrainer needed for single-GPU test steps."""
+    trainer = object.__new__(trainer_cls)
+    trainer.model_parts = [model]
+    trainer.loss_fn = cross_entropy_loss
+    trainer.parallel_dims = SimpleNamespace(pp_enabled=False, cp_enabled=False)
+    trainer.train_context = get_train_context(False)
+    trainer.model_config = model_config
+    trainer.device = torch.device("cuda")
+    trainer.tokenizer = tokenizer
+
+    if trainer_cls is GraphTrainer:
+        trainer.config = SimpleNamespace(
+            compile=SimpleNamespace(
+                mode="aot_fx_trace",
+                enable_passes=compile_enable_passes,
+                passes=[] if compile_passes is None else list(compile_passes),
+                joint_passes=[]
+                if compile_joint_passes is None
+                else list(compile_joint_passes),
+            ),
+            activation_checkpoint=ActivationCheckpointConfig(
+                mode=activation_checkpoint_mode
+            ),
+        )
+        trainer._fwd_bwd_step_module = None
+        trainer._traced_step = None
+    else:
+        trainer.config = SimpleNamespace()
+
+    return trainer

torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py

@@ -16,16 +16,13 @@
 import copy
 import unittest
 from collections.abc import Callable
-from types import SimpleNamespace
 
 import torch
 import torch.nn as nn
 from expecttest import assert_expected_inline
 from tests.utils import hash_gradient, hash_model
 
-from torchtitan.components.loss import cross_entropy_loss
 from torchtitan.components.tokenizer import HuggingFaceTokenizer
-from torchtitan.distributed.utils import get_train_context
 from torchtitan.experiments.graph_trainer.deepseek_v3 import (
     model_registry as dsv3_model_registry,
 )
@@ -36,6 +33,9 @@
     model_registry as llama3_model_registry,
 )
 from torchtitan.experiments.graph_trainer.llama3.parallelize import annotate_llama
+from torchtitan.experiments.graph_trainer.tests._trainer_test_utils import (
+    build_minimal_trainer,
+)
 from torchtitan.experiments.graph_trainer.trainer import GraphTrainer
 from torchtitan.tools.utils import has_cuda_capability
 from torchtitan.trainer import Trainer
@@ -57,42 +57,6 @@ def _set_deterministic(seed: int = SEED) -> None:
 _TOKENIZER_PATH = "./tests/assets/tokenizer"
 
 
-def _build_trainer(
-    model: nn.Module,
-    model_config,
-    trainer_cls: type,
-    *,
-    enable_passes: bool = True,
-) -> Trainer:
-    """Build a minimal Trainer/GraphTrainer for single-GPU non-distributed testing.
-
-    Uses object.__new__ to bypass __init__ because the full Trainer constructor
-    requires a distributed environment, job config, and checkpoint manager that
-    are unnecessary for single-GPU numerical verification. The attributes set
-    below are the minimal set required by forward_backward_step().
-    """
-    trainer = object.__new__(trainer_cls)
-    trainer.model_parts = [model]
-    trainer.loss_fn = cross_entropy_loss
-    trainer.parallel_dims = SimpleNamespace(pp_enabled=False, cp_enabled=False)
-    trainer.train_context = get_train_context(False)
-    trainer.model_config = model_config
-    trainer.device = torch.device("cuda")
-    trainer.tokenizer = HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH)
-
-    if trainer_cls is GraphTrainer:
-        trainer.config = SimpleNamespace(
-            compile=SimpleNamespace(
-                mode="aot_fx_trace",
-                enable_passes=enable_passes,
-            )
-        )
-        trainer._fwd_bwd_step_module = None
-        trainer._traced_step = None
-
-    return trainer
-
-
 @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
 class BitwiseDeterministicBase(unittest.TestCase):
     """Base class for bitwise determinism tests.
@@ -131,8 +95,12 @@ def _run_steps(
         # Annotate after deepcopy: annotate_fn wrappers capture bound methods
         # that don't rebind correctly through copy.deepcopy.
         self.annotate_model(model)
-        trainer = _build_trainer(
-            model, self.model_config, trainer_cls, enable_passes=enable_passes
+        trainer = build_minimal_trainer(
+            model,
+            self.model_config,
+            trainer_cls,
+            compile_enable_passes=enable_passes,
+            tokenizer=HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH),
         )
         global_valid_tokens = torch.tensor(
             BATCH_SIZE * SEQ_LEN, dtype=torch.float, device="cuda"
@@ -318,11 +286,6 @@ def test_eager_self_deterministic(self):
             """8bb6e647c3edaa229cc65872086ccc5c4e1b7f1647bb01da4506ab777a64a0db""",
         )
 
-    # TODO: OOMs during flex_attention compilation on A100 GPUs.
-    # Revisit when GraphTrainer addresses peak memory during compilation.
-    @unittest.skipUnless(
-        has_cuda_capability(9, 0), "OOMs during flex_attention compilation on A100"
-    )
     def test_aot_fx_trace_vs_eager(self):
         """aot_fx_trace with passes and eager produce bitwise identical results."""
         run_eager = self._run_steps(copy.deepcopy(self.model), Trainer)