From 3d04ab986d660814b84b079badaf53c4c3866a2b Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 30 Mar 2026 09:24:21 -0700
Subject: [PATCH 01/42] [graph_trainer] Replace trace_module/run_traced_module
 with aot_function API
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Authored-by: Claude

Redesign the graph trainer's tracing API based on the aot_function design doc.

Key changes:

make_fx_tracer.py:
- Rename trace_module -> aot_function. Takes any callable (not just nn.Module)
  with nn.Module instances auto-detected in args and their params/buffers
  lifted as graph inputs. When fn is an nn.Module, it is prepended to args
  and type(fn).__call__ is used as the callable.
- Delete run_traced_module. TracedResult is now directly callable — pass the
  same positional args (with live modules) to execute the graph. Fresh params
  are read from the modules automatically on each call.
- Store and restore output pytree spec so TracedResult.__call__ returns the
  same pytree structure as the original function (e.g. single tensor, list,
  tuple, dict), not a flat list.
- Add _ModuleParamsMeta with FQN storage. Parameter FQNs are recorded at trace
  time and validated at execute time to catch module structure mismatches.
- Add _collect_module_params helper for multi-module param extraction.
- Install TracingContext before make_fx so invoke_subgraph deduplication works.
- Validate that all pytree leaves in args are tensors or primitives
  (int/float/bool/str). Non-primitive values (callables, custom objects) must
  be captured in fn's closure or registered via pytree.register_constant /
  register_pytree_node.

trainer.py:
- Replace FwdBwdStepModule (nn.Module wrapper that only existed because the
  old trace_module required nn.Module as fn) with _make_fwd_bwd_step, a plain
  function factory. The model is now passed as an arg, loss_fn is captured in
  the closure.
- Remove manual params_and_buffers dict construction — TracedResult.__call__
  reads fresh params from the live module automatically.
- Add TODO for investigating loss_fn interaction with non-strict trace.

test_trace_module.py:
- Replace TrainStepModule with _make_train_step plain function factory.
- Remove _get_params_and_buffers helper (no longer needed).
- Update all callsites: trace_module -> aot_function, run_traced_module ->
  direct TracedResult.__call__.
- Register BlockMask as pytree node at module level so flex_attention tests
  pass the leaf validation.
- Add test_module_not_first_arg: module at position 1 in args.
- Add test_multiple_modules: two nn.Modules interleaved with a tensor.
- Add test_mismatched_module_raises: FQN validation catches wrong module.
- Add test_non_tensor_leaf_raises: callable leaf in args raises ValueError.

All 7 model tests pass (llama3, llama4, qwen3, qwen3_moe, deepseek_v3,
gpt_oss, flex_attention_annotations).

