[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,31 @@
 from torchtitan.tools.logging import logger
 
 _AC_REGION_ID = "ac_region_id"
+_IS_BWD = "is_bwd"
+_MODULE_FQN = "module_fqn"
+
+
+def _flatten_bucket_fqns(bucket_plans: list[list[str] | str]) -> list[str]:
+    """Flatten a possibly nested bucket plan list into a flat list of FQNs."""
+    result = []
+    for item in bucket_plans:
+        if isinstance(item, list):
+            result.extend(item)
+        else:
+            result.append(item)
+    return result
+
+
+def annotate_module_fqns(model: nn.Module) -> None:
+    """Annotate bucket-level modules' forward with their fully-qualified names.
+
+    Uses get_transformer_block_buckets to determine which modules to annotate,
+    ensuring the annotation granularity matches the bucketing granularity.
+    """
+    fqns = _flatten_bucket_fqns(get_transformer_block_buckets(model))
+    for fqn in fqns:
+        submodule = model.get_submodule(fqn)
+        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

@@ -252,8 +252,11 @@ def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
     Walks the graph in a single pass. The first node seen for each seq_nr is
     treated as the forward node.
     Subsequent nodes with the same seq_nr (typically backward nodes) receive
-    the forward node's custom metadata.
+    the forward node's custom metadata and are marked with ``is_bwd=True``
+    in their ``custom`` dict.
     """
+    from torchtitan.experiments.graph_trainer.common_utils import _IS_BWD
+
     seq_nr_to_fwd_node: dict[int, torch.fx.Node] = {}
 
     for node in fx_g.graph.nodes:
@@ -268,6 +271,7 @@ def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
             custom = fwd_node.meta.get("custom")
             if custom:
                 node.meta.setdefault("custom", {}).update(custom)
+            node.meta.setdefault("custom", {})[_IS_BWD] = True
             nn_module_stack = fwd_node.meta.get("nn_module_stack")
             if nn_module_stack is not None:
                 node.meta["nn_module_stack"] = nn_module_stack.copy()
@@ -416,7 +420,7 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
             traced = make_fx(
                 fn_with_subclass_handling,
                 record_stack_traces=True,
-                record_module_stack=False,  # don't need nn_module_stack for now
+                record_module_stack=True,
             )(*fake_args)
 
         # Copy forward annotations to backward nodes.

torchtitan/experiments/graph_trainer/passes.py

@@ -37,7 +37,11 @@
 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,
+    _IS_BWD,
+    _MODULE_FQN,
+)
 from torchtitan.experiments.graph_trainer.make_fx_tracer import TracedResult
 from torchtitan.experiments.graph_trainer.reshard_after_forward import (
     annotate_fsdp_all_gather,
@@ -184,6 +188,7 @@ def remove_identity_slice_pass(
 
 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.
 
@@ -193,15 +198,24 @@ 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,
+    )
+
     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),
+        ),
     ]
 
     # cudagraph should be the last pass.
@@ -252,18 +266,84 @@ 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]]:
+        custom = node.meta.get("custom", {})
+        is_bwd = custom.get(_IS_BWD, False)
+        if is_bwd != bwd:
+            return []
+        fqn = 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,
     *,
     fsdp_manual_buckets,
 ) -> torch.fx.GraphModule:
+    """Apply aten-level manual bucketing and reordering optimization.
+
+    Used by the AOT and JIT modes where forward and backward graphs are
+    already separate.
     """
-    Apply aten-level manual bucketing and reordering optimization.
+    manual_overlap_bucketing(
+        gm,
+        module_bucket_plans=fsdp_manual_buckets,
+        insert_overlap_deps=False,
+        module_stack_fn=_make_module_fqn_stack_fn(bwd=False),
+    )
+    gm.recompile()
+    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
+        gm,
+        module_bucket_plans=fsdp_manual_buckets,
+        insert_overlap_deps=False,
+        module_stack_fn=_make_module_fqn_stack_fn(bwd=False),
     )
+
+    # 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
 

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,