Torch._dynamo.run vs torch.compile

Abhishek_Ghosh: <p>I understand that if you want to use PyTorch 2.0’s torch.compile feature, you wrap your module with torch.compile and you shall get the benefits.</p>
<p>I was going through <a href="https://github.com/pytorch/pytorch/tree/main/benchmarks/dynamo" rel="noopener nofollow ugc">PyTorch Benchmark Suite</a>, and in the speedup experiments there I found a call to:<br>
<code>torch._dynamo.run()</code></p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L664C13-L665C1">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L664C13-L665C1" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L664C13-L665C1" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L664C13-L665C1</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="664" style="counter-reset: li-counter 663 ;">
          <li>frozen_model_iter_fn = torch._dynamo.run(model_iter_fn)</li>
          <li></li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>The definition of the <code>torch._dynamo.run()</code> function is as follows:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/torch/_dynamo/decorators.py#L27C1-L34C1">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/torch/_dynamo/decorators.py#L27C1-L34C1" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/torch/_dynamo/decorators.py#L27C1-L34C1" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/torch/_dynamo/decorators.py#L27C1-L34C1</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="27" style="counter-reset: li-counter 26 ;">
          <li>def run(fn=None):</li>
          <li>    """Don't do any dynamic compiles, just use prior optimizations"""</li>
          <li>    if fn is not None:</li>
          <li>        fn = innermost_fn(fn)</li>
          <li>        assert callable(fn)</li>
          <li>        return RunOnlyContext()(fn)</li>
          <li>    return RunOnlyContext()</li>
          <li></li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I find the doc string:</p>
<blockquote>
<p>Don’t do any dynamic compiles, just use prior optimizations</p>
</blockquote>
<p>Just want to clarify, is this API latest. And is it at par (performance wise) with torch.compile (after a torch compiled model say has run certain number of warmups?)</p>
<p>Please correct me if my understanding is mistaken or not, but what I feel is that, we can torch.compile a module, which shall do certain optimizations and keep it aside, and then when we use this <code>torch._dynamo.run</code> API, it somehow makes use of those optimizations which we did previously?</p>
<p>If so, I am unable to get, how is it done? Because the</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L624">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L624" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L624" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/common.py#L624</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="614" style="counter-reset: li-counter 613 ;">
          <li>        raise RuntimeError(</li>
          <li>            f"randomize_input can not handle input of type {type(inputs)}"</li>
          <li>        )</li>
          <li></li>
          <li></li>
          <li>def maybe_mark_step(args):</li>
          <li>    if args.trace_on_xla:</li>
          <li>        xm.mark_step()</li>
          <li></li>
          <li></li>
          <li class="selected">def speedup_experiment(args, model_iter_fn, model, example_inputs, **kwargs):</li>
          <li>    """</li>
          <li>    Measure speedups over eager.</li>
          <li></li>
          <li>    Writes to ./speedups.csv</li>
          <li>    """</li>
          <li>    # if args.dynamic_shapes:</li>
          <li>    #     return speedup_experiment_ds(args, model_iter_fn, model, example_inputs)</li>
          <li></li>
          <li>    timings = np.zeros((args.repeat, 2), np.float64)</li>
          <li>    # if we randomize the input, we should also check the result is correct</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
function, takes in a regular <code>model</code> handle, and the function which is wrappered with <code>torch._dynamo.run</code> is not <code>model</code> but rather <code>model_iter_fn</code> which happens to be like either of the following function:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/torchbench.py#L418C5-L432C20">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/torchbench.py#L418C5-L432C20" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/torchbench.py#L418C5-L432C20" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/d4a1b3e09349de3a6a2dd43020e265a88965e4d7/benchmarks/dynamo/torchbench.py#L418C5-L432C20</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="418" style="counter-reset: li-counter 417 ;">
          <li>def forward_pass(self, mod, inputs, collect_outputs=True):</li>
          <li>    with self.autocast(**self.autocast_arg):</li>
          <li>        return mod(*inputs)</li>
          <li></li>
          <li>def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):</li>
          <li>    cloned_inputs = clone_inputs(inputs)</li>
          <li>    self.optimizer_zero_grad(mod)</li>
          <li>    with self.autocast(**self.autocast_arg):</li>
          <li>        pred = mod(*cloned_inputs)</li>
          <li>        loss = self.compute_loss(pred)</li>
          <li>    self.grad_scaler.scale(loss).backward()</li>
          <li>    self.optimizer_step()</li>
          <li>    if collect_outputs:</li>
          <li>        return collect_results(mod, pred, loss, cloned_inputs)</li>
          <li>    return None</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>which defines the logic of how to run the model.</p>
<p>Can there be performance gap if we use <code>torch._dynamo.run</code> as used in the above example (where the model was previously <code>torch.compile</code>d). Or does <code>torch.compile</code>d module after initial warmups, use the same <code>torch._dynamo.run</code>? And both of them are equivalent after a certain number of warmups?</p>

