Enable CudaGraph Annotations

tests/unit_tests/cpu/test_cudagraph.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from contextlib import nullcontext
+from typing import cast
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from torchtitan.distributed.cudagraph import (
+    _manager,
+    CUDAGraphWrapper,
+    get_cudagraph_annotations,
+)
+
+
+def test_tensor_input_indices_control_replay_copies() -> None:
+    static_input = torch.tensor(1)
+    excluded_input = torch.tensor(2)
+    copied_input = torch.tensor(3)
+
+    with (
+        patch.object(_manager, "maybe_initialize"),
+        patch.object(_manager, "register"),
+    ):
+        wrapper = CUDAGraphWrapper(
+            lambda *args: args,
+            (static_input, excluded_input, copied_input),
+            static_input_indices=(0,),
+            tensor_input_indices=[0, 2],
+        )
+
+    wrapper._warmup_remaining = 0
+    wrapper._args = (static_input, excluded_input, copied_input)
+    graph = cast(torch.cuda.CUDAGraph, MagicMock())
+    wrapper._graph = graph
+    wrapper._output = "output"
+
+    result = wrapper(torch.tensor(4), torch.tensor(5), torch.tensor(6))
+
+    assert result == "output"
+    assert static_input.item() == 1
+    assert excluded_input.item() == 2
+    assert copied_input.item() == 6
+    cast(MagicMock, graph.replay).assert_called_once_with()
+
+
+def test_cudagraph_wrapper_collects_annotations() -> None:
+    graph = cast(torch.cuda.CUDAGraph, MagicMock())
+    annotations = {42: [{"module_fqn": "layers.0"}]}
+    graph_pool = object()
+    stream = MagicMock()
+
+    with (
+        patch.object(_manager, "maybe_initialize"),
+        patch.object(_manager, "register"),
+        patch.object(_manager, "_graph_pool", graph_pool),
+        patch.object(_manager, "_stream", stream),
+        patch.object(_manager, "all_annotations", {}),
+        patch("torch.cuda.CUDAGraph", return_value=graph),
+        patch("torch.cuda.graph", return_value=nullcontext()) as cuda_graph,
+        patch(
+            "torchtitan.distributed.cudagraph.get_kernel_annotations",
+            return_value=annotations,
+        ),
+    ):
+        wrapper = CUDAGraphWrapper(lambda x: x, (torch.tensor(1),))
+        wrapper._warmup_remaining = 0
+
+        output = wrapper(torch.tensor(2))
+
+        assert output.item() == 2
+        assert get_cudagraph_annotations() == annotations
+        cuda_graph.assert_called_once_with(
+            graph,
+            pool=graph_pool,
+            stream=stream,
+            enable_annotations=True,
+            capture_error_mode="thread_local",
+        )

torchtitan/distributed/cudagraph.py

@@ -4,17 +4,17 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""Lightweight CUDA graph wrapper for eager-mode training steps.
-
-Adapted from ``torchtitan/experiments/graph_trainer/cudagraph.py``.
-"""
+"""Lightweight CUDA graph wrapper for training steps."""
 
+import gzip
+import json
 import warnings
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from typing import Any, cast
 
 import torch
+from torch.cuda._graph_annotations import get_kernel_annotations
 from torch.nn.attention.flex_attention import BlockMask
 from torch.utils import _pytree as pytree
 
@@ -123,14 +123,15 @@ def unflatten(self, flat_leaves: Sequence[Any]) -> Any:
 
 
 class _CUDAGraphManager:
-    """Singleton that owns a shared graph pool and stream."""
+    """Singleton that owns a shared graph pool, stream, and annotations."""
 
     def __init__(self) -> None:
         self._initialized = False
         self._wrappers: list["CUDAGraphWrapper"] = []
         self._graph_pool: Any = None
         self._stream: torch.cuda.Stream | None = None
         self._dummy_graph: torch.cuda.CUDAGraph | None = None
+        self.all_annotations: dict[int, list[Any]] = {}
 
     @property
     def graph_pool(self) -> Any:
@@ -186,6 +187,39 @@ def cudagraph_teardown() -> None:
     _manager.teardown()
 
 
+def get_cudagraph_annotations() -> dict[int, list[Any]]:
+    """Return all kernel annotations accumulated across CUDA graph captures."""
+    return _manager.all_annotations
+
+
+def cudagraph_annotate_trace_post_processor(trace_path: str) -> None:
+    """Post-process a profiler trace with captured CUDA graph annotations."""
+    annotations = get_cudagraph_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
+
+    open_trace = gzip.open if trace_path.endswith(".gz") else open
+    with open_trace(trace_path, "rt") as trace_file:
+        trace = json.load(trace_file)
+
+    count = annotate_trace(trace, annotations)
+    if count > 0:
+        with open_trace(trace_path, "wt") as trace_file:
+            json.dump(trace, trace_file)
+        logger.info(f"Annotated {count} CUDA graph kernel events in profiler trace")
+
+
 class CUDAGraphWrapper:
     """Wrap a callable with CUDA graph capture and replay.
 
