[GraphTrainer] Support transformer_block_bucketing in aot_fx_trace mode

aot_fx_trace mode traces with record_module_stack=False, so nodes lack
nn_module_stack metadata that manual_overlap_bucketing uses to match
nodes to modules. This change uses the existing annotate_fn mechanism
(already used for AC region tagging) to tag module FQNs during tracing,
then provides a custom module_stack_fn to manual_overlap_bucketing.

Changes:
- Add annotate_module_fqns() and get_module_stack_from_annotation() to
  common_utils.py
- Call annotate_module_fqns() in both llama3 and deepseek_v3 annotate
  functions
- Add module_stack_fn parameter to
  transformer_block_bucketing_reordering_pass
- Update construct_default_graph_passes to accept compile_config, model,
  and parallel_dims, and wire up the bucketing pass with reassign_to_pg
- Pass config/model/parallel_dims from GraphTrainer to
  construct_default_graph_passes
- Add numerics tests for both llama3 and deepseek_v3 with aot_fx_trace
  + transformer_block_bucketing

Author: yiming0416

torchtitan/experiments/graph_trainer/common_utils.py

@@ -18,6 +18,18 @@
 from torchtitan.tools.logging import logger
 
 _AC_REGION_ID = "ac_region_id"
+_MODULE_FQN = "module_fqn"
+
+
+def annotate_module_fqns(model: nn.Module) -> None:
+    """Annotate all modules' forward with their fully-qualified names.
+
+    Every named submodule (excluding the root) gets its forward method wrapped
+    with annotate_fn so that FX nodes carry the module_fqn metadata.
+    """
+    for fqn, submodule in model.named_modules():
+        if fqn:  # skip root module
+            submodule.forward = annotate_fn({_MODULE_FQN: fqn})(submodule.forward)
 
 
 def annotate_ac_regions(model: nn.Module) -> None:

torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py

@@ -19,6 +19,7 @@
 from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
 from torchtitan.experiments.graph_trainer.common_utils import (
     annotate_ac_regions,
+    annotate_module_fqns,
     apply_graph_ac,
 )
 from torchtitan.experiments.graph_trainer.compile import apply_compile
@@ -55,6 +56,7 @@ def annotate_deepseekv3(model: GraphTrainerDeepSeekV3Model) -> None:
     )
     MoE.forward = annotate_fn({"EP": "compute"})(MoE.forward)
 
+    annotate_module_fqns(model)
     annotate_ac_regions(model)
 
 

torchtitan/experiments/graph_trainer/llama3/parallelize.py

@@ -16,6 +16,7 @@
 from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
 from torchtitan.experiments.graph_trainer.common_utils import (
     annotate_ac_regions,
+    annotate_module_fqns,
     apply_graph_ac,
 )
 from torchtitan.experiments.graph_trainer.compile import apply_compile
@@ -32,10 +33,14 @@
 def annotate_llama(model: GraphTrainerLlama3Model) -> None:
     """Attach annotations to FX graph nodes with ``torch.fx.traceback.annotate_fn``
 
+    - Module FQN annotation: Tags each submodule's forward with its
+      fully-qualified name so that the transformer_block_bucketing pass
+      can identify which graph nodes belong to which module.
     - AC region annotation: Tags each transformer block's forward with a unique
       ac_region_id so that apply_sac_pass can assign per-block ac_graph_id
       boundaries for the min-cut partitioner.
     """
+    annotate_module_fqns(model)
     annotate_ac_regions(model)
 
 

torchtitan/experiments/graph_trainer/make_fx_tracer.py

@@ -12,9 +12,6 @@
 import torch
 import torch.nn as nn
 import torch.utils._pytree as pytree
-from torch._functorch._aot_autograd.logging_utils import (
-    setup_stacktrace_preservation_hooks,
-)
 from torch._guards import tracing, TracingContext
 from torch._subclasses import FakeTensorMode
 from torch.fx.experimental.proxy_tensor import make_fx
@@ -200,52 +197,6 @@ def _remove_cpu_shadow_chains(gm: torch.fx.GraphModule) -> None:
     gm.recompile()
 
 
