Update on "[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.

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 
<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 
<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" />


 For Llama3 8B, from the step-30 rank0 memory snapshots:

  - No SAC:
      - end-of-step memory_allocated() = 14.99 GiB
        GiB
      - end-of-step memory_reserved() = 22.75 GiB
  - SAC:
      - end-of-step memory_allocated() = 14.99 GiB
        GiB
      - end-of-step memory_reserved() = 16.6 GiB
     
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" />



[ghstack-poisoned]

Author: tugsbayasgalan

.github/workflows/integration_test_8gpu_graph_trainer.yaml

@@ -61,7 +61,7 @@ jobs:
         sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded"
         sudo chown -R $(id -u):$(id -g) "$RUNNER_TEMP/artifacts-to-be-uploaded"
 
-        python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_default --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8 --collect_peak_memory
+        python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_default --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8
 
 
         # Run the numerics unit tests (dense models only; MoE tests run in the H100 workflow)

tests/integration_tests/run_tests.py

@@ -5,12 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 import argparse
-import json
-import logging
 import os
 import subprocess
 import time
-from pathlib import Path
 
 from torchtitan.tools.logging import logger
 
@@ -26,9 +23,6 @@
     "h100": build_h100_tests_list,
 }
 
-TB_MAX_RESERVED_TAG = "memory/max_reserved(GiB)"
-TB_MAX_ACTIVE_TAG = "memory/max_active(GiB)"
-
 
 def _run_cmd(cmd, timeout=None):
     return subprocess.run(
@@ -41,60 +35,11 @@ def _run_cmd(cmd, timeout=None):
     )
 
 
-def _read_peak_memory(tb_root: Path) -> dict[str, float | int]:
-    try:
-        from tensorboard.backend.event_processing.event_accumulator import (
-            EventAccumulator,
-        )
-    except ModuleNotFoundError as exc:
-        raise RuntimeError(
-            "Peak memory collection requires tensorboard to be installed."
-        ) from exc
-
-    if not tb_root.exists():
-        raise FileNotFoundError(f"TensorBoard directory not found: {tb_root}")
-    run_dirs = [path for path in tb_root.iterdir() if path.is_dir()]
-    if len(run_dirs) != 1:
-        raise RuntimeError(
-            f"Expected exactly one TensorBoard run directory under {tb_root}, found {run_dirs}"
-        )
-
-    tensorboard_logger = logging.getLogger("tensorboard")
-    original_level = tensorboard_logger.level
-    tensorboard_logger.setLevel(logging.WARNING)
-    event_accumulator = EventAccumulator(str(run_dirs[0]))
-    try:
-        event_accumulator.Reload()
-    finally:
-        tensorboard_logger.setLevel(original_level)
-
-    scalar_tags = event_accumulator.Tags().get("scalars", [])
-
-    def _peak_for_tag(tag: str) -> tuple[float, int]:
-        if tag not in scalar_tags:
-            raise KeyError(
-                f"Scalar tag {tag!r} not found in {run_dirs[0]}: {scalar_tags}"
-            )
-        peak = max(event_accumulator.Scalars(tag), key=lambda scalar: scalar.value)
-        return peak.value, peak.step
-
-    reserved_gib, reserved_step = _peak_for_tag(TB_MAX_RESERVED_TAG)
-    active_gib, active_step = _peak_for_tag(TB_MAX_ACTIVE_TAG)
-    return {
-        "max_reserved_gib": reserved_gib,
-        "max_reserved_step": reserved_step,
-        "max_active_gib": active_gib,
-        "max_active_step": active_step,
-    }
-
-
 def run_single_test(
     test_flavor: OverrideDefinitions,
     output_dir: str,
     module: str | None = None,
     config: str | None = None,
-    *,
-    collect_peak_memory: bool = False,
 ):
     # run_test supports sequence of tests.
     test_name = test_flavor.test_name
