Conda pytorch-cuda=12 for arm64 and Grace Hopper

henrique: <p>Is there any way, plan or timeline to support pytorch-cuda=12 for GH200?<br>
Following <a href="https://pytorch.org/get-started/locally/" class="inline-onebox" rel="noopener nofollow ugc">Start Locally | PyTorch</a> we just get PackagesNotFoundError<br>
and pytorch-cuda=11 installs cuda 11.8 plus pytorch_cpu from conda_forge, which is not really usable <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>We can use the NGC containers or install it from source but conda/mamba is still the go to solution for many users…</p>
<p>Cheers!</p>

ptrblck: <p>The binaries are already being build in CI runs, e.g. <a href="https://github.com/pytorch/pytorch/actions/runs/9115630888?pr=126174">here</a> (scroll down, download the corresponding Python version, and <code>pip install</code> the wheel locally), and the CI job integration allowing a <code>pip install</code> from the nightly index is WIP e.g. <a href="https://github.com/pytorch/pytorch/pull/126174">here</a> (latest update from 5 mins ago, so we are actively working on it).</p>

henrique: <p>Thanks a lot for the info and quick reply, Piotr!<br>
It works for me and it is great to know that it should soon be available in the official channels. Awesome work!</p>
<p>ps.: it seems the wheel has grown significantly in size; is that just a product of the dev build?</p>

ptrblck: <aside class="quote no-group" data-username="henrique" data-post="3" data-topic="203285">
<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/henrique/48/10098_2.png" class="avatar"> henrique:</div>
<blockquote>
<p>ps.: it seems the wheel has grown significantly in size; is that just a product of the dev build?</p>
</blockquote>
</aside>
<p>The current wheel is shipping as a “large” wheel, which packages every dependency (including cuDNN, NCCL, cuBLAS, etc.) into the wheel’s <code>lib</code> folder (we’ve used the same workflow for nightly builds before). One of our next steps would be to use the CUDA PyPI dependencies (as is done for x86 Linux wheels) to build a “small” wheel again. In the end the same data will be downloaded but instead of a gigantic wheel, users would download different wheels instead.</p>

dunlopj: <p>Hello ptrblck,</p>
<p>I know this topic is closed, but I didn’t want to create a new one. How exactly do you install pytorch 2.3.1 with Cuda 12.1 on an aarch64 system? I also have a new system with a grace hopper gpu. It appears that it’s still not yet available on conda. Only cuda 11.8 is available, however, it still downloads the cpu version of pytorch. Pip install does the same. I’m trying to install for Python 3.11. Any help would be greatly appreciated. Thanks.</p>

ptrblck: <p>You would need to install the current nightly pip wheels with CUDA 12.4. <code>torch==2.3.1+cu121</code> is not available for ARM+CUDA.</p>

youkaichao1: <p>Hi, ptrblck, I checked the pytorch 2.4.1 release on pypi, and it does not have aarch64 + cuda support yet. see <a href="https://pypi.org/project/torch/2.4.1/#files" class="inline-onebox" rel="noopener nofollow ugc">torch · PyPI</a> .</p>
<p>when will aarch64 + cuda available for public release?</p>

ptrblck: <p>I see the wheels were build e.g. <code>2.4.1+cu124</code> for Python 3.10 is located here: <a href="https://download.pytorch.org/whl/cu124/torch-2.4.1-cp310-cp310-linux_aarch64.whl#sha256=baa065a4fb7805c78f16841cfc4f3fc3c6823d1de726087e583c68abe553dad7">https://download.pytorch.org/whl/cu124/torch-2.4.1-cp310-cp310-linux_aarch64.whl#sha256=baa065a4fb7805c78f16841cfc4f3fc3c6823d1de726087e583c68abe553dad7</a></p>
<p>Did you forget to specify the <code>index-url</code>?</p>

dunlopj: <p>Hi ptrblck, is there a guide to installing the pip wheels? Thanks.</p>
<p>PS: We were trying to set up the NGC containers, but no luck with those as yet.</p>

ptrblck: <p>Select <code>2.4.1</code> with <code>CUDA 12.4</code> and install <code>torch</code>:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7aa9156801fd2c01966c288c434dcf2d5341ed9.png" data-download-href="https://discuss.pytorch.org/uploads/default/e7aa9156801fd2c01966c288c434dcf2d5341ed9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7aa9156801fd2c01966c288c434dcf2d5341ed9_2_690x278.png" alt="image" data-base62-sha1="x3pJAEuAw0zKovX0vOkmAxKBzrH" width="690" height="278" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7aa9156801fd2c01966c288c434dcf2d5341ed9_2_690x278.png, https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7aa9156801fd2c01966c288c434dcf2d5341ed9.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7aa9156801fd2c01966c288c434dcf2d5341ed9.png 2x" data-dominant-color="F5DCD8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">834×337 28.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">pip install torch --index-url https://download.pytorch.org/whl/cu124
...
Collecting torch
  Downloading https://download.pytorch.org/whl/cu124/torch-2.4.1-cp310-cp310-linux_aarch64.whl (2355.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.4/2.4 GB 43.2 MB/s eta 0:00:00
...
python -c "import torch; print(torch.__version__); print(torch.version.cuda); print(torch.randn(1).cuda())"
2.4.1
12.4
tensor([-1.6937], device='cuda:0')
</code></pre>
<p><code>vision</code> and <code>audio</code> are available as nightlies only for now.</p>

dunlopj: <p>Thank you. I’ll give it a try.</p>

dunlopj: <p>Ok this is only the CPU version. I mistakenly thought it was the cuda version. Pytorch is not recognizing the GPU.</p>

ptrblck: <p>No, it’s the PyTorch binary with CUDA support as seen in my example executed on a Grace Hopper node.</p>

dunlopj: <p>Ok. I miss that part of your previous message. I’ll try again, because I did install, but when I did the torch.cuda.is_available(), I got False.</p>

dunlopj: <p>I got through. I created a fresh environment and installed python 3.10. I had 3.11 installed and it was downloading torch 2.0.1. Thanks a lot. I really appreciate it.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/9/69bc9ccae8f9dd3790791982f4ae1310f8f13763.png" alt="Screenshot from 2024-09-23 22-00-52" data-base62-sha1="f5oeyoc532L0TpRzCzCwCfuGmIP" width="678" height="198"></p>

ptrblck: <p>Great, thanks for verifying! That’s still a bit weird as the Python==3.11 wheels are also available at the specified URL. In any case, good to hear it’s working.</p>

dunlopj: <p>Got it installed, but my code is not recognizing the GPU. It runs on the CPU.</p>

ptrblck: <p>Your code snippet shows you are able to create a tensor on the GPU so your GPU is recognized and usable.</p>

dunlopj: <p>But the actual code I’m trying to run doesn’t pick up the cuda device. Below is a snippet of the code to acquire the device.</p>
<p>if self.args.use_gpu:<br>
device = torch.device(‘cuda’)<br>
print(f’Use GPU: cuda: {torch.cuda.current_device()}')<br>
else:<br>
device = torch.device(‘cpu’)<br>
print(‘Use CPU’)</p>

ptrblck: <p>Check if <code>self.args.use_gpu</code> is set to <code>True</code> and if not, set it.<br>
We’ve now verified that ARM + CUDA binaries are available and are shipping since PyTorch 2.4.0 (as well as the nightly binaries). It seems you might have more questions related to your actual script, so I would recommend starting a new thread to keep this one for the actual binary support.</p>