Repository: pytorch
Issue: Bug when combining AMP with channels_last memory format
Labels: module: cuda,triaged,module: memory format,module: amp (automated mixed precision)

## 🐛 Bug when combining AMP with channels_last memory format

CUDA unexpectedly fails with different errors when training torchvision models with AMP and channels_last memory format.

## To Reproduce

Steps to reproduce the behavior:

1. Create new conda environment:
`conda create -n test_environment python=3.7 && conda activate test_environment`

2. Install torch, torchvision, cudatoolkit:
`conda install pytorch=1.7.1 torchvision=0.8.2 cudatoolkit=10.2 -c pytorch`
3. Run the following code snippet:

```
import torch
from torch.cuda import amp
from torchvision.models import mnasnet0_5 as model_class


device = torch.device('cuda:0')
torch.cuda.set_device(device)

model = model_class(pretrained=False)
model = model.to(device=device, memory_format=torch.channels_last)

model.train()

scaler = amp.GradScaler(enabled=True)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

x = torch.zeros((1, 3, 256, 256), dtype=torch.float, device=device)
x = x.contiguous(memory_format=torch.channels_last)

for i in range(100):

    optimizer.zero_grad()

    with amp.autocast(enabled=True):
        loss = model(x).contiguous(memory_format=torch.contiguous_format).mean()

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    print('done with batch ' + str(i + 1))
```

This script run results with:
```
(test_environment) naumenko@Cube02:~$ python amp_and_channels_last.py
done with batch 0
Traceback (most recent call last):
  File "amp_and_channels_last.py", line 28, in <module>
    scaler.step(optimizer)
  File "/home/naumenko/miniconda3/envs/test_environment/lib/python3.7/site-packages/torch/cuda/amp/grad_scaler.py", line 320, in step
    if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()):
  File "/home/naumenko/miniconda3/envs/test_environment/lib/python3.7/site-packages/torch/cuda/amp/grad_scaler.py", line 320, in <genexpr>
    if not sum(v.item() for v in optimizer_state["found_inf_per_device"].values()):
RuntimeError: CUDA error: an illegal memory access was encountered
```

When running this code with batch size=16 and CUDA_LAUNCH_BLOCKING=1, it produces the following stack trace:
```
(test_environment) naumenko@Cube02:~$ CUDA_LAUNCH_BLOCKING=1 CUDA_VISIBLE_DEVICES=3 python amp_and_channels_last.py
done with batch 0
Traceback (most recent call last):
  File "amp_and_channels_last.py", line 27, in <module>
    scaler.scale(loss).backward()
  File "/home/naumenko/miniconda3/envs/test_environment/lib/python3.7/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/naumenko/miniconda3/envs/test_environment/lib/python3.7/site-packages/torch/autograd/__init__.py", line 132,in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED
You can try to repro this exception using the following code snippet. If that doesn't trigger the error, please include your original repro script when reporting this issue.

import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.allow_tf32 = True
data = torch.randn([1, 24, 64, 64], dtype=torch.half, device='cuda', requires_grad=True).to(memory_format=torch.channels_last)
net = torch.nn.Conv2d(24, 16, kernel_size=[1, 1], padding=[0, 0], stride=[1, 1], dilation=[1, 1], groups=1)
net = net.cuda().half().to(memory_format=torch.channels_last)
out = net(data)
out.backward(torch.randn_like(out))
torch.cuda.synchronize()

ConvolutionParams
    data_type = CUDNN_DATA_HALF
    padding = [0, 0, 0]
    stride = [1, 1, 0]
    dilation = [1, 1, 0]
    groups = 1
    deterministic = false
    allow_tf32 = true
input: TensorDescriptor 0x7f4968075e10
    type = CUDNN_DATA_HALF
    nbDims = 4
    dimA = 1, 24, 64, 64,
    strideA = 98304, 1, 1536, 24,
output: TensorDescriptor 0x7f49680740c0
    type = CUDNN_DATA_HALF
    nbDims = 4
    dimA = 1, 16, 64, 64,
    strideA = 65536, 1, 1024, 16,
weight: FilterDescriptor 0x7f496807dd50
    type = CUDNN_DATA_HALF
    tensor_format = CUDNN_TENSOR_NHWC
    nbDims = 4
    dimA = 16, 24, 1, 1,
Pointer addresses:
    input: 0x7f497179f200
    output: 0x7f49709cc000
    weight: 0x7f49719ff800
Additional pointer addresses:
    grad_output: 0x7f49709cc000
    grad_input: 0x7f497179f200
Backward data algorithm: 1
```

## Expected behavior

Expected behavior: clear run without errors.

## Environment

Environment settings:
PyTorch version: 1.7.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.5 LTS (x86_64)
GCC version: (Ubuntu 5.5.0-12ubuntu1~16.04) 5.5.0 20171010
Clang version: Could not collect
CMake version: version 3.5.1

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration:
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti

Nvidia driver version: 430.26
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.7.1
[pip3] torchvision==0.8.2
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] mkl                       2020.2                      256
[conda] mkl-service               2.3.0            py37he8ac12f_0
[conda] mkl_fft                   1.2.0            py37h23d657b_0
[conda] mkl_random                1.1.1            py37h0573a6f_0
[conda] numpy                     1.19.2           py37h54aff64_0
[conda] numpy-base                1.19.2           py37hfa32c7d_0
[conda] pytorch                   1.7.1           py3.7_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] torchvision               0.8.2                py37_cu102    pytorch

cc @mcarilli @ngimel @VitalyFedyunin @jamesr66a @ppwwyyxx