Hi @leimao, thank you for your blogs!
I ran your example scripts but got similar results when running with and without cudagraph, and only got speedups in the partial cudagraph setting on an A100. What do you think could be the cause ? Details are below:
CUDA Graph Whole Network Capture Example
======================================================================
Using device: NVIDIA A100 PCIe
Model configuration:
Batch size: 6400
Input dim: 4096
Hidden dims: 2048 -> 1024 -> 512
Output dim: 256
======================================================================
SCENARIO 1: Training WITHOUT CUDA Graph
======================================================================
Training WITHOUT CUDA graph...
Completed 10 iterations.
Profiling trace saved to: traces/trace_without_make_graphed_callables.json
Top 10 operations by CUDA time (without CUDA graph):
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
torch/nn/modules/module.py(1782): _call_impl 0.10% 175.624us 6.32% 11.227ms 106.921us 0.000us 0.00% 178.934ms 1.704ms 0 B 0 B 4.83 GB 0 B 105
autograd::engine::evaluate_function: AddmmBackward0 0.24% 421.385us 1.77% 3.150ms 112.517us 0.000us 0.00% 75.765ms 2.706ms 0 B 0 B -365.65 MB -1.25 GB 28
AddmmBackward0 0.14% 254.272us 1.07% 1.902ms 67.933us 0.000us 0.00% 74.904ms 2.675ms 0 B 0 B 917.00 MB 0 B 28
aten::mm 0.47% 841.454us 0.68% 1.214ms 24.766us 74.904ms 44.11% 74.904ms 1.529ms 0 B 0 B 917.00 MB 917.00 MB 49
ampere_sgemm_32x32_sliced1x4_nt 0.00% 0.000us 0.00% 0.000us 0.000us 62.539ms 36.83% 62.539ms 2.316ms 0 B 0 B 0 B 0 B 27
torch_cuda_graph_make_graphed_callables.py(234): <mo... 0.00% 1.894us 96.92% 172.212ms 172.212ms 0.000us 0.00% 60.623ms 60.623ms 0 B 0 B 0 B 0 B 1
torch_cuda_graph_make_graphed_callables.py(139): mai... 0.00% 3.206us 96.92% 172.210ms 172.210ms 0.000us 0.00% 60.623ms 60.623ms 0 B 0 B 0 B 0 B 1
common.py(82): train_without_cuda_graph -0.05% -80.441us 96.92% 172.207ms 172.207ms 0.000us 0.00% 60.623ms 60.623ms 0 B 0 B 0 B 0 B 1
ProfilerStep* -0.21% -368.851us 8.18% 14.527ms 2.075ms 0.000us 0.00% 60.623ms 8.660ms 0 B 0 B 0 B 0 B 7
## forward_pass ## 0.00% 0.000us 0.00% 0.000us 0.000us 60.079ms 35.38% 60.079ms 8.583ms 0 B 0 B 0 B 0 B 7
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 177.676ms
Self CUDA time total: 169.809ms
======================================================================
SCENARIO 2: Training WITH CUDA Graph
======================================================================
Preparing CUDA graph (warmup + capture)...
Creating graphed model...
CUDA graph model ready.
CUDA graph ready.
Training with graph replay...
Completed 10 iterations.
Profiling trace saved to: traces/trace_with_make_graphed_callables.json
Top 10 operations by CUDA time (with CUDA graph):
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
autograd::engine::evaluate_function: GraphedBackward... 0.11% 208.352us 0.55% 1.073ms 153.343us 0.000us 0.00% 77.270ms 11.039ms 0 B 0 B -43.75 MB -43.75 MB 7
GraphedBackward 0.02% 32.351us 0.43% 831.230us 118.747us 77.228ms 43.81% 77.270ms 11.039ms 0 B 0 B 0 B 0 B 7
ampere_sgemm_32x32_sliced1x4_nt 0.00% 0.000us 0.00% 0.000us 0.000us 62.418ms 35.41% 62.418ms 2.312ms 0 B 0 B 0 B 0 B 27
torch_cuda_graph_make_graphed_callables.py(234): <mo... 0.00% 2.755us 99.01% 193.060ms 193.060ms 0.000us 0.00% 61.158ms 61.158ms 0 B 0 B 0 B 0 B 1
torch_cuda_graph_make_graphed_callables.py(171): mai... 0.00% 3.617us 99.01% 193.058ms 193.058ms 0.000us 0.00% 61.158ms 61.158ms 0 B 0 B 0 B 0 B 1
torch_cuda_graph_make_graphed_callables.py(91): trai... -0.04% -73.518us 99.01% 193.054ms 193.054ms 0.000us 0.00% 61.158ms 61.158ms 0 B 0 B 0 B 0 B 1
ProfilerStep* -0.17% -330.068us 3.67% 7.153ms 1.022ms 0.000us 0.00% 61.158ms 8.737ms 0 B 0 B 0 B 0 B 7
torch/nn/modules/module.py(1782): _call_impl 0.01% 29.009us 0.83% 1.626ms 116.109us 0.000us 0.00% 60.608ms 4.329ms 0 B 0 B 43.75 MB 0 B 14
## forward_pass_graphed ## 0.00% 0.000us 0.00% 0.000us 0.000us 60.592ms 34.37% 60.592ms 8.656ms 0 B 0 B 0 B 0 B 7
## forward_pass_graphed ## 0.07% 128.550us 0.58% 1.137ms 162.363us 0.000us 0.00% 60.455ms 8.636ms 0 B 0 B 0 B 0 B 7
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 194.990ms
Self CUDA time total: 176.283ms
======================================================================
Profiling completed successfully!
View traces in Chrome: chrome://tracing
- traces/trace_without_make_graphed_callables.json
- traces/trace_with_make_graphed_callables.json
======================================================================
======================================================================
SCENARIO 3: Training WITH PARTIAL CUDA Graph (only block2)
======================================================================
Preparing CUDA graph for block2 only (warmup + capture)...
Creating partially graphed model (only block2)...
CUDA graph for block2 ready.
CUDA graph for block2 ready.
Training with graph replay...
Completed 10 iterations.
Profiling trace saved to: traces/trace_with_partial_make_graphed_callables.json
Top 10 operations by CUDA time (with partial CUDA graph - block2 only):
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
torch/nn/modules/module.py(1782): _call_impl 0.14% 149.481us 9.30% 10.150ms 120.827us 0.000us 0.00% 158.798ms 1.890ms 0 B 0 B 3.12 GB 0 B 84
torch_cuda_graph_make_graphed_callables.py(234): <mo... 0.00% 3.817us 96.22% 104.966ms 104.966ms 0.000us 0.00% 57.518ms 57.518ms 0 B 0 B 0 B 0 B 1
torch_cuda_graph_make_graphed_callables.py(212): mai... 0.00% 3.898us 96.22% 104.962ms 104.962ms 0.000us 0.00% 57.518ms 57.518ms 0 B 0 B 0 B 0 B 1
torch_cuda_graph_make_graphed_callables.py(91): trai... -0.07% -76.975us 96.22% 104.958ms 104.958ms 0.000us 0.00% 57.518ms 57.518ms 0 B 0 B 0 B 0 B 1
ProfilerStep* -0.34% -371.782us 11.75% 12.817ms 1.831ms 0.000us 0.00% 57.518ms 8.217ms 0 B 0 B 0 B 0 B 7
## forward_pass_graphed ## 0.00% 0.000us 0.00% 0.000us 0.000us 57.329ms 58.02% 57.329ms 8.190ms 0 B 0 B 0 B 0 B 7
## forward_pass_graphed ## 0.15% 167.391us 3.65% 3.986ms 569.417us 0.000us 0.00% 57.156ms 8.165ms 0 B 0 B 641.38 MB -43.75 MB 7
nn.Module: MLPModel_0 0.01% 12.966us 3.50% 3.819ms 545.504us 0.000us 0.00% 57.156ms 8.165ms 0 B 0 B 685.12 MB 0 B 7
common.py(48): forward 0.09% 102.431us 3.47% 3.782ms 540.357us 0.000us 0.00% 57.156ms 8.165ms 0 B 0 B 685.12 MB -350.00 MB 7
ampere_sgemm_128x64_tn 0.00% 0.000us 0.00% 0.000us 0.000us 51.730ms 52.36% 51.730ms 5.748ms 0 B 0 B 0 B 0 B 9
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 109.085ms
Self CUDA time total: 98.801ms
======================================================================
All profiling completed successfully!
View traces in Chrome: chrome://tracing
- traces/trace_without_make_graphed_callables.json
- traces/trace_with_make_graphed_callables.json
- traces/trace_with_partial_make_graphed_callables.json
======================================================================
Hi @leimao, thank you for your blogs!
I ran your example scripts but got similar results when running with and without cudagraph, and only got speedups in the partial cudagraph setting on an A100. What do you think could be the cause ? Details are below: