[graph_trainer] Add remat pass and torch.no_grad() execution to minimal_fx_tracer

- Annotate backward FX nodes with {"remat_pass_tag": "is_backward"} during
  _patch_engine_run_backward so remat_using_tags_for_fwd_loss_bwd_graph can
  identify the forward/backward boundary.
- Apply remat_using_tags_for_fwd_loss_bwd_graph as a default post-trace pass.
  Nodes tagged PREFER_RECOMPUTE (from selective AC) are duplicated before
  backward and the forward copies are DCE'd, reducing peak memory.
- Execute traced graph under torch.no_grad() since the graph already contains
  explicit backward ops. Without this, PyTorch builds a redundant autograd
  graph keeping all forward intermediates alive via grad_fn references.
- Add test_llama_1b_peak_memory: verifies traced+AC peak memory is within 20%
  of eager+AC on Llama 1B (BS=2, seq=2048, bf16).

ghstack-source-id: b038c22
Pull Request resolved: #2766

Author: tugsbayasgalan

torchtitan/experiments/graph_trainer/make_fx_tracer.py

@@ -12,6 +12,9 @@
 import torch
 import torch.nn as nn
 import torch.utils._pytree as pytree
+from torch._functorch._activation_checkpointing.remat_using_tags_for_fwd_loss_bwd_graph_pass import (
+    remat_using_tags_for_fwd_loss_bwd_graph,
+)
 from torch._functorch._aot_autograd.logging_utils import (
     setup_stacktrace_preservation_hooks,
 )
@@ -202,7 +205,8 @@ def _remove_cpu_shadow_chains(gm: torch.fx.GraphModule) -> None:
 
 @contextmanager
 def _patch_engine_run_backward() -> Generator[None, None, None]:
-    """Patch _engine_run_backward to install stacktrace preservation hooks.
+    """Patch _engine_run_backward to install stacktrace preservation hooks and
+    annotate backward nodes for the activation checkpointing remat pass.
 
     Why this is needed:
     When make_fx traces a function that calls loss.backward(), the backward
@@ -215,6 +219,11 @@ def _patch_engine_run_backward() -> Generator[None, None, None]:
     nodes back to their forward counterparts (needed by
     ``_copy_fwd_metadata_to_bw_nodes``).
 
+    Additionally, we annotate all backward nodes with
+    ``{"remat_pass_tag": "is_backward"}`` so that
+    ``remat_using_tags_for_fwd_loss_bwd_graph`` can identify the backward
+    region boundary for activation checkpointing rematerialization.
+
     We must patch the name in both modules since ``torch.autograd.__init__``
     imports it via ``from .graph import``.
     """
@@ -231,7 +240,8 @@ def _patched(t_outputs, *args, **kwargs):  # type: ignore[no-untyped-def]
         ]
         if roots:
             setup_stacktrace_preservation_hooks(roots)
-        return _orig_fn(t_outputs, *args, **kwargs)
+        with torch.fx.traceback.annotate({"remat_pass_tag": "is_backward"}):
+            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]
@@ -327,6 +337,11 @@ def __call__(self, *args: Any) -> Any:
 
         The module must be the first argument (position 0), matching the
         convention enforced by :func:`minimal_fx_tracer`.
+
+        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.
         """
         mod = args[0]
         params_dict = _get_params_and_buffers(mod)
@@ -348,7 +363,8 @@ def __call__(self, *args: Any) -> Any:
         all_args = params_flat + list(user_args_flat)
         flat_inputs, _ = _unwrap_subclasses(all_args)
 
-        flat_outputs = self.gm(*flat_inputs)
+        with torch.no_grad():
+            flat_outputs = self.gm(*flat_inputs)
         wrapped = _wrap_subclasses(
             flat_outputs, self.num_flat_outputs, self.output_subclass_layouts
         )
@@ -478,6 +494,11 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
 
     _remove_cpu_shadow_chains(traced)
 
