Enable CudaGraph Annotations

[SherlockNoMad]

Human Note: this is for enabling the mark_kernels, so that we see annotations in profiling trace when cudagraph is used.
Also removed the duplicated implementation in GraphTrainer.

Summary

Graph Trainer maintained its own CUDA graph manager and wrapper even though core TorchTitan already provided the same capture, warmup, input-copy, replay, and teardown behavior. The two implementations had begun to diverge, which made CUDA graph fixes and profiler annotation support harder to maintain.

This PR:

• adds explicit tensor_input_indices support to the core CUDAGraphWrapper
• moves CUDA graph annotation collection and trace post-processing into the core implementation
• always enables CUDA graph annotations and removes the profiler trace_post_processor configuration field
• updates Graph Trainer to use the core CUDA graph wrapper and teardown path
• removes the duplicated Graph Trainer CUDA graph manager and wrapper
• temporarily disables test_precompile_vs_trace, because loaded precompile artifacts currently provide no example inputs for CUDA graph setup

The precompile test will be fixed and re-enabled in a follow-up by constructing the wrapper with live inputs from the training loop instead of serializing inputs in the artifact.

Test plan

• pytest -q tests/unit_tests/cpu/test_cudagraph.py torchtitan/experiments/graph_trainer/tests/test_precompile.py torchtitan/experiments/graph_trainer/tests/test_profiler.py
• pytest -q torchtitan/experiments/graph_trainer/tests/test_passes.py -k cudagraph
• pytest -q torchtitan/experiments/graph_trainer/tests/test_trace_module.py -k cudagraph
• pytest -q torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py::TestLlama3BitwiseDeterministic::test_aot_fx_trace_vs_eager
• targeted ufmt and flake8 checks

[tianyu-l]

if it's always enabled, don't need this function -- just always set with torch.cuda.graph(..., enable_annotations=True)

[SherlockNoMad]

done

[tianyu-l]

This probably should go to config, rather than being ad hoc-ly changed in trainer build. E.g. if there had been other processors, they'll be overwritten silently.

[SherlockNoMad]

I have removed trace_post_processor as a config. it was originally added to support annotation_post_processing. since we are have enable_annotations=True, we can just always run the post_processing.

it's a no-op when there is no annotation collected.

[SherlockNoMad]

will be fixed in the following PR.

[tianyu-l]

one comment

[tianyu-l]

need to make sure this doesn't fail on non-nvidia chips