Pytorch amp CUDA error with Transformer

username: thyeros
cooked: <p>Hi. I need some help on this error.</p>
<p>I’m encountering some CUDA error with a Transformer model. Here is a very simple example to reproduce the error. Can someone help with this error? This code works fine without AMP on and also for any typical ConvNets. I wonder if this is anything AMP-incompatible stuff in pytorch Transformer.</p>
<hr>
<p>NVIDIA-SMI 418.40.04    Driver Version: 418.40.04    CUDA Version: 10.1  Tesla V100-SXM2</p>
<p>Python 3.7.6 (default, Jan  8 2020, 19:59:22)<br>
[GCC 7.3.0] :: Anaconda, Inc. on linux<br>
Type “help”, “copyright”, “credits” or “license” for more information.</p>
<blockquote>
<blockquote>
<blockquote>
<p>import torch<br>
print(torch.<strong>version</strong>)<br>
1.6.0+cu101</p>
</blockquote>
</blockquote>
</blockquote>
<hr>
<p>import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F<br>
from torch.nn import TransformerEncoder, TransformerEncoderLayer</p>
<p>class SimpleVT(nn.Module):</p>
<pre><code>def __init__(self):
    super().__init__()
    self.enc_layers = TransformerEncoderLayer(40, 2, 20, 0.5)
    self.encoder = TransformerEncoder(self.enc_layers, 2)
    self.decoder = nn.Linear(40, 2)
    
def forward(self, x):
    x = self.enc_layers(x)
    x = self.encoder(x)
    x = self.decoder(x)
    return x
</code></pre>
<p>model = SimpleVT().cuda()</p>
<p>x = torch.rand([64, 102, 40]).cuda()<br>
y_hat = torch.rand([64, 102, 2]).cuda()</p>
<p>with torch.cuda.amp.autocast(enabled=True):<br>
y = model(x)<br>
loss = nn.MSELoss()(y_hat, y)<br>
loss.backward()</p>
<h2><a name="printysize-1" class="anchor" href="#printysize-1"></a>print(y.size())</h2>
<p>This code fails with the following error.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/f/8f7e64d9f85f8c90690daea1b96dc5c5400051db.png" data-download-href="https://discuss.pytorch.org/uploads/default/8f7e64d9f85f8c90690daea1b96dc5c5400051db" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f7e64d9f85f8c90690daea1b96dc5c5400051db_2_394x500.png" alt="image" data-base62-sha1="ktp4IotAR3Kr5JJF5KTpB5GrxyP" width="394" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f7e64d9f85f8c90690daea1b96dc5c5400051db_2_394x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f7e64d9f85f8c90690daea1b96dc5c5400051db_2_591x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8f7e64d9f85f8c90690daea1b96dc5c5400051db_2_788x1000.png 2x" data-dominant-color="E7E8E7"><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">1722×2180 707 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>

username: ptrblck
cooked: <p>I cannot reproduce it on a V100-SMX2 32GB using PyTorch <code>1.6.0</code> and the CUDA10.1 binaries and get a valid output of <code>torch.Size([64, 102, 2])</code>.</p>
<p>EDIT: the only difference seems to be the driver, so could you update it?</p>

username: thyeros
cooked: <p>Can you tell me which version worked for you?</p>

username: ptrblck
cooked: <p>I’m using <code>450.51.06</code>.</p>