[ghstack-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 293 ++++++++++++------
 .../graph_trainer/tests/test_trace_module.py  | 243 +++++++++------
 .../experiments/graph_trainer/trainer.py      |  62 ++--
 3 files changed, 390 insertions(+), 208 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 53b0379939..b8e8f87b60 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -4,8 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import itertools
-from collections.abc import Generator
+import contextlib
+from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any
@@ -16,6 +16,7 @@
 from torch._functorch._aot_autograd.logging_utils import (
     setup_stacktrace_preservation_hooks,
 )
+from torch._guards import TracingContext, tracing
 from torch._subclasses import FakeTensorMode
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.fx.traceback import preserve_node_meta
@@ -57,14 +58,11 @@ class SubclassLayout:
 
 
 @dataclass
-class TracedResult:
-    """Holds the traced graph and metadata needed to run it."""
+class _ModuleParamsMeta:
+    """Per-module parameter metadata captured during tracing."""
 
-    gm: torch.fx.GraphModule
-    params_len: int
-    params_spec: pytree.TreeSpec
-    input_subclass_layouts: list[SubclassLayout]
-    output_subclass_layouts: list[SubclassLayout]
+    fqns: list[str]
+    num_params: int
 
 
 def _unwrap_subclass(t: torch.Tensor) -> tuple[list[torch.Tensor], SubclassMeta | None]:
@@ -252,41 +250,152 @@ def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
                 node.meta["nn_module_stack"] = nn_module_stack.copy()
 
 
-def trace_module(
-    mod: nn.Module,
-    args: tuple,
-) -> TracedResult:
-    """Trace ``mod(*args)`` into a flat FX graph, unwrapping tensor subclasses.
+def _collect_module_params(
+    args: tuple, module_indices: list[int]
+) -> tuple[list[_ModuleParamsMeta], list[torch.Tensor]]:
+    """Extract flattened params/buffers for each nn.Module arg."""
+    per_module: list[_ModuleParamsMeta] = []
+    all_flat: list[torch.Tensor] = []
+    for i in module_indices:
+        mod = args[i]
+        params_dict = {
+            **dict(mod.named_parameters(remove_duplicate=False)),
+            **dict(mod.named_buffers(remove_duplicate=False)),
+        }
+        fqns = list(params_dict.keys())
+        flat = list(params_dict.values())
+        per_module.append(_ModuleParamsMeta(fqns=fqns, num_params=len(flat)))
+        all_flat.extend(flat)
+    return per_module, all_flat
 
-    Parameters and buffers are lifted as extra graph inputs so the returned
-    graph is a pure function.  Tensor subclasses (e.g. DTensor) are recursively
-    unwrapped into plain tensors for tracing, and the layouts needed to rewrap
-    them are recorded in the returned :class:`TracedResult`.
 
-    Args:
-        mod: The module to trace.
-        args: The user arguments to trace with.
+class TracedResult:
+    """Holds the traced graph and metadata needed to execute it.
+
+    Returned by :func:`aot_function`.  Call the instance directly to execute
+    the traced graph with fresh parameters read from the live modules::
+
+        traced = aot_function(train_step, (model, tokens, labels))
+        result = traced(model, tokens, labels)
     """
-    named_parameters = dict(mod.named_parameters(remove_duplicate=False))
-    named_buffers = dict(mod.named_buffers(remove_duplicate=False))
 
-    params_and_buffers = {**named_parameters, **named_buffers}
-    params_and_buffers_flat, params_spec = pytree.tree_flatten(params_and_buffers)
-    params_len = len(params_and_buffers_flat)
+    def __init__(
+        self,
+        gm: torch.fx.GraphModule,
+        module_indices: list[int],
+        per_module_params: list[_ModuleParamsMeta],
+        input_subclass_layouts: list[SubclassLayout],
+        output_subclass_layouts: list[SubclassLayout],
+        output_spec: pytree.TreeSpec,
+    ) -> None:
+        self.gm = gm
+        self.module_indices = module_indices
+        self.per_module_params = per_module_params
+        self.input_subclass_layouts = input_subclass_layouts
+        self.output_subclass_layouts = output_subclass_layouts
+        self.output_spec = output_spec
+
+    def __call__(self, *args: Any) -> list[torch.Tensor]:
+        """Execute the traced graph, reading fresh params from modules in ``args``."""
+        module_indices_set = set(self.module_indices)
+
+        # Read current params from live modules.
+        all_params_flat: list[torch.Tensor] = []
+        for i, pmp in zip(self.module_indices, self.per_module_params, strict=True):
+            mod = args[i]
+            params_dict = {
+                **dict(mod.named_parameters(remove_duplicate=False)),
+                **dict(mod.named_buffers(remove_duplicate=False)),
+            }
+            fqns = list(params_dict.keys())
+            if fqns != pmp.fqns:
+                raise ValueError(
+                    f"Module at arg position {i} has different parameter/buffer "
+                    f"names than during tracing.\n"
+                    f"  Traced: {pmp.fqns}\n"
+                    f"  Got:    {fqns}"
+                )
+            flat, _ = pytree.tree_flatten(params_dict)
+            all_params_flat.extend(flat)
+
+        user_args = [a for i, a in enumerate(args) if i not in module_indices_set]
+        user_args_flat, _ = pytree.tree_flatten(user_args)
+
+        all_args = all_params_flat + list(user_args_flat)
+        flat_inputs: list = []
+        for a in all_args:
+            if isinstance(a, torch.Tensor) and is_traceable_wrapper_subclass(a):
+                inner, _ = _unwrap_subclass(a)
+                flat_inputs.extend(inner)
+            else:
+                flat_inputs.append(a)
 
-    def functional_call(*all_args):
-        flat_params = all_args[:params_len]
-        user_args = all_args[params_len:]
-        params = pytree.tree_unflatten(list(flat_params), params_spec)
-        with stateless._reparametrize_module(mod, params):
-            return mod.forward(*user_args)
+        flat_outputs = self.gm(*flat_inputs)
+        wrapped = _wrap_to_subclasses(flat_outputs, self.output_subclass_layouts)
+        return pytree.tree_unflatten(wrapped, self.output_spec)
 
-    user_args_flat, user_args_spec = pytree.tree_flatten(args)
-    full_args = tuple(params_and_buffers_flat) + tuple(user_args_flat)
 
-    unwrapped_args = []
-    input_layouts: list[SubclassLayout] = []
+def aot_function(
+    fn: nn.Module | Callable,
+    args: tuple,
+) -> TracedResult:
+    """Trace ``fn(*args)`` into a flat FX graph, unwrapping tensor subclasses.
+
+    Parameters and buffers from ``nn.Module`` instances in ``args`` are lifted
+    as extra graph inputs so the returned graph is a pure function.  Tensor
+    subclasses (e.g. DTensor) are recursively unwrapped into plain tensors for
+    tracing, and the layouts needed to rewrap them are recorded in the returned
+    :class:`TracedResult`.
+
+    ``fn`` may be an ``nn.Module`` — in which case it is treated as though the
+    caller wrote ``aot_function(lambda m, *a: m(*a), (fn,) + args)``, i.e. the
+    module is prepended to ``args`` and its parameters are lifted automatically.
+
+    The returned :class:`TracedResult` is directly callable — pass the same
+    positional arguments (with live modules) to execute the graph::
 
+        traced = aot_function(train_step, (model, tokens, labels))
+        result = traced(model, tokens, labels)
+
+    Args:
+        fn: The callable (or ``nn.Module``) to trace.
+        args: The positional arguments to trace with.  ``nn.Module`` instances
+            are detected automatically and their parameters are lifted.
+    """
+    # When fn is an nn.Module, treat it as the first arg so its params get
+    # lifted like any other module arg.
+    if isinstance(fn, nn.Module):
+        args = (fn,) + args
+        fn = type(fn).__call__
+    module_indices = [i for i, a in enumerate(args) if isinstance(a, nn.Module)]
+    module_indices_set = set(module_indices)
+
+    per_module_params, all_params_flat = _collect_module_params(args, module_indices)
+    params_len = len(all_params_flat)
+
+    # User args: positional args that are not nn.Module instances.
+    user_args = [a for i, a in enumerate(args) if i not in module_indices_set]
+    user_args_flat, user_args_spec = pytree.tree_flatten(user_args)
+
+    # Validate leaves: tensors and make_fx-safe primitives (int, float, bool,
+    # str) are allowed.  Everything else (callables, custom objects) should be
+    # registered as pytree nodes/constants or captured in fn's closure.
+    _ALLOWED_LEAF_TYPES = (torch.Tensor, int, float, bool, str, type(None))
+    for leaf in user_args_flat:
+        if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
+            raise ValueError(
+                f"aot_function requires all pytree leaves in args to be tensors "
+                f"or primitives (int/float/bool/str), got {type(leaf).__name__}. "
+                f"Non-primitive values should either be registered as pytree "
+                f"nodes (register_pytree_node) or constants "
+                f"(pytree.register_constant), or captured in fn's closure."
+            )
+
+    # Combined flat input: [*params, *user_args] with subclasses unwrapped.
+    full_args = all_params_flat + list(user_args_flat)
+
+    unwrapped_args: list = []
+    input_layouts: list[SubclassLayout] = []
     for arg in full_args:
         if isinstance(arg, torch.Tensor) and is_traceable_wrapper_subclass(arg):
             inner_tensors, meta = _unwrap_subclass(arg)
@@ -300,46 +409,69 @@ def functional_call(*all_args):
         allow_non_fake_inputs=True,
         shape_env=torch.fx.experimental.symbolic_shapes.ShapeEnv(),
     )
-
-    def to_fake(t):
-        if isinstance(t, torch.Tensor):
-            return fake_mode.from_tensor(t, static_shapes=True)
-        return t
-
-    fake_args = tuple(to_fake(a) for a in unwrapped_args)
+    fake_args = tuple(
+        (
+            fake_mode.from_tensor(a, static_shapes=True)
+            if isinstance(a, torch.Tensor)
+            else a
+        )
+        for a in unwrapped_args
+    )
 
     output_layouts: list[SubclassLayout] = []
+    output_spec: pytree.TreeSpec | None = None
 
-    def fn_with_subclass_handling(*plain_args):
-        nonlocal output_layouts
+    def fn_with_subclass_handling(*plain_args: Any) -> list:
+        nonlocal output_layouts, output_spec
         output_layouts = []
 
-        wrapped_args = _wrap_to_subclasses(plain_args, input_layouts)
-
-        params_args = wrapped_args[:params_len]
-        user_args_wrapped = wrapped_args[params_len:]
-        user_args_restored = pytree.tree_unflatten(
-            list(user_args_wrapped), user_args_spec
-        )
-
-        with _patch_engine_run_backward():
-            outputs = functional_call(*params_args, *user_args_restored)
-
-        flat_outputs, _ = pytree.tree_flatten(outputs)
-        unwrapped_outputs = []
-        for out in flat_outputs:
+        wrapped = _wrap_to_subclasses(plain_args, input_layouts)
+        all_params = wrapped[:params_len]
+        user_flat = wrapped[params_len:]
+
+        # Reconstruct per-module param dicts.
+        offset = 0
+        module_param_dicts = []
+        for pmp in per_module_params:
+            flat = all_params[offset : offset + pmp.num_params]
+            module_param_dicts.append(dict(zip(pmp.fqns, flat, strict=True)))
+            offset += pmp.num_params
+
+        user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
+
+        # Reconstruct the original args: module positions keep the live module
+        # (already in args), non-module positions get the traced user tensors.
+        rebuilt: list = list(args)
+        user_idx = 0
+        for i in range(len(args)):
+            if i not in module_indices_set:
+                rebuilt[i] = user_list[user_idx]
+                user_idx += 1
+
+        with contextlib.ExitStack() as stack:
+            for i, pmp_dict in zip(module_indices, module_param_dicts, strict=True):
+                stack.enter_context(
+                    stateless._reparametrize_module(args[i], pmp_dict)
+                )
+
+            with _patch_engine_run_backward():
+                result = fn(*rebuilt)
+
+        flat_outs, output_spec = pytree.tree_flatten(result)
+        unwrapped_outs: list = []
+        for out in flat_outs:
             if isinstance(out, torch.Tensor) and is_traceable_wrapper_subclass(out):
                 inner, meta = _unwrap_subclass(out)
-                unwrapped_outputs.extend(inner)
+                unwrapped_outs.extend(inner)
                 output_layouts.append(SubclassLayout(len(inner), meta))
             else:
-                unwrapped_outputs.append(out)
+                unwrapped_outs.append(out)
                 output_layouts.append(SubclassLayout(1, None))
+        return unwrapped_outs
 
-        return unwrapped_outputs
-
+    ctx = TracingContext(fake_mode)
     # preserve_node_meta propagates fx.traceback.annotate metadata to traced nodes
-    with fake_mode, preserve_node_meta(), _skip_nested_compile():
+    with fake_mode, tracing(ctx), preserve_node_meta(), _skip_nested_compile():
         traced = make_fx(
             fn_with_subclass_handling,
             record_stack_traces=True,
@@ -349,39 +481,14 @@ def fn_with_subclass_handling(*plain_args):
     # Copy forward annotations to backward nodes.
     # Must run before DCE so that forward nodes used for matching aren't removed.
     _copy_fwd_metadata_to_bw_nodes(traced)
-
     _remove_cpu_shadow_chains(traced)
 
+    assert output_spec is not None
     return TracedResult(
         gm=traced,
-        params_len=params_len,
-        params_spec=params_spec,
+        module_indices=module_indices,
+        per_module_params=per_module_params,
         input_subclass_layouts=input_layouts,
         output_subclass_layouts=output_layouts,
+        output_spec=output_spec,
     )
-
-
-def run_traced_module(
-    traced_result: TracedResult,
-    params_and_buffers: dict[str, torch.Tensor],
-    args: tuple,
-) -> list[torch.Tensor]:
-    """Execute a traced graph and rewrap outputs into their original subclass types.
-
-    Accepts a ``params_and_buffers`` dict (from ``named_parameters`` /
-    ``named_buffers``) instead of the module itself, so callers control exactly
-    which parameter snapshot is used.
-    """
-    params_flat, _ = pytree.tree_flatten(params_and_buffers)
-    user_args_flat, _ = pytree.tree_flatten(args)
-
-    all_args = []
-    for a in itertools.chain(params_flat, user_args_flat):
-        if isinstance(a, torch.Tensor) and is_traceable_wrapper_subclass(a):
-            inner, _ = _unwrap_subclass(a)
-            all_args.extend(inner)
-        else:
-            all_args.append(a)
-
-    flat_outputs = traced_result.gm(*all_args)
-    return _wrap_to_subclasses(flat_outputs, traced_result.output_subclass_layouts)
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index a7ce654265..1bd8ed8780 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -12,16 +12,22 @@
 import torch.nn as nn
 from torch.testing._internal.common_fsdp import FSDPTest
 
+from torchtitan.experiments.graph_trainer.common_utils import (
+    register_blockmask_pytree_node,
+)
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
     _copy_fwd_metadata_to_bw_nodes,
     _patch_engine_run_backward,
-    run_traced_module,
-    trace_module,
+    aot_function,
 )
 from torchtitan.models.common.attention import (
     annotate_flex_attention_for_regional_inductor,
 )
 
+# BlockMask must be registered as a pytree node so its tensor children
+# are properly traced as graph inputs instead of opaque leaves.
+register_blockmask_pytree_node()
+
 
 def get_loss(logits, labels):
     return torch.nn.functional.cross_entropy(
@@ -31,28 +37,18 @@ def get_loss(logits, labels):
     )
 
 
-class TrainStepModule(nn.Module):
-    def __init__(self, model, loss_fn):
-        super().__init__()
-        self.model = model
-        self.loss_fn = loss_fn
+def _make_train_step(loss_fn):
+    """Return a plain function for aot_function tracing.  loss_fn is captured in closure."""
 
-    def forward(self, *args):
+    def train_step(model, *args):
         *fwd_args, labels = args
-        logits = self.model(*fwd_args)
-        loss = self.loss_fn(logits, labels)
-        # Must look up params in forward (not __init__) so that
-        # _reparametrize_module's swapped parameters are captured during tracing.
-        params = list(self.model.parameters())
+        logits = model(*fwd_args)
+        loss = loss_fn(logits, labels)
+        params = list(model.parameters())
         grads = torch.autograd.grad(loss, params)
         return [loss] + list(grads)
 
-
-def _get_params_and_buffers(mod):
-    return {
-        **dict(mod.named_parameters(remove_duplicate=False)),
-        **dict(mod.named_buffers(remove_duplicate=False)),
-    }
+    return train_step
 
 
 def create_model(config_cls, model_config, device="cuda", dtype=torch.float32):
@@ -123,28 +119,25 @@ def _make_mlp(self):
 
     def test_mlp_forward(self):
         model, tokens, labels, loss_fn = self._make_mlp()
-        traced_result = trace_module(model, (tokens,))
+        traced = aot_function(model, (tokens,))
         out_eager = model(tokens)
-        params_and_buffers = _get_params_and_buffers(model)
-        wrapped = run_traced_module(traced_result, params_and_buffers, (tokens,))
-        self.assertTrue(torch.equal(out_eager, wrapped[0]))
+        wrapped = traced(model, tokens)
+        self.assertTrue(torch.equal(out_eager, wrapped))
 
     def test_mlp_train_step(self):
         model_ref, tokens, labels, loss_fn = self._make_mlp()
         model_test = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
         model_test.load_state_dict(model_ref.state_dict())
 
-        train_step = TrainStepModule(model_ref, loss_fn)
-        traced_result = trace_module(train_step, (tokens, labels))
+        train_step = _make_train_step(loss_fn)
+        traced = aot_function(train_step, (model_ref, tokens, labels))
 
         logits_ref = model_ref(tokens)
         loss_ref = loss_fn(logits_ref, labels)
         loss_ref.backward()
         grads_ref = [p.grad.clone() for p in model_ref.parameters()]
 
-        train_step_copy = TrainStepModule(model_test, loss_fn)
-        params_and_buffers = _get_params_and_buffers(train_step_copy)
-        wrapped = run_traced_module(traced_result, params_and_buffers, (tokens, labels))
+        wrapped = traced(model_test, tokens, labels)
         loss_tr = wrapped[0]
         grads_tr = wrapped[1:]
 
@@ -157,9 +150,8 @@ def test_mlp_multistep_bitwise(self):
         model_test = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
         model_test.load_state_dict(model_ref.state_dict())
 
-        train_step_ref = TrainStepModule(model_ref, loss_fn)
-        train_step_copy = TrainStepModule(model_test, loss_fn)
-        traced_result = trace_module(train_step_ref, (tokens, labels))
+        train_step = _make_train_step(loss_fn)
+        traced = aot_function(train_step, (model_ref, tokens, labels))
 
         opt_ref = torch.optim.Adam(model_ref.parameters(), lr=self.LR)
         opt_copy = torch.optim.Adam(model_test.parameters(), lr=self.LR)
@@ -172,10 +164,7 @@ def test_mlp_multistep_bitwise(self):
             opt_ref.step()
             opt_ref.zero_grad()
 
-            params_and_buffers = _get_params_and_buffers(train_step_copy)
-            wrapped = run_traced_module(
-                traced_result, params_and_buffers, (tokens, labels)
-            )
+            wrapped = traced(model_test, tokens, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
@@ -189,6 +178,102 @@ def test_mlp_multistep_bitwise(self):
             for gr, gt in zip(grads_ref, grads_tr, strict=True):
                 self.assertTrue(torch.equal(gr, gt), f"Step {step}: grad mismatch")
 
+    def test_module_not_first_arg(self):
+        """Module in the middle of args: fn(tokens, model, labels)."""
+        model_ref, tokens, labels, loss_fn = self._make_mlp()
+        model_test = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
+        model_test.load_state_dict(model_ref.state_dict())
+
+        def train_step(tokens, model, labels):
+            logits = model(tokens)
+            loss = get_loss(logits, labels)
+            params = list(model.parameters())
+            grads = torch.autograd.grad(loss, params)
+            return [loss] + list(grads)
+
+        traced = aot_function(train_step, (tokens, model_ref, labels))
+
+        logits_ref = model_ref(tokens)
+        loss_ref = get_loss(logits_ref, labels)
+        loss_ref.backward()
+        grads_ref = [p.grad.clone() for p in model_ref.parameters()]
+
+        wrapped = traced(tokens, model_test, labels)
+        loss_tr = wrapped[0]
+        grads_tr = wrapped[1:]
+
+        self.assertTrue(torch.equal(loss_ref, loss_tr))
+        for gr, gt in zip(grads_ref, grads_tr, strict=True):
+            self.assertTrue(torch.equal(gr, gt))
+
+    def test_multiple_modules(self):
+        """Two modules interleaved with tensors: fn(model_a, x, model_b)."""
+        torch.manual_seed(42)
+        dim = 64
+        model_a = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
+        model_b = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
+        model_a_copy = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
+        model_b_copy = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
+        model_a_copy.load_state_dict(model_a.state_dict())
+        model_b_copy.load_state_dict(model_b.state_dict())
+
+        x = torch.randn(4, dim, device=self.DEVICE, dtype=self.DTYPE)
+
+        def two_model_step(model_a, x, model_b):
+            out = model_b(torch.relu(model_a(x)))
+            loss = out.sum()
+            params = list(model_a.parameters()) + list(model_b.parameters())
+            grads = torch.autograd.grad(loss, params)
+            return [loss] + list(grads)
+
+        traced = aot_function(two_model_step, (model_a, x, model_b))
+
+        # Eager reference.
+        out_ref = model_b(torch.relu(model_a(x)))
+        loss_ref = out_ref.sum()
+        loss_ref.backward()
+        grads_ref = [p.grad.clone() for p in model_a.parameters()] + [
+            p.grad.clone() for p in model_b.parameters()
+        ]
+
+        wrapped = traced(model_a_copy, x, model_b_copy)
+        loss_tr = wrapped[0]
+        grads_tr = wrapped[1:]
+
+        self.assertTrue(torch.equal(loss_ref, loss_tr))
+        for gr, gt in zip(grads_ref, grads_tr, strict=True):
+            self.assertTrue(torch.equal(gr, gt))
+
+    def test_mismatched_module_raises(self):
+        """Executing with a module that has different params than at trace time raises."""
+        model, tokens, labels, loss_fn = self._make_mlp()
+
+        def forward(model, tokens):
+            return model(tokens)
+
+        traced = aot_function(forward, (model, tokens))
+
+        # A model with a different architecture (extra layer → different FQNs).
+        different_model = nn.Sequential(
+            nn.Embedding(256, 64),
+            nn.Linear(64, 256),
+        ).to(device=self.DEVICE, dtype=self.DTYPE)
+
+        with self.assertRaises(ValueError, msg="different parameter/buffer names"):
+            traced(different_model, tokens)
+
+    def test_non_tensor_leaf_raises(self):
+        """Passing a callable leaf in args raises (should be in closure instead)."""
+
+        def fn(model, x, loss_fn):
+            return loss_fn(model(x))
+
+        model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
+        tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
+
+        with self.assertRaises(ValueError, msg="all pytree leaves"):
+            aot_function(fn, (model, tokens, lambda x: x.sum()))
+
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
 class TestTraceDTensor(unittest.TestCase):
@@ -238,15 +323,14 @@ def test_dtensor_forward(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         tokens_dt = DTensor.from_local(tokens, mesh, [Replicate()])
 
-        traced_result = trace_module(model, (tokens_dt,))
+        traced = aot_function(model, (tokens_dt,))
         has_subclass = any(
-            layout.meta is not None for layout in traced_result.input_subclass_layouts
+            layout.meta is not None for layout in traced.input_subclass_layouts
         )
         self.assertTrue(has_subclass)
 
         out_eager = model(tokens_dt)
-        params_and_buffers = _get_params_and_buffers(model)
-        wrapped = run_traced_module(traced_result, params_and_buffers, (tokens_dt,))
+        wrapped = traced(model, tokens_dt)
         self.assertTrue(torch.equal(out_eager.full_tensor(), wrapped[0].full_tensor()))
 
     def test_dtensor_train_step(self):
@@ -267,19 +351,15 @@ def test_dtensor_train_step(self):
         tokens_dt = DTensor.from_local(tokens, mesh, [Replicate()])
         labels_dt = DTensor.from_local(labels, mesh, [Replicate()])
 
-        train_step = TrainStepModule(model_ref, get_loss)
-        traced_result = trace_module(train_step, (tokens_dt, labels_dt))
+        train_step = _make_train_step(get_loss)
+        traced = aot_function(train_step, (model_ref, tokens_dt, labels_dt))
 
         logits_ref = model_ref(tokens_dt)
         loss_ref = get_loss(logits_ref, labels_dt)
         loss_ref.backward()
         grads_ref = [p.grad.clone() for p in model_ref.parameters()]
 
-        train_step_copy = TrainStepModule(model_test, get_loss)
-        params_and_buffers = _get_params_and_buffers(train_step_copy)
-        wrapped = run_traced_module(
-            traced_result, params_and_buffers, (tokens_dt, labels_dt)
-        )
+        wrapped = traced(model_test, tokens_dt, labels_dt)
         loss_tr = wrapped[0]
         grads_tr = wrapped[1:]
 
@@ -301,15 +381,15 @@ def setUp(self):
     def test_backward_nodes_have_seq_nr(self):
         """Verify that backward FX nodes get seq_nr metadata via the patched engine."""
         model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
-        train_step = TrainStepModule(model, get_loss)
+        train_step = _make_train_step(get_loss)
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced_result = trace_module(train_step, (tokens, labels))
+        traced = aot_function(train_step, (model, tokens, labels))
 
         # Collect seq_nr values from all call_function nodes
         seq_nrs = []
-        for node in traced_result.gm.graph.nodes:
+        for node in traced.gm.graph.nodes:
             if node.op == "call_function" and "seq_nr" in node.meta:
                 seq_nrs.append(node.meta["seq_nr"])
 
@@ -329,17 +409,13 @@ def test_copy_fwd_metadata_propagates_custom(self):
         """Verify _copy_fwd_metadata_to_bw_nodes copies custom metadata to bwd nodes."""
         model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
 
-        # Use annotate to set custom metadata on forward nodes, then trace
-        # with backward to verify it propagates
-        train_step = TrainStepModule(model, get_loss)
+        train_step = _make_train_step(get_loss)
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced_result = trace_module(train_step, (tokens, labels))
-        gm = traced_result.gm
+        traced = aot_function(train_step, (model, tokens, labels))
+        gm = traced.gm
 
-        # Manually set custom metadata on the first fwd node for each seq_nr
-        # to test that _copy_fwd_metadata_to_bw_nodes works
         seq_nr_first: dict[int, torch.fx.Node] = {}
         for node in gm.graph.nodes:
             if node.op == "call_function" and "seq_nr" in node.meta:
@@ -348,16 +424,13 @@ def test_copy_fwd_metadata_propagates_custom(self):
                     seq_nr_first[seq_nr] = node
                     node.meta["custom"] = {"test_key": "test_value"}
 
-        # Run the copy pass again
         _copy_fwd_metadata_to_bw_nodes(gm)
 
-        # Check that bwd nodes with shared seq_nr got the custom metadata
         for node in gm.graph.nodes:
             if node.op != "call_function" or "seq_nr" not in node.meta:
                 continue
             seq_nr = node.meta["seq_nr"]
             if node is not seq_nr_first.get(seq_nr):
-                # This is a backward node
                 custom = node.meta.get("custom")
                 self.assertIsNotNone(
                     custom,
@@ -373,11 +446,9 @@ def test_patch_engine_restores_original(self):
         orig_fn = torch.autograd.graph._engine_run_backward
 
         with _patch_engine_run_backward():
-            # Inside the context, it should be patched
             self.assertIsNot(torch.autograd.graph._engine_run_backward, orig_fn)
             self.assertIsNot(torch.autograd._engine_run_backward, orig_fn)
 
-        # After the context, it should be restored
         self.assertIs(torch.autograd.graph._engine_run_backward, orig_fn)
         self.assertIs(torch.autograd._engine_run_backward, orig_fn)
 
@@ -409,20 +480,24 @@ def _run_bitwise_test(
         num_steps=5,
         lr=1e-3,
     ):
-        train_step_ref = TrainStepModule(model_ref, get_loss)
+        train_step = _make_train_step(get_loss)
 
-        with annotate_flex_attention_for_regional_inductor() if use_regional_inductor else contextlib.nullcontext():
-            traced_result = trace_module(train_step_ref, (*fwd_args, labels))
+        with (
+            annotate_flex_attention_for_regional_inductor()
+            if use_regional_inductor
+            else contextlib.nullcontext()
+        ):
+            traced = aot_function(train_step, (model_ref, *fwd_args, labels))
 
         if check_collective_ops:
             ag = sum(
                 1
-                for n in traced_result.gm.graph.nodes
+                for n in traced.gm.graph.nodes
                 if "all_gather_into_tensor" in str(n.target)
             )
             rs = sum(
                 1
-                for n in traced_result.gm.graph.nodes
+                for n in traced.gm.graph.nodes
                 if "reduce_scatter_tensor" in str(n.target)
             )
             self.assertTrue(
@@ -431,7 +506,7 @@ def _run_bitwise_test(
             )
 
         if use_regional_inductor:
-            _apply_regional_inductor(traced_result)
+            _apply_regional_inductor(traced)
 
         opt_ref = torch.optim.Adam(model_ref.parameters(), lr=lr)
         opt_copy = torch.optim.Adam(model_test.parameters(), lr=lr)
@@ -444,11 +519,7 @@ def _run_bitwise_test(
             opt_ref.step()
             opt_ref.zero_grad()
 
-            train_step_copy = TrainStepModule(model_test, get_loss)
-            params_and_buffers = _get_params_and_buffers(train_step_copy)
-            wrapped = run_traced_module(
-                traced_result, params_and_buffers, (*fwd_args, labels)
-            )
+            wrapped = traced(model_test, *fwd_args, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
@@ -618,11 +689,11 @@ def test_flex_attention_annotations(self):
             "sliding_window_mask": sliding_window_mask,
         }
         with annotate_flex_attention_for_regional_inductor():
-            traced_result = trace_module(model, (tokens, attn_masks))
+            traced = aot_function(model, (tokens, attn_masks))
 
         flex_nodes = [
             n
-            for n in traced_result.gm.graph.nodes
+            for n in traced.gm.graph.nodes
             if "flex_attention" in str(n.target) and "backward" not in str(n.target)
         ]
         self.assertGreater(len(flex_nodes), 0, "No FlexAttentionHOP nodes found")
@@ -704,20 +775,22 @@ def _run_fsdp_model_test(
         else:
             fwd_args = (tokens,)
 
-        train_step_ref = TrainStepModule(model_ref, get_loss)
+        train_step = _make_train_step(get_loss)
 
-        with annotate_flex_attention_for_regional_inductor() if use_regional_inductor else contextlib.nullcontext():
-            traced_result = trace_module(train_step_ref, (*fwd_args, labels))
+        with (
+            annotate_flex_attention_for_regional_inductor()
+            if use_regional_inductor
+            else contextlib.nullcontext()
+        ):
+            traced = aot_function(train_step, (model_ref, *fwd_args, labels))
 
         ag = sum(
             1
-            for n in traced_result.gm.graph.nodes
+            for n in traced.gm.graph.nodes
             if "all_gather_into_tensor" in str(n.target)
         )
         rs = sum(
-            1
-            for n in traced_result.gm.graph.nodes
-            if "reduce_scatter_tensor" in str(n.target)
+            1 for n in traced.gm.graph.nodes if "reduce_scatter_tensor" in str(n.target)
         )
         self.assertTrue(
             ag > 0 and rs > 0,
@@ -725,7 +798,7 @@ def _run_fsdp_model_test(
         )
 
         if use_regional_inductor:
-            _apply_regional_inductor(traced_result)
+            _apply_regional_inductor(traced)
 
         opt_ref = torch.optim.Adam(model_ref.parameters(), lr=1e-3)
         opt_copy = torch.optim.Adam(model_test.parameters(), lr=1e-3)
@@ -738,11 +811,7 @@ def _run_fsdp_model_test(
             opt_ref.step()
             opt_ref.zero_grad()
 
-            train_step_copy = TrainStepModule(model_test, get_loss)
-            params_and_buffers = _get_params_and_buffers(train_step_copy)
-            wrapped = run_traced_module(
-                traced_result, params_and_buffers, (*fwd_args, labels)
-            )
+            wrapped = traced(model_test, *fwd_args, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index 5cc4d0ad54..926196170f 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -13,31 +13,33 @@
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
 from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
-    run_traced_module,
-    trace_module,
+    aot_function,
     TracedResult,
 )
 from torchtitan.trainer import Trainer
 
 
-class FwdBwdStepModule(nn.Module):
-    """Wraps model + loss_fn + autograd.grad into a single traceable forward.
+def _make_fwd_bwd_step(loss_fn):
+    """Return a plain function that traces the entire fwd+loss+bwd step.
 
-    This allows make_fx to trace through the entire fwd+loss+bwd as one graph.
-    """
+    ``loss_fn`` is captured in the closure so it is not a graph input.
 
-    def __init__(self, model, loss_fn):
-        super().__init__()
-        self.model = model
-        self.loss_fn = loss_fn
+    TODO: investigate how loss_fn interacts with non-strict trace. Currently
+    it is captured as a closure variable, but non-strict tracing may need it
+    registered via pytree.register_constant or passed differently.
+    """
 
-    def forward(self, inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs):
-        pred = self.model(inputs, **extra_inputs, **extra_kwargs)
-        loss = self.loss_fn(pred, labels) / global_valid_tokens
-        params = [p for p in self.model.parameters() if p.requires_grad]
+    def fwd_bwd_step(
+        model, inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs
+    ):
+        pred = model(inputs, **extra_inputs, **extra_kwargs)
+        loss = loss_fn(pred, labels) / global_valid_tokens
+        params = [p for p in model.parameters() if p.requires_grad]
         grads = torch.autograd.grad(loss, params)
         return [loss] + list(grads)
 
+    return fwd_bwd_step
+
 
 class GraphTrainer(Trainer):
     @dataclass(kw_only=True, slots=True)
@@ -55,7 +57,6 @@ def __init__(self, config):
             )
 
         # Lazy state for aot_fx_trace mode
-        self._fwd_bwd_step_module: FwdBwdStepModule | None = None
         self._traced_step: TracedResult | None = None
 
     def forward_backward_step(
@@ -101,23 +102,28 @@ def _make_fx_forward_backward_step(
         extra_kwargs: dict[str, Any],
     ) -> torch.Tensor:
         if self._traced_step is None:
-            self._fwd_bwd_step_module = FwdBwdStepModule(model, self.loss_fn)
-
+            fwd_bwd_fn = _make_fwd_bwd_step(self.loss_fn)
             with self.train_context(), self.maybe_enable_amp:
-                self._traced_step = trace_module(
-                    self._fwd_bwd_step_module,
-                    (inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs),
+                self._traced_step = aot_function(
+                    fwd_bwd_fn,
+                    (
+                        model,
+                        inputs,
+                        labels,
+                        global_valid_tokens,
+                        extra_inputs,
+                        extra_kwargs,
+                    ),
                 )
 
-        params_and_buffers = {
-            **dict(self._fwd_bwd_step_module.named_parameters(remove_duplicate=False)),
-            **dict(self._fwd_bwd_step_module.named_buffers(remove_duplicate=False)),
-        }
         with self.train_context(), self.maybe_enable_amp:
-            outputs = run_traced_module(
-                self._traced_step,
-                params_and_buffers,
-                (inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs),
+            outputs = self._traced_step(
+                model,
+                inputs,
+                labels,
+                global_valid_tokens,
+                extra_inputs,
+                extra_kwargs,
             )
         loss = outputs[0]
         grads = outputs[1:]

From da526a8e89aead992ca4ae63be80ecca50dbfb8e Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 30 Mar 2026 11:56:43 -0700
Subject: [PATCH 02/42] Update on "[graph_trainer] Replace
 trace_module/run_traced_module with aot_function API"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Authored-by: Claude

Redesign the graph trainer's tracing API based on the aot_function design doc.

Key changes:

make_fx_tracer.py:
- Rename trace_module -> aot_function. Takes any callable (not just nn.Module)
  with nn.Module instances auto-detected in args and their params/buffers
  lifted as graph inputs. When fn is an nn.Module, it is prepended to args
  and type(fn).__call__ is used as the callable.
- Delete run_traced_module. TracedResult is now directly callable — pass the
  same positional args (with live modules) to execute the graph. Fresh params
  are read from the modules automatically on each call.
- Store and restore output pytree spec so TracedResult.__call__ returns the
  same pytree structure as the original function (e.g. single tensor, list,
  tuple, dict), not a flat list.
- Add _ModuleParamsMeta with FQN storage. Parameter FQNs are recorded at trace
  time and validated at execute time to catch module structure mismatches.
- Add _collect_module_params helper for multi-module param extraction.
- Install TracingContext before make_fx so invoke_subgraph deduplication works.
- Validate that all pytree leaves in args are tensors or primitives
  (int/float/bool/str). Non-primitive values (callables, custom objects) must
  be captured in fn's closure or registered via pytree.register_constant /
  register_pytree_node.

trainer.py:
- Replace FwdBwdStepModule (nn.Module wrapper that only existed because the
  old trace_module required nn.Module as fn) with _make_fwd_bwd_step, a plain
  function factory. The model is now passed as an arg, loss_fn is captured in
  the closure.
- Remove manual params_and_buffers dict construction — TracedResult.__call__
  reads fresh params from the live module automatically.
- Add TODO for investigating loss_fn interaction with non-strict trace.

test_trace_module.py:
- Replace TrainStepModule with _make_train_step plain function factory.
- Remove _get_params_and_buffers helper (no longer needed).
- Update all callsites: trace_module -> aot_function, run_traced_module ->
  direct TracedResult.__call__.
- Register BlockMask as pytree node at module level so flex_attention tests
  pass the leaf validation.
- Add test_module_not_first_arg: module at position 1 in args.
- Add test_multiple_modules: two nn.Modules interleaved with a tensor.
- Add test_mismatched_module_raises: FQN validation catches wrong module.
- Add test_non_tensor_leaf_raises: callable leaf in args raises ValueError.

All 7 model tests pass (llama3, llama4, qwen3, qwen3_moe, deepseek_v3,
gpt_oss, flex_attention_annotations).

[ghstack-poisoned]
---
 .../graph_trainer/tests/test_trace_module.py     | 16 ++++++++--------
 torchtitan/experiments/graph_trainer/trainer.py  |  4 ++--
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 1bd8ed8780..bd4dbf3921 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -37,7 +37,7 @@ def get_loss(logits, labels):
     )
 
 
-def _make_train_step(loss_fn):
+def make_train_step(loss_fn):
     """Return a plain function for aot_function tracing.  loss_fn is captured in closure."""
 
     def train_step(model, *args):
@@ -129,7 +129,7 @@ def test_mlp_train_step(self):
         model_test = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
         model_test.load_state_dict(model_ref.state_dict())
 
-        train_step = _make_train_step(loss_fn)
+        train_step = make_train_step(loss_fn)
         traced = aot_function(train_step, (model_ref, tokens, labels))
 
         logits_ref = model_ref(tokens)
@@ -150,7 +150,7 @@ def test_mlp_multistep_bitwise(self):
         model_test = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
         model_test.load_state_dict(model_ref.state_dict())
 
-        train_step = _make_train_step(loss_fn)
+        train_step = make_train_step(loss_fn)
         traced = aot_function(train_step, (model_ref, tokens, labels))
 
         opt_ref = torch.optim.Adam(model_ref.parameters(), lr=self.LR)
@@ -351,7 +351,7 @@ def test_dtensor_train_step(self):
         tokens_dt = DTensor.from_local(tokens, mesh, [Replicate()])
         labels_dt = DTensor.from_local(labels, mesh, [Replicate()])
 
-        train_step = _make_train_step(get_loss)
+        train_step = make_train_step(get_loss)
         traced = aot_function(train_step, (model_ref, tokens_dt, labels_dt))
 
         logits_ref = model_ref(tokens_dt)
@@ -381,7 +381,7 @@ def setUp(self):
     def test_backward_nodes_have_seq_nr(self):
         """Verify that backward FX nodes get seq_nr metadata via the patched engine."""
         model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
-        train_step = _make_train_step(get_loss)
+        train_step = make_train_step(get_loss)
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
@@ -409,7 +409,7 @@ def test_copy_fwd_metadata_propagates_custom(self):
         """Verify _copy_fwd_metadata_to_bw_nodes copies custom metadata to bwd nodes."""
         model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
 
-        train_step = _make_train_step(get_loss)
+        train_step = make_train_step(get_loss)
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
@@ -480,7 +480,7 @@ def _run_bitwise_test(
         num_steps=5,
         lr=1e-3,
     ):
-        train_step = _make_train_step(get_loss)
+        train_step = make_train_step(get_loss)
 
         with (
             annotate_flex_attention_for_regional_inductor()
@@ -775,7 +775,7 @@ def _run_fsdp_model_test(
         else:
             fwd_args = (tokens,)
 
-        train_step = _make_train_step(get_loss)
+        train_step = make_train_step(get_loss)
 
         with (
             annotate_flex_attention_for_regional_inductor()
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index 926196170f..f7835f7011 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -19,7 +19,7 @@
 from torchtitan.trainer import Trainer
 
 
-def _make_fwd_bwd_step(loss_fn):
+def make_fwd_bwd_step(loss_fn):
     """Return a plain function that traces the entire fwd+loss+bwd step.
 
     ``loss_fn`` is captured in the closure so it is not a graph input.
@@ -102,7 +102,7 @@ def _make_fx_forward_backward_step(
         extra_kwargs: dict[str, Any],
     ) -> torch.Tensor:
         if self._traced_step is None:
-            fwd_bwd_fn = _make_fwd_bwd_step(self.loss_fn)
+            fwd_bwd_fn = make_fwd_bwd_step(self.loss_fn)
             with self.train_context(), self.maybe_enable_amp:
                 self._traced_step = aot_function(
                     fwd_bwd_fn,

From 3394aaf54b65854e0e1239cca3139860f70d6ac5 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 30 Mar 2026 20:08:28 -0700
Subject: [PATCH 03/42] Update on "[graph_trainer] Replace
 trace_module/run_traced_module with aot_function API"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Authored-by: Claude

Redesign the graph trainer's tracing API based on the aot_function design doc.

Key changes:

make_fx_tracer.py:
- Rename trace_module -> aot_function. Takes any callable (not just nn.Module)
  with nn.Module instances auto-detected in args and their params/buffers
  lifted as graph inputs. When fn is an nn.Module, it is prepended to args
  and type(fn).__call__ is used as the callable.
- Delete run_traced_module. TracedResult is now directly callable — pass the
  same positional args (with live modules) to execute the graph. Fresh params
  are read from the modules automatically on each call.
- Store and restore output pytree spec so TracedResult.__call__ returns the
  same pytree structure as the original function (e.g. single tensor, list,
  tuple, dict), not a flat list.
- Add _ModuleParamsMeta with FQN storage. Parameter FQNs are recorded at trace
  time and validated at execute time to catch module structure mismatches.
- Add _collect_module_params helper for multi-module param extraction.
- Install TracingContext before make_fx so invoke_subgraph deduplication works.
- Validate that all pytree leaves in args are tensors or primitives
  (int/float/bool/str). Non-primitive values (callables, custom objects) must
  be captured in fn's closure or registered via pytree.register_constant /
  register_pytree_node.

trainer.py:
- Replace FwdBwdStepModule (nn.Module wrapper that only existed because the
  old trace_module required nn.Module as fn) with _make_fwd_bwd_step, a plain
  function factory. The model is now passed as an arg, loss_fn is captured in
  the closure.
- Remove manual params_and_buffers dict construction — TracedResult.__call__
  reads fresh params from the live module automatically.
- Add TODO for investigating loss_fn interaction with non-strict trace.

test_trace_module.py:
- Replace TrainStepModule with _make_train_step plain function factory.
- Remove _get_params_and_buffers helper (no longer needed).
- Update all callsites: trace_module -> aot_function, run_traced_module ->
  direct TracedResult.__call__.
- Register BlockMask as pytree node at module level so flex_attention tests
  pass the leaf validation.
- Add test_module_not_first_arg: module at position 1 in args.
- Add test_multiple_modules: two nn.Modules interleaved with a tensor.
- Add test_mismatched_module_raises: FQN validation catches wrong module.
- Add test_non_tensor_leaf_raises: callable leaf in args raises ValueError.

All 7 model tests pass (llama3, llama4, qwen3, qwen3_moe, deepseek_v3,
gpt_oss, flex_attention_annotations).

[ghstack-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 296 +++++++++---------
 .../graph_trainer/tests/test_trace_module.py  |  76 +----
 2 files changed, 159 insertions(+), 213 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index b8e8f87b60..9a9daa6b44 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -4,7 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import contextlib
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -57,14 +56,6 @@ class SubclassLayout:
     meta: SubclassMeta | None
 
 
-@dataclass
-class _ModuleParamsMeta:
-    """Per-module parameter metadata captured during tracing."""
-
-    fqns: list[str]
-    num_params: int
-
-
 def _unwrap_subclass(t: torch.Tensor) -> tuple[list[torch.Tensor], SubclassMeta | None]:
     if not is_traceable_wrapper_subclass(t):
         return [t], None
@@ -105,19 +96,46 @@ def _wrap_to_subclass(
     )
 
 
-def _wrap_to_subclasses(
-    flat_tensors: tuple[torch.Tensor, ...] | list[torch.Tensor],
-    layouts: list[SubclassLayout],
-) -> list[torch.Tensor]:
+def _unwrap_subclasses(
+    args: list,
+) -> tuple[list, dict[int, SubclassLayout]]:
+    """Unwrap tensor subclasses into plain tensors.
+
+    Returns the flattened plain tensors and a dict mapping original arg index
+    to its SubclassLayout.  Plain tensors have no entry.
+    """
+    flat: list = []
+    layouts: dict[int, SubclassLayout] = {}
+    for i, arg in enumerate(args):
+        if isinstance(arg, torch.Tensor) and is_traceable_wrapper_subclass(arg):
+            inner_tensors, meta = _unwrap_subclass(arg)
+            layouts[i] = SubclassLayout(len(inner_tensors), meta)
+            flat.extend(inner_tensors)
+        else:
+            flat.append(arg)
+    return flat, layouts
+
+
+def _wrap_subclasses(
+    flat_tensors: tuple | list,
+    num_args: int,
+    layouts: dict[int, SubclassLayout],
+) -> list:
+    """Rewrap plain tensors back into their original subclass types.
+
+    Positions not in ``layouts`` are plain tensors (taken one-to-one).
+    """
     wrapped = []
     idx = 0
-    for layout in layouts:
-        tensors = flat_tensors[idx : idx + layout.num_tensors]
-        idx += layout.num_tensors
-        if layout.meta is None:
-            wrapped.append(tensors[0])
-        else:
+    for i in range(num_args):
+        if i in layouts:
+            layout = layouts[i]
+            tensors = flat_tensors[idx : idx + layout.num_tensors]
+            idx += layout.num_tensors
             wrapped.append(_wrap_to_subclass(list(tensors), layout.meta))
+        else:
+            wrapped.append(flat_tensors[idx])
+            idx += 1
     return wrapped
 
 
@@ -192,10 +210,6 @@ def _patch_engine_run_backward() -> Generator[None, None, None]:
     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``.
     """
@@ -250,30 +264,11 @@ def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
                 node.meta["nn_module_stack"] = nn_module_stack.copy()
 
 
-def _collect_module_params(
-    args: tuple, module_indices: list[int]
-) -> tuple[list[_ModuleParamsMeta], list[torch.Tensor]]:
-    """Extract flattened params/buffers for each nn.Module arg."""
-    per_module: list[_ModuleParamsMeta] = []
-    all_flat: list[torch.Tensor] = []
-    for i in module_indices:
-        mod = args[i]
-        params_dict = {
-            **dict(mod.named_parameters(remove_duplicate=False)),
-            **dict(mod.named_buffers(remove_duplicate=False)),
-        }
-        fqns = list(params_dict.keys())
-        flat = list(params_dict.values())
-        per_module.append(_ModuleParamsMeta(fqns=fqns, num_params=len(flat)))
-        all_flat.extend(flat)
-    return per_module, all_flat
-
-
 class TracedResult:
     """Holds the traced graph and metadata needed to execute it.
 
     Returned by :func:`aot_function`.  Call the instance directly to execute
-    the traced graph with fresh parameters read from the live modules::
+    the traced graph with fresh parameters read from the live module::
 
         traced = aot_function(train_step, (model, tokens, labels))
         result = traced(model, tokens, labels)
@@ -282,56 +277,78 @@ class TracedResult:
     def __init__(
         self,
         gm: torch.fx.GraphModule,
-        module_indices: list[int],
-        per_module_params: list[_ModuleParamsMeta],
-        input_subclass_layouts: list[SubclassLayout],
-        output_subclass_layouts: list[SubclassLayout],
+        module_index: int,
+        param_fqns: list[str],
+        num_params: int,
+        num_flat_inputs: int,
+        input_subclass_layouts: dict[int, SubclassLayout],
+        num_flat_outputs: int,
+        output_subclass_layouts: dict[int, SubclassLayout],
         output_spec: pytree.TreeSpec,
     ) -> None:
+        """
+        Args:
+            gm: The traced FX graph (a pure function of flat tensors).
+            module_index: Position of the ``nn.Module`` in the original ``args``.
+            param_fqns: Fully qualified names of the module's parameters and
+                buffers, recorded at trace time for validation.
+            num_params: Number of parameters + buffers (flat count).
+            num_flat_inputs: Number of original args (before subclass unwrapping).
+            input_subclass_layouts: Maps arg positions that are tensor subclasses
+                to their unwrap/rewrap metadata.  Plain tensors have no entry.
+            num_flat_outputs: Number of original outputs (before subclass unwrapping).
+            output_subclass_layouts: Maps output positions that are tensor subclasses
+                to their rewrap metadata.  Plain tensors have no entry.
+            output_spec: Pytree spec of the original function's return value,
+                used to reconstruct the output structure.
+        """
         self.gm = gm
-        self.module_indices = module_indices
-        self.per_module_params = per_module_params
+        self.module_index = module_index
+        self.param_fqns = param_fqns
+        self.num_params = num_params
+        self.num_flat_inputs = num_flat_inputs
         self.input_subclass_layouts = input_subclass_layouts
+        self.num_flat_outputs = num_flat_outputs
         self.output_subclass_layouts = output_subclass_layouts
         self.output_spec = output_spec
+        self._validated = False
+
+    def __call__(self, *args: Any) -> Any:
+        """Execute the traced graph, reading fresh params from the module in ``args``.
 
-    def __call__(self, *args: Any) -> list[torch.Tensor]:
-        """Execute the traced graph, reading fresh params from modules in ``args``."""
-        module_indices_set = set(self.module_indices)
-
-        # Read current params from live modules.
-        all_params_flat: list[torch.Tensor] = []
-        for i, pmp in zip(self.module_indices, self.per_module_params, strict=True):
-            mod = args[i]
-            params_dict = {
-                **dict(mod.named_parameters(remove_duplicate=False)),
-                **dict(mod.named_buffers(remove_duplicate=False)),
-            }
+        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[self.module_index]
+        params_dict = {
+            **dict(mod.named_parameters(remove_duplicate=False)),
+            **dict(mod.named_buffers(remove_duplicate=False)),
+        }
+        if not self._validated:
             fqns = list(params_dict.keys())
-            if fqns != pmp.fqns:
+            if fqns != self.param_fqns:
                 raise ValueError(
-                    f"Module at arg position {i} has different parameter/buffer "
-                    f"names than during tracing.\n"
-                    f"  Traced: {pmp.fqns}\n"
+                    f"Module at arg position {self.module_index} has different "
+                    f"parameter/buffer names than during tracing.\n"
+                    f"  Traced: {self.param_fqns}\n"
                     f"  Got:    {fqns}"
                 )
-            flat, _ = pytree.tree_flatten(params_dict)
-            all_params_flat.extend(flat)
+            self._validated = True
+        params_flat = list(params_dict.values())
 
-        user_args = [a for i, a in enumerate(args) if i not in module_indices_set]
+        user_args = [a for i, a in enumerate(args) if i != self.module_index]
         user_args_flat, _ = pytree.tree_flatten(user_args)
 
-        all_args = all_params_flat + list(user_args_flat)
-        flat_inputs: list = []
-        for a in all_args:
-            if isinstance(a, torch.Tensor) and is_traceable_wrapper_subclass(a):
-                inner, _ = _unwrap_subclass(a)
-                flat_inputs.extend(inner)
-            else:
-                flat_inputs.append(a)
-
-        flat_outputs = self.gm(*flat_inputs)
-        wrapped = _wrap_to_subclasses(flat_outputs, self.output_subclass_layouts)
+        all_args = params_flat + list(user_args_flat)
+        flat_inputs, _ = _unwrap_subclasses(all_args)
+
+        with torch.no_grad():
+            flat_outputs = self.gm(*flat_inputs)
+        wrapped = _wrap_subclasses(
+            flat_outputs, self.num_flat_outputs, self.output_subclass_layouts
+        )
         return pytree.tree_unflatten(wrapped, self.output_spec)
 
 
@@ -341,45 +358,62 @@ def aot_function(
 ) -> TracedResult:
     """Trace ``fn(*args)`` into a flat FX graph, unwrapping tensor subclasses.
 
-    Parameters and buffers from ``nn.Module`` instances in ``args`` are lifted
-    as extra graph inputs so the returned graph is a pure function.  Tensor
-    subclasses (e.g. DTensor) are recursively unwrapped into plain tensors for
-    tracing, and the layouts needed to rewrap them are recorded in the returned
-    :class:`TracedResult`.
+    Exactly one ``nn.Module`` must be present in ``args``.  Its parameters and
+    buffers are lifted as extra graph inputs so the returned graph is a pure
+    function.  Tensor subclasses (e.g. DTensor) are recursively unwrapped into
+    plain tensors for tracing, and the layouts needed to rewrap them are
+    recorded in the returned :class:`TracedResult`.
 
     ``fn`` may be an ``nn.Module`` — in which case it is treated as though the
     caller wrote ``aot_function(lambda m, *a: m(*a), (fn,) + args)``, i.e. the
     module is prepended to ``args`` and its parameters are lifted automatically.
 
     The returned :class:`TracedResult` is directly callable — pass the same
-    positional arguments (with live modules) to execute the graph::
+    positional arguments (with the live module) to execute the graph::
 
         traced = aot_function(train_step, (model, tokens, labels))
         result = traced(model, tokens, labels)
 
     Args:
         fn: The callable (or ``nn.Module``) to trace.
-        args: The positional arguments to trace with.  ``nn.Module`` instances
-            are detected automatically and their parameters are lifted.
+        args: The positional arguments to trace with.  Must contain exactly one
+            ``nn.Module`` whose parameters will be lifted.
     """
     # When fn is an nn.Module, treat it as the first arg so its params get
     # lifted like any other module arg.
     if isinstance(fn, nn.Module):
         args = (fn,) + args
         fn = type(fn).__call__
-    module_indices = [i for i, a in enumerate(args) if isinstance(a, nn.Module)]
-    module_indices_set = set(module_indices)
-
-    per_module_params, all_params_flat = _collect_module_params(args, module_indices)
-    params_len = len(all_params_flat)
 
-    # User args: positional args that are not nn.Module instances.
-    user_args = [a for i, a in enumerate(args) if i not in module_indices_set]
+    # Find the single nn.Module in args — must be at position 0.
+    module_indices = [i for i, a in enumerate(args) if isinstance(a, nn.Module)]
+    if len(module_indices) != 1:
+        raise ValueError(
+            f"aot_function expects exactly one nn.Module in args, "
+            f"got {len(module_indices)} at positions {module_indices}."
+        )
+    if module_indices[0] != 0:
+        raise ValueError(
+            f"The nn.Module must be the first argument (position 0), "
+            f"got it at position {module_indices[0]}."
+        )
+    module_index = 0
+    mod = args[module_index]
+
+    # Extract params/buffers from the module.
+    params_dict = {
+        **dict(mod.named_parameters(remove_duplicate=False)),
+        **dict(mod.named_buffers(remove_duplicate=False)),
+    }
+    param_fqns = list(params_dict.keys())
+    params_flat = list(params_dict.values())
+    num_params = len(params_flat)
+
+    # User args: everything except the module.
+    user_args = [a for i, a in enumerate(args) if i != module_index]
     user_args_flat, user_args_spec = pytree.tree_flatten(user_args)
 
-    # Validate leaves: tensors and make_fx-safe primitives (int, float, bool,
-    # str) are allowed.  Everything else (callables, custom objects) should be
-    # registered as pytree nodes/constants or captured in fn's closure.
+    # Validate leaves: tensors and make_fx-safe primitives are allowed.
     _ALLOWED_LEAF_TYPES = (torch.Tensor, int, float, bool, str, type(None))
     for leaf in user_args_flat:
         if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
@@ -392,18 +426,9 @@ def aot_function(
             )
 
     # Combined flat input: [*params, *user_args] with subclasses unwrapped.
-    full_args = all_params_flat + list(user_args_flat)
-
-    unwrapped_args: list = []
-    input_layouts: list[SubclassLayout] = []
-    for arg in full_args:
-        if isinstance(arg, torch.Tensor) and is_traceable_wrapper_subclass(arg):
-            inner_tensors, meta = _unwrap_subclass(arg)
-            unwrapped_args.extend(inner_tensors)
-            input_layouts.append(SubclassLayout(len(inner_tensors), meta))
-        else:
-            unwrapped_args.append(arg)
-            input_layouts.append(SubclassLayout(1, None))
+    full_args = params_flat + list(user_args_flat)
+    num_full_args = len(full_args)
+    unwrapped_args, input_layouts = _unwrap_subclasses(full_args)
 
     fake_mode = FakeTensorMode(
         allow_non_fake_inputs=True,
@@ -418,55 +443,37 @@ def aot_function(
         for a in unwrapped_args
     )
 
-    output_layouts: list[SubclassLayout] = []
+    output_layouts: dict[int, SubclassLayout] = {}
+    num_flat_outputs: int = 0
     output_spec: pytree.TreeSpec | None = None
 
     def fn_with_subclass_handling(*plain_args: Any) -> list:
-        nonlocal output_layouts, output_spec
-        output_layouts = []
-
-        wrapped = _wrap_to_subclasses(plain_args, input_layouts)
-        all_params = wrapped[:params_len]
-        user_flat = wrapped[params_len:]
+        nonlocal output_layouts, output_spec, num_flat_outputs
+        output_layouts = {}
 
-        # Reconstruct per-module param dicts.
-        offset = 0
-        module_param_dicts = []
-        for pmp in per_module_params:
-            flat = all_params[offset : offset + pmp.num_params]
-            module_param_dicts.append(dict(zip(pmp.fqns, flat, strict=True)))
-            offset += pmp.num_params
+        wrapped = _wrap_subclasses(plain_args, num_full_args, input_layouts)
+        params_wrapped = wrapped[:num_params]
+        user_flat = wrapped[num_params:]
 
+        params_for_mod = dict(zip(param_fqns, params_wrapped, strict=True))
         user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
 
-        # Reconstruct the original args: module positions keep the live module
-        # (already in args), non-module positions get the traced user tensors.
+        # Reconstruct the original args: module position keeps the live module,
+        # other positions get the traced user tensors.
         rebuilt: list = list(args)
         user_idx = 0
         for i in range(len(args)):
-            if i not in module_indices_set:
+            if i != module_index:
                 rebuilt[i] = user_list[user_idx]
                 user_idx += 1
 
-        with contextlib.ExitStack() as stack:
-            for i, pmp_dict in zip(module_indices, module_param_dicts, strict=True):
-                stack.enter_context(
-                    stateless._reparametrize_module(args[i], pmp_dict)
-                )
-
+        with stateless._reparametrize_module(mod, params_for_mod):
             with _patch_engine_run_backward():
                 result = fn(*rebuilt)
 
         flat_outs, output_spec = pytree.tree_flatten(result)
-        unwrapped_outs: list = []
-        for out in flat_outs:
-            if isinstance(out, torch.Tensor) and is_traceable_wrapper_subclass(out):
-                inner, meta = _unwrap_subclass(out)
-                unwrapped_outs.extend(inner)
-                output_layouts.append(SubclassLayout(len(inner), meta))
-            else:
-                unwrapped_outs.append(out)
-                output_layouts.append(SubclassLayout(1, None))
+        num_flat_outputs = len(flat_outs)
+        unwrapped_outs, output_layouts = _unwrap_subclasses(flat_outs)
         return unwrapped_outs
 
     ctx = TracingContext(fake_mode)
@@ -486,9 +493,12 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
     assert output_spec is not None
     return TracedResult(
         gm=traced,
-        module_indices=module_indices,
-        per_module_params=per_module_params,
+        module_index=module_index,
+        param_fqns=param_fqns,
+        num_params=num_params,
+        num_flat_inputs=num_full_args,
         input_subclass_layouts=input_layouts,
+        num_flat_outputs=num_flat_outputs,
         output_subclass_layouts=output_layouts,
         output_spec=output_spec,
     )
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index bd4dbf3921..2ffbdb067d 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -178,72 +178,6 @@ def test_mlp_multistep_bitwise(self):
             for gr, gt in zip(grads_ref, grads_tr, strict=True):
                 self.assertTrue(torch.equal(gr, gt), f"Step {step}: grad mismatch")
 
-    def test_module_not_first_arg(self):
-        """Module in the middle of args: fn(tokens, model, labels)."""
-        model_ref, tokens, labels, loss_fn = self._make_mlp()
-        model_test = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
-        model_test.load_state_dict(model_ref.state_dict())
-
-        def train_step(tokens, model, labels):
-            logits = model(tokens)
-            loss = get_loss(logits, labels)
-            params = list(model.parameters())
-            grads = torch.autograd.grad(loss, params)
-            return [loss] + list(grads)
-
-        traced = aot_function(train_step, (tokens, model_ref, labels))
-
-        logits_ref = model_ref(tokens)
-        loss_ref = get_loss(logits_ref, labels)
-        loss_ref.backward()
-        grads_ref = [p.grad.clone() for p in model_ref.parameters()]
-
-        wrapped = traced(tokens, model_test, labels)
-        loss_tr = wrapped[0]
-        grads_tr = wrapped[1:]
-
-        self.assertTrue(torch.equal(loss_ref, loss_tr))
-        for gr, gt in zip(grads_ref, grads_tr, strict=True):
-            self.assertTrue(torch.equal(gr, gt))
-
-    def test_multiple_modules(self):
-        """Two modules interleaved with tensors: fn(model_a, x, model_b)."""
-        torch.manual_seed(42)
-        dim = 64
-        model_a = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
-        model_b = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
-        model_a_copy = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
-        model_b_copy = nn.Linear(dim, dim).to(device=self.DEVICE, dtype=self.DTYPE)
-        model_a_copy.load_state_dict(model_a.state_dict())
-        model_b_copy.load_state_dict(model_b.state_dict())
-
-        x = torch.randn(4, dim, device=self.DEVICE, dtype=self.DTYPE)
-
-        def two_model_step(model_a, x, model_b):
-            out = model_b(torch.relu(model_a(x)))
-            loss = out.sum()
-            params = list(model_a.parameters()) + list(model_b.parameters())
-            grads = torch.autograd.grad(loss, params)
-            return [loss] + list(grads)
-
-        traced = aot_function(two_model_step, (model_a, x, model_b))
-
-        # Eager reference.
-        out_ref = model_b(torch.relu(model_a(x)))
-        loss_ref = out_ref.sum()
-        loss_ref.backward()
-        grads_ref = [p.grad.clone() for p in model_a.parameters()] + [
-            p.grad.clone() for p in model_b.parameters()
-        ]
-
-        wrapped = traced(model_a_copy, x, model_b_copy)
-        loss_tr = wrapped[0]
-        grads_tr = wrapped[1:]
-
-        self.assertTrue(torch.equal(loss_ref, loss_tr))
-        for gr, gt in zip(grads_ref, grads_tr, strict=True):
-            self.assertTrue(torch.equal(gr, gt))
-
     def test_mismatched_module_raises(self):
         """Executing with a module that has different params than at trace time raises."""
         model, tokens, labels, loss_fn = self._make_mlp()
@@ -482,11 +416,12 @@ def _run_bitwise_test(
     ):
         train_step = make_train_step(get_loss)
 
-        with (
+        maybe_regional_inductor = (
             annotate_flex_attention_for_regional_inductor()
             if use_regional_inductor
             else contextlib.nullcontext()
-        ):
+        )
+        with maybe_regional_inductor:
             traced = aot_function(train_step, (model_ref, *fwd_args, labels))
 
         if check_collective_ops:
@@ -777,11 +712,12 @@ def _run_fsdp_model_test(
 
         train_step = make_train_step(get_loss)
 
-        with (
+        maybe_regional_inductor = (
             annotate_flex_attention_for_regional_inductor()
             if use_regional_inductor
             else contextlib.nullcontext()
-        ):
+        )
+        with maybe_regional_inductor:
             traced = aot_function(train_step, (model_ref, *fwd_args, labels))
 
         ag = sum(

From 8f276bafecb1f3d85ca59c75fad7d0ba7dd954d9 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 30 Mar 2026 20:21:42 -0700
Subject: [PATCH 04/42] Update on "[graph_trainer] Replace
 trace_module/run_traced_module with aot_function API"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Authored-by: Claude

Redesign the graph trainer's tracing API based on the aot_function design doc.

Key changes:

make_fx_tracer.py:
- Rename trace_module -> aot_function. Takes any callable (not just nn.Module)
  with nn.Module instances auto-detected in args and their params/buffers
  lifted as graph inputs. When fn is an nn.Module, it is prepended to args
  and type(fn).__call__ is used as the callable.
- Delete run_traced_module. TracedResult is now directly callable — pass the
  same positional args (with live modules) to execute the graph. Fresh params
  are read from the modules automatically on each call.
- Store and restore output pytree spec so TracedResult.__call__ returns the
  same pytree structure as the original function (e.g. single tensor, list,
  tuple, dict), not a flat list.
- Add _ModuleParamsMeta with FQN storage. Parameter FQNs are recorded at trace
  time and validated at execute time to catch module structure mismatches.
- Add _collect_module_params helper for multi-module param extraction.
- Install TracingContext before make_fx so invoke_subgraph deduplication works.
- Validate that all pytree leaves in args are tensors or primitives
  (int/float/bool/str). Non-primitive values (callables, custom objects) must
  be captured in fn's closure or registered via pytree.register_constant /
  register_pytree_node.

trainer.py:
- Replace FwdBwdStepModule (nn.Module wrapper that only existed because the
  old trace_module required nn.Module as fn) with _make_fwd_bwd_step, a plain
  function factory. The model is now passed as an arg, loss_fn is captured in
  the closure.
- Remove manual params_and_buffers dict construction — TracedResult.__call__
  reads fresh params from the live module automatically.
- Add TODO for investigating loss_fn interaction with non-strict trace.

test_trace_module.py:
- Replace TrainStepModule with _make_train_step plain function factory.
- Remove _get_params_and_buffers helper (no longer needed).
- Update all callsites: trace_module -> aot_function, run_traced_module ->
  direct TracedResult.__call__.
- Register BlockMask as pytree node at module level so flex_attention tests
  pass the leaf validation.
- Add test_module_not_first_arg: module at position 1 in args.
- Add test_multiple_modules: two nn.Modules interleaved with a tensor.
- Add test_mismatched_module_raises: FQN validation catches wrong module.
- Add test_non_tensor_leaf_raises: callable leaf in args raises ValueError.

All 7 model tests pass (llama3, llama4, qwen3, qwen3_moe, deepseek_v3,
gpt_oss, flex_attention_annotations).

[ghstack-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 39 ++++++++-----------
 1 file changed, 16 insertions(+), 23 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 9a9daa6b44..c997632d18 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -277,7 +277,6 @@ class TracedResult:
     def __init__(
         self,
         gm: torch.fx.GraphModule,
-        module_index: int,
         param_fqns: list[str],
         num_params: int,
         num_flat_inputs: int,
@@ -289,7 +288,6 @@ def __init__(
         """
         Args:
             gm: The traced FX graph (a pure function of flat tensors).
-            module_index: Position of the ``nn.Module`` in the original ``args``.
             param_fqns: Fully qualified names of the module's parameters and
                 buffers, recorded at trace time for validation.
             num_params: Number of parameters + buffers (flat count).
@@ -303,7 +301,6 @@ def __init__(
                 used to reconstruct the output structure.
         """
         self.gm = gm
-        self.module_index = module_index
         self.param_fqns = param_fqns
         self.num_params = num_params
         self.num_flat_inputs = num_flat_inputs
@@ -316,12 +313,15 @@ def __init__(
     def __call__(self, *args: Any) -> Any:
         """Execute the traced graph, reading fresh params from the module in ``args``.
 
+        The module must be the first argument (position 0), matching the
+        convention enforced by :func:`aot_function`.
+
         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[self.module_index]
+        mod = args[0]
         params_dict = {
             **dict(mod.named_parameters(remove_duplicate=False)),
             **dict(mod.named_buffers(remove_duplicate=False)),
@@ -330,15 +330,15 @@ def __call__(self, *args: Any) -> Any:
             fqns = list(params_dict.keys())
             if fqns != self.param_fqns:
                 raise ValueError(
-                    f"Module at arg position {self.module_index} has different "
-                    f"parameter/buffer names than during tracing.\n"
+                    f"Module at args[0] has different parameter/buffer "
+                    f"names than during tracing.\n"
                     f"  Traced: {self.param_fqns}\n"
                     f"  Got:    {fqns}"
                 )
             self._validated = True
         params_flat = list(params_dict.values())
 
-        user_args = [a for i, a in enumerate(args) if i != self.module_index]
+        user_args = list(args[1:])
         user_args_flat, _ = pytree.tree_flatten(user_args)
 
         all_args = params_flat + list(user_args_flat)
@@ -358,7 +358,7 @@ def aot_function(
 ) -> TracedResult:
     """Trace ``fn(*args)`` into a flat FX graph, unwrapping tensor subclasses.
 
-    Exactly one ``nn.Module`` must be present in ``args``.  Its parameters and
+    The first element of ``args`` must be an ``nn.Module``.  Its parameters and
     buffers are lifted as extra graph inputs so the returned graph is a pure
     function.  Tensor subclasses (e.g. DTensor) are recursively unwrapped into
     plain tensors for tracing, and the layouts needed to rewrap them are
@@ -376,8 +376,8 @@ def aot_function(
 
     Args:
         fn: The callable (or ``nn.Module``) to trace.
-        args: The positional arguments to trace with.  Must contain exactly one
-            ``nn.Module`` whose parameters will be lifted.
+        args: The positional arguments to trace with.  The first element must
+            be an ``nn.Module`` whose parameters will be lifted.
     """
     # When fn is an nn.Module, treat it as the first arg so its params get
     # lifted like any other module arg.
@@ -397,8 +397,7 @@ def aot_function(
             f"The nn.Module must be the first argument (position 0), "
             f"got it at position {module_indices[0]}."
         )
-    module_index = 0
-    mod = args[module_index]
+    mod = args[0]
 
     # Extract params/buffers from the module.
     params_dict = {
@@ -409,8 +408,8 @@ def aot_function(
     params_flat = list(params_dict.values())
     num_params = len(params_flat)
 
-    # User args: everything except the module.
-    user_args = [a for i, a in enumerate(args) if i != module_index]
+    # User args: everything after the module.
+    user_args = list(args[1:])
     user_args_flat, user_args_spec = pytree.tree_flatten(user_args)
 
     # Validate leaves: tensors and make_fx-safe primitives are allowed.
@@ -458,14 +457,9 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
         params_for_mod = dict(zip(param_fqns, params_wrapped, strict=True))
         user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
 
-        # Reconstruct the original args: module position keeps the live module,
-        # other positions get the traced user tensors.
-        rebuilt: list = list(args)
-        user_idx = 0
-        for i in range(len(args)):
-            if i != module_index:
-                rebuilt[i] = user_list[user_idx]
-                user_idx += 1
+        # Reconstruct the original args: module at position 0 keeps the live
+        # module, remaining positions get the traced user tensors.
+        rebuilt = [mod] + user_list
 
         with stateless._reparametrize_module(mod, params_for_mod):
             with _patch_engine_run_backward():
@@ -493,7 +487,6 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
     assert output_spec is not None
     return TracedResult(
         gm=traced,
-        module_index=module_index,
         param_fqns=param_fqns,
         num_params=num_params,
         num_flat_inputs=num_full_args,

From 4f56847ab4c2a1b8112884f8bd517266d7f884a4 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 30 Mar 2026 20:53:42 -0700
Subject: [PATCH 05/42] Update on "[graph_trainer] Replace
 trace_module/run_traced_module with aot_function API"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Authored-by: Claude

Redesign the graph trainer's tracing API based on the aot_function design doc.

Key changes:

make_fx_tracer.py:
- Rename trace_module -> aot_function. Takes any callable (not just nn.Module)
  with nn.Module instances auto-detected in args and their params/buffers
  lifted as graph inputs. When fn is an nn.Module, it is prepended to args
  and type(fn).__call__ is used as the callable.
- Delete run_traced_module. TracedResult is now directly callable — pass the
  same positional args (with live modules) to execute the graph. Fresh params
  are read from the modules automatically on each call.
- Store and restore output pytree spec so TracedResult.__call__ returns the
  same pytree structure as the original function (e.g. single tensor, list,
  tuple, dict), not a flat list.
- Add _ModuleParamsMeta with FQN storage. Parameter FQNs are recorded at trace
  time and validated at execute time to catch module structure mismatches.
- Add _collect_module_params helper for multi-module param extraction.
- Install TracingContext before make_fx so invoke_subgraph deduplication works.
- Validate that all pytree leaves in args are tensors or primitives
  (int/float/bool/str). Non-primitive values (callables, custom objects) must
  be captured in fn's closure or registered via pytree.register_constant /
  register_pytree_node.

trainer.py:
- Replace FwdBwdStepModule (nn.Module wrapper that only existed because the
  old trace_module required nn.Module as fn) with _make_fwd_bwd_step, a plain
  function factory. The model is now passed as an arg, loss_fn is captured in
  the closure.
- Remove manual params_and_buffers dict construction — TracedResult.__call__
  reads fresh params from the live module automatically.
- Add TODO for investigating loss_fn interaction with non-strict trace.

test_trace_module.py:
- Replace TrainStepModule with _make_train_step plain function factory.
- Remove _get_params_and_buffers helper (no longer needed).
- Update all callsites: trace_module -> aot_function, run_traced_module ->
  direct TracedResult.__call__.
- Register BlockMask as pytree node at module level so flex_attention tests
  pass the leaf validation.
- Add test_module_not_first_arg: module at position 1 in args.
- Add test_multiple_modules: two nn.Modules interleaved with a tensor.
- Add test_mismatched_module_raises: FQN validation catches wrong module.
- Add test_non_tensor_leaf_raises: callable leaf in args raises ValueError.

All 7 model tests pass (llama3, llama4, qwen3, qwen3_moe, deepseek_v3,
gpt_oss, flex_attention_annotations).

[ghstack-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 30 ++++++++-----------
 .../graph_trainer/tests/test_trace_module.py  | 17 +++++++++--
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index c997632d18..e79ce54a87 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -353,38 +353,34 @@ def __call__(self, *args: Any) -> Any:
 
 
 def aot_function(
-    fn: nn.Module | Callable,
+    fn: Callable,
     args: tuple,
 ) -> TracedResult:
     """Trace ``fn(*args)`` into a flat FX graph, unwrapping tensor subclasses.
 
-    The first element of ``args`` must be an ``nn.Module``.  Its parameters and
-    buffers are lifted as extra graph inputs so the returned graph is a pure
-    function.  Tensor subclasses (e.g. DTensor) are recursively unwrapped into
-    plain tensors for tracing, and the layouts needed to rewrap them are
-    recorded in the returned :class:`TracedResult`.
+    ``args[0]`` must be an ``nn.Module``.  Its parameters and buffers are
+    lifted as extra graph inputs so the returned graph is a pure function.
+    Tensor subclasses (e.g. DTensor) are recursively unwrapped into plain
+    tensors for tracing, and the layouts needed to rewrap them are recorded
+    in the returned :class:`TracedResult`.
 
-    ``fn`` may be an ``nn.Module`` — in which case it is treated as though the
-    caller wrote ``aot_function(lambda m, *a: m(*a), (fn,) + args)``, i.e. the
-    module is prepended to ``args`` and its parameters are lifted automatically.
+    ``fn`` must be a plain callable (not an ``nn.Module``).  This keeps the
+    trace and execute calling conventions identical — the same ``args`` are
+    passed at both trace time and execution time, with no hidden arg
+    prepending.  Non-tensor, non-module values like ``loss_fn`` should be
+    captured in ``fn``'s closure rather than passed as args.
 
     The returned :class:`TracedResult` is directly callable — pass the same
-    positional arguments (with the live module) to execute the graph::
+    positional arguments (with the live module first) to execute the graph::
 
         traced = aot_function(train_step, (model, tokens, labels))
         result = traced(model, tokens, labels)
 
     Args:
-        fn: The callable (or ``nn.Module``) to trace.
+        fn: The callable to trace.
         args: The positional arguments to trace with.  The first element must
             be an ``nn.Module`` whose parameters will be lifted.
     """
-    # When fn is an nn.Module, treat it as the first arg so its params get
-    # lifted like any other module arg.
-    if isinstance(fn, nn.Module):
-        args = (fn,) + args
-        fn = type(fn).__call__
-
     # Find the single nn.Module in args — must be at position 0.
     module_indices = [i for i, a in enumerate(args) if isinstance(a, nn.Module)]
     if len(module_indices) != 1:
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 2ffbdb067d..0b5d8940a2 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -119,7 +119,11 @@ def _make_mlp(self):
 
     def test_mlp_forward(self):
         model, tokens, labels, loss_fn = self._make_mlp()
-        traced = aot_function(model, (tokens,))
+
+        def forward(model, tokens):
+            return model(tokens)
+
+        traced = aot_function(forward, (model, tokens))
         out_eager = model(tokens)
         wrapped = traced(model, tokens)
         self.assertTrue(torch.equal(out_eager, wrapped))
@@ -257,7 +261,10 @@ def test_dtensor_forward(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         tokens_dt = DTensor.from_local(tokens, mesh, [Replicate()])
 
-        traced = aot_function(model, (tokens_dt,))
+        def forward(model, tokens):
+            return model(tokens)
+
+        traced = aot_function(forward, (model, tokens_dt))
         has_subclass = any(
             layout.meta is not None for layout in traced.input_subclass_layouts
         )
@@ -624,7 +631,11 @@ def test_flex_attention_annotations(self):
             "sliding_window_mask": sliding_window_mask,
         }
         with annotate_flex_attention_for_regional_inductor():
-            traced = aot_function(model, (tokens, attn_masks))
+
+            def forward(model, tokens, attn_masks):
+                return model(tokens, attn_masks)
+
+            traced = aot_function(forward, (model, tokens, attn_masks))
 
         flex_nodes = [
             n

From 155d120a5da4183a5e0b84d0d2cd7a93fad3a6b2 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 08:00:32 -0700
Subject: [PATCH 06/42] Update on "[graph_trainer] Replace
 trace_module/run_traced_module with aot_function API"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Authored-by: Claude

Redesign the graph trainer's tracing API based on the aot_function design doc.

Key changes:

make_fx_tracer.py:
- Rename trace_module -> aot_function. Takes any callable (not just nn.Module)
  with nn.Module instances auto-detected in args and their params/buffers
  lifted as graph inputs. When fn is an nn.Module, it is prepended to args
  and type(fn).__call__ is used as the callable.
- Delete run_traced_module. TracedResult is now directly callable — pass the
  same positional args (with live modules) to execute the graph. Fresh params
  are read from the modules automatically on each call.
- Store and restore output pytree spec so TracedResult.__call__ returns the
  same pytree structure as the original function (e.g. single tensor, list,
  tuple, dict), not a flat list.
- Add _ModuleParamsMeta with FQN storage. Parameter FQNs are recorded at trace
  time and validated at execute time to catch module structure mismatches.
- Add _collect_module_params helper for multi-module param extraction.
- Install TracingContext before make_fx so invoke_subgraph deduplication works.
- Validate that all pytree leaves in args are tensors or primitives
  (int/float/bool/str). Non-primitive values (callables, custom objects) must
  be captured in fn's closure or registered via pytree.register_constant /
  register_pytree_node.

trainer.py:
- Replace FwdBwdStepModule (nn.Module wrapper that only existed because the
  old trace_module required nn.Module as fn) with _make_fwd_bwd_step, a plain
  function factory. The model is now passed as an arg, loss_fn is captured in
  the closure.
- Remove manual params_and_buffers dict construction — TracedResult.__call__
  reads fresh params from the live module automatically.
- Add TODO for investigating loss_fn interaction with non-strict trace.

test_trace_module.py:
- Replace TrainStepModule with _make_train_step plain function factory.
- Remove _get_params_and_buffers helper (no longer needed).
- Update all callsites: trace_module -> aot_function, run_traced_module ->
  direct TracedResult.__call__.
- Register BlockMask as pytree node at module level so flex_attention tests
  pass the leaf validation.
- Add test_module_not_first_arg: module at position 1 in args.
- Add test_multiple_modules: two nn.Modules interleaved with a tensor.
- Add test_mismatched_module_raises: FQN validation catches wrong module.
- Add test_non_tensor_leaf_raises: callable leaf in args raises ValueError.

All 7 model tests pass (llama3, llama4, qwen3, qwen3_moe, deepseek_v3,
gpt_oss, flex_attention_annotations).

[ghstack-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 60 +++++++++----------
 .../graph_trainer/tests/test_trace_module.py  | 35 ++++++-----
 .../experiments/graph_trainer/trainer.py      |  8 +--
 3 files changed, 52 insertions(+), 51 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index e79ce54a87..b1ce10e599 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -22,6 +22,11 @@
 from torch.nn.utils import stateless
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
 
+# Tensors and make_fx-safe primitives are allowed as pytree leaves in args.
+# Everything else (callables, custom objects) should be registered as pytree
+# nodes/constants or captured in fn's closure.
+_ALLOWED_LEAF_TYPES = (torch.Tensor, int, float, bool, str, type(None))
+
 
 @contextmanager
 def _skip_nested_compile() -> Generator[None, None, None]:
@@ -267,11 +272,25 @@ def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
 class TracedResult:
     """Holds the traced graph and metadata needed to execute it.
 
-    Returned by :func:`aot_function`.  Call the instance directly to execute
+    Returned by :func:`minimal_fx_tracer`.  Call the instance directly to execute
     the traced graph with fresh parameters read from the live module::
 
-        traced = aot_function(train_step, (model, tokens, labels))
+        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
         result = traced(model, tokens, labels)
+
+    Args:
+        gm: The traced FX graph (a pure function of flat tensors).
+        param_fqns: Fully qualified names of the module's parameters and
+            buffers, recorded at trace time for validation.
+        num_params: Number of parameters + buffers (flat count).
+        num_flat_inputs: Number of original args (before subclass unwrapping).
+        input_subclass_layouts: Maps arg positions that are tensor subclasses
+            to their unwrap/rewrap metadata.  Plain tensors have no entry.
+        num_flat_outputs: Number of original outputs (before subclass unwrapping).
+        output_subclass_layouts: Maps output positions that are tensor subclasses
+            to their rewrap metadata.  Plain tensors have no entry.
+        output_spec: Pytree spec of the original function's return value,
+            used to reconstruct the output structure.
     """
 
     def __init__(
@@ -285,21 +304,6 @@ def __init__(
         output_subclass_layouts: dict[int, SubclassLayout],
         output_spec: pytree.TreeSpec,
     ) -> None:
-        """
-        Args:
-            gm: The traced FX graph (a pure function of flat tensors).
-            param_fqns: Fully qualified names of the module's parameters and
-                buffers, recorded at trace time for validation.
-            num_params: Number of parameters + buffers (flat count).
-            num_flat_inputs: Number of original args (before subclass unwrapping).
-            input_subclass_layouts: Maps arg positions that are tensor subclasses
-                to their unwrap/rewrap metadata.  Plain tensors have no entry.
-            num_flat_outputs: Number of original outputs (before subclass unwrapping).
-            output_subclass_layouts: Maps output positions that are tensor subclasses
-                to their rewrap metadata.  Plain tensors have no entry.
-            output_spec: Pytree spec of the original function's return value,
-                used to reconstruct the output structure.
-        """
         self.gm = gm
         self.param_fqns = param_fqns
         self.num_params = num_params
@@ -314,12 +318,7 @@ def __call__(self, *args: Any) -> Any:
         """Execute the traced graph, reading fresh params from the module in ``args``.
 
         The module must be the first argument (position 0), matching the
-        convention enforced by :func:`aot_function`.
-
-        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.
+        convention enforced by :func:`minimal_fx_tracer`.
         """
         mod = args[0]
         params_dict = {
@@ -344,15 +343,14 @@ def __call__(self, *args: Any) -> Any:
         all_args = params_flat + list(user_args_flat)
         flat_inputs, _ = _unwrap_subclasses(all_args)
 
-        with torch.no_grad():
-            flat_outputs = self.gm(*flat_inputs)
+        flat_outputs = self.gm(*flat_inputs)
         wrapped = _wrap_subclasses(
             flat_outputs, self.num_flat_outputs, self.output_subclass_layouts
         )
         return pytree.tree_unflatten(wrapped, self.output_spec)
 
 
-def aot_function(
+def minimal_fx_tracer(
     fn: Callable,
     args: tuple,
 ) -> TracedResult:
@@ -373,7 +371,7 @@ def aot_function(
     The returned :class:`TracedResult` is directly callable — pass the same
     positional arguments (with the live module first) to execute the graph::
 
-        traced = aot_function(train_step, (model, tokens, labels))
+        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
         result = traced(model, tokens, labels)
 
     Args:
@@ -385,7 +383,7 @@ def aot_function(
     module_indices = [i for i, a in enumerate(args) if isinstance(a, nn.Module)]
     if len(module_indices) != 1:
         raise ValueError(
-            f"aot_function expects exactly one nn.Module in args, "
+            f"minimal_fx_tracer expects exactly one nn.Module in args, "
             f"got {len(module_indices)} at positions {module_indices}."
         )
     if module_indices[0] != 0:
@@ -408,12 +406,11 @@ def aot_function(
     user_args = list(args[1:])
     user_args_flat, user_args_spec = pytree.tree_flatten(user_args)
 
-    # Validate leaves: tensors and make_fx-safe primitives are allowed.
-    _ALLOWED_LEAF_TYPES = (torch.Tensor, int, float, bool, str, type(None))
+    # Validate leaves.
     for leaf in user_args_flat:
         if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
             raise ValueError(
-                f"aot_function requires all pytree leaves in args to be tensors "
+                f"minimal_fx_tracer requires all pytree leaves in args to be tensors "
                 f"or primitives (int/float/bool/str), got {type(leaf).__name__}. "
                 f"Non-primitive values should either be registered as pytree "
                 f"nodes (register_pytree_node) or constants "
@@ -478,6 +475,7 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
     # Copy forward annotations to backward nodes.
     # Must run before DCE so that forward nodes used for matching aren't removed.
     _copy_fwd_metadata_to_bw_nodes(traced)
+
     _remove_cpu_shadow_chains(traced)
 
     assert output_spec is not None
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 0b5d8940a2..572501c238 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -18,7 +18,7 @@
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
     _copy_fwd_metadata_to_bw_nodes,
     _patch_engine_run_backward,
-    aot_function,
+    minimal_fx_tracer,
 )
 from torchtitan.models.common.attention import (
     annotate_flex_attention_for_regional_inductor,
@@ -38,7 +38,7 @@ def get_loss(logits, labels):
 
 
 def make_train_step(loss_fn):
-    """Return a plain function for aot_function tracing.  loss_fn is captured in closure."""
+    """Return a plain function for minimal_fx_tracer tracing.  loss_fn is captured in closure."""
 
     def train_step(model, *args):
         *fwd_args, labels = args
@@ -123,7 +123,7 @@ def test_mlp_forward(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = aot_function(forward, (model, tokens))
+        traced = minimal_fx_tracer(forward, (model, tokens))
         out_eager = model(tokens)
         wrapped = traced(model, tokens)
         self.assertTrue(torch.equal(out_eager, wrapped))
@@ -134,7 +134,7 @@ def test_mlp_train_step(self):
         model_test.load_state_dict(model_ref.state_dict())
 
         train_step = make_train_step(loss_fn)
-        traced = aot_function(train_step, (model_ref, tokens, labels))
+        traced = minimal_fx_tracer(train_step, (model_ref, tokens, labels))
 
         logits_ref = model_ref(tokens)
         loss_ref = loss_fn(logits_ref, labels)
@@ -155,7 +155,7 @@ def test_mlp_multistep_bitwise(self):
         model_test.load_state_dict(model_ref.state_dict())
 
         train_step = make_train_step(loss_fn)
-        traced = aot_function(train_step, (model_ref, tokens, labels))
+        traced = minimal_fx_tracer(train_step, (model_ref, tokens, labels))
 
         opt_ref = torch.optim.Adam(model_ref.parameters(), lr=self.LR)
         opt_copy = torch.optim.Adam(model_test.parameters(), lr=self.LR)
@@ -189,7 +189,7 @@ def test_mismatched_module_raises(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = aot_function(forward, (model, tokens))
+        traced = minimal_fx_tracer(forward, (model, tokens))
 
         # A model with a different architecture (extra layer → different FQNs).
         different_model = nn.Sequential(
@@ -210,7 +210,7 @@ def fn(model, x, loss_fn):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
         with self.assertRaises(ValueError, msg="all pytree leaves"):
-            aot_function(fn, (model, tokens, lambda x: x.sum()))
+            minimal_fx_tracer(fn, (model, tokens, lambda x: x.sum()))
 
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
@@ -264,7 +264,7 @@ def test_dtensor_forward(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = aot_function(forward, (model, tokens_dt))
+        traced = minimal_fx_tracer(forward, (model, tokens_dt))
         has_subclass = any(
             layout.meta is not None for layout in traced.input_subclass_layouts
         )
@@ -293,7 +293,7 @@ def test_dtensor_train_step(self):
         labels_dt = DTensor.from_local(labels, mesh, [Replicate()])
 
         train_step = make_train_step(get_loss)
-        traced = aot_function(train_step, (model_ref, tokens_dt, labels_dt))
+        traced = minimal_fx_tracer(train_step, (model_ref, tokens_dt, labels_dt))
 
         logits_ref = model_ref(tokens_dt)
         loss_ref = get_loss(logits_ref, labels_dt)
@@ -326,7 +326,7 @@ def test_backward_nodes_have_seq_nr(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced = aot_function(train_step, (model, tokens, labels))
+        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
 
         # Collect seq_nr values from all call_function nodes
         seq_nrs = []
@@ -354,9 +354,11 @@ def test_copy_fwd_metadata_propagates_custom(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced = aot_function(train_step, (model, tokens, labels))
+        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
         gm = traced.gm
 
+        # Manually set custom metadata on the first fwd node for each seq_nr
+        # to test that _copy_fwd_metadata_to_bw_nodes works
         seq_nr_first: dict[int, torch.fx.Node] = {}
         for node in gm.graph.nodes:
             if node.op == "call_function" and "seq_nr" in node.meta:
@@ -365,13 +367,16 @@ def test_copy_fwd_metadata_propagates_custom(self):
                     seq_nr_first[seq_nr] = node
                     node.meta["custom"] = {"test_key": "test_value"}
 
+        # Run the copy pass again
         _copy_fwd_metadata_to_bw_nodes(gm)
 
+        # Check that bwd nodes with shared seq_nr got the custom metadata
         for node in gm.graph.nodes:
             if node.op != "call_function" or "seq_nr" not in node.meta:
                 continue
             seq_nr = node.meta["seq_nr"]
             if node is not seq_nr_first.get(seq_nr):
+                # This is a backward node
                 custom = node.meta.get("custom")
                 self.assertIsNotNone(
                     custom,
@@ -387,9 +392,11 @@ def test_patch_engine_restores_original(self):
         orig_fn = torch.autograd.graph._engine_run_backward
 
         with _patch_engine_run_backward():
+            # Inside the context, it should be patched
             self.assertIsNot(torch.autograd.graph._engine_run_backward, orig_fn)
             self.assertIsNot(torch.autograd._engine_run_backward, orig_fn)
 
+        # After the context, it should be restored
         self.assertIs(torch.autograd.graph._engine_run_backward, orig_fn)
         self.assertIs(torch.autograd._engine_run_backward, orig_fn)
 
@@ -429,7 +436,7 @@ def _run_bitwise_test(
             else contextlib.nullcontext()
         )
         with maybe_regional_inductor:
-            traced = aot_function(train_step, (model_ref, *fwd_args, labels))
+            traced = minimal_fx_tracer(train_step, (model_ref, *fwd_args, labels))
 
         if check_collective_ops:
             ag = sum(
@@ -635,7 +642,7 @@ def test_flex_attention_annotations(self):
             def forward(model, tokens, attn_masks):
                 return model(tokens, attn_masks)
 
-            traced = aot_function(forward, (model, tokens, attn_masks))
+            traced = minimal_fx_tracer(forward, (model, tokens, attn_masks))
 
         flex_nodes = [
             n
@@ -729,7 +736,7 @@ def _run_fsdp_model_test(
             else contextlib.nullcontext()
         )
         with maybe_regional_inductor:
-            traced = aot_function(train_step, (model_ref, *fwd_args, labels))
+            traced = minimal_fx_tracer(train_step, (model_ref, *fwd_args, labels))
 
         ag = sum(
             1
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index f7835f7011..1df158814e 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -13,7 +13,7 @@
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
 from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
-    aot_function,
+    minimal_fx_tracer,
     TracedResult,
 )
 from torchtitan.trainer import Trainer
@@ -23,10 +23,6 @@ def make_fwd_bwd_step(loss_fn):
     """Return a plain function that traces the entire fwd+loss+bwd step.
 
     ``loss_fn`` is captured in the closure so it is not a graph input.
-
-    TODO: investigate how loss_fn interacts with non-strict trace. Currently
-    it is captured as a closure variable, but non-strict tracing may need it
-    registered via pytree.register_constant or passed differently.
     """
 
     def fwd_bwd_step(
@@ -104,7 +100,7 @@ def _make_fx_forward_backward_step(
         if self._traced_step is None:
             fwd_bwd_fn = make_fwd_bwd_step(self.loss_fn)
             with self.train_context(), self.maybe_enable_amp:
-                self._traced_step = aot_function(
+                self._traced_step = minimal_fx_tracer(
                     fwd_bwd_fn,
                     (
                         model,

From 49bd670b518756f998ef8dade8841e92875926e7 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 08:00:35 -0700
Subject: [PATCH 07/42] [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-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 27 +++++-
 .../graph_trainer/tests/test_trace_module.py  | 95 +++++++++++++++++++
 2 files changed, 119 insertions(+), 3 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index b1ce10e599..42589ecd2a 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/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]
@@ -319,6 +329,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 = {
@@ -343,7 +358,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,
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 572501c238..fa312cf637 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -664,6 +664,101 @@ 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
+        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)
+
+        # Create reference weights on CPU.
+        ref_model = Llama3Model(config).to(device=self.DEVICE, dtype=dtype)
+        with torch.no_grad():
+            ref_model.init_weights(buffer_device=torch.device(self.DEVICE))
+        state = {k: v.cpu() for k, v in ref_model.state_dict().items()}
+        del ref_model
+        torch.cuda.empty_cache()
+
+        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))
+            model.load_state_dict(state)
+            apply_ac(model, ActivationCheckpointConfig(mode="selective"))
+            return model
+
+        def run_steps(step_fn, model):
+            """Run num_steps training steps, return losses and peak memory."""
+            opt = torch.optim.Adam(model.parameters(), lr=1e-4)
+            losses = []
+            torch.cuda.reset_peak_memory_stats()
+            for _ in range(num_steps):
+                loss, grads = step_fn(model)
+                losses.append(loss.detach().cpu())
+                for p, g in zip(model.parameters(), grads, strict=True):
+                    p.grad = g
+                opt.step()
+                opt.zero_grad()
+            peak = torch.cuda.max_memory_allocated() / 1e9
+            del opt
+            return losses, peak
+
+        def eager_step(model):
+            logits = model(tokens)
+            loss = get_loss(logits, labels)
+            loss.backward()
+            grads = [p.grad.clone() for p in model.parameters()]
+            return loss, grads
+
+        def traced_step_factory(model):
+            train_step = make_train_step(get_loss)
+            traced = minimal_fx_tracer(train_step, (model, tokens, labels))
+
+            def step(model):
+                result = traced(model, tokens, labels)
+                return result[0], result[1:]
+
+            return step
+
+        # Eager + AC
+        model_eager = make_model()
+        eager_losses, peak_eager = run_steps(eager_step, model_eager)
+        del model_eager
+        torch.cuda.empty_cache()
+
+        # Traced + AC
+        model_traced = make_model()
+        traced_step = traced_step_factory(model_traced)
+        traced_losses, peak_traced = run_steps(traced_step, model_traced)
+        del model_traced
+        torch.cuda.empty_cache()
+
+        torch.use_deterministic_algorithms(False)
+
+        # Verify bitwise-identical loss across all steps.
+        for step, (el, tl) in enumerate(
+            zip(eager_losses, traced_losses, strict=True)
+        ):
+            self.assertTrue(
+                torch.equal(el, tl),
+                f"Step {step}: eager loss={el.item():.6f} vs traced loss={tl.item():.6f}",
+            )
+
+        # Verify 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

From d7e80e08ad638769b7860e0d3f4d5ffe39ed0bb1 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 08:17:19 -0700
Subject: [PATCH 08/42] Update on "[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-poisoned]
---
 .../graph_trainer/tests/test_trace_module.py  | 103 ++++++++----------
 1 file changed, 48 insertions(+), 55 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index fa312cf637..e53dc4466e 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -674,83 +674,76 @@ def test_llama_1b_peak_memory(self):
         batch_size, seq_len = 2, 2048
         dtype = torch.bfloat16
         num_steps = 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)
-
-        # Create reference weights on CPU.
-        ref_model = Llama3Model(config).to(device=self.DEVICE, dtype=dtype)
-        with torch.no_grad():
-            ref_model.init_weights(buffer_device=torch.device(self.DEVICE))
-        state = {k: v.cpu() for k, v in ref_model.state_dict().items()}
-        del ref_model
-        torch.cuda.empty_cache()
+        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))
-            model.load_state_dict(state)
             apply_ac(model, ActivationCheckpointConfig(mode="selective"))
             return model
 
-        def run_steps(step_fn, model):
-            """Run num_steps training steps, return losses and peak memory."""
-            opt = torch.optim.Adam(model.parameters(), lr=1e-4)
-            losses = []
+        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):
-                loss, grads = step_fn(model)
-                losses.append(loss.detach().cpu())
-                for p, g in zip(model.parameters(), grads, strict=True):
-                    p.grad = g
+                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 opt
-            return losses, peak
-
-        def eager_step(model):
-            logits = model(tokens)
-            loss = get_loss(logits, labels)
-            loss.backward()
-            grads = [p.grad.clone() for p in model.parameters()]
-            return loss, grads
-
-        def traced_step_factory(model):
-            train_step = make_train_step(get_loss)
-            traced = minimal_fx_tracer(train_step, (model, tokens, labels))
-
-            def step(model):
-                result = traced(model, tokens, labels)
-                return result[0], result[1:]
-
-            return step
-
-        # Eager + AC
-        model_eager = make_model()
-        eager_losses, peak_eager = run_steps(eager_step, model_eager)
-        del model_eager
-        torch.cuda.empty_cache()
+            del model, opt
+            torch.cuda.empty_cache()
+            return step_results, peak
 
-        # Traced + AC
-        model_traced = make_model()
-        traced_step = traced_step_factory(model_traced)
-        traced_losses, peak_traced = run_steps(traced_step, model_traced)
-        del model_traced
+        ref = make_model()
+        state = ref.state_dict()
+        del ref
         torch.cuda.empty_cache()
 
-        torch.use_deterministic_algorithms(False)
+        eager_results, peak_eager = run("eager")
+        traced_results, peak_traced = run("traced")
 
-        # Verify bitwise-identical loss across all steps.
-        for step, (el, tl) in enumerate(
-            zip(eager_losses, traced_losses, strict=True)
+        # 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}: eager loss={el.item():.6f} vs traced loss={tl.item():.6f}",
+                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}",
+                )
 
-        # Verify peak memory within 10%.
+        # Peak memory within 10%.
         ratio = peak_traced / peak_eager
         self.assertLess(
             ratio,

From 2ebcc9482dd9726b0e9399173e9765f874fcfc95 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 09:37:07 -0700
Subject: [PATCH 09/42] Update base for Update on "[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-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 32 +++++++++----------
 .../graph_trainer/tests/test_trace_module.py  |  3 +-
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index b1ce10e599..70b80ebf46 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -269,6 +269,14 @@ def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
                 node.meta["nn_module_stack"] = nn_module_stack.copy()
 
 
+def _get_params_and_buffers(mod: nn.Module) -> dict[str, torch.Tensor]:
+    """Return a merged dict of the module's named parameters and buffers."""
+    return {
+        **dict(mod.named_parameters(remove_duplicate=False)),
+        **dict(mod.named_buffers(remove_duplicate=False)),
+    }
+
+
 class TracedResult:
     """Holds the traced graph and metadata needed to execute it.
 
@@ -321,10 +329,7 @@ def __call__(self, *args: Any) -> Any:
         convention enforced by :func:`minimal_fx_tracer`.
         """
         mod = args[0]
-        params_dict = {
-            **dict(mod.named_parameters(remove_duplicate=False)),
-            **dict(mod.named_buffers(remove_duplicate=False)),
-        }
+        params_dict = _get_params_and_buffers(mod)
         if not self._validated:
             fqns = list(params_dict.keys())
             if fqns != self.param_fqns:
@@ -379,25 +384,20 @@ def minimal_fx_tracer(
         args: The positional arguments to trace with.  The first element must
             be an ``nn.Module`` whose parameters will be lifted.
     """
-    # Find the single nn.Module in args — must be at position 0.
-    module_indices = [i for i, a in enumerate(args) if isinstance(a, nn.Module)]
-    if len(module_indices) != 1:
+    if not isinstance(args[0], nn.Module):
         raise ValueError(
-            f"minimal_fx_tracer expects exactly one nn.Module in args, "
-            f"got {len(module_indices)} at positions {module_indices}."
+            "minimal_fx_tracer requires args[0] to be an nn.Module, "
+            f"got {type(args[0]).__name__}."
         )
-    if module_indices[0] != 0:
+    if any(isinstance(a, nn.Module) for a in args[1:]):
         raise ValueError(
-            f"The nn.Module must be the first argument (position 0), "
-            f"got it at position {module_indices[0]}."
+            "minimal_fx_tracer supports exactly one nn.Module at args[0]. "
+            "Additional nn.Module instances found in args[1:]."
         )
     mod = args[0]
 
     # Extract params/buffers from the module.
-    params_dict = {
-        **dict(mod.named_parameters(remove_duplicate=False)),
-        **dict(mod.named_buffers(remove_duplicate=False)),
-    }
+    params_dict = _get_params_and_buffers(mod)
     param_fqns = list(params_dict.keys())
     params_flat = list(params_dict.values())
     num_params = len(params_flat)
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 572501c238..f6fa635a44 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -266,7 +266,8 @@ def forward(model, tokens):
 
         traced = minimal_fx_tracer(forward, (model, tokens_dt))
         has_subclass = any(
-            layout.meta is not None for layout in traced.input_subclass_layouts
+            layout.meta is not None
+            for layout in traced.input_subclass_layouts.values()
         )
         self.assertTrue(has_subclass)
 

From 6960a36d9ddc17eaf3feb3b0ed49d46c961851eb Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 10:00:25 -0700
Subject: [PATCH 10/42] Update base for Update on "[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-poisoned]
---
 torchtitan/experiments/graph_trainer/make_fx_tracer.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 70b80ebf46..90663cf256 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -215,6 +215,10 @@ def _patch_engine_run_backward() -> Generator[None, None, None]:
     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``.
     """

From 0b66ad1478b817a33f05f5bde2ddb53efe3cc796 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 10:08:50 -0700
Subject: [PATCH 11/42] Update base for Update on "[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-poisoned]

From cc9798314a5f2555a630aac3c044c16f8328f9ff Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 14:54:22 -0700
Subject: [PATCH 12/42] Update base for Update on "[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-poisoned]

From 0783db347b5c5ec25f1a1f7cccef84ee6349cdde Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 15:03:33 -0700
Subject: [PATCH 13/42] Update base for Update on "[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-poisoned]

From 4b48a21e8c929cbb783fc9d48aa96e4aa4d77408 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 19:42:17 -0700
Subject: [PATCH 14/42] Update base for Update on "[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-poisoned]

From 84c908a960dab611d3562ac673a062344831188e Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 31 Mar 2026 19:51:46 -0700
Subject: [PATCH 15/42] Update base for Update on "[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-poisoned]

From d5382494c3851b5e7283fcf14c0190ddbb06d354 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Wed, 1 Apr 2026 13:43:41 -0700
Subject: [PATCH 16/42] Update base for Update on "[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-poisoned]

From 201ed869a9c7dcfe8600e74594fb208b1d52b15d Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Wed, 1 Apr 2026 14:03:15 -0700
Subject: [PATCH 17/42] Update base for Update on "[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-poisoned]

From fd4bd7791359eddb8b7a60621c635cfcbfad0c75 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Wed, 1 Apr 2026 20:01:20 -0700
Subject: [PATCH 18/42] Update base for Update on "[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-poisoned]
---
 torchtitan/experiments/graph_trainer/make_fx_tracer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 90663cf256..075e7f82a7 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -15,7 +15,7 @@
 from torch._functorch._aot_autograd.logging_utils import (
     setup_stacktrace_preservation_hooks,
 )
-from torch._guards import TracingContext, tracing
+from torch._guards import tracing, TracingContext
 from torch._subclasses import FakeTensorMode
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.fx.traceback import preserve_node_meta

From 803cc90ee7668d917172b38c35cea177be42132a Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Thu, 2 Apr 2026 08:10:11 -0700
Subject: [PATCH 19/42] Update base for Update on "[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-poisoned]

From 6559ba69a969802032ebd9dec2a6b07ba86e1bd5 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Thu, 2 Apr 2026 19:42:58 -0700
Subject: [PATCH 20/42] Update base for Update on "[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-poisoned]

From 3036baed003edef3f96e2c0fafa93c376da82b1f Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Sat, 4 Apr 2026 05:40:47 -0700
Subject: [PATCH 21/42] Update base for Update on "[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-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 110 +++++++++++-------
 .../graph_trainer/tests/test_trace_module.py  |  43 +++----
 .../experiments/graph_trainer/trainer.py      |  21 ++--
 3 files changed, 100 insertions(+), 74 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 075e7f82a7..7e812e3b4e 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -284,11 +284,12 @@ def _get_params_and_buffers(mod: nn.Module) -> dict[str, torch.Tensor]:
 class TracedResult:
     """Holds the traced graph and metadata needed to execute it.
 
-    Returned by :func:`minimal_fx_tracer`.  Call the instance directly to execute
-    the traced graph with fresh parameters read from the live module::
+    Returned by :func:`minimal_fx_tracer`. Call the tracer returned by
+    :func:`minimal_fx_tracer` with trace-time args, then execute it with
+    :func:`run_traced`::
 
-        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
-        result = traced(model, tokens, labels)
+        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
+        result = run_traced(traced, model, tokens, labels)
 
     Args:
         gm: The traced FX graph (a pure function of flat tensors).
@@ -324,42 +325,9 @@ def __init__(
         self.num_flat_outputs = num_flat_outputs
         self.output_subclass_layouts = output_subclass_layouts
         self.output_spec = output_spec
-        self._validated = False
-
-    def __call__(self, *args: Any) -> Any:
-        """Execute the traced graph, reading fresh params from the module in ``args``.
-
-        The module must be the first argument (position 0), matching the
-        convention enforced by :func:`minimal_fx_tracer`.
-        """
-        mod = args[0]
-        params_dict = _get_params_and_buffers(mod)
-        if not self._validated:
-            fqns = list(params_dict.keys())
-            if fqns != self.param_fqns:
-                raise ValueError(
-                    f"Module at args[0] has different parameter/buffer "
-                    f"names than during tracing.\n"
-                    f"  Traced: {self.param_fqns}\n"
-                    f"  Got:    {fqns}"
-                )
-            self._validated = True
-        params_flat = list(params_dict.values())
-
-        user_args = list(args[1:])
-        user_args_flat, _ = pytree.tree_flatten(user_args)
-
-        all_args = params_flat + list(user_args_flat)
-        flat_inputs, _ = _unwrap_subclasses(all_args)
-
-        flat_outputs = self.gm(*flat_inputs)
-        wrapped = _wrap_subclasses(
-            flat_outputs, self.num_flat_outputs, self.output_subclass_layouts
-        )
-        return pytree.tree_unflatten(wrapped, self.output_spec)
 
 
-def minimal_fx_tracer(
+def _trace_with_args(
     fn: Callable,
     args: tuple,
 ) -> TracedResult:
@@ -377,11 +345,11 @@ def minimal_fx_tracer(
     prepending.  Non-tensor, non-module values like ``loss_fn`` should be
     captured in ``fn``'s closure rather than passed as args.
 
-    The returned :class:`TracedResult` is directly callable — pass the same
-    positional arguments (with the live module first) to execute the graph::
+    Execute the returned :class:`TracedResult` with :func:`run_traced`, passing
+    the same positional arguments (with the live module first)::
 
-        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
-        result = traced(model, tokens, labels)
+        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
+        result = run_traced(traced, model, tokens, labels)
 
     Args:
         fn: The callable to trace.
@@ -493,3 +461,61 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
         output_subclass_layouts=output_layouts,
         output_spec=output_spec,
     )
+
+
+def minimal_fx_tracer(fn: Callable) -> Callable[..., TracedResult]:
+    """Return a tracer for ``fn`` that traces a concrete module/input invocation.
+
+    ``fn`` must be a plain callable (not an ``nn.Module``). The returned
+    callable expects the trace-time positional arguments, with the live module
+    at position 0::
+
+        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
+        result = run_traced(traced, model, tokens, labels)
+    """
+
+    def trace_with_args(*args: Any) -> TracedResult:
+        return _trace_with_args(fn, args)
+
+    return trace_with_args
+
+
+def run_traced(traced_result: TracedResult, *args: Any) -> Any:
+    """Execute a traced graph with fresh parameters read from the live module.
+
+    This is a reference implementation of traced-graph execution. It keeps the
+    parameter lookup, subclass unwrapping, and output reconstruction logic
+    explicit instead of baking those semantics into ``TracedResult`` itself.
+
+    The module must be the first argument (position 0), matching the
+    convention enforced by :func:`minimal_fx_tracer`.
+    """
+
+    mod = args[0]
+    params_dict = _get_params_and_buffers(mod)
+    fqns = list(params_dict.keys())
+    if fqns != traced_result.param_fqns:
+        raise ValueError(
+            f"Module at args[0] has different parameter/buffer "
+            f"names than during tracing.\n"
+            f"  Traced: {traced_result.param_fqns}\n"
+            f"  Got:    {fqns}"
+        )
+    params_flat = list(params_dict.values())
+
+    user_args = list(args[1:])
+    user_args_flat, _ = pytree.tree_flatten(user_args)
+
+    all_args = params_flat + list(user_args_flat)
+    flat_inputs, _ = _unwrap_subclasses(all_args)
+
+    flat_outputs = traced_result.gm(*flat_inputs)
+    wrapped = _wrap_subclasses(
+        flat_outputs,
+        traced_result.num_flat_outputs,
+        traced_result.output_subclass_layouts,
+    )
+    return pytree.tree_unflatten(wrapped, traced_result.output_spec)
+
+
+run_traced_module = run_traced
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index f6fa635a44..48bf34363b 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -19,6 +19,7 @@
     _copy_fwd_metadata_to_bw_nodes,
     _patch_engine_run_backward,
     minimal_fx_tracer,
+    run_traced,
 )
 from torchtitan.models.common.attention import (
     annotate_flex_attention_for_regional_inductor,
@@ -123,9 +124,9 @@ def test_mlp_forward(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = minimal_fx_tracer(forward, (model, tokens))
+        traced = minimal_fx_tracer(forward)(model, tokens)
         out_eager = model(tokens)
-        wrapped = traced(model, tokens)
+        wrapped = run_traced(traced, model, tokens)
         self.assertTrue(torch.equal(out_eager, wrapped))
 
     def test_mlp_train_step(self):
@@ -134,14 +135,14 @@ def test_mlp_train_step(self):
         model_test.load_state_dict(model_ref.state_dict())
 
         train_step = make_train_step(loss_fn)
-        traced = minimal_fx_tracer(train_step, (model_ref, tokens, labels))
+        traced = minimal_fx_tracer(train_step)(model_ref, tokens, labels)
 
         logits_ref = model_ref(tokens)
         loss_ref = loss_fn(logits_ref, labels)
         loss_ref.backward()
         grads_ref = [p.grad.clone() for p in model_ref.parameters()]
 
-        wrapped = traced(model_test, tokens, labels)
+        wrapped = run_traced(traced, model_test, tokens, labels)
         loss_tr = wrapped[0]
         grads_tr = wrapped[1:]
 
@@ -155,7 +156,7 @@ def test_mlp_multistep_bitwise(self):
         model_test.load_state_dict(model_ref.state_dict())
 
         train_step = make_train_step(loss_fn)
-        traced = minimal_fx_tracer(train_step, (model_ref, tokens, labels))
+        traced = minimal_fx_tracer(train_step)(model_ref, tokens, labels)
 
         opt_ref = torch.optim.Adam(model_ref.parameters(), lr=self.LR)
         opt_copy = torch.optim.Adam(model_test.parameters(), lr=self.LR)
@@ -168,7 +169,7 @@ def test_mlp_multistep_bitwise(self):
             opt_ref.step()
             opt_ref.zero_grad()
 
-            wrapped = traced(model_test, tokens, labels)
+            wrapped = run_traced(traced, model_test, tokens, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
@@ -189,7 +190,7 @@ def test_mismatched_module_raises(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = minimal_fx_tracer(forward, (model, tokens))
+        traced = minimal_fx_tracer(forward)(model, tokens)
 
         # A model with a different architecture (extra layer → different FQNs).
         different_model = nn.Sequential(
@@ -198,7 +199,7 @@ def forward(model, tokens):
         ).to(device=self.DEVICE, dtype=self.DTYPE)
 
         with self.assertRaises(ValueError, msg="different parameter/buffer names"):
-            traced(different_model, tokens)
+            run_traced(traced, different_model, tokens)
 
     def test_non_tensor_leaf_raises(self):
         """Passing a callable leaf in args raises (should be in closure instead)."""
@@ -210,7 +211,7 @@ def fn(model, x, loss_fn):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
         with self.assertRaises(ValueError, msg="all pytree leaves"):
-            minimal_fx_tracer(fn, (model, tokens, lambda x: x.sum()))
+            minimal_fx_tracer(fn)(model, tokens, lambda x: x.sum())
 
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
@@ -264,7 +265,7 @@ def test_dtensor_forward(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = minimal_fx_tracer(forward, (model, tokens_dt))
+        traced = minimal_fx_tracer(forward)(model, tokens_dt)
         has_subclass = any(
             layout.meta is not None
             for layout in traced.input_subclass_layouts.values()
@@ -272,8 +273,8 @@ def forward(model, tokens):
         self.assertTrue(has_subclass)
 
         out_eager = model(tokens_dt)
-        wrapped = traced(model, tokens_dt)
-        self.assertTrue(torch.equal(out_eager.full_tensor(), wrapped[0].full_tensor()))
+        wrapped = run_traced(traced, model, tokens_dt)
+        self.assertTrue(torch.equal(out_eager.full_tensor(), wrapped.full_tensor()))
 
     def test_dtensor_train_step(self):
         from torch.distributed._tensor import DTensor, Replicate
@@ -294,14 +295,14 @@ def test_dtensor_train_step(self):
         labels_dt = DTensor.from_local(labels, mesh, [Replicate()])
 
         train_step = make_train_step(get_loss)
-        traced = minimal_fx_tracer(train_step, (model_ref, tokens_dt, labels_dt))
+        traced = minimal_fx_tracer(train_step)(model_ref, tokens_dt, labels_dt)
 
         logits_ref = model_ref(tokens_dt)
         loss_ref = get_loss(logits_ref, labels_dt)
         loss_ref.backward()
         grads_ref = [p.grad.clone() for p in model_ref.parameters()]
 
-        wrapped = traced(model_test, tokens_dt, labels_dt)
+        wrapped = run_traced(traced, model_test, tokens_dt, labels_dt)
         loss_tr = wrapped[0]
         grads_tr = wrapped[1:]
 
@@ -327,7 +328,7 @@ def test_backward_nodes_have_seq_nr(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
+        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
 
         # Collect seq_nr values from all call_function nodes
         seq_nrs = []
@@ -355,7 +356,7 @@ def test_copy_fwd_metadata_propagates_custom(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced = minimal_fx_tracer(train_step, (model, tokens, labels))
+        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
         gm = traced.gm
 
         # Manually set custom metadata on the first fwd node for each seq_nr
@@ -437,7 +438,7 @@ def _run_bitwise_test(
             else contextlib.nullcontext()
         )
         with maybe_regional_inductor:
-            traced = minimal_fx_tracer(train_step, (model_ref, *fwd_args, labels))
+            traced = minimal_fx_tracer(train_step)(model_ref, *fwd_args, labels)
 
         if check_collective_ops:
             ag = sum(
@@ -469,7 +470,7 @@ def _run_bitwise_test(
             opt_ref.step()
             opt_ref.zero_grad()
 
-            wrapped = traced(model_test, *fwd_args, labels)
+            wrapped = run_traced(traced, model_test, *fwd_args, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
@@ -643,7 +644,7 @@ def test_flex_attention_annotations(self):
             def forward(model, tokens, attn_masks):
                 return model(tokens, attn_masks)
 
-            traced = minimal_fx_tracer(forward, (model, tokens, attn_masks))
+            traced = minimal_fx_tracer(forward)(model, tokens, attn_masks)
 
         flex_nodes = [
             n
@@ -737,7 +738,7 @@ def _run_fsdp_model_test(
             else contextlib.nullcontext()
         )
         with maybe_regional_inductor:
-            traced = minimal_fx_tracer(train_step, (model_ref, *fwd_args, labels))
+            traced = minimal_fx_tracer(train_step)(model_ref, *fwd_args, labels)
 
         ag = sum(
             1
@@ -766,7 +767,7 @@ def _run_fsdp_model_test(
             opt_ref.step()
             opt_ref.zero_grad()
 
-            wrapped = traced(model_test, *fwd_args, labels)
+            wrapped = run_traced(traced, model_test, *fwd_args, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index 1df158814e..f3ccd0c130 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -14,6 +14,7 @@
 from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
     minimal_fx_tracer,
+    run_traced,
     TracedResult,
 )
 from torchtitan.trainer import Trainer
@@ -100,20 +101,18 @@ def _make_fx_forward_backward_step(
         if self._traced_step is None:
             fwd_bwd_fn = make_fwd_bwd_step(self.loss_fn)
             with self.train_context(), self.maybe_enable_amp:
-                self._traced_step = minimal_fx_tracer(
-                    fwd_bwd_fn,
-                    (
-                        model,
-                        inputs,
-                        labels,
-                        global_valid_tokens,
-                        extra_inputs,
-                        extra_kwargs,
-                    ),
+                self._traced_step = minimal_fx_tracer(fwd_bwd_fn)(
+                    model,
+                    inputs,
+                    labels,
+                    global_valid_tokens,
+                    extra_inputs,
+                    extra_kwargs,
                 )
 
         with self.train_context(), self.maybe_enable_amp:
-            outputs = self._traced_step(
+            outputs = run_traced(
+                self._traced_step,
                 model,
                 inputs,
                 labels,

From 386810a432a756b17a05b27fed9922acd06e2ea3 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Sat, 4 Apr 2026 06:12:47 -0700
Subject: [PATCH 22/42] Update base for Update on "[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-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 323 ++++++++----------
 1 file changed, 142 insertions(+), 181 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 7e812e3b4e..6a41ac7f5e 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -281,203 +281,164 @@ def _get_params_and_buffers(mod: nn.Module) -> dict[str, torch.Tensor]:
     }
 
 
+@dataclass
 class TracedResult:
-    """Holds the traced graph and metadata needed to execute it.
-
-    Returned by :func:`minimal_fx_tracer`. Call the tracer returned by
-    :func:`minimal_fx_tracer` with trace-time args, then execute it with
-    :func:`run_traced`::
-
-        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
-        result = run_traced(traced, model, tokens, labels)
-
-    Args:
-        gm: The traced FX graph (a pure function of flat tensors).
-        param_fqns: Fully qualified names of the module's parameters and
-            buffers, recorded at trace time for validation.
-        num_params: Number of parameters + buffers (flat count).
-        num_flat_inputs: Number of original args (before subclass unwrapping).
-        input_subclass_layouts: Maps arg positions that are tensor subclasses
-            to their unwrap/rewrap metadata.  Plain tensors have no entry.
-        num_flat_outputs: Number of original outputs (before subclass unwrapping).
-        output_subclass_layouts: Maps output positions that are tensor subclasses
-            to their rewrap metadata.  Plain tensors have no entry.
-        output_spec: Pytree spec of the original function's return value,
-            used to reconstruct the output structure.
+    """Holds the traced graph and execution metadata.
+
+    Attributes:
+        gm: The traced FX graph as a pure function of flat tensors.
+        param_fqns: Trace-time parameter/buffer FQNs for execution-time validation.
+        num_params: Number of lifted parameters and buffers.
+        num_flat_inputs: Number of flat graph inputs before subclass unwrapping.
+        input_subclass_layouts: Subclass unwrap/rewrap metadata for inputs.
+        num_flat_outputs: Number of flat graph outputs before subclass rewrapping.
+        output_subclass_layouts: Subclass unwrap/rewrap metadata for outputs.
+        output_spec: Original output pytree spec used during reconstruction.
     """
 
-    def __init__(
-        self,
-        gm: torch.fx.GraphModule,
-        param_fqns: list[str],
-        num_params: int,
-        num_flat_inputs: int,
-        input_subclass_layouts: dict[int, SubclassLayout],
-        num_flat_outputs: int,
-        output_subclass_layouts: dict[int, SubclassLayout],
-        output_spec: pytree.TreeSpec,
-    ) -> None:
-        self.gm = gm
-        self.param_fqns = param_fqns
-        self.num_params = num_params
-        self.num_flat_inputs = num_flat_inputs
-        self.input_subclass_layouts = input_subclass_layouts
-        self.num_flat_outputs = num_flat_outputs
-        self.output_subclass_layouts = output_subclass_layouts
-        self.output_spec = output_spec
-
-
-def _trace_with_args(
-    fn: Callable,
-    args: tuple,
-) -> TracedResult:
-    """Trace ``fn(*args)`` into a flat FX graph, unwrapping tensor subclasses.
-
-    ``args[0]`` must be an ``nn.Module``.  Its parameters and buffers are
-    lifted as extra graph inputs so the returned graph is a pure function.
-    Tensor subclasses (e.g. DTensor) are recursively unwrapped into plain
-    tensors for tracing, and the layouts needed to rewrap them are recorded
-    in the returned :class:`TracedResult`.
+    gm: torch.fx.GraphModule
+    param_fqns: list[str]
+    num_params: int
+    num_flat_inputs: int
+    input_subclass_layouts: dict[int, SubclassLayout]
+    num_flat_outputs: int
+    output_subclass_layouts: dict[int, SubclassLayout]
+    output_spec: pytree.TreeSpec
+
+
+def minimal_fx_tracer(fn: Callable) -> Callable[..., TracedResult]:
+    """Return a tracer for ``fn`` that traces a concrete module/input invocation.
 
-    ``fn`` must be a plain callable (not an ``nn.Module``).  This keeps the
-    trace and execute calling conventions identical — the same ``args`` are
-    passed at both trace time and execution time, with no hidden arg
-    prepending.  Non-tensor, non-module values like ``loss_fn`` should be
-    captured in ``fn``'s closure rather than passed as args.
+    ``fn`` must be a plain callable (not an ``nn.Module``). The returned
+    callable expects the trace-time positional arguments, with the live module
+    at position 0.
 
     Execute the returned :class:`TracedResult` with :func:`run_traced`, passing
     the same positional arguments (with the live module first)::
 
-        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
-        result = run_traced(traced, model, tokens, labels)
-
-    Args:
-        fn: The callable to trace.
-        args: The positional arguments to trace with.  The first element must
-            be an ``nn.Module`` whose parameters will be lifted.
-    """
-    if not isinstance(args[0], nn.Module):
-        raise ValueError(
-            "minimal_fx_tracer requires args[0] to be an nn.Module, "
-            f"got {type(args[0]).__name__}."
-        )
-    if any(isinstance(a, nn.Module) for a in args[1:]):
-        raise ValueError(
-            "minimal_fx_tracer supports exactly one nn.Module at args[0]. "
-            "Additional nn.Module instances found in args[1:]."
-        )
-    mod = args[0]
+        traced_result = minimal_fx_tracer(train_step)(model, tokens, labels)
+        result = run_traced(traced_result, model, tokens, labels)
 
-    # Extract params/buffers from the module.
-    params_dict = _get_params_and_buffers(mod)
-    param_fqns = list(params_dict.keys())
-    params_flat = list(params_dict.values())
-    num_params = len(params_flat)
+    The trace-time args must satisfy these constraints:
+    - ``args[0]`` must be an ``nn.Module`` whose parameters and buffers are
+      lifted as extra graph inputs
+    - there must be no additional ``nn.Module`` instances in ``args[1:]``
+    - all remaining pytree leaves must be tensors or make_fx-safe primitives
+      (``int``, ``float``, ``bool``, ``str``, ``None``)
 
-    # User args: everything after the module.
-    user_args = list(args[1:])
-    user_args_flat, user_args_spec = pytree.tree_flatten(user_args)
+    Tensor subclasses (for example ``DTensor``) are recursively unwrapped into
+    plain tensors for tracing, and the layouts needed to rewrap them are stored
+    in the returned :class:`TracedResult`.
+    """
 
-    # Validate leaves.
-    for leaf in user_args_flat:
-        if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
+    def _trace_with_args(*args: Any) -> TracedResult:
+        if not isinstance(args[0], nn.Module):
             raise ValueError(
-                f"minimal_fx_tracer requires all pytree leaves in args to be tensors "
-                f"or primitives (int/float/bool/str), got {type(leaf).__name__}. "
-                f"Non-primitive values should either be registered as pytree "
-                f"nodes (register_pytree_node) or constants "
-                f"(pytree.register_constant), or captured in fn's closure."
+                "minimal_fx_tracer requires args[0] to be an nn.Module, "
+                f"got {type(args[0]).__name__}."
             )
-
-    # Combined flat input: [*params, *user_args] with subclasses unwrapped.
-    full_args = params_flat + list(user_args_flat)
-    num_full_args = len(full_args)
-    unwrapped_args, input_layouts = _unwrap_subclasses(full_args)
-
-    fake_mode = FakeTensorMode(
-        allow_non_fake_inputs=True,
-        shape_env=torch.fx.experimental.symbolic_shapes.ShapeEnv(),
-    )
-    fake_args = tuple(
-        (
-            fake_mode.from_tensor(a, static_shapes=True)
-            if isinstance(a, torch.Tensor)
-            else a
+        if any(isinstance(a, nn.Module) for a in args[1:]):
+            raise ValueError(
+                "minimal_fx_tracer supports exactly one nn.Module at args[0]. "
+                "Additional nn.Module instances found in args[1:]."
+            )
+        mod = args[0]
+
+        # Extract params/buffers from the module.
+        params_dict = _get_params_and_buffers(mod)
+        param_fqns = list(params_dict.keys())
+        params_flat = list(params_dict.values())
+        num_params = len(params_flat)
+
+        # User args: everything after the module.
+        user_args = list(args[1:])
+        user_args_flat, user_args_spec = pytree.tree_flatten(user_args)
+
+        # Validate leaves.
+        for leaf in user_args_flat:
+            if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
+                raise ValueError(
+                    f"minimal_fx_tracer requires all pytree leaves in args to be tensors "
+                    f"or primitives (int/float/bool/str), got {type(leaf).__name__}. "
+                    f"Non-primitive values should either be registered as pytree "
+                    f"nodes (register_pytree_node) or constants "
+                    f"(pytree.register_constant), or captured in fn's closure."
+                )
+
+        # Combined flat input: [*params, *user_args] with subclasses unwrapped.
+        full_args = params_flat + list(user_args_flat)
+        num_full_args = len(full_args)
+        unwrapped_args, input_layouts = _unwrap_subclasses(full_args)
+
+        fake_mode = FakeTensorMode(
+            allow_non_fake_inputs=True,
+            shape_env=torch.fx.experimental.symbolic_shapes.ShapeEnv(),
+        )
+        fake_args = tuple(
+            (
+                fake_mode.from_tensor(a, static_shapes=True)
+                if isinstance(a, torch.Tensor)
+                else a
+            )
+            for a in unwrapped_args
         )
-        for a in unwrapped_args
-    )
-
-    output_layouts: dict[int, SubclassLayout] = {}
-    num_flat_outputs: int = 0
-    output_spec: pytree.TreeSpec | None = None
-
-    def fn_with_subclass_handling(*plain_args: Any) -> list:
-        nonlocal output_layouts, output_spec, num_flat_outputs
-        output_layouts = {}
-
-        wrapped = _wrap_subclasses(plain_args, num_full_args, input_layouts)
-        params_wrapped = wrapped[:num_params]
-        user_flat = wrapped[num_params:]
-
-        params_for_mod = dict(zip(param_fqns, params_wrapped, strict=True))
-        user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
-
-        # Reconstruct the original args: module at position 0 keeps the live
-        # module, remaining positions get the traced user tensors.
-        rebuilt = [mod] + user_list
-
-        with stateless._reparametrize_module(mod, params_for_mod):
-            with _patch_engine_run_backward():
-                result = fn(*rebuilt)
-
-        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():
-        traced = make_fx(
-            fn_with_subclass_handling,
-            record_stack_traces=True,
-            record_module_stack=False,  # don't need nn_module_stack for now
-        )(*fake_args)
-
-    # Copy forward annotations to backward nodes.
-    # Must run before DCE so that forward nodes used for matching aren't removed.
-    _copy_fwd_metadata_to_bw_nodes(traced)
-
-    _remove_cpu_shadow_chains(traced)
-
-    assert output_spec is not None
-    return TracedResult(
-        gm=traced,
-        param_fqns=param_fqns,
-        num_params=num_params,
-        num_flat_inputs=num_full_args,
-        input_subclass_layouts=input_layouts,
-        num_flat_outputs=num_flat_outputs,
-        output_subclass_layouts=output_layouts,
-        output_spec=output_spec,
-    )
-
-
-def minimal_fx_tracer(fn: Callable) -> Callable[..., TracedResult]:
-    """Return a tracer for ``fn`` that traces a concrete module/input invocation.
-
-    ``fn`` must be a plain callable (not an ``nn.Module``). The returned
-    callable expects the trace-time positional arguments, with the live module
-    at position 0::
-
-        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
-        result = run_traced(traced, model, tokens, labels)
-    """
 
-    def trace_with_args(*args: Any) -> TracedResult:
-        return _trace_with_args(fn, args)
+        output_layouts: dict[int, SubclassLayout] = {}
+        num_flat_outputs: int = 0
+        output_spec: pytree.TreeSpec | None = None
+
+        def fn_with_subclass_handling(*plain_args: Any) -> list:
+            nonlocal output_layouts, output_spec, num_flat_outputs
+            output_layouts = {}
+
+            wrapped = _wrap_subclasses(plain_args, num_full_args, input_layouts)
+            params_wrapped = wrapped[:num_params]
+            user_flat = wrapped[num_params:]
+
+            params_for_mod = dict(zip(param_fqns, params_wrapped, strict=True))
+            user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
+
+            # Reconstruct the original args: module at position 0 keeps the live
+            # module, remaining positions get the traced user tensors.
+            rebuilt = [mod] + user_list
+
+            with stateless._reparametrize_module(mod, params_for_mod):
+                with _patch_engine_run_backward():
+                    result = fn(*rebuilt)
+
+            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():
+            traced = make_fx(
+                fn_with_subclass_handling,
+                record_stack_traces=True,
+                record_module_stack=False,  # don't need nn_module_stack for now
+            )(*fake_args)
+
+        # Copy forward annotations to backward nodes.
+        # Must run before DCE so that forward nodes used for matching aren't removed.
+        _copy_fwd_metadata_to_bw_nodes(traced)
+
+        _remove_cpu_shadow_chains(traced)
+
+        assert output_spec is not None
+        return TracedResult(
+            gm=traced,
+            param_fqns=param_fqns,
+            num_params=num_params,
+            num_flat_inputs=num_full_args,
+            input_subclass_layouts=input_layouts,
+            num_flat_outputs=num_flat_outputs,
+            output_subclass_layouts=output_layouts,
+            output_spec=output_spec,
+        )
 
-    return trace_with_args
+    return _trace_with_args
 
 
 def run_traced(traced_result: TracedResult, *args: Any) -> Any:

From 6580931220916094b618b96324cfbf94f69f0ea2 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 6 Apr 2026 07:13:26 -0700
Subject: [PATCH 23/42] Update base for Update on "[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-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 32 ++--------------
 .../graph_trainer/tests/test_trace_module.py  | 38 ++++++++++++++++---
 2 files changed, 35 insertions(+), 35 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index cd820b41c5..0308f4cd9a 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -61,31 +61,6 @@ class SubclassLayout:
     meta: SubclassMeta | None
 
 
-@dataclass
-class TracedResult:
-    """Holds the traced graph and execution metadata.
-
-    Attributes:
-        gm: The traced FX graph as a pure function of flat tensors.
-        example_inputs: Trace-time fake flat inputs used by downstream graph passes.
-        param_fqns: Trace-time parameter/buffer FQNs for execution-time validation.
-        num_params: Number of lifted parameters and buffers.
-        num_flat_inputs: Number of flat graph inputs before subclass unwrapping.
-        input_subclass_layouts: Subclass unwrap/rewrap metadata for inputs.
-        num_flat_outputs: Number of flat graph outputs before subclass rewrapping.
-        output_subclass_layouts: Subclass unwrap/rewrap metadata for outputs.
-        output_spec: Original output pytree spec used during reconstruction.
-    """
-
-    gm: torch.fx.GraphModule
-    example_inputs: tuple[Any, ...]
-    param_fqns: list[str]
-    num_params: int
-    num_flat_inputs: int
-    input_subclass_layouts: dict[int, SubclassLayout]
-    num_flat_outputs: int
-    output_subclass_layouts: dict[int, SubclassLayout]
-    output_spec: pytree.TreeSpec
 def _unwrap_subclass(t: torch.Tensor) -> tuple[list[torch.Tensor], SubclassMeta | None]:
     if not is_traceable_wrapper_subclass(t):
         return [t], None
@@ -311,10 +286,11 @@ def _get_params_and_buffers(mod: nn.Module) -> dict[str, torch.Tensor]:
 
 @dataclass
 class TracedResult:
-    """Holds the traced graph and execution metadata.
+    """Execution metadata returned by :func:`minimal_fx_tracer`.
 
     Attributes:
         gm: The traced FX graph as a pure function of flat tensors.
+        example_inputs: Trace-time fake flat inputs used by downstream graph passes.
         param_fqns: Trace-time parameter/buffer FQNs for execution-time validation.
         num_params: Number of lifted parameters and buffers.
         num_flat_inputs: Number of flat graph inputs before subclass unwrapping.
@@ -325,6 +301,7 @@ class TracedResult:
     """
 
     gm: torch.fx.GraphModule
+    example_inputs: tuple[Any, ...]
     param_fqns: list[str]
     num_params: int
     num_flat_inputs: int
@@ -506,6 +483,3 @@ def run_traced(traced_result: TracedResult, *args: Any) -> Any:
         traced_result.output_subclass_layouts,
     )
     return pytree.tree_unflatten(wrapped, traced_result.output_spec)
-
-
-run_traced_module = run_traced
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 0aa12f4f68..99333fb403 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -13,6 +13,7 @@
 from torch.testing._internal.common_fsdp import FSDPTest
 
 from torchtitan.experiments.graph_trainer.common_utils import (
+    annotate_flex_attention_for_regional_inductor,
     register_blockmask_pytree_node,
 )
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
@@ -21,9 +22,6 @@
     minimal_fx_tracer,
     run_traced,
 )
-from torchtitan.models.common.attention import (
-    annotate_flex_attention_for_regional_inductor,
-)
 
 # BlockMask must be registered as a pytree node so its tensor children
 # are properly traced as graph inputs instead of opaque leaves.
@@ -213,6 +211,28 @@ def fn(model, x, loss_fn):
         with self.assertRaises(ValueError, msg="all pytree leaves"):
             minimal_fx_tracer(fn)(model, tokens, lambda x: x.sum())
 
+    def test_module_must_be_first_arg(self):
+        def forward(tokens, model):
+            return model(tokens)
+
+        model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
+        tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
+
+        with self.assertRaises(ValueError, msg="args[0]"):
+            minimal_fx_tracer(forward)(tokens, model)
+
+    def test_additional_module_arg_raises(self):
+        def forward(model, other_model, tokens):
+            del other_model
+            return model(tokens)
+
+        model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
+        other_model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
+        tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
+
+        with self.assertRaises(ValueError, msg="Additional nn.Module"):
+            minimal_fx_tracer(forward)(model, other_model, tokens)
+
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
 class TestTraceDTensor(unittest.TestCase):
@@ -389,16 +409,17 @@ def test_copy_fwd_metadata_propagates_custom(self):
     def test_backward_nodes_have_stack_trace(self):
         """Verify that backward nodes get stack_trace from their forward counterpart."""
         model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
-        train_step = TrainStepModule(model, get_loss)
+        train_step = make_train_step(get_loss)
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced_result = trace_module(train_step, (tokens, labels))
+        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
 
         # Find backward nodes: nodes sharing a seq_nr with an earlier (forward) node
         seq_nr_first: dict[int, torch.fx.Node] = {}
         bwd_nodes_missing_stack_trace = []
-        for node in traced_result.gm.graph.nodes:
+        num_checked = 0
+        for node in traced.gm.graph.nodes:
             if node.op != "call_function" or "seq_nr" not in node.meta:
                 continue
             seq_nr = node.meta["seq_nr"]
@@ -406,9 +427,14 @@ def test_backward_nodes_have_stack_trace(self):
                 seq_nr_first[seq_nr] = node
             else:
                 # This is a backward node
+                fwd_node = seq_nr_first[seq_nr]
+                if not fwd_node.stack_trace:
+                    continue
+                num_checked += 1
                 if not node.stack_trace:
                     bwd_nodes_missing_stack_trace.append((node.name, seq_nr))
 
+        self.assertEqual(num_checked, 24)
         self.assertEqual(
             bwd_nodes_missing_stack_trace,
             [],

From 8bd77b7de426e7cbfd1b95859f38d01737f337ca Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 6 Apr 2026 07:43:14 -0700
Subject: [PATCH 24/42] Update base for Update on "[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-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 23 +++++++-----
 .../graph_trainer/tests/test_trace_module.py  | 36 +++++++++----------
 2 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 0308f4cd9a..7e6c3183e2 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -447,7 +447,11 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
     return _trace_with_args
 
 
-def run_traced(traced_result: TracedResult, *args: Any) -> Any:
+def run_traced(
+    traced_result: TracedResult,
+    *args: Any,
+    validate_module_fqns: bool = False,
+) -> Any:
     """Execute a traced graph with fresh parameters read from the live module.
 
     This is a reference implementation of traced-graph execution. It keeps the
@@ -460,14 +464,15 @@ def run_traced(traced_result: TracedResult, *args: Any) -> Any:
 
     mod = args[0]
     params_dict = _get_params_and_buffers(mod)
-    fqns = list(params_dict.keys())
-    if fqns != traced_result.param_fqns:
-        raise ValueError(
-            f"Module at args[0] has different parameter/buffer "
-            f"names than during tracing.\n"
-            f"  Traced: {traced_result.param_fqns}\n"
-            f"  Got:    {fqns}"
-        )
+    if validate_module_fqns:
+        fqns = list(params_dict.keys())
+        if fqns != traced_result.param_fqns:
+            raise ValueError(
+                f"Module at args[0] has different parameter/buffer "
+                f"names than during tracing.\n"
+                f"  Traced: {traced_result.param_fqns}\n"
+                f"  Got:    {fqns}"
+            )
     params_flat = list(params_dict.values())
 
     user_args = list(args[1:])
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 99333fb403..6fa54e6cea 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -181,8 +181,20 @@ def test_mlp_multistep_bitwise(self):
             for gr, gt in zip(grads_ref, grads_tr, strict=True):
                 self.assertTrue(torch.equal(gr, gt), f"Step {step}: grad mismatch")
 
-    def test_mismatched_module_raises(self):
-        """Executing with a module that has different params than at trace time raises."""
+    def test_non_tensor_leaf_raises(self):
+        """Passing a callable leaf in args raises (should be in closure instead)."""
+
+        def fn(model, x, loss_fn):
+            return loss_fn(model(x))
+
+        model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
+        tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
+
+        with self.assertRaises(ValueError, msg="all pytree leaves"):
+            minimal_fx_tracer(fn)(model, tokens, lambda x: x.sum())
+
+    def test_mismatched_module_raises_when_validation_enabled(self):
+        """Opt-in module FQN validation catches execution with the wrong module."""
         model, tokens, labels, loss_fn = self._make_mlp()
 
         def forward(model, tokens):
@@ -190,26 +202,13 @@ def forward(model, tokens):
 
         traced = minimal_fx_tracer(forward)(model, tokens)
 
-        # A model with a different architecture (extra layer → different FQNs).
         different_model = nn.Sequential(
             nn.Embedding(256, 64),
             nn.Linear(64, 256),
         ).to(device=self.DEVICE, dtype=self.DTYPE)
 
         with self.assertRaises(ValueError, msg="different parameter/buffer names"):
-            run_traced(traced, different_model, tokens)
-
-    def test_non_tensor_leaf_raises(self):
-        """Passing a callable leaf in args raises (should be in closure instead)."""
-
-        def fn(model, x, loss_fn):
-            return loss_fn(model(x))
-
-        model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
-        tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
-
-        with self.assertRaises(ValueError, msg="all pytree leaves"):
-            minimal_fx_tracer(fn)(model, tokens, lambda x: x.sum())
+            run_traced(traced, different_model, tokens, validate_module_fqns=True)
 
     def test_module_must_be_first_arg(self):
         def forward(tokens, model):
@@ -287,8 +286,7 @@ def forward(model, tokens):
 
         traced = minimal_fx_tracer(forward)(model, tokens_dt)
         has_subclass = any(
-            layout.meta is not None
-            for layout in traced.input_subclass_layouts.values()
+            layout.meta is not None for layout in traced.input_subclass_layouts.values()
         )
         self.assertTrue(has_subclass)
 
@@ -962,7 +960,7 @@ def run_grad(model):
                 loss = get_loss(logits, labels)
                 params = [p for p in model.parameters() if p.requires_grad]
                 grads = torch.autograd.grad(loss, params)
-                for p, g in zip(params, grads):
+                for p, g in zip(params, grads, strict=True):
                     p.grad = g
 
             # Warmup

From 9693ffeea557a63bd98d71d634fc8f3bb2261555 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 6 Apr 2026 08:49:58 -0700
Subject: [PATCH 25/42] Update base for Update on "[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-poisoned]
---
 torchtitan/experiments/graph_trainer/trainer.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index f2190bd004..7c508f9b8c 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -10,6 +10,9 @@
 import torch
 import torch.nn as nn
 
+from torchtitan.experiments.graph_trainer.common_utils import (
+    register_blockmask_pytree_node,
+)
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
 from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
@@ -101,6 +104,11 @@ def _make_fx_forward_backward_step(
     ) -> torch.Tensor:
         if self._traced_step is None:
             fwd_bwd_fn = make_fwd_bwd_step(self.loss_fn)
+            # Flex attention masks can flow through extra_inputs / extra_kwargs.
+            from torch.nn.attention.flex_attention import BlockMask
+
+            if BlockMask not in torch.utils._pytree.SUPPORTED_NODES:
+                register_blockmask_pytree_node()
             with self.train_context(), self.maybe_enable_amp:
                 self._traced_step = minimal_fx_tracer(fwd_bwd_fn)(
                     model,

From eba50380d2a31bae4a0640654f9c1726a9b2ba82 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 6 Apr 2026 11:16:31 -0700
Subject: [PATCH 26/42] Update base for Update on "[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-poisoned]
---
 .../graph_trainer/make_fx_tracer.py           | 180 ++++++++++--------
 .../graph_trainer/tests/test_trace_module.py  | 110 ++++++++---
 .../experiments/graph_trainer/trainer.py      |  14 +-
 3 files changed, 188 insertions(+), 116 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index 7e6c3183e2..20ab4cc7b1 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -276,7 +276,7 @@ def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
                 node.meta["stack_trace"] = stack_trace
 
 
-def _get_params_and_buffers(mod: nn.Module) -> dict[str, torch.Tensor]:
+def extract_module_state(mod: nn.Module) -> dict[str, torch.Tensor]:
     """Return a merged dict of the module's named parameters and buffers."""
     return {
         **dict(mod.named_parameters(remove_duplicate=False)),
@@ -291,8 +291,7 @@ class TracedResult:
     Attributes:
         gm: The traced FX graph as a pure function of flat tensors.
         example_inputs: Trace-time fake flat inputs used by downstream graph passes.
-        param_fqns: Trace-time parameter/buffer FQNs for execution-time validation.
-        num_params: Number of lifted parameters and buffers.
+        state_fqns: Trace-time state keys.
         num_flat_inputs: Number of flat graph inputs before subclass unwrapping.
         input_subclass_layouts: Subclass unwrap/rewrap metadata for inputs.
         num_flat_outputs: Number of flat graph outputs before subclass rewrapping.
@@ -302,8 +301,7 @@ class TracedResult:
 
     gm: torch.fx.GraphModule
     example_inputs: tuple[Any, ...]
-    param_fqns: list[str]
-    num_params: int
+    state_fqns: list[str]
     num_flat_inputs: int
     input_subclass_layouts: dict[int, SubclassLayout]
     num_flat_outputs: int
@@ -312,66 +310,51 @@ class TracedResult:
 
 
 def minimal_fx_tracer(fn: Callable) -> Callable[..., TracedResult]:
-    """Return a tracer for ``fn`` that traces a concrete module/input invocation.
+    """Return a tracer for a stateless ``fn`` with explicit ``state`` input.
 
     ``fn`` must be a plain callable (not an ``nn.Module``). The returned
-    callable expects the trace-time positional arguments, with the live module
-    at position 0.
+    callable expects ``state`` as the first positional argument, followed by
+    the traced user inputs::
 
-    Execute the returned :class:`TracedResult` with :func:`run_traced`, passing
-    the same positional arguments (with the live module first)::
+        traced_result = minimal_fx_tracer(train_step)(state, tokens, labels)
+        result = run_traced(traced_result, state, tokens, labels)
 
-        traced_result = minimal_fx_tracer(train_step)(model, tokens, labels)
-        result = run_traced(traced_result, model, tokens, labels)
-
-    The trace-time args must satisfy these constraints:
-    - ``args[0]`` must be an ``nn.Module`` whose parameters and buffers are
-      lifted as extra graph inputs
-    - there must be no additional ``nn.Module`` instances in ``args[1:]``
-    - all remaining pytree leaves must be tensors or make_fx-safe primitives
+    The trace-time ``state`` and ``args`` must satisfy these constraints:
+    - ``state`` must be a ``dict[str, Tensor]`` of parameters/buffers
+    - all pytree leaves must be tensors or make_fx-safe primitives
       (``int``, ``float``, ``bool``, ``str``, ``None``)
+    - there must be no ``nn.Module`` instances in ``state`` or ``args``
 
     Tensor subclasses (for example ``DTensor``) are recursively unwrapped into
     plain tensors for tracing, and the layouts needed to rewrap them are stored
     in the returned :class:`TracedResult`.
     """
 
-    def _trace_with_args(*args: Any) -> TracedResult:
-        if not isinstance(args[0], nn.Module):
-            raise ValueError(
-                "minimal_fx_tracer requires args[0] to be an nn.Module, "
-                f"got {type(args[0]).__name__}."
-            )
-        if any(isinstance(a, nn.Module) for a in args[1:]):
-            raise ValueError(
-                "minimal_fx_tracer supports exactly one nn.Module at args[0]. "
-                "Additional nn.Module instances found in args[1:]."
-            )
-        mod = args[0]
-
-        # Extract params/buffers from the module.
-        params_dict = _get_params_and_buffers(mod)
-        param_fqns = list(params_dict.keys())
-        params_flat = list(params_dict.values())
-        num_params = len(params_flat)
-
-        # User args: everything after the module.
-        user_args = list(args[1:])
+    def _trace_with_args(state: Any, *args: Any) -> TracedResult:
+        state_fqns = list(state.keys())
+        state_flat = list(state.values())
+        user_args = list(args)
         user_args_flat, user_args_spec = pytree.tree_flatten(user_args)
 
         # Validate leaves.
-        for leaf in user_args_flat:
+        for leaf in [*state_flat, *user_args_flat]:
+            if isinstance(leaf, nn.Module):
+                raise ValueError(
+                    "minimal_fx_tracer requires explicit tensor state, not nn.Module "
+                    "instances. Use trace_train_step(...) for the reference "
+                    "train-step wrapper."
+                )
             if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
                 raise ValueError(
-                    f"minimal_fx_tracer requires all pytree leaves in args to be tensors "
-                    f"or primitives (int/float/bool/str), got {type(leaf).__name__}. "
-                    f"Non-primitive values should either be registered as pytree "
-                    f"nodes (register_pytree_node) or constants "
+                    "minimal_fx_tracer requires all pytree leaves in state/args to "
+                    f"be tensors or primitives (int/float/bool/str), got "
+                    f"{type(leaf).__name__}. Non-primitive values should either be "
+                    "registered as pytree nodes (register_pytree_node) or constants "
                     f"(pytree.register_constant), or captured in fn's closure."
                 )
 
-        # Combined flat input: [*params, *user_args] with subclasses unwrapped.
-        full_args = params_flat + list(user_args_flat)
+        # Combined flat input: [*state, *user_args] with subclasses unwrapped.
+        full_args = list(state_flat) + list(user_args_flat)
         num_full_args = len(full_args)
         unwrapped_args, input_layouts = _unwrap_subclasses(full_args)
 
@@ -397,19 +380,14 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
             output_layouts = {}
 
             wrapped = _wrap_subclasses(plain_args, num_full_args, input_layouts)
-            params_wrapped = wrapped[:num_params]
-            user_flat = wrapped[num_params:]
+            state_wrapped = wrapped[: len(state_flat)]
+            user_flat = wrapped[len(state_flat) :]
 
-            params_for_mod = dict(zip(param_fqns, params_wrapped, strict=True))
+            state_for_fn = dict(zip(state_fqns, state_wrapped, strict=True))
             user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
 
-            # Reconstruct the original args: module at position 0 keeps the live
-            # module, remaining positions get the traced user tensors.
-            rebuilt = [mod] + user_list
-
-            with stateless._reparametrize_module(mod, params_for_mod):
-                with _patch_engine_run_backward():
-                    result = fn(*rebuilt)
+            with _patch_engine_run_backward():
+                result = fn(state_for_fn, *user_list)
 
             flat_outs, output_spec = pytree.tree_flatten(result)
             num_flat_outputs = len(flat_outs)
@@ -435,8 +413,7 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
         return TracedResult(
             gm=traced,
             example_inputs=fake_args,
-            param_fqns=param_fqns,
-            num_params=num_params,
+            state_fqns=state_fqns,
             num_flat_inputs=num_full_args,
             input_subclass_layouts=input_layouts,
             num_flat_outputs=num_flat_outputs,
@@ -449,36 +426,23 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
 
 def run_traced(
     traced_result: TracedResult,
+    state: Any,
     *args: Any,
-    validate_module_fqns: bool = False,
 ) -> Any:
     """Execute a traced graph with fresh parameters read from the live module.
 
     This is a reference implementation of traced-graph execution. It keeps the
-    parameter lookup, subclass unwrapping, and output reconstruction logic
+    state handling, subclass unwrapping, and output reconstruction logic
     explicit instead of baking those semantics into ``TracedResult`` itself.
-
-    The module must be the first argument (position 0), matching the
-    convention enforced by :func:`minimal_fx_tracer`.
     """
-
-    mod = args[0]
-    params_dict = _get_params_and_buffers(mod)
-    if validate_module_fqns:
-        fqns = list(params_dict.keys())
-        if fqns != traced_result.param_fqns:
-            raise ValueError(
-                f"Module at args[0] has different parameter/buffer "
-                f"names than during tracing.\n"
-                f"  Traced: {traced_result.param_fqns}\n"
-                f"  Got:    {fqns}"
-            )
-    params_flat = list(params_dict.values())
-
-    user_args = list(args[1:])
-    user_args_flat, _ = pytree.tree_flatten(user_args)
-
-    all_args = params_flat + list(user_args_flat)
+    state_flat = list(state.values())
+    user_args_flat, _ = pytree.tree_flatten(list(args))
+    if any(isinstance(leaf, nn.Module) for leaf in [*state_flat, *user_args_flat]):
+        raise ValueError(
+            "run_traced requires explicit tensor state, not nn.Module instances. "
+            "Use run_traced_train_step(...) for the reference train-step wrapper."
+        )
+    all_args = list(state_flat) + list(user_args_flat)
     flat_inputs, _ = _unwrap_subclasses(all_args)
 
     flat_outputs = traced_result.gm(*flat_inputs)
@@ -488,3 +452,57 @@ def run_traced(
         traced_result.output_subclass_layouts,
     )
     return pytree.tree_unflatten(wrapped, traced_result.output_spec)
+
+
+def trace_train_step(fn: Callable) -> Callable[..., TracedResult]:
+    """Reference implementation for capturing a whole train step via the core API."""
+
+    def _trace_with_module(module: nn.Module, *args: Any) -> TracedResult:
+        if not isinstance(module, nn.Module):
+            raise ValueError(
+                "trace_train_step requires args[0] to be an nn.Module, "
+                f"got {type(module).__name__}."
+            )
+        if any(isinstance(arg, nn.Module) for arg in args):
+            raise ValueError(
+                "trace_train_step supports exactly one nn.Module at args[0]. "
+                "Additional nn.Module instances found in args[1:]."
+            )
+
+        def _stateless_fn(state: dict[str, torch.Tensor], *user_args: Any) -> Any:
+            with stateless._reparametrize_module(module, state):
+                return fn(module, *user_args)
+
+        return minimal_fx_tracer(_stateless_fn)(extract_module_state(module), *args)
+
+    return _trace_with_module
+
+
+def run_traced_train_step(
+    traced_result: TracedResult,
+    module: nn.Module,
+    *args: Any,
+    validate_module_fqns: bool = False,
+) -> Any:
+    """Reference implementation for executing a traced whole train step."""
+
+    if not isinstance(module, nn.Module):
+        raise ValueError(
+            "run_traced_train_step requires args[0] to be an nn.Module, "
+            f"got {type(module).__name__}."
+        )
+    if any(isinstance(arg, nn.Module) for arg in args):
+        raise ValueError(
+            "run_traced_train_step supports exactly one nn.Module at args[0]. "
+            "Additional nn.Module instances found in args[1:]."
+        )
+
+    # TODO: Consider stronger state validation once the long-term state API settles.
+    state = extract_module_state(module)
+    if validate_module_fqns and list(state.keys()) != traced_result.state_fqns:
+        raise ValueError(
+            "module has different parameter/buffer names than during tracing.\n"
+            f"  Traced: {traced_result.state_fqns}\n"
+            f"  Got:    {list(state.keys())}"
+        )
+    return run_traced(traced_result, state, *args)
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 6fa54e6cea..0dad75d915 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -19,8 +19,11 @@
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
     _copy_fwd_metadata_to_bw_nodes,
     _patch_engine_run_backward,
+    extract_module_state,
     minimal_fx_tracer,
     run_traced,
+    run_traced_train_step,
+    trace_train_step,
 )
 
 # BlockMask must be registered as a pytree node so its tensor children
@@ -37,7 +40,7 @@ def get_loss(logits, labels):
 
 
 def make_train_step(loss_fn):
-    """Return a plain function for minimal_fx_tracer tracing.  loss_fn is captured in closure."""
+    """Return a plain function for module-first tracing. loss_fn is captured in closure."""
 
     def train_step(model, *args):
         *fwd_args, labels = args
@@ -50,6 +53,20 @@ def train_step(model, *args):
     return train_step
 
 
+def make_stateless_train_step(model, loss_fn):
+    """Return a state-first function for the minimal_fx_tracer core API."""
+
+    def train_step(state, *args):
+        *fwd_args, labels = args
+        with torch.nn.utils.stateless._reparametrize_module(model, state):
+            logits = model(*fwd_args)
+        loss = loss_fn(logits, labels)
+        grads = torch.autograd.grad(loss, list(state.values()))
+        return [loss] + list(grads)
+
+    return train_step
+
+
 def create_model(config_cls, model_config, device="cuda", dtype=torch.float32):
     model = config_cls(model_config)
     model.to(device=device, dtype=dtype)
@@ -122,9 +139,9 @@ def test_mlp_forward(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = minimal_fx_tracer(forward)(model, tokens)
+        traced = trace_train_step(forward)(model, tokens)
         out_eager = model(tokens)
-        wrapped = run_traced(traced, model, tokens)
+        wrapped = run_traced_train_step(traced, model, tokens)
         self.assertTrue(torch.equal(out_eager, wrapped))
 
     def test_mlp_train_step(self):
@@ -133,14 +150,14 @@ def test_mlp_train_step(self):
         model_test.load_state_dict(model_ref.state_dict())
 
         train_step = make_train_step(loss_fn)
-        traced = minimal_fx_tracer(train_step)(model_ref, tokens, labels)
+        traced = trace_train_step(train_step)(model_ref, tokens, labels)
 
         logits_ref = model_ref(tokens)
         loss_ref = loss_fn(logits_ref, labels)
         loss_ref.backward()
         grads_ref = [p.grad.clone() for p in model_ref.parameters()]
 
-        wrapped = run_traced(traced, model_test, tokens, labels)
+        wrapped = run_traced_train_step(traced, model_test, tokens, labels)
         loss_tr = wrapped[0]
         grads_tr = wrapped[1:]
 
@@ -154,7 +171,7 @@ def test_mlp_multistep_bitwise(self):
         model_test.load_state_dict(model_ref.state_dict())
 
         train_step = make_train_step(loss_fn)
-        traced = minimal_fx_tracer(train_step)(model_ref, tokens, labels)
+        traced = trace_train_step(train_step)(model_ref, tokens, labels)
 
         opt_ref = torch.optim.Adam(model_ref.parameters(), lr=self.LR)
         opt_copy = torch.optim.Adam(model_test.parameters(), lr=self.LR)
@@ -167,7 +184,7 @@ def test_mlp_multistep_bitwise(self):
             opt_ref.step()
             opt_ref.zero_grad()
 
-            wrapped = run_traced(traced, model_test, tokens, labels)
+            wrapped = run_traced_train_step(traced, model_test, tokens, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
@@ -191,7 +208,7 @@ def fn(model, x, loss_fn):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
         with self.assertRaises(ValueError, msg="all pytree leaves"):
-            minimal_fx_tracer(fn)(model, tokens, lambda x: x.sum())
+            trace_train_step(fn)(model, tokens, lambda x: x.sum())
 
     def test_mismatched_module_raises_when_validation_enabled(self):
         """Opt-in module FQN validation catches execution with the wrong module."""
@@ -200,7 +217,7 @@ def test_mismatched_module_raises_when_validation_enabled(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = minimal_fx_tracer(forward)(model, tokens)
+        traced = trace_train_step(forward)(model, tokens)
 
         different_model = nn.Sequential(
             nn.Embedding(256, 64),
@@ -208,17 +225,56 @@ def forward(model, tokens):
         ).to(device=self.DEVICE, dtype=self.DTYPE)
 
         with self.assertRaises(ValueError, msg="different parameter/buffer names"):
-            run_traced(traced, different_model, tokens, validate_module_fqns=True)
+            run_traced_train_step(
+                traced, different_model, tokens, validate_module_fqns=True
+            )
 
-    def test_module_must_be_first_arg(self):
-        def forward(tokens, model):
+    def test_trace_train_step_requires_module_first_arg(self):
+        def forward(model, tokens):
             return model(tokens)
 
+        tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
+
+        with self.assertRaises(ValueError, msg="args\\[0\\]"):
+            trace_train_step(forward)(tokens)
+
+    def test_core_explicit_state_executes(self):
         model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        with self.assertRaises(ValueError, msg="args[0]"):
-            minimal_fx_tracer(forward)(tokens, model)
+        def forward(state, tokens):
+            with torch.nn.utils.stateless._reparametrize_module(model, state):
+                return model(tokens)
+
+        state = extract_module_state(model)
+        traced = minimal_fx_tracer(forward)(state, tokens)
+        out_ref = forward(state, tokens)
+        out_traced = run_traced(traced, state, tokens)
+
+        self.assertTrue(torch.equal(out_ref, out_traced))
+
+    def test_core_explicit_state_train_step(self):
+        model_ref, tokens, labels, loss_fn = self._make_mlp()
+        model_test = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE)
+        model_test.load_state_dict(model_ref.state_dict())
+
+        state_ref = extract_module_state(model_ref)
+        state_test = extract_module_state(model_test)
+        train_step = make_stateless_train_step(model_ref, loss_fn)
+        traced = minimal_fx_tracer(train_step)(state_ref, tokens, labels)
+
+        logits_ref = model_ref(tokens)
+        loss_ref = loss_fn(logits_ref, labels)
+        loss_ref.backward()
+        grads_ref = [p.grad.clone() for p in model_ref.parameters()]
+
+        wrapped = run_traced(traced, state_test, tokens, labels)
+        loss_tr = wrapped[0]
+        grads_tr = wrapped[1:]
+
+        self.assertTrue(torch.equal(loss_ref, loss_tr))
+        for gr, gt in zip(grads_ref, grads_tr, strict=True):
+            self.assertTrue(torch.equal(gr, gt))
 
     def test_additional_module_arg_raises(self):
         def forward(model, other_model, tokens):
@@ -230,7 +286,7 @@ def forward(model, other_model, tokens):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
         with self.assertRaises(ValueError, msg="Additional nn.Module"):
-            minimal_fx_tracer(forward)(model, other_model, tokens)
+            trace_train_step(forward)(model, other_model, tokens)
 
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
@@ -284,14 +340,14 @@ def test_dtensor_forward(self):
         def forward(model, tokens):
             return model(tokens)
 
-        traced = minimal_fx_tracer(forward)(model, tokens_dt)
+        traced = trace_train_step(forward)(model, tokens_dt)
         has_subclass = any(
             layout.meta is not None for layout in traced.input_subclass_layouts.values()
         )
         self.assertTrue(has_subclass)
 
         out_eager = model(tokens_dt)
-        wrapped = run_traced(traced, model, tokens_dt)
+        wrapped = run_traced_train_step(traced, model, tokens_dt)
         self.assertTrue(torch.equal(out_eager.full_tensor(), wrapped.full_tensor()))
 
     def test_dtensor_train_step(self):
@@ -313,14 +369,14 @@ def test_dtensor_train_step(self):
         labels_dt = DTensor.from_local(labels, mesh, [Replicate()])
 
         train_step = make_train_step(get_loss)
-        traced = minimal_fx_tracer(train_step)(model_ref, tokens_dt, labels_dt)
+        traced = trace_train_step(train_step)(model_ref, tokens_dt, labels_dt)
 
         logits_ref = model_ref(tokens_dt)
         loss_ref = get_loss(logits_ref, labels_dt)
         loss_ref.backward()
         grads_ref = [p.grad.clone() for p in model_ref.parameters()]
 
-        wrapped = run_traced(traced, model_test, tokens_dt, labels_dt)
+        wrapped = run_traced_train_step(traced, model_test, tokens_dt, labels_dt)
         loss_tr = wrapped[0]
         grads_tr = wrapped[1:]
 
@@ -346,7 +402,7 @@ def test_backward_nodes_have_seq_nr(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
+        traced = trace_train_step(train_step)(model, tokens, labels)
 
         # Collect seq_nr values from all call_function nodes
         seq_nrs = []
@@ -374,7 +430,7 @@ def test_copy_fwd_metadata_propagates_custom(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
+        traced = trace_train_step(train_step)(model, tokens, labels)
         gm = traced.gm
 
         # Manually set custom metadata on the first fwd node for each seq_nr
@@ -411,7 +467,7 @@ def test_backward_nodes_have_stack_trace(self):
         tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE)
         labels = torch.randint(0, 256, (2, 32), device=self.DEVICE)
 
-        traced = minimal_fx_tracer(train_step)(model, tokens, labels)
+        traced = trace_train_step(train_step)(model, tokens, labels)
 
         # Find backward nodes: nodes sharing a seq_nr with an earlier (forward) node
         seq_nr_first: dict[int, torch.fx.Node] = {}
@@ -491,7 +547,7 @@ def _run_bitwise_test(
             else contextlib.nullcontext()
         )
         with maybe_regional_inductor:
-            traced = minimal_fx_tracer(train_step)(model_ref, *fwd_args, labels)
+            traced = trace_train_step(train_step)(model_ref, *fwd_args, labels)
 
         if check_collective_ops:
             ag = sum(
@@ -523,7 +579,7 @@ def _run_bitwise_test(
             opt_ref.step()
             opt_ref.zero_grad()
 
-            wrapped = run_traced(traced, model_test, *fwd_args, labels)
+            wrapped = run_traced_train_step(traced, model_test, *fwd_args, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
@@ -702,7 +758,7 @@ def test_flex_attention_annotations(self):
             def forward(model, tokens, attn_masks):
                 return model(tokens, attn_masks)
 
-            traced = minimal_fx_tracer(forward)(model, tokens, attn_masks)
+            traced = trace_train_step(forward)(model, tokens, attn_masks)
 
         flex_nodes = [
             n
@@ -796,7 +852,7 @@ def _run_fsdp_model_test(
             else contextlib.nullcontext()
         )
         with maybe_regional_inductor:
-            traced = minimal_fx_tracer(train_step)(model_ref, *fwd_args, labels)
+            traced = trace_train_step(train_step)(model_ref, *fwd_args, labels)
 
         ag = sum(
             1
@@ -825,7 +881,7 @@ def _run_fsdp_model_test(
             opt_ref.step()
             opt_ref.zero_grad()
 
-            wrapped = run_traced(traced, model_test, *fwd_args, labels)
+            wrapped = run_traced_train_step(traced, model_test, *fwd_args, labels)
             loss_tr = wrapped[0]
             grads_tr = wrapped[1:]
             for p, g in zip(model_test.parameters(), grads_tr, strict=True):
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index 7c508f9b8c..992edb9dba 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -16,9 +16,9 @@
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
 from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
-    minimal_fx_tracer,
-    run_traced,
     TracedResult,
+    run_traced_train_step,
+    trace_train_step,
 )
 from torchtitan.experiments.graph_trainer.passes import apply_default_graph_passes
 from torchtitan.trainer import Trainer
@@ -30,9 +30,7 @@ def make_fwd_bwd_step(loss_fn):
     ``loss_fn`` is captured in the closure so it is not a graph input.
     """
 
-    def fwd_bwd_step(
-        model, inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs
-    ):
+    def fwd_bwd_step(model, inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs):
         pred = model(inputs, **extra_inputs, **extra_kwargs)
         loss = loss_fn(pred, labels) / global_valid_tokens
         params = [p for p in model.parameters() if p.requires_grad]
@@ -110,7 +108,7 @@ def _make_fx_forward_backward_step(
             if BlockMask not in torch.utils._pytree.SUPPORTED_NODES:
                 register_blockmask_pytree_node()
             with self.train_context(), self.maybe_enable_amp:
-                self._traced_step = minimal_fx_tracer(fwd_bwd_fn)(
+                self._traced_step = trace_train_step(fwd_bwd_fn)(
                     model,
                     inputs,
                     labels,
@@ -124,7 +122,7 @@ def _make_fx_forward_backward_step(
                 self._traced_step.example_inputs,
             )
         with self.train_context(), self.maybe_enable_amp:
-            outputs = run_traced(
+            outputs = run_traced_train_step(
                 self._traced_step,
                 model,
                 inputs,
@@ -136,7 +134,7 @@ def _make_fx_forward_backward_step(
         loss = outputs[0]
         grads = outputs[1:]
 
-        for param, grad in zip(params, grads):
+        for param, grad in zip(params, grads, strict=True):
             if param.grad is None:
                 param.grad = grad
             else:

From f1ae755d98cdace28b8a7c58839ab54299c49ca1 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 6 Apr 2026 11:38:59 -0700
Subject: [PATCH 27/42] Update base for Update on "[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-poisoned]
---
 torchtitan/experiments/graph_trainer/common_utils.py   | 10 +++++++++-
 .../graph_trainer/tests/test_trace_module.py           | 10 ++++------
 torchtitan/experiments/graph_trainer/trainer.py        |  8 ++------
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py
index c0e647e6e6..f85473da76 100644
--- a/torchtitan/experiments/graph_trainer/common_utils.py
+++ b/torchtitan/experiments/graph_trainer/common_utils.py
@@ -70,6 +70,14 @@ def register_blockmask_pytree_node():
         )
 
 
+def maybe_register_blockmask_pytree_node() -> None:
+    """Register BlockMask if it is not already in the global pytree registry."""
+    from torch.nn.attention.flex_attention import BlockMask
+
+    if BlockMask not in torch.utils._pytree.SUPPORTED_NODES:
+        register_blockmask_pytree_node()
+
+
 def end_with_pass(passes: list[Callable], names: list[str]) -> bool:
     return (
         len(passes) > 0
@@ -177,7 +185,7 @@ def apply_graph_ac(
     if ac_config.mode != "selective":
         raise ValueError(
             f"graph_trainer only supports activation_checkpoint.mode 'selective' or "
-            f"'none', got '{ac_config.mode}'. Use 'selective' for graph-based SAC."
+            f"'none', got {ac_config.mode!r}. Use 'selective' for graph-based SAC."
         )
 
     joint_pass_names = getattr(compile_config, "joint_passes", [])
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index 0dad75d915..ea56782ab6 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -14,7 +14,7 @@
 
 from torchtitan.experiments.graph_trainer.common_utils import (
     annotate_flex_attention_for_regional_inductor,
-    register_blockmask_pytree_node,
+    maybe_register_blockmask_pytree_node,
 )
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
     _copy_fwd_metadata_to_bw_nodes,
@@ -26,11 +26,6 @@
     trace_train_step,
 )
 
-# BlockMask must be registered as a pytree node so its tensor children
-# are properly traced as graph inputs instead of opaque leaves.
-register_blockmask_pytree_node()
-
-
 def get_loss(logits, labels):
     return torch.nn.functional.cross_entropy(
         logits.flatten(0, 1).float(),
@@ -546,6 +541,7 @@ def _run_bitwise_test(
             if use_regional_inductor
             else contextlib.nullcontext()
         )
+        maybe_register_blockmask_pytree_node()
         with maybe_regional_inductor:
             traced = trace_train_step(train_step)(model_ref, *fwd_args, labels)
 
@@ -753,6 +749,7 @@ def test_flex_attention_annotations(self):
             "basic_mask": basic_mask,
             "sliding_window_mask": sliding_window_mask,
         }
+        maybe_register_blockmask_pytree_node()
         with annotate_flex_attention_for_regional_inductor():
 
             def forward(model, tokens, attn_masks):
@@ -851,6 +848,7 @@ def _run_fsdp_model_test(
             if use_regional_inductor
             else contextlib.nullcontext()
         )
+        maybe_register_blockmask_pytree_node()
         with maybe_regional_inductor:
             traced = trace_train_step(train_step)(model_ref, *fwd_args, labels)
 
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index 992edb9dba..8610acabd4 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -11,7 +11,7 @@
 import torch.nn as nn
 
 from torchtitan.experiments.graph_trainer.common_utils import (
-    register_blockmask_pytree_node,
+    maybe_register_blockmask_pytree_node,
 )
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
 from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown
@@ -102,11 +102,7 @@ def _make_fx_forward_backward_step(
     ) -> torch.Tensor:
         if self._traced_step is None:
             fwd_bwd_fn = make_fwd_bwd_step(self.loss_fn)
-            # Flex attention masks can flow through extra_inputs / extra_kwargs.
-            from torch.nn.attention.flex_attention import BlockMask
-
-            if BlockMask not in torch.utils._pytree.SUPPORTED_NODES:
-                register_blockmask_pytree_node()
+            maybe_register_blockmask_pytree_node()
             with self.train_context(), self.maybe_enable_amp:
                 self._traced_step = trace_train_step(fwd_bwd_fn)(
                     model,

From 3bcc83815e1ae34865cb1a9870855a83578cfcff Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Mon, 6 Apr 2026 13:31:45 -0700
Subject: [PATCH 28/42] Update base for Update on "[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-poisoned]
---
 torchtitan/experiments/graph_trainer/common_utils.py      | 8 +++++---
 .../experiments/graph_trainer/tests/test_trace_module.py  | 1 +
 torchtitan/experiments/graph_trainer/trainer.py           | 6 ++++--
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py
index f85473da76..6d5dc95be4 100644
--- a/torchtitan/experiments/graph_trainer/common_utils.py
+++ b/torchtitan/experiments/graph_trainer/common_utils.py
@@ -12,7 +12,7 @@
 import torch.nn as nn
 from torch.distributed.tensor import DTensor, Replicate
 from torch.fx.traceback import annotate_fn
-from torch.utils._pytree import register_pytree_node, tree_map
+from torch.utils._pytree import register_constant, register_pytree_node, tree_map
 
 from torchtitan.config import CompileConfig
 from torchtitan.distributed import ParallelDims
@@ -71,11 +71,13 @@ def register_blockmask_pytree_node():
 
 
 def maybe_register_blockmask_pytree_node() -> None:
-    """Register BlockMask if it is not already in the global pytree registry."""
-    from torch.nn.attention.flex_attention import BlockMask
+    """Register flex-attention pytree helpers if they are missing."""
+    from torch.nn.attention.flex_attention import BlockMask, _MaskModWrapper
 
     if BlockMask not in torch.utils._pytree.SUPPORTED_NODES:
         register_blockmask_pytree_node()
+    if _MaskModWrapper not in torch.utils._pytree.SUPPORTED_NODES:
+        register_constant(_MaskModWrapper)
 
 
 def end_with_pass(passes: list[Callable], names: list[str]) -> bool:
diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
index ea56782ab6..9d774dc502 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py
@@ -26,6 +26,7 @@
     trace_train_step,
 )
 
+
 def get_loss(logits, labels):
     return torch.nn.functional.cross_entropy(
         logits.flatten(0, 1).float(),
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index 8610acabd4..377e5f3afd 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -16,9 +16,9 @@
 from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig
 from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
-    TracedResult,
     run_traced_train_step,
     trace_train_step,
+    TracedResult,
 )
 from torchtitan.experiments.graph_trainer.passes import apply_default_graph_passes
 from torchtitan.trainer import Trainer
@@ -30,7 +30,9 @@ def make_fwd_bwd_step(loss_fn):
     ``loss_fn`` is captured in the closure so it is not a graph input.
     """
 
-    def fwd_bwd_step(model, inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs):
+    def fwd_bwd_step(
+        model, inputs, labels, global_valid_tokens, extra_inputs, extra_kwargs
+    ):
         pred = model(inputs, **extra_inputs, **extra_kwargs)
         loss = loss_fn(pred, labels) / global_valid_tokens
         params = [p for p in model.parameters() if p.requires_grad]

From 9b92650c47818fb3ffa3802af77194e4994d2dde Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 7 Apr 2026 15:18:48 -0700
Subject: [PATCH 29/42] Update base for Update on "[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-poisoned]

From 01d002a63208e2f2275a8e318e254a1a1b8bd0c0 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 7 Apr 2026 20:05:51 -0700
Subject: [PATCH 30/42] Update base for Update on "[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-poisoned]

From fa359987e170684ee3ebb1409c1d9b8385e59d44 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Wed, 8 Apr 2026 11:17:16 -0700
Subject: [PATCH 31/42] Update base for 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.

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-poisoned]

From 0ea91a0ab3c6eec470ef83d71df3739a813c29bb Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Thu, 9 Apr 2026 11:23:29 -0700
Subject: [PATCH 32/42] Update base for 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.

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-poisoned]

From 32942cdc6be2c035258cf21a5d12d5d644d72b33 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Sun, 12 Apr 2026 15:17:29 -0700
Subject: [PATCH 33/42] Update base for 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.

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-poisoned]

From a838dc2d89a0599debe07aa02d6b8646d5f6e7f0 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Sun, 12 Apr 2026 20:51:31 -0700
Subject: [PATCH 34/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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 aot and aot_fx_trace

<img width="918" height="750" alt="Screenshot 2026-04-12 at 11 39 01 PM" src="https://github.com/user-attachments/assets/059188c5-fc14-4a34-876a-966d938bef3b" />


[ghstack-poisoned]

From 7bb5d546a1ff8170db9b696cf22f4887cba1071e Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Sun, 12 Apr 2026 21:43:31 -0700
Subject: [PATCH 35/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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 aot and aot_fx_trace

<img width="918" height="750" alt="Screenshot 2026-04-12 at 11 39 01 PM" src="https://github.com/user-attachments/assets/059188c5-fc14-4a34-876a-966d938bef3b" />


[ghstack-poisoned]

From 404b1a3d275fa319ef0375580f12a755a5ef6f99 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 14 Apr 2026 10:10:59 -0700
Subject: [PATCH 36/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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 aot and aot_fx_trace

<img width="918" height="750" alt="Screenshot 2026-04-12 at 11 39 01 PM" src="https://github.com/user-attachments/assets/059188c5-fc14-4a34-876a-966d938bef3b" />


[ghstack-poisoned]

From 5a02e95117d0901366401578fd9c859a9090af9b Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 14 Apr 2026 10:24:34 -0700
Subject: [PATCH 37/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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





[ghstack-poisoned]

From b28baa4b334439ecc3aa8c8e022239fab5d5fc73 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 14 Apr 2026 10:42:13 -0700
Subject: [PATCH 38/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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





[ghstack-poisoned]

From e7e251cb4cfa5cedaf363520606b181cf61f6ae5 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 14 Apr 2026 11:10:00 -0700
Subject: [PATCH 39/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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





[ghstack-poisoned]

From 64bf2c7805362f6de375673a083d2687af60f06f Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Tue, 14 Apr 2026 21:12:02 -0700
Subject: [PATCH 40/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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]

From 7f57f921bdd5a31ab65bb225e360b4a23e21a6ac Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Wed, 15 Apr 2026 11:09:59 -0700
Subject: [PATCH 41/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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]

From 723e0e468e18db936530b741e69d84da099a4921 Mon Sep 17 00:00:00 2001
From: Tugsbayasgalan Manlaibaatar <tmanlaibaatar@fb.com>
Date: Wed, 15 Apr 2026 11:50:48 -0700
Subject: [PATCH 42/42] Update base for Update on "[graph_trainer] Add
 torch.no_grad() and graph-based SAC to traced execution"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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]