CUDA Call Failed with torch.cuda.current_device(), A100, Pytorch 2.2.2, CUDA 12.3, what do I do with env variables to use jupyter notebook?

DKchemistry: <p>Hi,</p>
<p>Apologies if this is solved somewhere, I’ve googled around but I can’t figure out what to do.</p>
<p>I have <code>2.2.2</code> in a conda environment on a linux machine, installed via</p>
<p><code>conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia</code></p>
<p>I have following nvidia-smi:</p>
<pre><code class="lang-auto">Tue Apr  9 03:35:46 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03              Driver Version: 535.54.03    CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA A100 80GB PCIe          Off | 00000000:31:00.0 Off |                   On |
| N/A   67C    P0             252W / 300W |    791MiB / 81920MiB |     N/A      Default |
|                                         |                      |              Enabled |
+-----------------------------------------+----------------------+----------------------+
|   1  NVIDIA A100 80GB PCIe          Off | 00000000:4B:00.0 Off |                   On |
| N/A   67C    P0             251W / 300W |    739MiB / 81920MiB |     N/A      Default |
|                                         |                      |              Enabled |
+-----------------------------------------+----------------------+----------------------+
|   2  NVIDIA A100 80GB PCIe          Off | 00000000:B1:00.0 Off |                   On |
| N/A   34C    P0              45W / 300W |     87MiB / 81920MiB |     N/A      Default |
|                                         |                      |              Enabled |
+-----------------------------------------+----------------------+----------------------+
|   3  NVIDIA A100 80GB PCIe          Off | 00000000:CA:00.0 Off |                   On |
| N/A   34C    P0              42W / 300W |     87MiB / 81920MiB |     N/A      Default |
|                                         |                      |              Enabled |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| MIG devices:                                                                          |
+------------------+--------------------------------+-----------+-----------------------+
| GPU  GI  CI  MIG |                   Memory-Usage |        Vol|      Shared           |
|      ID  ID  Dev |                     BAR1-Usage | SM     Unc| CE ENC DEC OFA JPG    |
|                  |                                |        ECC|                       |
|==================+================================+===========+=======================|
|  0    0   0   0  |             751MiB / 81050MiB  | 98      0 |  7   0    5    1    1 |
|                  |               5MiB / 131072MiB |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  1    0   0   0  |             725MiB / 81050MiB  | 98      0 |  7   0    5    1    1 |
|                  |               5MiB / 131072MiB |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  2    7   0   0  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  2    8   0   1  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  2    9   0   2  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  2   10   0   3  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  2   11   0   4  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  2   12   0   5  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  2   13   0   6  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  3    7   0   0  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  3    8   0   1  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  3    9   0   2  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  3   11   0   3  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  3   12   0   4  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  3   13   0   5  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
|  3   14   0   6  |              12MiB /  9728MiB  | 14      0 |  1   0    0    0    0 |
|                  |               0MiB / 16383MiB  |           |                       |
+------------------+--------------------------------+-----------+-----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0    0    0    1877752      C   ...odinger2024-1/internal/bin/gdesmond      686MiB |
|    1    0    0    1877801      C   ...odinger2024-1/internal/bin/gdesmond      686MiB |
+---------------------------------------------------------------------------------------+
</code></pre>
<p>In my python interperter I get:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; torch.__version__
'2.2.2'
&gt;&gt;&gt; torch.cuda.is_available()
True
</code></pre>
<p>But, if I run:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; torch.cuda.current_device()
# long traceback
RuntimeError: device &gt;= 0 &amp;&amp; device &lt; num_gpus INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1711403380481/work/aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch. device=1, num_gpus=
# long traceback
</code></pre>
<p>I will put the full error below. I am not sure what to do here? I used to run some models with tensorflow that required me to do something like</p>
<pre><code class="lang-auto">export CUDA_VISIBLE_DEVICES=MIG-e343b8ed-9408-52cc-90b7-683cf8ec07f0
</code></pre>
<p>Before I could run any jobs, usually in a format like this, inside a <code>.sh</code> script:</p>
<pre><code class="lang-auto">export CUDA_VISIBLE_DEVICES=MIG-8ad5411d-a2b7-5a4c-9d7e-d973c6c03b08
source activate dd-nvi
python -u progressive_docking_2.py 
</code></pre>
<p>I tried to do the same thing here, in terminal I passed:</p>
<pre><code class="lang-auto">export CUDA_VISIBLE_DEVICES=MIG-e343b8ed-9408-52cc-90b7-683cf8ec07f0
</code></pre>
<p>Then echo:</p>
<pre><code class="lang-auto">echo $CUDA_VISIBLE_DEVICES
MIG-e343b8ed-9408-52cc-90b7-683cf8ec07f0
</code></pre>
<p>Now I could run</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.cuda.is_available()
True
&gt;&gt;&gt; torch.cuda.current_device()
0
&gt;&gt;&gt; torch.cuda.get_device_name(0)
'NVIDIA A100 80GB PCIe MIG 1g.10gb'
</code></pre>
<p>This all looks correct. I can now move a tensor to device:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; X_train = torch.FloatTensor([0., 1., 2.])
&gt;&gt;&gt; X_train = X_train.to("cuda")
</code></pre>
<p>However, I can seem to actually use this environment in a jupyter notebook on this server, because I get the error regarding CUDA call failed. I’ve tried setting stuff like this</p>
<pre><code class="lang-auto">!export CUDA_VISIBLE_DEVICES=MIG-e343b8ed-9408-52cc-90b7-683cf8ec07f0
</code></pre>
<p>But it doesn’t seem to do anything as I can’t echo it back (but I can echo test, see the space between them in the output, indicating $CUDA_VISIBLE_DEVICES is empty.)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/a/5a2ad2aa2e87057c84c52af3e96dbd47e6cd7b45.png" data-download-href="https://discuss.pytorch.org/uploads/default/5a2ad2aa2e87057c84c52af3e96dbd47e6cd7b45" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/a/5a2ad2aa2e87057c84c52af3e96dbd47e6cd7b45_2_588x500.png" alt="image" data-base62-sha1="cREIDGfBgQB82u0CGwJ3jF4JeNT" width="588" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/a/5a2ad2aa2e87057c84c52af3e96dbd47e6cd7b45_2_588x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/a/5a2ad2aa2e87057c84c52af3e96dbd47e6cd7b45_2_882x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/a/5a2ad2aa2e87057c84c52af3e96dbd47e6cd7b45_2_1176x1000.png 2x" data-dominant-color="2A2335"><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">1632×1386 210 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>
<p>And so now when I get to the part of my code that does something with CUDA, it fails with that same CUDA call failed error</p>
<p>(edit: I can only embed one image so I guess I can’t share that line, I am doing this):</p>
<pre><code class="lang-auto">X_train = X_train.to(device)

