Repository: pytorch
Issue: nn.DataParallel not working on A100 with nvidia-driver 470.42.01, Cuda 11.1, Debian 10
Labels: module: multi-gpu,triaged,module: data parallel

## 🐛 Bug

Using a single GPU works fine, but just changing the model to model = nn.DataParallel(model) creates issues. Either it halts like in some infinite loop. Or fails with an error `RuntimeError: CUDA error: unspecified launch failure
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.`
<!-- . -->

## To Reproduce

Steps to reproduce the behavior:

1. Copy the following code in a file
2. Run `python <filename>`

```
import torch
import torch.nn as nn
import torch.nn.functional as F
import random
from torch import optim

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Test(nn.Module):
	def __init__(self):
		super(Test, self).__init__()
		self.W = [nn.Linear(200, 100, bias=True), nn.Linear(100, 10, bias=True)]
		self.W = nn.ModuleList(self.W)
		self.g = [nn.ReLU(), nn.ReLU()]
		self.g = nn.ModuleList(self.g)
		self.g = nn.ModuleList(self.g)
		self.dropout = nn.Dropout(0.3)

	def forward(self, x):
		temp = self.dropout(self.g[0](self.W[0](x)))
		out = self.W[1](temp)
		return out

X = torch.randn(100000,200).to(device)
y = torch.tensor([random.randint(0,9) for i in range(100000)]).to(device)
print("Initialising the model now...")
model = Test().to(device)
optimizer = optim.Adam(list(model.parameters()), lr=0.008)
model = nn.DataParallel(model)
print("Done. Did Data Parallel as well")

ce = nn.CrossEntropyLoss()

out = model(X)
loss = ce(out, y) # THIS TAKES EITHER QUITE SOME TIME OR FAILS
print("Loss calculated")
loss.backward() # THIS TAKES EITHER QUITE SOME TIME OR FAILS
print("Backprop done")
optimizer.step()  # THIS TAKES EITHER QUITE SOME TIME OR FAILS
print("loss was: " + str(loss))
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior
Should run normally as expected

<!-- 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
```

Collecting environment information...
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (Debian 8.3.0-6) 8.3.0
Clang version: Could not collect
CMake version: version 3.13.4
Libc version: glibc-2.10

Python version: 3.7.10 | packaged by conda-forge | (default, Feb 19 2021, 16:07:37)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-4.19.0-17-cloud-amd64-x86_64-with-debian-10.10
Is CUDA available: True
CUDA runtime version: 11.4.48
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB
GPU 8: NVIDIA A100-SXM4-40GB
GPU 9: NVIDIA A100-SXM4-40GB
GPU 10: NVIDIA A100-SXM4-40GB
GPU 11: NVIDIA A100-SXM4-40GB
GPU 12: NVIDIA A100-SXM4-40GB
GPU 13: NVIDIA A100-SXM4-40GB
GPU 14: NVIDIA A100-SXM4-40GB
GPU 15: NVIDIA A100-SXM4-40GB

Nvidia driver version: 470.42.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] gnumpy==0.2
[pip3] numpy==1.19.5
[pip3] torch==1.9.0
[pip3] torchaudio==0.9.0a0+33b2469
[pip3] torchvision==0.2.2
[conda] blas                      2.106                       mkl    conda-forge
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] gnumpy                    0.2                      pypi_0    pypi
[conda] libblas                   3.9.0                     6_mkl    conda-forge
[conda] libcblas                  3.9.0                     6_mkl    conda-forge
[conda] liblapack                 3.9.0                     6_mkl    conda-forge
[conda] liblapacke                3.9.0                     6_mkl    conda-forge
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.19.5           py37haa41c4c_1    conda-forge
[conda] pytorch                   1.9.0           py3.7_cuda11.1_cudnn8.0.5_0    pytorch
[conda] torchaudio                0.9.0                      py37    pytorch
[conda] torchvision               0.2.2                      py_3    pytorch

## Additional context
The above provided code is just an example where it's taking quite some time to run. On other slightly different models I tested nn.DataParallel and got similar issues. In fact, on the above code it's still making some progress. On the other runs (that is using other files I was working on), it gave runtime error `RuntimeError: CUDA error: unspecified launch failure
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`

Note that not having nn.DataParallel() works pretty fine.

Can someone please tell if I'm making a mistake somewhere?
<!-- Add any other context about the problem here. -->