How to use `torch.compile` with CUDA graphs when using gradient activation checkpointing

veritas: <p>Hello! I would like to ask about using the latest <code>torch.compile</code> with the somewhat older <code>torch.cuda.make_graphed_callables</code> function to create CUDA graphs when gradient activation checkpointing and multi-GPU training is enabled.<br>
It is currently my understanding that <code>torch.compile</code> reduces CUDA overhead while CUDA graphs reduce CPU overhead by fusing calls to the CUDA device. However, I have not seen any explanation or tutorial on whether they can be used together. Is it possible to capture the CUDA graph of a model that has undergone <code>torch.compile</code>? I do not think that there is a fundamental reason why this could not be done, though it may be technically time-consuming.<br>
Also, activation checkpointing is a common optimization in Transformer architectures, which are also frequently trained on multiple GPUs. Is it possible to combine the <code>torch.compile</code> and creation of CUDA graphs with multi-GPU training, possibly only with DDP but not with FSDP?<br>
Though these conditions may appear complicated, they are quite common in training large language models (LLMs), which have gained enormous popularity. I believe that making those optimizations would be very beneficial and I would like to thank any answers or informative discussion in advance.</p>

eqy: <p>If you wish to use <code>torch.compile</code> with CUDA graphs, the preferred method to do so would probably be via the option <code>mode="reduce-overhead"</code> which should use CUDA graphs according to the argument documentation:</p>
<pre><code class="lang-auto">       mode (str): Can be either "default", "reduce-overhead" or "max-autotune"
        - "default" is the default mode, which is a good balance between performance and overhead
        - "reduce-overhead" is a mode that reduces the overhead of python with CUDA graphs, useful for small batches
        - "max-autotune" is a mode that that leverages Triton based matrix multiplications and convolutions
        - To see the exact configs that each mode sets you can call `torch._inductor.list_mode_options()`
</code></pre>
<p>I believe <code>torch.compile</code> should be compatible with DDP, and compatible with FSDP as some issues were recently addressed e.g., <a href="https://github.com/pytorch/pytorch/issues/98494" class="inline-onebox" rel="noopener nofollow ugc">[FSDP] `use_orig_params=True` with CPU offloading and Gradient Accumulation: RuntimeError · Issue #98494 · pytorch/pytorch · GitHub</a>.</p>

ipe-zhangyz: <p>I tried adding <code>torch.compile()</code> to the sample in this blog: <a href="https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/" class="inline-onebox" rel="noopener nofollow ugc">Accelerating PyTorch with CUDA Graphs | PyTorch</a><br>
The output of <code>nsys profile</code> shows that kernels executed before and after the model gotten compiled are different. After the model is compiled, some kernels like<code> triton__0d1d2d3</code> are executed, while there is no such kernel without the <code>torch.compile()</code>.<br>
I’m doubting if there is no benefit to use torch.comple() together with CUDA grahp, as in the mentioned case the walltime of graph replay is almost the same with and without torch.comple().<br>
By the way, if using the reduce-overhead mode, the compiled model can not be used during the capture of the graph, or there will be an error:</p>
<blockquote>
<p>RuntimeError: Cannot call CUDAGeneratorImpl::current_seed during CUDA graph capture. If you need this call to be captured, please file an issue. Current cudaStreamCaptureStatus: cudaStreamCaptureStatusActive</p>
</blockquote>
<p>Howerer, the default mode does not encounter this error.<br>
Hope there will be some official tutorial about using torch.compile() with CUDA graph.</p>

ptrblck: <p>I don’t believe manually capturing a compiled graph would be a supported use case and would recommend sticking to <a class="mention" href="/u/eqy">@eqy</a> 's suggestion in using the <code>reduce-overhead</code> mode instead.</p>