X_test = X_test.to(device)

y_train = y_train.float().to(device)

y_test = y_test.float().to(device)
</code></pre>
<pre><code class="lang-auto">RuntimeError: device &gt;= 0 &amp;&amp; device &lt; num_gpus INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1711403380481/work/aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch. device=1, num_gpus= The above exception was the direct cause of the following exception: DeferredCudaCallError Traceback (most recent call last) Cell In[27], [line 3](vscode-notebook-cell:?execution_count=27&amp;line=3) [1](vscode-notebook-cell:?execution_count=27&amp;line=1) # %% ----&gt; [3](vscode-notebook-cell:?execution_count=27&amp;line=3) X_train = X_train.to(device) [4](vscode-notebook-cell:?execution_count=27&amp;line=4) X_test = X_test.to(device) [5](vscode-notebook-cell:?execution_count=27&amp;line=5) y_train = y_train.float().to(device)
</code></pre>
<p>I am really stuck as to what I am supposed to do here. I tried to be thorough about following a bunch of tutorials before getting here and my models are running (slowly) on my local Mac using MPS. I’ve gotten TensorFlow to work on this server before so I have a little experience but out of ideas.</p>

ptrblck: <p>You could either <code>export CUDA_VISIBLE_DEVICES</code> pointing to the desired MIG slice globally or disable MIG allowing PyTorch to see all devices.</p>