Torch.compile seems to hang

murphyk: <p>As a followup to <a href="https://discuss.pytorch.org/t/is-there-an-equivalent-of-jax-lax-scan-eg-in-torch-func/177088" class="inline-onebox">Is there an equivalent of jax.lax.scan (eg in torch.func)?</a>, I have been trying to compile my kalman filter code using torch.compile.<br>
However it seems to hang (even though the simple example of torch.compile in the documentation works). I am running torch 2.0 on a macbook with Intel CPU.</p>
<p>Here is the key code snippet:</p>
<pre><code class="lang-auto">def kf_pt(params, emissions, return_covs=False, compile=False):
    F, Q, R = params['F'], params['Q'], params['R']
    def step(carry, t):
        ...
    
    if compile:
        **step = torch.compile(step)**
    num_timesteps = len(emissions)
    D = len(params['mu0'])
    ll = 0
    carry = (ll, params['mu0'], params['Sigma0'])
    for t in range(num_timesteps):
        if return_covs:
            carry, (filtered_means[t], filtered_covs[t]) = step(carry, t)
        else:
            carry, filtered_means[t] = step(carry, t)
    return ll, filtered_means, filtered_covs
</code></pre>
<p>Full code is at <a href="https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf-linreg-pt.ipynb" class="inline-onebox" rel="noopener nofollow ugc">dynamax/kf-linreg-pt.ipynb at main · probml/dynamax · GitHub</a></p>
<p>When I call <code>ll_pt, kf_means_pt, kf_covs_pt = kf_pt(param_dict_pt, Y_pt, return_covs, **compile=True**) </code> it just hangs (for several minutes). If I interupt it I see</p>
<pre><code class="lang-auto">--------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
&lt;timed exec&gt; in &lt;module&gt;

&lt;ipython-input-42-ada65422be98&gt; in kf_pt(params, emissions, return_covs, compile)
     39     for t in range(num_timesteps):
     40         if return_covs:
---&gt; 41             carry, (filtered_means[t], filtered_covs[t]) = step(carry, t)
     42         else:
     43             carry, filtered_means[t] = step(carry, t)

/opt/anaconda3/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py in _fn(*args, **kwargs)
    207             dynamic_ctx.__enter__()
    208             try:
--&gt; 209                 return fn(*args, **kwargs)
    210             finally:
    211                 set_eval_frame(prior)

/opt/anaconda3/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py in catch_errors(frame, cache_size)
    335 
    336         with compile_lock:
--&gt; 337             return callback(frame, cache_size, hooks)
    338 
    339     catch_errors._torchdynamo_orig_callable = callback  # type: ignore[attr-defined]

...

/opt/anaconda3/lib/python3.8/site-packages/torch/_inductor/graph.py in compile_to_module(self)
    573             print(code)
    574 
--&gt; 575         mod = PyCodeCache.load(code)
    576         for name, value in self.constants.items():
    577             setattr(mod, name, value)

/opt/anaconda3/lib/python3.8/site-packages/torch/_inductor/codecache.py in load(cls, source_code)
    526                 mod.__file__ = path
    527                 mod.key = key
--&gt; 528                 exec(code, mod.__dict__, mod.__dict__)
    529                 # another thread might set this first
    530                 cls.cache.setdefault(key, mod)

/var/folders/mn/vt7cgfsx6zs9vblhvbbk7pf8003xtr/T/torchinductor_kpmurphy/ym/cymnh75e3azdeay6wzkazx6yl3p457udrtmnkd3pyeo5vzgsls7k.py in &lt;module&gt;
    109 
    110 
--&gt; 111 async_compile.wait(globals())
    112 del async_compile
    113 

/opt/anaconda3/lib/python3.8/site-packages/torch/_inductor/codecache.py in wait(self, scope)
    713                     pbar.set_postfix_str(key)
    714                 if isinstance(result, (Future, TritonFuture)):
--&gt; 715                     scope[key] = result.result()
    716                     pbar.update(1)
    717 

/opt/anaconda3/lib/python3.8/concurrent/futures/_base.py in result(self, timeout)
    432                 return self.__get_result()
    433 
--&gt; 434             self._condition.wait(timeout)
    435 
    436             if self._state in [CANCELLED, CANCELLED_AND_NOTIFIED]:

/opt/anaconda3/lib/python3.8/threading.py in wait(self, timeout)
    300         try:    # restore state no matter what (e.g., KeyboardInterrupt)
    301             if timeout is None:
--&gt; 302                 waiter.acquire()
    303                 gotit = True
    304             else:

KeyboardInterrupt: 
</code></pre>

marksaroufim: <p>This was a bug which was fixed in the nightlies, try instead <code>torch.compile(.., backend="aot_eager")</code> - inductor the default backend is not supported on Mac</p>
<p><code>pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu</code></p>

murphyk: <aside class="quote no-group quote-modified" data-username="marksaroufim" data-post="2" data-topic="177089">
<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>backend=“aot_eager”</p>
</blockquote>
</aside>
<p>Thanks. That fixes the problem. However, I noticed by code is 5x slower after compiling.<br>
I get the warning</p>
<pre><code class="lang-auto">2023-04-10 17:59:00,831] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: 'step' (kf_linreg_jax_vs_pt.py:171)
   reasons:  L['t'] == 63
