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

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

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.

Adds SAC option on aot_fx_trace on GraphTrainer. 

Adds option to dump the peak memory from CUDA caching allocator. 

Adds similar test as test_bitwise_equivalency.py for peak memory 

Here is the run comparison between eager and aot_fx_trace with SAC
<img width="1114" height="455" alt="Screenshot 2026-04-14 at 1 18 41 PM"
src="https://github.com/user-attachments/assets/bf4ffdd3-8783-47be-ac64-32587787591f"
/>

Llama3B 8b with dp=4 tp=2 on 10 step run before SAC with cuda graph
seq_len=1024
<img width="1728" height="585" alt="Screenshot 2026-04-14 at 10 39
41 PM"
src="https://github.com/user-attachments/assets/d8a1fe31-bd91-41a4-8e6f-a0e738cc2b76"
/>

Llama3B 8b with dp=4 tp=2 on 10 step run after SAC with cuda graph
seq_len=1024
<img width="1728" height="552" alt="Screenshot 2026-04-14 at 10 39
26 PM"
src="https://github.com/user-attachments/assets/273b6e65-c2a3-473b-a2db-53b67a95be3a"
/>

Llama3B 8b with dp=4 tp=2 on 10 step run before SAC without cuda graph
with seq_len=8192
<img width="1728" height="577" alt="Screenshot 2026-04-15 at 4 07 50 PM"
src="https://github.com/user-attachments/assets/ca7108eb-6dc6-476c-b9d9-fa320f96945a"
/>

Llama3B 8b with dp=4 tp=2 on 10 step run after SAC without cuda graph
with seq_len=8192
<img width="1721" height="592" alt="Screenshot 2026-04-15 at 4 07 43 PM"
src="https://github.com/user-attachments/assets/80fcdcb0-ca38-4012-b71d-2b091adbfddd"
/>

Deepseek3 16B before SAC 
<img width="1728" height="600" alt="Screenshot 2026-04-14 at 10 54
31 PM"
src="https://github.com/user-attachments/assets/61180fa4-debd-4c10-834e-e4936cfbd67e"
/>

Deepseek3 16B after SAC 

<img width="1728" height="578" alt="Screenshot 2026-04-14 at 10 58
30 PM"
src="https://github.com/user-attachments/assets/3e3b4e71-e13f-43b7-9404-ad9a0a04f75c"
/>

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

@@ -93,7 +93,6 @@ def run_single_test(
 
 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(",")}

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.
 
@@ -402,33 +353,36 @@ 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
-
         # Disable autograd multithreading so that backward tracing
-        # runs on the calling thread.  Without this, the C++ autograd
+        # 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.
+        # _skip_nested_compile lets the current make_fx trace inline through
+        # torch.compile'd FlexAttention kernels instead of erroring.
+        # _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.autograd.set_multithreading_enabled(False),
+            torch.compiler._non_strict_tracing_context(),
         ):
             traced = make_fx(
                 fn_with_subclass_handling,
@@ -470,6 +424,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))
@@ -481,7 +439,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

@@ -50,6 +50,10 @@
 from torchtitan.tools.logging import logger
 
 
+def _is_backward_node(node: torch.fx.Node) -> bool:
+    return node.meta.get("autograd_backward", False)
+
+
 def compile_time_passes(
     traced_result: "TracedResult",
 ) -> list[Callable]:
@@ -69,6 +73,7 @@ def compile_time_passes(
         remove_detach_pass,
         remove_identity_view_pass,
         remove_identity_slice_pass,
+        selective_activation_remat_pass,
         # FlexAttention HOPs must be compiled (via regional_inductor) to
         # produce bitwise identical results to the eager Trainer path.
         # When left uncompiled, flex_attention still runs correctly but
@@ -452,6 +457,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,
@@ -469,7 +479,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
@@ -504,6 +513,27 @@ def apply_sac_pass(
     return gm
 
 
+def selective_activation_remat_pass(
+    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,
@@ -604,7 +634,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,60 @@
+# 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),
+                precompile_artifact_dir="",
+            ),
+            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

@@ -17,7 +17,6 @@
 import tempfile
 import unittest
 from collections.abc import Callable
-from types import SimpleNamespace
 
 import torch
 import torch.nn as nn
@@ -27,7 +26,6 @@
 
 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,
 )
@@ -38,6 +36,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.models.common.attention import FlexAttention
 from torchtitan.tools.utils import has_cuda_capability
@@ -60,43 +61,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,
-                precompile_artifact_dir="",
-            )
-        )
-        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.
@@ -149,8 +113,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"
@@ -446,11 +414,6 @@ def test_eager_self_deterministic(self):
             """b3f5b911dea6c9d36f508b08300220d7f39f142a7c34a49b7ef2543abb2065dc""",
         )
 
-    # 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)