Torch.compile() segfaults on CUDA 11.6

cakeislife100: <p>When running with a <code>torch.compile()</code> model with PyTorch 2.0 and CUDA 11.6, my code is segfaulting with just some sample code. When I remove <code>torch.compile()</code>, the code executes just fine. Any insight into what might be going on would be greatly appreciated <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20">.</p>
<p>Here is my CUDA environment:</p>
<pre><code class="lang-auto">NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.6 
</code></pre>
<p>The exact python packages:</p>
<pre><code class="lang-auto">pytorch-triton           2.0.0+0d7e753227
torch                    2.0.0.dev20230202+cu116
torchaudio               2.0.0.dev20230201+cu116
torchvision              0.15.0.dev20230201+cu116
</code></pre>
<p>And the code that I’m trying to execute.</p>
<pre><code class="lang-auto">import torch
import torchvision.models as models

if __name__ == "__main__":
    model = models.resnet18().cuda()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    # compiled_model = model # Works fine when not actually compiled
    compiled_model = torch.compile(model)

    x = torch.randn(16, 3, 224, 224).cuda()
    optimizer.zero_grad()
    out = compiled_model(x)
    out.sum().backward()
    optimizer.step()
</code></pre>

ptrblck: <p>Could you update PyTorch to the latest nightly as the <code>pytorch-triton</code> module was recently updated to check if you would still see the segfault?</p>

cakeislife100: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> As far as I can see from the cu116 index, the latest version of torch for cu116 is the one from 2023-02-02 which is what I already have installed: <a href="https://download.pytorch.org/whl/nightly/torch/" rel="noopener nofollow ugc">https://download.pytorch.org/whl/nightly/torch/</a>. (And that’s the one that’s installed when I do <code>pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu116</code>)</p>
<p>Is there a newer one that you know of?</p>
<p>The strange part is, when I tried this about a month ago, this code worked (no segfault) with the same CUDA + nvidia driver set up. But I’m not sure which versions of I had then of the different Torch libraries…</p>

ptrblck: <p>You are right and the binaries using CUDA 11.6 were deprecated in favor of 11.7 and 11.8.<br>
Could you install one of these newer nightly releases and rerun your code, please?</p>

cakeislife100: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Ah okay. Is PyTorch 2.0 no longer meant to be compatible with CUDA 11.6?</p>
<p>Also, when I tried it with 11.7 and 11.8, I got <a href="https://discuss.pytorch.org/t/runtimeerror-cannot-find-ptxas/172958">this “RuntimeError: Cannot find ptxas” error</a>. I’m guessing I should just hold tight to wait on a resolution to that issue?</p>

ptrblck: <p>Yes, the binaries with CUDA 11.6 will be deprecated soon and the PyTorch 2.0 binary release will support CUDA 11.7 and 11.8 as described <a href="https://pytorch.org/blog/deprecation-cuda-python-support/">here</a>.<br>
Note that you will still be able to build PyTorch from source with other CUDA versions and the “deprecation” only means the pip wheels and conda binaries will ship with CUDA 11.7 and 11.8 in the next release.</p>
<p>The <code>ptxas</code> issue seems to be a regression and is tracked <a href="https://github.com/pytorch/pytorch/issues/95082">here</a> and <a href="https://github.com/pytorch/builder/issues/1318">here</a>.</p>

cakeislife100: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Just tried the new 0228 versions with CUDA 11.7 and that worked. Thanks for your help!</p>