Repository: pytorch
Issue: RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED when batch size is too large
Labels: module: cudnn,triaged

## 🐛 Bug

Got the RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED when execute a CNN model. It only appears when I have multiple GPU and the batch size is too large ( for my task it is ok for 128 but not 256).



And the GPU memory still has some space

```
[dunan@nvl-003 ~]$ nvidia-smi
Thu Mar  7 15:51:19 2019
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 396.44                 Driver Version: 396.44                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  On   | 00000000:15:00.0 Off |                    0 |
| N/A   43C    P0    76W / 300W |  30731MiB / 32510MiB |     74%   E. Process |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2...  On   | 00000000:16:00.0 Off |                    0 |
| N/A   43C    P0    73W / 300W |  27283MiB / 32510MiB |     78%   E. Process |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     17565      C   python                                     30720MiB |
|    1     17565      C   python                                     27272MiB |
+-----------------------------------------------------------------------------+

```

## To Reproduce

```
Traceback (most recent call last):
  File "/mnt/home/dunan/Job/2018/DNA_classification/DeepDFam/predict.py", line 60, in <module>
    predict(args)
  File "/mnt/home/dunan/Job/2018/DNA_classification/DeepDFam/predict.py", line 31, in predict
    logit = data_parallel(model, feature)
  File "/mnt/home/dunan/anaconda/envs/python3/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py",
 line 188, in data_parallel
    outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
  File "/mnt/home/dunan/anaconda/envs/python3/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py"
, line 83, in parallel_apply
    raise output
  File "/mnt/home/dunan/anaconda/envs/python3/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py"
, line 59, in _worker
    output = module(*input, **kwargs)
  File "/mnt/home/dunan/anaconda/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 48
9, in __call__
    result = self.forward(*input, **kwargs)
  File "/mnt/home/dunan/Job/2018/DNA_classification/DeepDFam/models/SingleFrameCNN.py", line 354, in forward
    x = [F.relu(conv(x)).squeeze(3) for i, conv in enumerate(self.convs1)]
  File "/mnt/home/dunan/Job/2018/DNA_classification/DeepDFam/models/SingleFrameCNN.py", line 354, in <listcomp
>
    x = [F.relu(conv(x)).squeeze(3) for i, conv in enumerate(self.convs1)]
  File "/mnt/home/dunan/anaconda/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 48
9, in __call__
    result = self.forward(*input, **kwargs)
  File "/mnt/home/dunan/Job/2018/DNA_classification/DeepDFam/models/SingleFrameCNN.py", line 307, in forward
    x1 = [conv(x1[:, :, i:seq_len - self.kernel_size[0] + i + 1, :]) for i, conv in enumerate(self.convs1)]
  File "/mnt/home/dunan/Job/2018/DNA_classification/DeepDFam/models/SingleFrameCNN.py", line 307, in <listcomp
>
    x1 = [conv(x1[:, :, i:seq_len - self.kernel_size[0] + i + 1, :]) for i, conv in enumerate(self.convs1)]
  File "/mnt/home/dunan/anaconda/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 48
9, in __call__
    result = self.forward(*input, **kwargs)
  File "/mnt/home/dunan/anaconda/envs/python3/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 320,
 in forward
    self.padding, self.dilation, self.groups)
RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED
```

And here is the modified convolution layer I used

```
class PosWiseConv(nn.Module):
    def __init__(self, channel_in, channel_out, kernel_size, conv_layer=None):
        super(PosWiseConv, self).__init__()
        self.channel_in = channel_in
        self.channel_out = channel_out
        self.kernel_size = kernel_size  # (length, num_token)
        self.convs1 = nn.ModuleList(
            [nn.Conv2d(self.channel_in, self.channel_out, (1, self.kernel_size[1]), bias=False) for i in
             range(self.kernel_size[0])])
        self.convs2 = nn.ModuleList(
            [nn.Conv2d(self.channel_in, self.channel_out, (1, self.kernel_size[1]), bias=False) for i in
             range(self.kernel_size[0])])
        self.convs3 = nn.ModuleList(
            [nn.Conv2d(self.channel_in, self.channel_out, (1, self.kernel_size[1]), bias=False) for i in
             range(self.kernel_size[0])])

        # weight shape is (channel_out, channel_in, kernel_size[0], kernel_size[1])
        # bias shape is (channel_out)
        if conv_layer is not None:
            for i, conv in enumerate(self.convs1):
                conv.weight = nn.Parameter(conv_layer.weight[:, :, i, :].unsqueeze(2))
                self.convs2[i].weight = nn.Parameter(conv_layer.weight[:, :, i, :].unsqueeze(2))
                self.convs3[i].weight = nn.Parameter(conv_layer.weight[:, :, i, :].unsqueeze(2))

        self.bias = conv_layer.bias

    def forward(self, input):
        # input shape: batch_size, channel(3), seq_len, num_vocab
        seq_len = input.shape[2]
        x1 = input[:, 0, :, :].unsqueeze(1)
        x2 = input[:, 1, :, :].unsqueeze(1)
        x3 = input[:, 2, :, :].unsqueeze(1)

        x1 = [conv(x1[:, :, i:seq_len - self.kernel_size[0] + i + 1, :]) for i, conv in enumerate(self.convs1)]
        x2 = [conv(x2[:, :, i:seq_len - self.kernel_size[0] + i + 1, :]) for i, conv in enumerate(self.convs2)]
        x3 = [conv(x3[:, :, i:seq_len - self.kernel_size[0] + i + 1, :]) for i, conv in enumerate(self.convs3)]

        # x[i].shape: batch, channel_out, seq_len - kenel_size + 1, 1
        x = [F.max_pool2d(torch.cat([out, x2[i], x3[i]], 3), (1, 3)) for i, out in
             enumerate(x1)]

        x = torch.sum(torch.cat(x, 3), 3, keepdim=True)
        # expect x shape: (batch, channel_out, seq_len - kenel_size + 1, 1)

        # add bias: shape (channel_out)
        bias = self.bias.unsqueeze(0).unsqueeze(2).unsqueeze(3)
        bias = bias.expand(x.shape)

        x = torch.sum(torch.cat([x, bias], 3), 3, keepdim=True)

        return x
```

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

PyTorch version: 1.0.0
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: CentOS Linux release 7.4.1708 (Core)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-16)
CMake version: version 3.9.4

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB

Nvidia driver version: 396.44
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] msgpack-numpy==0.4.3.2
[pip3] numpy==1.16.1
[pip3] numpydoc==0.7.0
[pip3] torch==1.0.0
[pip3] torchtext==0.2.3
[pip3] torchvision==0.2.1
[conda] blas                      1.0                         mkl
[conda] magma-cuda80              2.3.0                         1    pytorch
[conda] mkl                       2018.0.3                      1
[conda] mkl-dnn                   0.14                    intel_1  [intel]  intel
[conda] mkl-include               2018.0.2                      1
[conda] mkl-service               1.1.2                    py36_3
[conda] mkl_fft                   1.0.10                   py36_0    conda-forge
[conda] mkl_random                1.0.2                    py36_0    conda-forge
[conda] pytorch                   1.0.0           py3.6_cuda9.0.176_cudnn7.4.1_1    pytorch
[conda] torchtext                 0.2.3                    pypi_0    pypi
[conda] torchvision               0.2.1                    py36_1    pytorch

## Additional context

<!-- Add any other context about the problem here. -->


cc @csarofeen @ptrblck