+    # Rematerialize activations tagged PREFER_RECOMPUTE by selective AC.
+    # Duplicates recomputable forward ops before the backward region and
+    # DCEs the original copies, reducing peak memory.
+    traced = remat_using_tags_for_fwd_loss_bwd_graph(traced)
+
     assert output_spec is not None
     return TracedResult(
         gm=traced,

torchtitan/experiments/graph_trainer/tests/test_trace_module.py

@@ -665,6 +665,94 @@ def forward(model, tokens, attn_masks):
                 f"{node.name} missing ac_region_id annotation",
             )
 
+    def test_llama_1b_peak_memory(self):
+        """Traced+AC peak memory within 10% of eager+AC, with bitwise-identical numerics."""
+        from torchtitan.config import ActivationCheckpointConfig
+        from torchtitan.distributed.activation_checkpoint import apply_ac
+        from torchtitan.models.llama3 import llama3_configs, Llama3Model
+
+        config = llama3_configs["1B"]
+        batch_size, seq_len = 2, 2048
+        dtype = torch.bfloat16
+        num_steps = 3
+        lr = 1e-3
+
+        tokens = torch.randint(
+            0, config.vocab_size, (batch_size, seq_len), device=self.DEVICE
+        )
+        labels = torch.randint(
+            0, config.vocab_size, (batch_size, seq_len), device=self.DEVICE
+        )
+
+        def make_model():
+            model = Llama3Model(config).to(device=self.DEVICE, dtype=dtype)
+            with torch.no_grad():
+                model.init_weights(buffer_device=torch.device(self.DEVICE))
+            apply_ac(model, ActivationCheckpointConfig(mode="selective"))
+            return model
+
+        def run(mode):
+            model = make_model()
+            model.load_state_dict(state)
+            if mode == "traced":
+                train_step = make_train_step(get_loss)
+                traced = minimal_fx_tracer(train_step, (model, tokens, labels))
+            opt = torch.optim.Adam(model.parameters(), lr=lr)
+            step_results = []
+            torch.cuda.reset_peak_memory_stats()
+            for _ in range(num_steps):
+                if mode == "traced":
+                    result = traced(model, tokens, labels)
+                    loss, grads = result[0], result[1:]
+                    for p, g in zip(model.parameters(), grads, strict=True):
+                        p.grad = g
+                else:
+                    logits = model(tokens)
+                    loss = get_loss(logits, labels)
+                    loss.backward()
+                    grads = [p.grad for p in model.parameters()]
+                step_results.append((
+                    loss.detach().cpu(),
+                    [g.clone().cpu() for g in grads],
+                ))
+                opt.step()
+                opt.zero_grad()
+            peak = torch.cuda.max_memory_allocated() / 1e9
+            del model, opt
+            torch.cuda.empty_cache()
+            return step_results, peak
+
+        ref = make_model()
+        state = ref.state_dict()
+        del ref
+        torch.cuda.empty_cache()
+
+        eager_results, peak_eager = run("eager")
+        traced_results, peak_traced = run("traced")
+
+        # Bitwise-identical loss and grads across all steps.
+        for step, ((el, eg), (tl, tg)) in enumerate(
+            zip(eager_results, traced_results, strict=True)
+        ):
+            self.assertTrue(
+                torch.equal(el, tl),
+                f"Step {step}: loss mismatch — eager={el.item():.6f} vs traced={tl.item():.6f}",
+            )
+            for i, (ge, gt) in enumerate(zip(eg, tg, strict=True)):
+                self.assertTrue(
+                    torch.equal(ge, gt),
+                    f"Step {step}: grad[{i}] mismatch — max_diff={( ge - gt).abs().max().item():.2e}",
+                )
+
+        # Peak memory within 10%.
+        ratio = peak_traced / peak_eager
+        self.assertLess(
+            ratio,
+            1.1,
+            f"Traced+AC peak memory ({peak_traced:.2f} GB) is more than "
+            f"10% above eager+AC ({peak_eager:.2f} GB), ratio={ratio:.2f}",
+        )
+
 
 class TestTraceFSDP(FSDPTest):
     @property