-@contextmanager
-def _patch_engine_run_backward() -> Generator[None, None, None]:
-    """Patch _engine_run_backward to install stacktrace preservation hooks.
-
-    Why this is needed:
-    When make_fx traces a function that calls loss.backward(), the backward
-    pass is decomposed into primitive ATen ops. Normally (in eager autograd),
-    ``setup_stacktrace_preservation_hooks`` is called by the autograd engine
-    to propagate ``seq_nr`` from forward ops to their corresponding backward
-    ops. Under make_fx tracing, this hook setup doesn't happen automatically
-    because the engine path differs, so backward FX nodes end up without
-    ``seq_nr`` metadata. Without ``seq_nr``, we can't correlate backward
-    nodes back to their forward counterparts (needed by
-    ``_copy_fwd_metadata_to_bw_nodes``).
-
-    This context manager patches ``_engine_run_backward`` to call
-    ``setup_stacktrace_preservation_hooks`` before the autograd engine runs,
-    restoring ``seq_nr`` propagation during tracing.
-
-    We must patch the name in both modules since ``torch.autograd.__init__``
-    imports it via ``from .graph import``.
-    """
-    import torch.autograd
-    import torch.autograd.graph
-
-    _orig_fn = torch.autograd.graph._engine_run_backward
-
-    def _patched(t_outputs, *args, **kwargs):  # type: ignore[no-untyped-def]
-        roots = [
-            t.grad_fn
-            for t in t_outputs
-            if isinstance(t, torch.Tensor) and t.grad_fn is not None
-        ]
-        if roots:
-            setup_stacktrace_preservation_hooks(roots)
-        return _orig_fn(t_outputs, *args, **kwargs)
-
-    torch.autograd.graph._engine_run_backward = _patched  # type: ignore[assignment]
-    torch.autograd._engine_run_backward = _patched  # type: ignore[assignment]
-    try:
-        yield
-    finally:
-        torch.autograd.graph._engine_run_backward = _orig_fn  # type: ignore[assignment]
-        torch.autograd._engine_run_backward = _orig_fn  # type: ignore[assignment]
-
-
 def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
     """Copy forward node metadata (custom) to later nodes sharing the same seq_nr.
 
@@ -402,17 +353,25 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
             state_for_fn = dict(zip(state_fqns, state_wrapped, strict=True))
             user_list = pytree.tree_unflatten(list(user_flat), user_args_spec)
 
-            with _patch_engine_run_backward():
-                result = fn(state_for_fn, *user_list)
+            result = fn(state_for_fn, *user_list)
 
             flat_outs, output_spec = pytree.tree_flatten(result)
             num_flat_outputs = len(flat_outs)
             unwrapped_outs, output_layouts = _unwrap_subclasses(flat_outs)
             return unwrapped_outs
 
         ctx = TracingContext(fake_mode)
-        # preserve_node_meta propagates fx.traceback.annotate metadata to traced nodes
-        with fake_mode, tracing(ctx), preserve_node_meta(), _skip_nested_compile():
+        # preserve_node_meta propagates fx.traceback.annotate metadata to traced nodes.
+        # _non_strict_tracing_context + _patch_autograd_grad install autograd
+        # hooks so backward nodes get seq_nr and autograd_backward metadata.
+        with (
+            fake_mode,
+            tracing(ctx),
+            preserve_node_meta(),
+            _skip_nested_compile(),
+            torch.compiler._non_strict_tracing_context(),
+            torch.compiler._patch_autograd_grad(),
+        ):
             traced = make_fx(
                 fn_with_subclass_handling,
                 record_stack_traces=True,

torchtitan/experiments/graph_trainer/passes.py

@@ -36,7 +36,7 @@
 from torch.utils.checkpoint import CheckpointPolicy
 
 from torchtitan.distributed.activation_checkpoint import _get_save_ops
-from torchtitan.experiments.graph_trainer.common_utils import _AC_REGION_ID
+from torchtitan.experiments.graph_trainer.common_utils import _AC_REGION_ID, _MODULE_FQN
 from torchtitan.experiments.graph_trainer.custom_codegen import custom_codegen_pass
 from torchtitan.experiments.graph_trainer.make_fx_tracer import TracedResult
 from torchtitan.experiments.graph_trainer.remove_noop_passes import (
@@ -52,6 +52,7 @@
 
 def construct_default_graph_passes(
     traced_result: "TracedResult",
+    model: torch.nn.Module,
 ) -> list[Callable]:
     """Build the default pass list for the aot_fx_trace compile path.
 