ptrblck: <p>I don’t know the exact difference and why this call is used in the benchmarks (<a class="mention" href="/u/marksaroufim">@marksaroufim</a> might know), but wanted to point out that <code>torch._dynamo</code> is an internal module and should thus not be directly used (note the underscore at the beginning of the module).</p>

marksaroufim: <p>Yeah so as Piotr mentioned this is an internal function you want to instead be using functions that are in the <code>torch.compiler</code> namespace</p>
<p>First off <code>torch.compile()</code> just calls another function called <code>torch._dynamo.optimize()</code> I would not use <code>run</code> in your code because it disables dynamic recompiles which I believe would have some silent correctness issues</p>
<p>Although I’m still not sure where your interest in <code>run()</code> stems from, are you trying to do serialization?</p>

Abhishek_Ghosh: <p>I was just trying to reuse some of the performance benchmark scripts from pytorch benchmark suite (the one which is used in the <a href="https://hud.pytorch.org/benchmark/compilers" rel="noopener nofollow ugc">TorchInductor Performance Dashboard</a>).</p>
<p>I started from</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/641ec2115f300a3e3b39c75f6a32ee3f64afcf30/.ci/pytorch/test.sh#L343-L418">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/641ec2115f300a3e3b39c75f6a32ee3f64afcf30/.ci/pytorch/test.sh#L343-L418" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/641ec2115f300a3e3b39c75f6a32ee3f64afcf30/.ci/pytorch/test.sh#L343-L418" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/641ec2115f300a3e3b39c75f6a32ee3f64afcf30/.ci/pytorch/test.sh#L343-L418</a></h4>



    <pre class="onebox"><code class="lang-sh">
      <ol class="start lines" start="343" style="counter-reset: li-counter 342 ;">
          <li>test_perf_for_dashboard() {</li>
          <li>  TEST_REPORTS_DIR=$(pwd)/test/test-reports</li>
          <li>  mkdir -p "$TEST_REPORTS_DIR"</li>
          <li></li>
          <li>  local suite="$1"</li>
          <li>  shift</li>
          <li></li>
          <li>  local backend=inductor</li>
          <li>  local modes=()</li>
          <li>  if [[ "$DASHBOARD_TAG" == *training-true* ]]; then</li>
          <li>    modes+=(training)</li>
          <li>  fi</li>
          <li>  if [[ "$DASHBOARD_TAG" == *inference-true* ]]; then</li>
          <li>    modes+=(inference)</li>
          <li>  fi</li>
          <li>  # TODO: All the accuracy tests can be skipped once the CI accuracy checking is stable enough</li>
          <li>  local targets=(accuracy performance)</li>
          <li></li>
          <li>  for mode in "${modes[@]}"; do</li>
          <li>    if [[ "$mode" == "inference" ]]; then</li>
      </ol>
    </code></pre>


  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/641ec2115f300a3e3b39c75f6a32ee3f64afcf30/.ci/pytorch/test.sh#L343-L418" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I thought that reusing most of the code from PyTorch itself shall be a good start.<br>
The above script does basically something like:</p>
<pre><code class="lang-auto">./benchmarks/dynamo/torchbench.py --performance --training --amp --backend=inductor --output=torchbench_training.csv
</code></pre>
<p>This <code>torchbench.py</code> during the speedup calculation uses the <code>torch._dynamo.run()</code> function in the <code>common.py</code> file.</p>
<p>That is why I am curious about it. (I am not doing serialization)</p>
<aside class="quote no-group" data-username="marksaroufim" data-post="3" data-topic="201696">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/marksaroufim/48/36747_2.png" class="avatar"> marksaroufim:</div>
<blockquote>
<p>I would not use <code>run</code> in your code because it disables dynamic recompiles which I believe would have some silent correctness issues</p>
</blockquote>
</aside>
<p>Since you said that, is the same <code>torch._dynamo.run</code> function still being used in the latest PyTorch CI Dashboard workflow? Like is it safe to use the run() as in the common.py file — given the situation in that file? Why or why not?</p>
<p>I guess with changing inputs, we might hit a different branch (say), and we might have to compile some part dynamically, but if inputs remain the same (as in the common.py file) then this function is alright in that situation…</p>

b-gran: <p><code>torch._dynamo.run</code> is recommended in the <a href="https://pytorch.org/docs/stable/torch.compiler_faq.html#why-are-you-recompiling-in-production" rel="noopener nofollow ugc">torch.compiler FAQ</a> for production use-cases where recompilation is not needed.</p>
<p>Is this still the recommendation for disabling recompilation? Is there better way to disable recompilation?</p>

soulitzer: <p>You can set <code>torch._dynamo.config.error_on_recompile = True</code> to validate that your code is not recompiling.</p>