</code></pre>
<p>Here is the code:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf_linreg_jax_vs_pt.py#L184">
  <header class="source">

      <a href="https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf_linreg_jax_vs_pt.py#L184" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf_linreg_jax_vs_pt.py#L184" target="_blank" rel="noopener nofollow ugc">probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf_linreg_jax_vs_pt.py#L184</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="174" style="counter-reset: li-counter 173 ;">
          <li>    y = emissions[t]</li>
          <li>    #ll += MVN(H @ pred_mean, H @ pred_cov @ H.T + R).log_prob(y)</li>
          <li>    filtered_mean, filtered_cov = condition_on_pt(pred_mean, pred_cov, H, R, y)</li>
          <li>    pred_mean, pred_cov = predict_pt(filtered_mean, filtered_cov, F, Q)</li>
          <li>    carry = (ll, pred_mean, pred_cov)</li>
          <li>    if return_covs:</li>
          <li>        return carry, (filtered_mean, filtered_cov)</li>
          <li>    else:</li>
          <li>        return carry, filtered_mean</li>
          <li></li>
          <li class="selected">if compile:</li>
          <li>    if platform.system() == 'Darwin':</li>
          <li>        # https://discuss.pytorch.org/t/torch-compile-seems-to-hang/177089</li>
          <li>        step = torch.compile(step, backend="aot_eager")</li>
          <li>    else:</li>
          <li>        step = torch.compile(step)</li>
          <li>num_timesteps = len(emissions)</li>
          <li>D = len(params['mu0'])</li>
          <li>filtered_means = torch.zeros((num_timesteps, D))</li>
          <li>if return_covs:</li>
          <li>    filtered_covs = torch.zeros((num_timesteps, D, D))</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>Any ideas?</p>

marksaroufim: <aside class="quote no-group" data-username="murphyk" data-post="3" data-topic="177089">
<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/murphyk/48/59291_2.png" class="avatar"> murphyk:</div>
<blockquote>
<pre><code class="lang-auto">2023-04-10 17:59:00,831] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: 'step' (kf_linreg_jax_vs_pt.py:171)
   reasons:  L['t'] == 63
</code></pre>
</blockquote>
</aside>
<p>Apologies for the delay but there are a few issues</p>
<ol>
<li>AOT eager won’t have substantial speedups relative to inductor</li>
<li>Your model is dynamic so the recompilations are what is giving the appearance of a hang</li>
<li>For some reason <code>dynamic=True</code> has no impact</li>
<li>After it recompiles your model OOMs</li>
<li>Your benchmark suite doesn’t factor in warmup<br>
6 <code>torch.compile</code> speedups will be most significant on newer GPUs like A100</li>
</ol>
<p>For now I’m running experiments on an A10G which means I default to <code>inductor</code> which is the compiler with the most significant speedups</p>
<p>For the recompiles I ran <code>TORCH_LOGS="recompiles" python kf_linreg_jax_vs_pt.py </code> and it showed that the culprit was the <code>t</code> variable here <a href="https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf_linreg_jax_vs_pt.py#L86">https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf_linreg_jax_vs_pt.py#L86</a> if I fix it to 1 then the hang goes away</p>
<p>To handle dynamic models we have an argument to <code>torch.compile(m, dynamic=True)</code> but for some reason I don’t see its impact on your model</p>
<p>Even with <code>t</code> fixed your model is OOM’ing, I was able to make the OOM go away with this line and I’ve opened an issue about this on github since pattern matcher is on by default <a href="https://github.com/pytorch/pytorch/issues/99084" class="inline-onebox">OOM in fuse_attention inductor pass · Issue #99084 · pytorch/pytorch · GitHub</a></p>
<pre><code class="lang-auto">import torch._inductor.config
torch._inductor.config.pattern_matcher = False
</code></pre>
<p>As far as benchmarking is concerned you call <code>torch.compile(m)</code> the compilation only happens on the first inference so for benchmarks you need to remove the first inference time since the assumption is you’ll amortize it over a long enough experiment</p>
<p>Finally I added a <code>torch.set_default_device('cuda')</code> in your script which defaults all tensors to GPU on an A100 and the benchmarks are not anything to write home about quite yet but the speedups are there</p>
<pre><code class="lang-auto">sm/demos$ python kf_linreg_jax_vs_pt.py 
torch, time=0.604 compile False N 100 D 500
[2023-04-13 21:36:10,010] torch._inductor.utils: [WARNING] make_fallback(aten.cumprod): a decomposition exists, we should switch to it
torch, time=4.884 compile True N 100 D 500
torch, time=0.552 compile True N 100 D 500
</code></pre>

murphyk: <p>Thanks. I know that jax.jit does tracing, and as long as the shapes don’t change when the args change, it doesn’t recompile the function. Here <code>t</code> just specifies which row to take out of a global variable X, so the shape of the variables is constant across iterations. I assumed torch.compile would do something similar to jax.</p>

marksaroufim: <p>It should have been the case as well, PR for fix is out <a href="https://github.com/pytorch/pytorch/pull/99123" class="inline-onebox" rel="noopener nofollow ugc">Don't specialize when indexing by SymInt by ezyang · Pull Request #99123 · pytorch/pytorch · GitHub</a></p>
<p>And the OOM issue has a PR out too <a href="https://github.com/pytorch/pytorch/pull/99128" class="inline-onebox" rel="noopener nofollow ugc">[inductor] Use FakeTensorMode() when creating patterns by jansel · Pull Request #99128 · pytorch/pytorch · GitHub</a></p>