@@ -197,14 +231,17 @@ class CUDAGraphWrapper:
             are stable across calls (e.g. model weights/buffers).
         should_check_address: Whether to verify static input tensor addresses
             before each replay. This should only be enabled for debugging.
+        tensor_input_indices: Indices of inputs that should be copied before
+            replay. When omitted, these are inferred from ``example_inputs``.
     """
 
     def __init__(
         self,
         fn: Callable,
         example_inputs: Sequence[Any],
-        static_input_indices: tuple[int, ...] | None = None,
+        static_input_indices: Sequence[int] | None = None,
         should_check_address: bool = False,
+        tensor_input_indices: Sequence[int] | None = None,
     ):
         self._fn = fn
         self._num_inputs = len(example_inputs)
@@ -218,11 +255,16 @@ def __init__(
                 f"{sorted(invalid_static_indices)}"
             )
 
-        self._input_indices_to_copy = [
-            i
-            for i, inp in enumerate(example_inputs)
-            if isinstance(inp, torch.Tensor) and i not in self._static_input_indices
-        ]
+        if tensor_input_indices is not None:
+            self._input_indices_to_copy = [
+                i for i in tensor_input_indices if i not in self._static_input_indices
+            ]
+        else:
+            self._input_indices_to_copy = [
+                i
+                for i, inp in enumerate(example_inputs)
+                if isinstance(inp, torch.Tensor) and i not in self._static_input_indices
+            ]
         self._tensor_metadata = {
             i: (inp.shape, inp.dtype, inp.device)
             for i, inp in enumerate(example_inputs)
@@ -309,9 +351,11 @@ def __call__(self, *args):
                 self._graph,
                 pool=_manager.graph_pool,
                 stream=_manager.stream,
+                enable_annotations=True,
                 capture_error_mode="thread_local",
             ):
                 self._output = self._fn(*args)
+            _manager.all_annotations.update(get_kernel_annotations())
             logger.info("Recorded CUDA graph")
 
         if self._should_check_address:

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

@@ -340,10 +340,10 @@ NGPU=4 MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_8b_c4_test ./run_
 
 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.
+`aot_fx_trace` path (bundled with the cudagraph pass). The profiler always runs
+``cudagraph_annotate_trace_post_processor`` after exporting a trace, so CUDA
+graph annotations are merged automatically and 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

torchtitan/experiments/graph_trainer/configs.py

@@ -246,7 +246,6 @@ def to_graph_trainer_config(
     from the graph_trainer model_registry. The compile field is removed and
     left as the GraphTrainer.Config 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)}
@@ -290,11 +289,4 @@ def to_graph_trainer_config(
             **{f.name: getattr(loss_cfg, f.name) for f in fields(loss_cfg)}
         )
 
-    # 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_processor = cudagraph_annotate_trace_post_processor()
-
     return GraphTrainer.Config(**d)