[graph_trainer] Add CUDA graph kernel annotation pass#2926
Merged
Conversation
meta-cla
Bot
added
the
CLA Signed
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
3 times, most recently
from
5ed6f0d to
7cdf3e2
Compare
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
from
7cdf3e2 to
8faf83e
Compare
yushangdi
changed the title
yushangdi
requested review from
SherlockNoMad,
aditvenk,
fegin,
tianyu-l,
wconstab,
wwwjn,
xmfan and
yiming0416
as code owners
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
from
8faf83e to
b28d5ca
Compare
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
from
b28d5ca to
92dc19d
Compare
yiming0416
reviewed
Apr 14, 2026
yiming0416
left a comment
yiming0416
left a comment
Contributor
There was a problem hiding this comment.
What happens when cuda<13.1? Would this be a no-op?
FWIW currently torchtitan's CI runs with pytorch cuda-13.0
Contributor
Author
yeah, this will be a no-op |
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
2 times, most recently
from
3b5c4f5 to
040f22d
Compare
| ] | ||
|
|
||
| if pass_names and "insert_kernel_annotations" in pass_names: | ||
| passes.append(insert_kernel_annotations_pass) |
Contributor
There was a problem hiding this comment.
add comment on it's ordering requirement, does this need to run before cudagraph?
Contributor
Author
There was a problem hiding this comment.
added. it needs to run before cudagraph pass
There was a problem hiding this comment.
✅ commit 9ddd440bd — CRITICAL 진짜 버그 확정. store.upsertCrawlPayload로 source_url/source_type/source_contents 3개 컬럼만 명시적으로 INSERT INTO ... ON DUPLICATE KEY UPDATE 처리. 기존 row의 extracted_result_json/inspection_status 보존됨.
Comment on lines
+292
to
+293
| Requires ``cuda-python`` package and CUDA toolkit/driver >= 13.1 | ||
| (or cuda-compat >= 13.1). Returns the graph unchanged when unavailable. |
Contributor
There was a problem hiding this comment.
what happens if the pass is run on env doesn't meet the requirement ?
can it be a no-op pass in such case?
Contributor
Author
There was a problem hiding this comment.
yeah, it will be a no-op pass
yiming0416
approved these changes
Apr 15, 2026
yiming0416
reviewed
Apr 15, 2026
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
from
77eeaf6 to
10d8db5
Compare
tianyu-l
reviewed
Apr 15, 2026
yiming0416
reviewed
Apr 16, 2026
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
5 times, most recently
from
af19dbb to
49777a1
Compare
Contributor
Author
|
Seeing some test failures like below on 8 GPU Model test. Likely OOM-style CUBLAS allocation failure, not caused by this PR. They're also present on a PR that only changes .md files: #3026 |
tianyu-l
approved these changes
Apr 20, 2026
tianyu-l
left a comment
tianyu-l
left a comment
Contributor
There was a problem hiding this comment.
sgtm
@felipemello1 related to your observability work
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
3 times, most recently
from
dac49c5 to
7fb2787
Compare
| if is_cudagraph_compatible(traced_result.gm): | ||
| if not precompiled: | ||
| pre_passes = compile_time_passes(traced_result) | ||
| if cudagraph_compatible: |
Contributor
There was a problem hiding this comment.
hmm... does this have to be a compile time pass?
Contributor
Author
There was a problem hiding this comment.
technically no, but it needs to happen before the custom_codegen_pass, and custom_codegen_pass is in compile_time_passes
| save_memory_snapshot_folder: str = "memory_snapshot" | ||
| """Memory snapshot files location.""" | ||
|
|
||
| trace_post_processors: list[Function.Config] = field(default_factory=list) |
Contributor
There was a problem hiding this comment.
do you expect this to actually grow into a list of post processors?
Contributor
There was a problem hiding this comment.
also, I am wondering if we actually need to introduce this callback?
I think ppl would always want processed trace file, as long as enable_cudagraph_annotations() is true
Contributor
Author
There was a problem hiding this comment.
not sure, but could be possible?
for example, if we want custom_cdegen_pass to attach more info to kernels other than just file name and line number, we could tag each region with an id, and then post-process the trace and attach full stack traces/node names/etc.
Contributor
There was a problem hiding this comment.
ok, but i don't think it need to be a list of processor, hook with single entry seems better.
if you need to do multiple thing in the callback, just do it inside the hook
Contributor
Author
There was a problem hiding this comment.
makes sense. updated to a single entry now.
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
from
7fb2787 to
306411e
Compare
SherlockNoMad
approved these changes
Apr 20, 2026
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
3 times, most recently
from
dbea916 to
4d2b30e
Compare
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.
yushangdi
force-pushed
the
sy_cudagraph_annotation
branch
from
4d2b30e to
56ceea7
Compare
2 tasks
tianyu-l
pushed a commit
that referenced
this pull request
Apr 21, 2026
## Summary Fixes a crash introduced by #2926 where `trace_post_processor: Function.Config | None` in `Profiler.Config` causes `tyro.cli()` to fail because tyro cannot parse `Callable[..., Any]` types from CLI arguments. **Error:** ``` Invalid input to tyro.cli() Unsupported type annotation for field profiler.trace-post-processor.fn with type collections.abc.Callable[..., typing.Any] ``` **Fix:** Annotate `trace_post_processor` with `tyro.conf.Suppress`, matching the existing pattern used for: - `model_spec` in `trainer.py:62` - `sample_processor` in `text_datasets.py:502` The field defaults to `None` and is only set programmatically by `CUDAGraphWrapper`, never via CLI. ## Test plan - [ ] Verify `python -m torchtitan.train ...` no longer crashes at config parsing - [ ] Verify graph_trainer still works (trace_post_processor set programmatically)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Need pytorch/pytorch#179867 and cuda 13.1 driver/compat
[graph_trainer] Add CUDA graph kernel annotation pass
Summary
Adds a compiler pass that automatically labels CUDA graph kernels with their originating
nn.Modulepath in profiler traces. When enabled, Perfetto/chrome traces show which component each kernel belongs to (e.g.,layers.0.attention.wq,layers.0.feed_forward.w1,norm). Both forward and backward kernels are annotated.No model code changes required — module paths are derived from the existing
nn.Modulehierarchy.How it works
annotate_module_fqns(model)walks thenn.Moduletree and wraps each submodule'sforwardwithtorch.fx.traceback.annotate_fn({"module_fqn": fqn}). This setsnode.meta["custom"]["module_fqn"]on all FX nodes during tracing — surviving dynamo,make_fx, and export tracers. Backward nodes receive the metadata via_copy_fwd_metadata_to_bw_nodes.insert_kernel_annotations_passreads themodule_fqnmetadata from FX nodes and insertsmark_kernels()enter/exit calls at module boundaries in the graph. Bundled withcudagraph_pass— runs automatically whenever cudagraph is used. No-op whencuda-pythonorcuda-compat >= 13.1is unavailable.CUDAGraphWrappercaptures the annotations during CUDA graph recording when the pass is active, and registers a trace post-processor with the core profiler.Profiler integration —
profiling.pyexposesregister_trace_post_processor()so experiment code can hook into trace export without core importing from experiments. The annotation post-processor merges kernel annotations into the trace automatically.Usage
The annotation pass runs automatically in
aot_fx_tracemode when cudagraph is compatible. Profiler traces come out pre-annotated:NGPU=1 MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_debugmodel ./run_train.sh \ --compile.mode aot_fx_trace \ --profiling.enable_profiling \ --profiling.profile_freq 4Open the exported trace at
outputs/.../rank0_trace.jsonin https://ui.perfetto.dev to see module labels on each kernel.Example output (Llama3 debug model, 6 layers)
Requirements
Requires PyTorch with
torch.cuda._graph_annotationsmodule andenable_annotationskwarg ontorch.cuda.graph(). Falls back gracefully when unavailable (pass is a no-op).Files changed
common_utils.pyannotate_module_fqns()and_MODULE_FQNconstantpasses.pyinsert_kernel_annotations_pass(bundled with cudagraph pass), with alternative approaches documentedcudagraph.pyenable_annotationsflag,enable_cudagraph_annotations(),get_cudagraph_annotations(), trace post-processor registrationllama3/parallelize.pyannotate_module_fqns(model)inannotate_llama()deepseek_v3/parallelize.pyannotate_module_fqns(model)inannotate_deepseekv3()qwen3/parallelize.pyannotate_module_fqns(model)inannotate_qwen3()tools/profiling.pyregister_trace_post_processor()callback API (no experiment imports)models/common/attention.pyimport os+DISABLE_LLVM_OPTtests/test_passes.pytrace_train_step, pass insertion, no-op when unavailable, no-op without metadata, xfail for same-class instancestests/profiler_tests.pytrace_train_stepfwd+bwd → pass → cudagraph → profile → verifymodule_fqnin trace for both fwd and bwd kernels.github/workflows/integration_test_8gpu_graph_trainer.yamlTest plan
Authored with Claude.
Example:
Result:
https://www.internalfb.com/intern/perfetto/open_trace/?manifold_path=perfetto_internal_traces%2Ftree%2Fshared_trace%2Fshangdiy_22d012fd-21e5-412f-aa06-d17a32426030_rank0_trace.json