Change to use annotate_module_components

Author: yushangdi

.github/workflows/integration_test_8gpu_graph_trainer.yaml

@@ -71,6 +71,10 @@ jobs:
         # Run the graph passes unit tests
         torchrun --nproc-per-node=8 -m pytest torchtitan/experiments/graph_trainer/tests/test_passes.py::TestReassignToPgPass -v
         pytest torchtitan/experiments/graph_trainer/tests/test_passes.py::TestApplySACPass -v
+        pytest torchtitan/experiments/graph_trainer/tests/test_passes.py::TestAnnotateModuleFqns -v
+
+        # Run kernel annotation profiler tests (skips if cuda-compat < 13.1)
+        pytest torchtitan/experiments/graph_trainer/tests/profiler_tests.py -v
 
         # Run the make_fx tracer unit tests
         pytest torchtitan/experiments/graph_trainer/tests/test_trace_module.py -v -k "TestTraceModule or TestTraceDTensor or TestMetadataPropagation"

.gitignore

@@ -49,4 +49,3 @@ CLAUDE.local.md
 /debug_*.py
 CLAUDE_CONTEXT/
 /.claude/settings.local.json
-agent_space/

torchtitan/experiments/graph_trainer/.claude/CLAUDE.md

@@ -109,3 +109,18 @@ This verifies that the aot_fx_trace path produces bitwise identical losses
 and gradients across runs, and matches eager numerics exactly. Any change
 that breaks this test must be investigated and fixed before proceeding with
 other tests.
+
+### CUDA Graph Kernel Annotations
+
+The `insert_kernel_annotations_pass` labels CUDA graph kernels with their
+originating `nn.Module` path in profiler traces. It runs automatically in the
+`aot_fx_trace` path (bundled with the cudagraph pass). The profiler
+post-processes the trace automatically via `_maybe_annotate_trace` in
+`torchtitan/tools/profiling.py` — no manual post-processing is needed.
+
+Requirements: `cuda-python` package and CUDA toolkit/driver >= 13.1
+(or `cuda-compat >= 13.1` on `LD_LIBRARY_PATH`). The pass is a no-op when
+these are unavailable.
+
+To view annotated traces, open the exported JSON in https://ui.perfetto.dev.
+Kernel events will have `module_fqn` fields like `layers.0.attention.wq`.

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/compile.py

