[graph_trainer] Add CUDA graph kernel annotation pass

Adds `insert_kernel_annotations_pass` that labels CUDA graph kernels with
their originating nn.Module path in profiler traces.

How it works:
1. `annotate_module_fqns(model)` wraps each submodule's forward with
   `torch.fx.traceback.annotate_fn({"module_fqn": fqn})`, setting
   `node.meta["custom"]["module_fqn"]` on FX nodes during tracing.
2. `insert_kernel_annotations_pass` reads the `module_fqn` metadata and
   inserts `mark_kernels` enter/exit calls at module boundaries.  Bundled
   with the cudagraph pass.
3. `CUDAGraphWrapper` captures annotations during graph recording when the
   pass is active.
4. `Profiler.Config.trace_post_processors` (new field) runs
   `cudagraph_annotate_trace_post_processor()` on each exported trace so
   profiler traces automatically carry `module_fqn` fields on graphed
   kernel events.

Requires PyTorch with `torch.cuda._graph_annotations` and the
`enable_annotations` kwarg on `torch.cuda.graph()`.  Falls back gracefully
when unavailable (pass is a no-op).

Authored with Claude.

Author: yushangdi

.github/workflows/integration_test_8gpu_graph_trainer.yaml

@@ -69,6 +69,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/test_profiler.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"

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

@@ -208,3 +208,19 @@ 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 post-processor
+is attached via ``Profiler.Config.trace_post_processors`` (see
+``cudagraph_annotate_trace_post_processor``) so exported traces are
+annotated automatically — 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,29 @@
 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 ``module_fqn`` in
+    ``node.meta["custom"]``.
+
+    Call once after model construction, before tracing/compilation.
+
+    .. note::
+
+        When traced via ``trace_train_step``, multiple instances of the
+        **same class with no parameters** may all receive the first instance's
+        FQN.  This is because ``_reparametrize_module`` with an empty state
+        dict causes ``make_fx`` to collapse the ``annotate_fn`` contexts for
+        parameterless same-class instances.
+    """
+    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/configs.py

@@ -69,6 +69,7 @@ def to_graph_trainer_config(
     from the graph_trainer model_registry. The compile field is removed and
     left as the GraphTrainerConfig default; callers should explicitly set it.
     """
+    from .cudagraph import cudagraph_annotate_trace_post_processor
     from .trainer import GraphTrainer
 
     d = {f.name: getattr(base_config, f.name) for f in fields(base_config)}
