Mixed precision(autocast) error for a simple network

username: scarecrow1123
cooked: <p>I’m trying out a simple model setup with <code>autocast</code> using the latest nightly version 1.6.0.dev20200406. The model has a single <code>Conv2d</code> and a <code>GRU</code> layer with <code>float32</code> input wrapped inside the <code>autocast</code> context as below:</p>
<pre><code class="lang-auto">class Net(torch.nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 3, (3, 3), padding=(1, 1))
        self.gru = torch.nn.GRU(dim, 1024, batch_first=True)

    def forward(self, x):
        with torch.cuda.amp.autocast():
            out = self.conv(x)
            out = out.reshape(-1, out.size(1) * out.size(2), out.size(3))
            out, _ = self.gru(out)
            return out

model = Net(512).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=.9)
scaler = torch.cuda.amp.GradScaler()
optimizer.zero_grad()

x = torch.rand((32, 1, 1600, 512), dtype=torch.float32).cuda()
y = model(x)
</code></pre>
<p>This doesn’t work and throws the following trace:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-11-8fb506ecd9e7&gt; in &lt;module&gt;
     17 
     18 x = torch.rand((32, 1, 1600, 512), dtype=torch.float32).cuda()
---&gt; 19 y = model(x)

~/miniconda3/envs/torch-test/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    556             result = self._slow_forward(*input, **kwargs)
    557         else:
--&gt; 558             result = self.forward(*input, **kwargs)
    559         for hook in self._forward_hooks.values():
    560             hook_result = hook(self, input, result)

&lt;ipython-input-11-8fb506ecd9e7&gt; in forward(self, x)
      8             out = self.conv(x)
      9             out = out.reshape(-1, out.size(1) * out.size(2), out.size(3))
---&gt; 10             out, _ = self.gru(out)
     11             return out
     12 

~/miniconda3/envs/torch-test/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    556             result = self._slow_forward(*input, **kwargs)
    557         else:
--&gt; 558             result = self.forward(*input, **kwargs)
    559         for hook in self._forward_hooks.values():
    560             hook_result = hook(self, input, result)

~/miniconda3/envs/torch-test/lib/python3.7/site-packages/torch/nn/modules/rnn.py in forward(self, input, hx)
    725         if batch_sizes is None:
    726             result = _VF.gru(input, hx, self._flat_weights, self.bias, self.num_layers,
--&gt; 727                              self.dropout, self.training, self.bidirectional, self.batch_first)
    728         else:
    729             result = _VF.gru(input, batch_sizes, hx, self._flat_weights, self.bias,

RuntimeError: cuDNN error: CUDNN_STATUS_BAD_PARAM
</code></pre>
<p><code>GRU</code> expects a <code>float32</code> input and so explicitly typecasting the <code>Conv2d</code> output to <code>float32</code> resolves this. I’m migrating from Apex, and I know that this shouldn’t be the case with <code>autocast</code> API as well. I’d like to know if I’m missing something here or if it is a limitation with the current version. <a class="mention" href="/u/mcarilli">@mcarilli</a></p>
<p>PyTorch version = 1.6.0.dev20200406+cu101<br>
CuDNN enabled<br>
CUDA = 10.1<br>
Os = Ubuntu 19.10</p>

username: mcarilli
cooked: <p>Glad you’re trying the native API, and thanks for the minimal repro.  I think this indicates my wrappers aren’t reflattening weights in the way cudnn rnns require.  Cudnn rnn support is a high priority, I’ll look at it soon.  In the meantime, the cell-based RNN API should work.</p>

username: Ognjen_Kocic
cooked: <p><a class="mention" href="/u/mcarilli">@mcarilli</a> I can confirm this also happens with <code>1.6.0.dev20200430</code>. Also it affects LSTM as well (my case), while the manually permuting inputs and changing <code>batch_first</code> doesn’t help (in the past some things were broken with <code>batch_first=True</code>).<br>
Also, not sure how to reach you regarding specific use case with variable batch size. So I read materials on AMP and to squeeze out max performance all dimensions on layers like <code>Linear</code> or <code>Conv2D</code> must be divisible by 8 and <code>batch_size</code> as well. Is there any kind of workaround to get decent performance with variable batch size? I can ask completely new question if that is preferred.</p>

username: mcarilli
cooked: <p>Fix is underway <a href="https://github.com/pytorch/pytorch/issues/36428#issuecomment-665748638" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/36428#issuecomment-665748638</a></p>