@@ -61,17 +62,25 @@ def construct_default_graph_passes(
 
     Args:
         traced_result: The traced graph and metadata from ``trace_train_step``.
+        model: The model for extracting transformer block buckets.
 
     Returns:
         An ordered list of graph passes ready to apply.
     """
+    from torchtitan.experiments.graph_trainer.common_utils import (
+        get_transformer_block_buckets,
+    )
     from torchtitan.models.common.attention import FlexAttention
 
     passes: list[Callable] = [
         functools.partial(tlparse_log_graph_pass, graph_name="make_fx_graph_traced"),
         remove_detach_pass,
         remove_identity_view_pass,
         remove_identity_slice_pass,
+        functools.partial(
+            joint_transformer_block_bucketing_reordering_pass,
+            fsdp_manual_buckets=get_transformer_block_buckets(model),
+        ),
         # FlexAttention HOPs must be compiled (via regional_inductor) to
         # produce bitwise identical results to the eager Trainer path.
         # When left uncompiled, flex_attention still runs correctly but
@@ -141,6 +150,31 @@ def autobucketing_reordering_pass(
     return gm
 
 
+def _make_module_fqn_stack_fn(
+    *, bwd: bool
+) -> Callable[[torch.fx.Node], list[tuple[str, type]]]:
+    """Create a module_stack_fn that filters nodes by forward/backward direction.
+
+    Args:
+        bwd: If True, only return FQNs for backward nodes. If False, only
+            return FQNs for forward nodes.
+
+    Returns:
+        A callable compatible with make_graph_view's module_stack_fn parameter.
+    """
+
+    def _stack_fn(node: torch.fx.Node) -> list[tuple[str, type]]:
+        is_bwd = node.meta.get("autograd_backward", False)
+        if is_bwd != bwd:
+            return []
+        fqn = node.meta.get("custom", {}).get(_MODULE_FQN)
+        if not fqn:
+            return []
+        return [(fqn, torch.nn.Module)]
+
+    return _stack_fn
+
+
 def transformer_block_bucketing_reordering_pass(
     gm: torch.fx.GraphModule,
     example_inputs: tuple | None = None,
@@ -157,6 +191,48 @@ def transformer_block_bucketing_reordering_pass(
     return gm
 
 
+def joint_transformer_block_bucketing_reordering_pass(
+    gm: torch.fx.GraphModule,
+    example_inputs: tuple | None = None,
+    *,
+    fsdp_manual_buckets,
+) -> torch.fx.GraphModule:
+    """Apply manual bucketing and reordering on a joint fwd+bwd graph.
+
+    The joint graph is processed in two passes to ensure each direction's
+    collectives are bucketed independently and reordering uses the correct
+    execution order (forward order for fwd, reversed order for bwd).
+
+    Used by the aot_fx_trace mode where forward and backward are in a
+    single graph.
+    """
+    # Forward pass: bucket and reorder in forward execution order
+    manual_overlap_bucketing(
+        gm,
+        module_bucket_plans=fsdp_manual_buckets,
+        insert_overlap_deps=False,
+        module_stack_fn=_make_module_fqn_stack_fn(bwd=False),
+    )
+
+    # Clear manual_bucket_node_type metadata left by the forward pass.
+    # The reordering is already applied via topological sort, but stale
+    # metadata would confuse the second call's _manual_reorder_graph
+    # which iterates all graph nodes regardless of direction.
+    for node in gm.graph.nodes:
+        node.meta.pop("manual_bucket_node_type", None)
+
+    # Backward pass: bucket and reorder backward collectives separately
+    manual_overlap_bucketing(
+        gm,
+        module_bucket_plans=fsdp_manual_buckets,
+        insert_overlap_deps=False,
+        module_stack_fn=_make_module_fqn_stack_fn(bwd=True),
+    )
+
+    gm.recompile()
+    return gm
+
+
 def _ops_filter_with_distributed(name: str) -> bool:
     """Ops filter that allows distributed collective ops for serialization.
 

torchtitan/experiments/graph_trainer/qwen3/parallelize.py

@@ -19,6 +19,7 @@
 from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
 from torchtitan.experiments.graph_trainer.common_utils import (
     annotate_ac_regions,
+    annotate_module_fqns,
     apply_graph_ac,
 )
 from torchtitan.experiments.graph_trainer.compile import apply_compile
@@ -66,6 +67,7 @@ def annotate_qwen3(model: GraphTrainerQwen3Model) -> None:
         FlexAttention.forward
     )
 
+    annotate_module_fqns(model)
     annotate_ac_regions(model)
 
 

torchtitan/experiments/graph_trainer/trainer.py

@@ -128,7 +128,10 @@ def _make_fx_forward_backward_step(
                 )
 
             if self.config.compile.enable_passes:
-                passes = construct_default_graph_passes(self._traced_step)
+                passes = construct_default_graph_passes(
+                    self._traced_step,
+                    model=model,
+                )
                 self._traced_step.gm = apply_graph_passes(
                     self._traced_step.gm,
                     self._traced_step.example_inputs,