Torch.compile - what is the best scope of compilation?

chedatomasz: <p>I am looking for clarification on the best point to wrap a model in torch.compile, relative to other wrappers / calls. Some previous answers stated it was subject to change in newer releases, so I am looking for information about the current state of things.<br>
Here’s a summary of what I’ve seen so far:</p>
<ol>
<li>
<p>DDP:<br>
According to <a href="https://pytorch.org/docs/main/notes/ddp#example" class="inline-onebox" rel="noopener nofollow ugc">Distributed Data Parallel — PyTorch main documentation</a>, the DDP wrapper should definitely be applied first, as a requirement for DDPOptimizer to work properly.<br>
In contrast, according to <a href="https://discuss.pytorch.org/t/torch-compile-before-or-after-cuda/176031" class="inline-onebox">Torch.compile() before or after .cuda()</a> answer by <a class="mention" href="/u/marksaroufim">@marksaroufim</a>, compile() should be called first on the inner module.<br>
This is reiterated in <a href="https://discuss.pytorch.org/t/how-should-i-use-torch-compile-properly/179021/2" class="inline-onebox">How should I use torch.compile properly? - #2 by marksaroufim</a><br>
Which of these is correct and how big is the impact?</p>
</li>
<li>
<p>.cuda()<br>
According to the answers by Mark linked above, .cuda() should be called first before compile(cuda_model). Is the impact noticeable? The second thread states it shouldn’t be a huge deal, but I’d like to make sure.</p>
</li>
<li>
<p>AMP<br>
I can see two main ways:</p>
</li>
</ol>
<pre><code class="lang-auto">@torch.compile
def uses_autocast(model, input):
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        output=model(input)
</code></pre>
<p>and</p>
<pre><code class="lang-auto">model = torch.compile(model)
with torch.autocast(device_type='cuda', dtype=torch.float16):
    output=model(input)
</code></pre>
<p>Per <a href="https://github.com/pytorch/pytorch/issues/100241" class="inline-onebox" rel="noopener nofollow ugc">Interaction of torch.no_grad and torch.autocast context managers with torch.compile · Issue #100241 · pytorch/pytorch · GitHub</a>, the second one seems to be recommended, as the graph breaks on context manager entry/exit.</p>
<p>Is this still valid for new versions? Are any changes expected here?</p>
<ol start="4">
<li>General extent of compilation<br>
The two scenarios here are</li>
</ol>
<pre><code class="lang-auto">model = torch.compile(model)

for input, target in data:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()
</code></pre>
<p>vs</p>
<pre><code class="lang-auto">@torch.compile
def train_step(...):
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

for input, target in data:
    train_step(...)
</code></pre>
<p>As the issue linked above mentions, theoretically the bigger the extent of compilation the better, but the real results may vary.<br>
In my experiments, compiling the whole train step functions sometimes yields improved performance.<br>
However, the docs and tutorials mainly refer to compiling the module, or just a standalone function (not a training step).<br>
Is compiling the whole step recommended / supported? Is there official guidance on the best extent of compilation? (In my experience, including dataloading gets messy).</p>
<p>Thanks in advance!</p>

ptrblck: <p>These are interesting and good questions! I’m unaware of a “best practice” guide, but I’m sure <a class="mention" href="/u/marksaroufim">@marksaroufim</a> would know more about the proper workflow to use <code>torch.compile</code>.</p>

marksaroufim: <p>Your sense is correct, basically the more things dynamo has to worry about the more likely things will be slower</p>
<p>Answers to your specific question</p>
<ol>
<li>Distributed: Ideally <code>torch.compile()</code> should optimize communication passes but because it doesn’t do so quite yet and because AOT autograd unrolls both the forward and backward pass, you can’t overlap communication effectively. This was fixed for DDP but not yet for FSDP for example so my suggestion in the past was outdated for DDP but still correct for FSDP</li>
<li><code>.cuda()</code> calls: any calls that changes devices is something torch.compile needs to reason about so it’s best to do these before compilation but it’s not necessary. This today messes up cudagraphs and you can read more here <a href="https://stackoverflow.com/questions/68196191/is-changing-the-device-in-a-cuda-graph-node-unavailable" class="inline-onebox" rel="noopener nofollow ugc">multi gpu - Is changing the device in a CUDA Graph node unavailable? - Stack Overflow</a></li>
<li>AMP: You’re right ideally compile within context managers. Support for this will generally be better in nightlies and things move too quickly for me to have an authorative view as to what to recommend, basically try the context manager and if things are too slow or break then compile at amore granular level</li>
<li>Compiling functions vs modules: When 2.0 was released support for optimizers wasn’t quite there but now it is so yes you should optimize the training loop. It’s also critical to compile funcitons instead of modules if your module doesnt call the <code>forward()</code> or <code>__call__()</code> functions, for example HuggingFace models tend to use generate in which case you want to <code>torch.compile(model.generate))</code>. I haven’t personally looked at compiling data loading yet, it’s interesting so if you find specific issues please lmk</li>
</ol>