@@ -103,7 +48,6 @@ def run_single_test(
     all_ranks = ",".join(map(str, range(test_flavor.ngpu)))
 
     for idx, override_arg in enumerate(test_flavor.override_args):
-        tb_folder = f"tb_{idx}"
         cmd = ""
         if module is not None:
             cmd += f"MODULE={module} "
@@ -115,11 +59,6 @@ def run_single_test(
         cmd = f'TORCH_TRACE="{output_dir}/{test_name}/compile_trace" ' + cmd
 
         cmd += " " + dump_folder_arg
-        if collect_peak_memory:
-            cmd += (
-                f" --metrics.enable_tensorboard --metrics.log_freq=1 "
-                f"--metrics.save_tb_folder={tb_folder}"
-            )
         if override_arg:
             cmd += " " + " ".join(override_arg)
         logger.info(
@@ -150,18 +89,6 @@ def run_single_test(
                 f"Command: {cmd}\n"
                 f"stderr: {result.stderr}\n"
             )
-        if collect_peak_memory:
-            tb_root = Path(output_dir) / test_name / tb_folder
-            peak_memory = _read_peak_memory(tb_root)
-            summary_path = Path(output_dir) / test_name / f"peak_memory_{idx}.json"
-            summary_path.write_text(json.dumps(peak_memory, indent=2))
-            logger.info(
-                f"Peak memory for {test_name}[{idx}]: "
-                f"reserved={peak_memory['max_reserved_gib']:.3f} GiB "
-                f"at step {peak_memory['max_reserved_step']}, "
-                f"active={peak_memory['max_active_gib']:.3f} GiB "
-                f"at step {peak_memory['max_active_step']}"
-            )
 
 
 def run_tests(
@@ -171,8 +98,6 @@ def run_tests(
     config=None,
 ):
     """Run all integration tests to test the core features of TorchTitan"""
-    collect_peak_memory = getattr(args, "collect_peak_memory", False)
-
     exclude_set = set()
     if hasattr(args, "exclude") and args.exclude:
         exclude_set = {name.strip() for name in args.exclude.split(",")}
@@ -207,7 +132,6 @@ def run_tests(
                     args.output_dir,
                     module,
                     config,
-                    collect_peak_memory=collect_peak_memory,
                 )
             except Exception as e:
                 logger.error(str(e))
@@ -278,12 +202,6 @@ def main():
         default=None,
         help="Comma-separated list of test names to skip",
     )
-    parser.add_argument(
-        "--collect_peak_memory",
-        default=False,
-        action="store_true",
-        help="Collect peak reserved/active CUDA memory from TensorBoard logs.",
-    )
     args = parser.parse_args()
 
     if not os.path.exists(args.output_dir):

tests/utils.py

@@ -11,11 +11,6 @@
 from torch.distributed.tensor import DTensor
 
 
-def _canonicalize_tensor_name(name: str) -> str:
-    """Normalize transparent wrapper segments so hashes reflect tensor contents."""
-    return name.replace("._checkpoint_wrapped_module", "")
-
-
 def _hash_model_impl(
     model: torch.nn.Module,
     algo: str,
@@ -34,7 +29,6 @@ def _hash_model_impl(
 
     def hash_named_tensor(name: str, obj) -> None:
         if isinstance(obj, torch.Tensor):
-            name = _canonicalize_tensor_name(name)
             if isinstance(obj, DTensor):
                 t = obj.to_local().cpu().contiguous()
             else:

torchtitan/experiments/graph_trainer/common_utils.py

@@ -155,29 +155,20 @@ def convert_modules_to_fqns(modules, module_to_fqn_mapping):
     return module_fqns
 
 
-def enable_graph_ac_for_mode(ac_mode: str) -> bool:
-    if ac_mode == "none":
-        return False
-    if ac_mode == "selective":
-        return True
-    raise ValueError(
-        "graph_trainer only supports activation_checkpoint.mode "
-        f"'selective' or 'none', got {ac_mode!r}. Use 'selective' for "
-        "graph-based SAC."
-    )
-
-
 def apply_graph_ac(
     compile_config: CompileConfig,
     ac_config: "ActivationCheckpointConfig",
 ) -> None:
     """Add apply_sac to compile joint passes for graph-based selective AC.
 
-    ``selective`` enables graph SAC, ``none`` is a no-op, and other modes
-    raise ValueError.
+    Must be called only when ac_config.mode != "none". Only "selective" mode
+    is supported; other modes raise ValueError.
     """
-    if not enable_graph_ac_for_mode(ac_config.mode):
-        return
+    if ac_config.mode != "selective":
+        raise ValueError(
+            f"graph_trainer only supports activation_checkpoint.mode 'selective' or "
+            f"'none', got {ac_config.mode!r}. Use 'selective' for graph-based SAC."
+        )
 
     joint_pass_names = getattr(compile_config, "joint_passes", [])
     if "apply_sac" not in joint_pass_names:

torchtitan/experiments/graph_trainer/make_fx_tracer.py

@@ -361,6 +361,11 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
 
         ctx = TracingContext(fake_mode)
         # preserve_node_meta propagates fx.traceback.annotate metadata to traced nodes
+        # _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),

torchtitan/experiments/graph_trainer/passes.py

@@ -187,7 +187,7 @@ def _is_backward_node(node: torch.fx.Node) -> bool:
 
 
 def construct_default_graph_passes(
-    traced_result: TracedResult,
+    traced_result: "TracedResult",
 ) -> list[Callable]:
     """Build the default pass list for the aot_fx_trace compile path.
 
@@ -205,6 +205,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,
@@ -253,13 +254,6 @@ def apply_graph_passes(
     return gm
 
 
-def graph_ac_pass(
-    gm: torch.fx.GraphModule, example_inputs: tuple | None = None
-) -> torch.fx.GraphModule:
-    """Apply graph-based SAC to a traced fwd+loss+bwd graph."""
-    return apply_ac_on_fwd_bwd_graph(gm)
-
-
 def autobucketing_reordering_pass(
     gm: torch.fx.GraphModule, example_inputs: tuple | None = None
 ) -> torch.fx.GraphModule:
@@ -553,8 +547,6 @@ def apply_sac_pass(
         if node.op != "call_function":
             continue
 
-        custom_meta = node.meta.get("custom", {})
-
         # Skip backward nodes — they must not carry recompute tags,
         # otherwise the remat pass would try to duplicate backward ops.
         if _is_backward_node(node):
@@ -577,6 +569,7 @@ 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
 
@@ -610,7 +603,9 @@ def apply_sac_pass(
     return gm
 
 
-def apply_ac_on_fwd_bwd_graph(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
+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

torchtitan/experiments/graph_trainer/tests/integration_tests.py

@@ -547,11 +547,6 @@ def main():
         help="test to run, acceptable values: `test_name` in `build_test_list` (default: all)",
     )
     parser.add_argument("--ngpu", default=8, type=int)
-    parser.add_argument(
-        "--collect_peak_memory",
-        action="store_true",
-        help="Collect peak reserved/active CUDA memory from TensorBoard logs.",
-    )
     args = parser.parse_args()
 
     if not os.path.exists(args.output_dir):

torchtitan/experiments/graph_trainer/trainer.py

@@ -11,8 +11,6 @@
 import torch.nn as nn
 
 from torchtitan.experiments.graph_trainer.common_utils import (
-    annotate_ac_regions,
-    enable_graph_ac_for_mode,
     maybe_register_blockmask_pytree_node,
 )
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
@@ -25,7 +23,6 @@
 from torchtitan.experiments.graph_trainer.passes import (
     apply_graph_passes,
     construct_default_graph_passes,
-    graph_ac_pass,
 )
 from torchtitan.trainer import Trainer
 
@@ -119,10 +116,6 @@ def _make_fx_forward_backward_step(
     ) -> torch.Tensor:
         if self._traced_step is None:
             fwd_bwd_fn = make_fwd_bwd_step(self.loss_fn)
-            ac_mode = self.config.activation_checkpoint.mode
-            enable_graph_ac = enable_graph_ac_for_mode(ac_mode)
-            if enable_graph_ac:
-                annotate_ac_regions(model)
             maybe_register_blockmask_pytree_node()
             with self.train_context():
                 self._traced_step = trace_train_step(fwd_bwd_fn)(
@@ -135,10 +128,7 @@ def _make_fx_forward_backward_step(
                 )
 
             if self.config.compile.enable_passes:
-                passes = []
-                if enable_graph_ac:
-                    passes.append(graph_ac_pass)
-                passes.extend(construct_default_graph_passes(self._traced_step))
+                passes = construct_default_graph_passes(self._traced_step)
                 self._traced_step.gm = apply_graph_passes(
                     self._traced_step.gm,
                     self._traced_step.example_inputs,