[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.
 
@@ -297,6 +248,9 @@ class TracedResult:
         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.
+        model: The original nn.Module, cached for downstream passes that need
+            module structure (e.g. transformer block bucketing). Only set when
+            traced via :func:`trace_train_step`.
     """
 
     gm: torch.fx.GraphModule
@@ -307,6 +261,7 @@ class TracedResult:
     num_flat_outputs: int
     output_subclass_layouts: dict[int, SubclassLayout]
     output_spec: pytree.TreeSpec
+    model: nn.Module | None = None
 
     @property
     def num_static_inputs(self) -> int:
@@ -402,17 +357,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,
@@ -489,7 +452,9 @@ 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)
+        result = minimal_fx_tracer(_stateless_fn)(extract_module_state(module), *args)
+        result.model = module
+        return result
 
     return _trace_with_module
 

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 (
@@ -65,13 +65,20 @@ def construct_default_graph_passes(
     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(traced_result.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
@@ -157,6 +164,62 @@ 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.
+    """
+
+    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."""
+
+        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
+
+    # 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),
+    )
+
+    # keep it until https://github.com/pytorch/pytorch/pull/180579 lands
+    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)