@@ -240,7 +240,6 @@ def apply_compile(
         return model
 
     torch._inductor.config.reorder_for_peak_memory = False
-    torch._inductor.config.triton.cudagraph_kernel_annotations = True
     torch._dynamo.config.capture_scalar_outputs = True
 
     fsdp_reshard_after_forward = get_fsdp_reshard_after_forward_policy(

torchtitan/experiments/graph_trainer/cudagraph.py

@@ -31,7 +31,10 @@ def __init__(self) -> None:
         self._initialized = False
         self._cudagraph_wrappers: list["CUDAGraphWrapper"] = []
         self._teardown_called = False
+        # toolsId (graph_id << 32 | node_id) -> list of annotation dicts
+        # (e.g. [{"module_fqn": "layers.0.attention.wq"}]).
         self.all_annotations: dict[int, list] = {}
+        self.enable_annotations: bool = False
 
     def maybe_initialize(self) -> None:
         if self._initialized:
@@ -106,6 +109,11 @@ def get_cudagraph_annotations() -> dict[int, list]:
     return _cg_manager.all_annotations
 
 
+def enable_cudagraph_annotations() -> None:
+    """Enable kernel annotation capture on subsequent CUDA graph recordings."""
+    _cg_manager.enable_annotations = True
+
+
 class CUDAGraphWrapper:
     """Wraps a callable with cudagraph. It warms up the callable, records cudagraph,
     and replays cudagraph during runtime. It also handles static input tensors, which
@@ -212,15 +220,15 @@ def __call__(self, *args):
                 self._cudagraph,
                 pool=_cg_manager.graph_pool,
                 stream=_cg_manager.stream,
-                enable_annotations=True,
+                enable_annotations=_cg_manager.enable_annotations,
             ):
                 # `output` is managed by pytorch's cudagraph pool
                 self._output = self._runnable(*args)
 
-            # Save kernel annotations for trace post-processing.
-            from torch.cuda._graph_annotations import get_kernel_annotations
+            if _cg_manager.enable_annotations:
+                from torch.cuda._graph_annotations import get_kernel_annotations
 
-            _cg_manager.all_annotations.update(get_kernel_annotations())
+                _cg_manager.all_annotations.update(get_kernel_annotations())
 
         if self._should_check_address:
             self._check_static_inputs_address()

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
@@ -42,6 +43,9 @@ def annotate_deepseekv3(model: GraphTrainerDeepSeekV3Model) -> None:
     - 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.
+    - Module FQN annotation: Tags each submodule's forward with its
+      fully-qualified name so that insert_kernel_annotations_pass can label
+      CUDA graph kernels with their originating module.
 
     """
     from torchtitan.distributed.expert_parallel import ExpertParallel
@@ -56,6 +60,7 @@ def annotate_deepseekv3(model: GraphTrainerDeepSeekV3Model) -> None:
     MoE.forward = annotate_fn({"EP": "compute"})(MoE.forward)
 
     annotate_ac_regions(model)
+    annotate_module_fqns(model)
 
 
 # Adapted from llama4/infra/parallelize.py

torchtitan/experiments/graph_trainer/graph_utils.py

@@ -519,11 +519,6 @@ def get_compiler_passes_from_config(
                 )
             )
         else:
-            # Auto-insert annotation pass before cudagraph capture.
-            if pass_name == "cudagraph":
-                compiler_passes.append(
-                    AVAILABLE_COMPILER_PASSES["insert_kernel_annotations"]
-                )
             compiler_passes.append(AVAILABLE_COMPILER_PASSES[pass_name])
 
     if pass_names:

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
@@ -35,8 +36,12 @@ def annotate_llama(model: GraphTrainerLlama3Model) -> None:
     - 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.
+    - Module FQN annotation: Tags each submodule's forward with its
+      fully-qualified name so that insert_kernel_annotations_pass can label
+      CUDA graph kernels with their originating module.
     """
     annotate_ac_regions(model)
+    annotate_module_fqns(model)
 
 
 def parallelize_llama(

torchtitan/experiments/graph_trainer/passes.py

@@ -80,11 +80,13 @@ def construct_default_graph_passes(
             flex_compile_config=FlexAttention.inductor_configs,
         ),
         regional_inductor_pass,
+        # insert_kernel_annotations must run before cudagraph: it inserts
+        # mark_kernels enter/exit calls that execute during CUDA graph
+        # capture.  It is a no-op when cuda-python or cuda-compat >= 13.1
+        # is not available.
+        insert_kernel_annotations_pass,
     ]
 
-    # Insert kernel annotation markers before cudagraph capture.
-    passes.append(insert_kernel_annotations_pass)
-
     # cudagraph should be the last pass.
     from torchtitan.experiments.graph_trainer.cudagraph import is_cudagraph_compatible
 
@@ -263,67 +265,77 @@ def _get_fake_mode_from_gm(gm: torch.fx.GraphModule):
     return gm
 
 
-def _mark_kernels_enter(annotation: dict) -> object:
-    """Helper called from FX graph to enter a mark_kernels scope."""
-    from torch.cuda._graph_annotations import mark_kernels
-
-    ctx = mark_kernels(annotation)
-    ctx.__enter__()
-    return ctx
-
-
-def _mark_kernels_exit(ctx: object) -> None:
-    """Helper called from FX graph to exit a mark_kernels scope."""
-    ctx.__exit__(None, None, None)  # type: ignore[union-attr]
-
-
 def insert_kernel_annotations_pass(
     gm: torch.fx.GraphModule,
     example_inputs: tuple | None = None,
 ) -> torch.fx.GraphModule:
-    """Insert mark_kernels() calls at component boundaries in the FX graph.
+    """Insert mark_kernels() calls at module boundaries in the FX graph.
 
-    Reads ``node.meta["custom"]["component"]`` (set via
-    ``torch.fx.traceback.annotate``) and inserts enter/exit calls so that
+    Reads ``node.meta["custom"]["module_fqn"]`` (set via
+    ``annotate_module_fqns``) and inserts enter/exit calls so that
     CUDA graph capture records the annotations.
+
+    Requires ``cuda-python`` package and CUDA toolkit/driver >= 13.1
+    (or cuda-compat >= 13.1).  Returns the graph unchanged when unavailable.
+
+    Also enables annotation capture on :class:`CUDAGraphWrapper` so that
+    ``enable_annotations=True`` is passed to ``torch.cuda.graph()``.
     """
+    from torch.cuda._graph_annotations import _is_tools_id_unavailable
+
+    from torchtitan.experiments.graph_trainer.common_utils import _MODULE_FQN
+    from torchtitan.experiments.graph_trainer.cudagraph import (
+        enable_cudagraph_annotations,
+    )
+
+    def _enter(annotation: dict) -> object:
+        from torch.cuda._graph_annotations import mark_kernels
+
+        ctx = mark_kernels(annotation)
+        ctx.__enter__()
+        return ctx
+
+    def _exit(ctx: object) -> None:
+        ctx.__exit__(None, None, None)  # type: ignore[union-attr]
+
+    if _is_tools_id_unavailable():
+        return gm
+
+    enable_cudagraph_annotations()
+
     graph = gm.graph
-    current_component: str | None = None
+    current_fqn: str | None = None
     current_ctx_node = None
 
     for node in list(graph.nodes):
-        component = (node.meta.get("custom") or {}).get("component")
+        fqn = (node.meta.get("custom") or {}).get(_MODULE_FQN)
 
-        if component != current_component:
+        if fqn != current_fqn:
             # Close previous scope
             if current_ctx_node is not None:
                 with graph.inserting_before(node):
-                    exit_node = graph.call_function(
-                        _mark_kernels_exit, (current_ctx_node,)
-                    )
+                    exit_node = graph.call_function(_exit, (current_ctx_node,))
                     exit_node.meta["custom"] = {}
                 current_ctx_node = None
 
             # Open new scope
-            if component is not None:
+            if fqn is not None:
                 with graph.inserting_before(node):
                     enter_node = graph.call_function(
-                        _mark_kernels_enter,
-                        ({"component": component},),
+                        _enter,
+                        ({_MODULE_FQN: fqn},),
                     )
                     enter_node.meta["custom"] = {}
                 current_ctx_node = enter_node
 
-            current_component = component
+            current_fqn = fqn
 
     # Close any trailing scope (before output/return)
     if current_ctx_node is not None:
         output_nodes = [n for n in graph.nodes if n.op == "output"]
         if output_nodes:
             with graph.inserting_before(output_nodes[0]):
-                exit_node = graph.call_function(
-                    _mark_kernels_exit, (current_ctx_node,)
-                )
+                exit_node = graph.call_function(_exit, (current_ctx_node,))
                 exit_node.meta["custom"] = {}
 
     graph.lint()
@@ -751,7 +763,6 @@ def tlparse_log_graph_pass(
     "auto_bucketing": autobucketing_reordering_pass,
     "transformer_block_bucketing": transformer_block_bucketing_reordering_pass,
     "regional_inductor": regional_inductor_pass,
-    "insert_kernel_annotations": insert_kernel_annotations_pass,
     "cudagraph": cudagraph_pass,
     "full_inductor_compilation": full_inductor_compilation_pass,
 }