Abhishek_Ghosh: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I do not get the concept of using <code>torch.compile</code> and CUDA Graph. I did not find any detailed documentation of the same.</p>
<p>But what I understand is the following — given a plain PyTorch model (a model which does not make use of CUDA Graph on its own), we can use <code>torch.compile</code> to compile it, and if we pass the <code>reduce-overhead</code> option in <code>mode</code> field of <code>torch.compile</code>, then the compiler shall be smart enough to understand which of the regions are “CUDA Graph capture safe” and accordingly graph those regions <em>automatically</em>. I say automatic based on the example code in the docs (about CUDAGraph Trees Integration)</p>
<pre data-code-wrap="python"><code class="lang-python">@torch.compile
def foo(x):
    # GRAPH 1
    y = x * x * x
    # graph break triggered here
    if y.sum() &gt; 0:
        # GRAPH 2
        z = y ** y
    else:
        # GRAPH 3
        z = (y.abs() ** y.abs())
    torch._dynamo.graph_break()
    # GRAPH 4
    return z * torch.rand_like(z)

# the first run warms up each graph, which does things like CuBlas or Triton benchmarking
foo(torch.arange(0, 10, device="cuda"))
# The second run does a CUDA Graph recording, and replays it
foo(torch.arange(0, 10, device="cuda"))
# Finally we hit the optimized, CUDA Graph replay path
foo(torch.arange(0, 10, device="cuda"))
</code></pre>
<p>where there is no explicit use of CUDA Graph, but the comments make it look like CUDA Graphs are being used in the background <em>automatically</em>.</p>
<p>Is my understanding correct, or is it flawed?</p>
<p>From PyTorch doc I got the following:</p>
<blockquote>
<p><strong>mode</strong> (<em>str</em>) –</p>
<p>Can be either “default”, “reduce-overhead”, “max-autotune” or “max-autotune-no-cudagraphs”</p>
<ul>
<li>”reduce-overhead” is a mode that reduces the overhead of python with CUDA graphs, useful for small batches. Reduction of overhead can come at the cost of more memory usage, as we will cache the workspace memory required for the invocation so that we do not have to reallocate it on subsequent runs. Reduction of overhead is not guaranteed to work; today, we only reduce overhead for CUDA only graphs which do not mutate inputs. There are other circumstances where CUDA graphs are not applicable; use TORCH_LOG=perf_hints to debug.</li>
<li>”max-autotune” is a mode that leverages Triton based matrix multiplications and convolutions It enables CUDA graphs by default.</li>
</ul>
</blockquote>
<p>Regarding this automatic use of CUDA Graph by the compiler, I guess this is what <a class="mention" href="/u/ezyang">@ezyang</a> was hinting at in this comment <a href="https://github.com/pytorch/pytorch/issues/71896#issuecomment-1169241054" rel="noopener nofollow ugc">here</a>.</p>
<p>Also, I guess <a class="mention" href="/u/ezyang">@ezyang</a> 's comment <a href="https://github.com/pytorch/pytorch/issues/119873#issuecomment-1945332993" rel="noopener nofollow ugc">here</a></p>
<blockquote>
<p>If you’re willing to do a detour in making your code torch.compile’able, it might tell you about what might be the problem.</p>
</blockquote>
<p>In the context of Dynamo and CUDA Graph, and the tool letting us know “what might be the problem”, I guess he is hinting at the use of <code>TORCH_LOG=perf_hints</code> for debug, as mentioned in the “reduce-overhead” documentation above.</p>
<p>Actually, I find the concept of the PyTorch compiler very interesting, and the fact that it can handle CUDA Graph (automatically?) is just provoking me to dig deeper into its details.</p>
<p>So, as a follow-up question, is there any detailed documentation/tutorial on the same?</p>

ptrblck: <p>Yes, your explanation is correct and <code>torch.compile</code> with <code>reduce-overhead</code> should apply CUDA Graphs automatically where applicable.<br>
However, there might still be a few rough edges as described in <a href="https://github.com/pytorch/pytorch/issues/118228">this</a> and <a href="https://github.com/pytorch/pytorch/issues/113180">this issue</a> explaining convergence issues and <a href="https://github.com/pytorch/pytorch/issues/119607">this</a> as well as <a href="https://github.com/pytorch/pytorch/issues/105181">this issue</a> describing memory leaks.</p>