YunfanZhang42: <p>Hi <a class="mention" href="/u/marksaroufim">@marksaroufim</a>, thank you for your reply! This basically explains things I am looking for. I have a  few questions regarding question <span class="hashtag-raw">#4</span>, compiling models vs compiling the training step:</p>
<blockquote>
<p>Compiling functions vs modules: When 2.0 was released support for optimizers wasn’t quite there but now it is so yes you should optimize the training loop.</p>
</blockquote>
<ol>
<li>I think 2.0 is the latest stable version as of now? Are support for optimizers improved in nightly?</li>
<li>If I compiled the model rather than the entire training step, would the loss.backward() also be JIT compiled? I guess the backward pass is just a reverse of the forward pass, so it should benefit from the graph optimization done in the forward pass as well?</li>
</ol>
<p>Thanks in advance for answering my questions!</p>

marksaroufim: <ol>
<li>Stable is 2.0.1 now, optimizer support is better there but tbh still best in nightlies</li>
<li>Yes that’s the intent</li>
</ol>

YunfanZhang42: <p>Thanks a lot <a class="mention" href="/u/marksaroufim">@marksaroufim</a>! Really appreciate the answer.</p>

divinho: <p><a class="mention" href="/u/marksaroufim">@marksaroufim</a> what’s the state of things now. I feel like this should go into the <a href="https://pytorch.org/docs/stable/torch.compiler_faq.html" rel="noopener nofollow ugc">FAQ</a>, which talks about DDP but never makes a suggestion for what to actually do</p>

bdhirsh: <p>For DDP and compile:</p>
<p>(1) running with them together by default should work (if not, let us know!). The “default” integration basically takes the graph that we capture, splits it into buckets/regions, and compiles each bucket separately - the idea being that we can issue all-reduces in the backward that overlap with each (compiled) bucket.</p>
<p>The main benefit of this integration is that it’s pretty simple (if you are seeing speedups compared to eager, then great). One potential drawback, though is that compile won’t be able to fuse things across buckets. In practice, this may or may not actually matter for getting good performance.</p>
<p>(2) There is an experimental config you can toggle to try to get better performance, if you’re willing to try out some experimental features that are still being hardened: <code>torch._dynamo.config.optimize_ddp = "python_reducer"</code> (<a href="https://github.com/pytorch/pytorch/blob/bf8709b08aaa1d46512790283f43188cf2c496c5/torch/_dynamo/config.py#L275C1-L275C13" rel="noopener nofollow ugc">link</a>).</p>
<p>That config basically ensures that the all-reduces in the backward show up inside of the compiled graph, so compile can get one giant backward graph to optimize - containing all of the backward ops, and the all-reduce communications. The main reason it’s not on by default is that it requires “compiled autograd” (Ed has a great podcast on it <a href="https://pytorch-dev-podcast.simplecast.com/episodes/compiled-autograd-TCcEyBRZ" rel="noopener nofollow ugc">here</a>). Compiled autograd is another feature in compile that effectively allows us to compile more of the backward than we do ordinarily - it can capture stuff like <code>AccumulateGrad</code> nodes in the backward - although it is currently undergoing a re-write and being further hardened. If you’re willing to try it out, you can do so with:</p>
<pre><code class="lang-auto"># compile your model
m = torch.compile(DDP(m))
# run the forward/backward under compiled autograd
with torch._dynamo.utils.maybe_enable_compiled_autograd(True):
    fw_out = m(inp)
    loss = loss_fn(fw_out, ..)
    loss.backward()

</code></pre>