@@ -96,4 +97,11 @@ def to_graph_trainer_config(
     if ac is not None and ac.mode != "none":
         d["activation_checkpoint"] = ActivationCheckpointConfig(mode="selective")
 
+    # Merge CUDA graph kernel annotations into profiler traces when profiling
+    # is active.  No-op otherwise (and no-op when requirements aren't met).
+    # It's also a no-op if there is CUDA graph is not enabled.
+    profiler = d.get("profiler")
+    if profiler is not None:
+        profiler.trace_post_processors.append(cudagraph_annotate_trace_post_processor())
+
     return GraphTrainer.Config(**d)

torchtitan/experiments/graph_trainer/cudagraph.py

@@ -11,7 +11,7 @@
 during compilation.
 """
 
-import logging
+import json
 import warnings
 from collections.abc import Callable, Sequence
 from typing import Any
@@ -21,7 +21,8 @@
 from torch._library.opaque_object import is_opaque_value
 from torch.utils._ordered_set import OrderedSet
 
-logger = logging.getLogger(__name__)
+from torchtitan.config.function import Function
+from torchtitan.tools.logging import logger
 
 
 class _CUDAGraphManager:
@@ -31,6 +32,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:
@@ -100,6 +105,54 @@ def cudagraph_teardown() -> None:
     _cg_manager.teardown()
 
 
+def get_cudagraph_annotations() -> dict[int, list]:
+    """Return all kernel annotations accumulated across CUDA graph captures."""
+    return _cg_manager.all_annotations
+
+
+def enable_cudagraph_annotations() -> None:
+    """Enable kernel annotation capture on subsequent CUDA graph recordings."""
+    _cg_manager.enable_annotations = True
+
+
+def cudagraph_annotate_trace_post_processor() -> Function.Config:
+    """Return a ``Function.Config`` that merges captured CUDA graph kernel
+    annotations into a profiler trace file.
+
+    Attach this to ``ProfilingConfig.trace_post_processors`` so that exported
+    profiler traces automatically carry ``module_fqn`` fields on graphed kernel
+    events.
+    """
+    return Function.Config(fn=_cudagraph_annotate_trace_file)
+
+
+def _cudagraph_annotate_trace_file(trace_path: str) -> None:
+    """Post-process a profiler trace with CUDA graph kernel annotations."""
+    annotations = _cg_manager.all_annotations
+    if not annotations:
+        return
+
+    try:
+        from torch.cuda._annotate_cuda_graph_trace import (  # pyrefly: ignore[missing-import]
+            annotate_trace,
+        )
+    except ImportError:
+        logger.warning(
+            "torch.cuda._annotate_cuda_graph_trace not available. "
+            "Upgrade PyTorch to enable trace CUDA graph kernel annotation."
+        )
+        return
+
+    with open(trace_path) as f:
+        trace = json.load(f)
+
+    count = annotate_trace(trace, annotations)
+    if count > 0:
+        with open(trace_path, "w") as f:
+            json.dump(trace, f)
+        logger.info(f"Annotated {count} CUDAGraph kernel event(s) in profiler trace")
+
+
 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
@@ -219,10 +272,16 @@ def __call__(self, *args):
                 self._cudagraph,
                 pool=_cg_manager.graph_pool,
                 stream=_cg_manager.stream,
+                enable_annotations=_cg_manager.enable_annotations,
             ):
                 # `output` is managed by pytorch's cudagraph pool
                 self._output = self._runnable(*args)
 
+            if _cg_manager.enable_annotations:
+                from torch.cuda._graph_annotations import 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.models.common.moe import MoE
@@ -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/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

@@ -116,11 +116,22 @@ def construct_default_graph_passes(
     from torchtitan.experiments.graph_trainer.cudagraph import is_cudagraph_compatible
 
     passes: list[Callable] = []
+    cudagraph_compatible = is_cudagraph_compatible(traced_result.gm)
+
     if not precompiled:
-        passes.extend(compile_time_passes(traced_result))
+        pre_passes = compile_time_passes(traced_result)
+        # insert_kernel_annotations must run before custom_codegen_pass
+        # (which saves code to disk and replaces forward with a loaded
+        # module) so that the inserted mark_kernels enter/exit calls end
+        # up in the saved graph.
+        if cudagraph_compatible:
+            pre_passes = _insert_before_custom_codegen(
+                pre_passes, insert_kernel_annotations_pass
+            )
+        passes.extend(pre_passes)
 
     # cudagraph should be the last pass.
-    if is_cudagraph_compatible(traced_result.gm):
+    if cudagraph_compatible:
         static_input_indices = list(range(traced_result.num_static_inputs))
         passes.append(
             functools.partial(
@@ -133,6 +144,26 @@ def construct_default_graph_passes(
     return passes
 
 
+def _insert_before_custom_codegen(
+    passes: list[Callable], new_pass: Callable
+) -> list[Callable]:
+    """Return a new pass list with ``new_pass`` inserted before custom_codegen_pass.
+
+    If custom_codegen_pass is not present, appends ``new_pass`` at the end.
+    """
+    result: list[Callable] = []
+    inserted = False
+    for p in passes:
+        name = getattr(p, "__name__", "")
+        if not inserted and name == "custom_codegen_pass":
+            result.append(new_pass)
+            inserted = True
+        result.append(p)
+    if not inserted:
+        result.append(new_pass)
+    return result
+
+
 def apply_graph_passes(
     gm: torch.fx.GraphModule,
     example_inputs: tuple,
@@ -320,6 +351,96 @@ def _get_fake_mode_from_gm(gm: torch.fx.GraphModule):
     return gm
 
 
+def insert_kernel_annotations_pass(
+    gm: torch.fx.GraphModule,
+    example_inputs: tuple | None = None,
+) -> torch.fx.GraphModule:
+    """Insert mark_kernels() calls at module boundaries in the FX graph.
+
+    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()``.
+
+    Alternative approaches:
+
+    1. **fx.Interpreter**: During cudagraph capture, run the graph via an
+       ``fx.Interpreter`` subclass that reads ``module_fqn`` metadata and
+       calls ``mark_kernels`` enter/exit around each node — avoids mutating
+       the graph.
+    2. **Custom CodeGen**: Use a custom ``torch.fx.graph.CodeGen`` to emit
+       enter/exit lines (or ``with`` blocks) directly in the generated
+       Python code.
+
+    The current graph-pass approach is the least invasive.
+    """
+    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_fqn: str | None = None
+    current_ctx_node = None
+
+    for node in list(graph.nodes):
+        fqn = (node.meta.get("custom") or {}).get(_MODULE_FQN)
+
+        if fqn != current_fqn:
+            # Close previous scope
+            if current_ctx_node is not None:
+                with graph.inserting_before(node):
+                    exit_node = graph.call_function(_exit, (current_ctx_node,))
+                    exit_node.meta["custom"] = {}
+                current_ctx_node = None
+
+            # Open new scope
+            if fqn is not None:
+                with graph.inserting_before(node):
+                    enter_node = graph.call_function(
+                        _enter,
+                        ({_MODULE_FQN: fqn},),
+                    )
+                    enter_node.meta["custom"] = {}
+                current_ctx_node = enter_node
+
+            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(_exit, (current_ctx_node,))
+                exit_node.meta["custom"] = {}
+
+    graph.lint()
+    gm.recompile()
+    return gm
+
+
 def cudagraph_pass(
     gm: torch.fx.GraphModule,
     example_inputs: tuple,