PC crashes and turns off when using large hidden size

mrdbourke: <p>Hello,</p>
<p>I’m trying to replicate the ViT paper: <a href="https://arxiv.org/abs/2010.11929" class="inline-onebox" rel="noopener nofollow ugc">[2010.11929] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale</a></p>
<p>I’ve reproduced the architecture, however, some strange things happen when increasing the size of the <code>nn.Linear()</code> layers.</p>
<p>For example, when trying to run a <code>nn.Linear()</code> layer with over <code>3000</code> hidden units, my PC crashes and turns off immediately.</p>
<p>It works with lower multiples such as:</p>
<ul>
<li>768 - works</li>
<li>1024 - works</li>
<li>2048 - works</li>
</ul>
<p>But anything over 3000, seems to crash my PC without any kind of error warning.</p>
<p>Sometimes it’ll work with 3001 but then won’t work with 3002 hidden units.</p>
<p>The code I’m using is replicating the <code>MLP</code> block in Table 1 of the paper linked above:</p>
<pre><code class="lang-python">import torch
from torch import nn

# Could also call this "FeedForward"
class MLPBlock(nn.Module):
    """Creates an MLPBlock of the Vision Transformer architecture."""
    def __init__(self,
                 embedding_dim, # embedding dimension (Hidden Size D in Table 1)
                 mlp_size, # MLP size in Table 1
                 dropout=0): # "Dropout... is applied to every dense layer... (Appendix B.1)"
        super().__init__()
        
        self.mlp = nn.Sequential(
            nn.Linear(in_features=embedding_dim,
                      out_features=mlp_size),
            nn.GELU(), # "The MLP contains two layers with a GELU non-linearity (section 3.1)."
            nn.Dropout(p=dropout),
            nn.Linear(in_features=mlp_size, # needs to take same in_features as out_features of layer above
                      out_features=embedding_dim), # take back to embedding_dim
            nn.Dropout(p=dropout)
        )

    def forward(self, x):
        return self.mlp(x)

# Create random tensor (same shape as paper)
z = torch.randn((1, 196, 768))
print(z.shape)

# Set MLP size
mlp_size = 1024

# No CUDA
cpu_device = "cpu"
print(f"\nUsing device: {cpu_device}")
print(f"Using MLP size: {mlp_size}")
mlp_block = MLPBlock(embedding_dim=768,
                     mlp_size=mlp_size).to(cpu_device) 
z_through_mlp_block = mlp_block(z.to(cpu_device))
print(z_through_mlp_block.shape)

# With CUDA
cuda_device = "cuda"
print(f"\nUsing device: {cuda_device}")
print(f"Using MLP size: {mlp_size}")
mlp_block = MLPBlock(embedding_dim=768,
                     mlp_size=mlp_size).to(cuda_device) 
z_through_mlp_block_cuda = mlp_block(z.to(cuda_device))
print(z_through_mlp_block_cuda.shape)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">torch.Size([1, 196, 768])

Using device: cpu
Using MLP size: 1024
torch.Size([1, 196, 768])

Using device: cuda
Using MLP size: 1024
torch.Size([1, 196, 768])
</code></pre>
<p>If I set <code>mlp_size</code> to be anything over 3000 in the code above, it crashes my whole PC.</p>
<h2>
<a name="tests-ive-done-1" class="anchor" href="#tests-ive-done-1"></a>Tests I’ve done</h2>
<p>I’m not quite sure what’s going on because I’ve tested the same code on Google Colab with a P100 GPU (~16GB memory) and it works fine at various <code>mlp_size</code> values, including values of 5000+.</p>
<p>But if I run the same code on my local machine with a NVIDIA TITAN RTX (~24GB memory), it crashes immediately.</p>
<h2>
<a name="my-hardware-2" class="anchor" href="#my-hardware-2"></a>My hardware</h2>
<p>I ran the script from the <a href="https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py" rel="noopener nofollow ugc">PyTorch GitHub</a> to show the various hardware/software I’ve got:</p>
<pre><code class="lang-auto">Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.11.0-27-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA TITAN RTX
Nvidia driver version: 470.57.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.3
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] torch==1.11.0
[pip3] torch-tb-profiler==0.3.1
[pip3] torchaudio==0.11.0
[pip3] torchinfo==1.7.0
[pip3] torchmetrics==0.7.2
[pip3] torchvision==0.12.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.21.2           py39h20f2e39_0  
[conda] numpy-base                1.21.2           py39h79a1101_0  
[conda] pytorch                   1.11.0          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-tb-profiler         0.3.1                    pypi_0    pypi
[conda] torchaudio                0.11.0               py39_cu113    pytorch
[conda] torchinfo                 1.7.0                    pypi_0    pypi
[conda] torchmetrics              0.7.2                    pypi_0    pypi
[conda] torchvision               0.12.0               py39_cu113    pytorch
</code></pre>
<h2>
<a name="my-thoughts-3" class="anchor" href="#my-thoughts-3"></a>My thoughts</h2>
<p>Is this potentially something to do with the max memory values I’ve set on my GPU?</p>
<p>I’m not sure where I’d look to find those.</p>
<p>And temperature wise, this happens regardless of warm start or cold start.</p>

ptrblck: <aside class="quote no-group" data-username="mrdbourke" data-post="1" data-topic="154563">
<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/mrdbourke/48/12134_2.png" class="avatar"> mrdbourke:</div>
<blockquote>
<p>Is this potentially something to do with the max memory values I’ve set on my GPU?</p>
</blockquote>
</aside>
<p>No, I don’t think so and based on your description I would guess your system shuts down due to a power spike if the GPU workload exceeds your PSU max. output.<br>
You could try to decrease the clock frequencies of your GPU and see if this throttling would work.</p>

mrdbourke: <p>I think you may be right about the PSU…</p>
<p>It’s now happening for (previously working) older scripts using less compute-heavy networks/data.</p>
<p>Also starting to get a CUDA issue about invalid device ordinals despite only having one CUDA device:</p>
<pre><code class="lang-auto">RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
</code></pre>

ptrblck: <p>The <code>CUDA error: invalid device ordinal</code> might be raised, if your system “dropped” the GPU, which could also indicate a shutdown e.g. due to a weak/faulty PSU. You could check <code>dmesg</code> for <code>Xid</code> error codes and see if some are indicating a bus error etc.</p>

mrdbourke: <p>Thank you!</p>
<p>Turns out you were right.</p>
<p>It was the PSU (or at least so far so good).</p>
<p>My PC started switching off randomly despite the activity (even just browsing the web).</p>
<p>I replaced the PSU today and it’s able to run PyTorch code flawlessly + do many other tasks without a shutdown.</p>
<p>Thank you again <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>