Repository: pytorch
Issue: `layer_norm` needs to be done in fp32 for fp16 inputs
Labels: module: numerical-stability,triaged,module: half,actionable,module: norms and normalization

## 🐛 Bug

As discussed at https://pytorch.slack.com/archives/C3PDTEV8E/p1634151917293100

layer_norm needs to be done in fp32 for fp16 inputs, otherwise overflow happens and there is a significant divergence that starts to add up over multiple chained uses.

This of course can be worked around on the user side, but chances are that it's almost never a desirable behavior to run layer norm in fp16, as overflow is pretty much guaranteed for large tensors.

## To Reproduce

```
import torch

width = 128
input = torch.rand((1,5,width))*0.1
input = input.cuda().half()
normalized_shape = (width,)
weight = torch.ones(width).cuda().half()
bias = torch.zeros(width).cuda().half()
eps = 1e-5
torch.backends.cudnn.enabled = True

output_fp16 = torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)
output_fp32 = torch.layer_norm(input.float(), normalized_shape, weight.float(), bias.float(), eps, torch.backends.cudnn.enabled).half()

torch.testing.assert_equal(output_fp16, output_fp32)
```

Currently it gives:
```
AssertionError: Tensors are not equal!

Mismatched elements: 290 / 640 (45.3%)
Greatest absolute difference: 0.0009765625 at (0, 0, 21)
Greatest relative difference: 0.13069908814589665 at (0, 1, 93)
```

## Expected behavior

output tensors should be either equal or at least much much closer to each other.

## Environment

```
PyTorch version: 1.9.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-88-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration: 
GPU 0: NVIDIA GeForce GTX 1070 Ti
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 465.19.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.1
[pip3] torch==1.9.1+cu111
[pip3] torch-scatter==2.0.8
[pip3] torchaudio==0.9.1
[pip3] torchvision==0.10.1+cu111
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0                 pypi_0    pypi
[conda] mkl-service               2.4.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.0            py38h42c9631_2
[conda] mkl_random                1.2.1            py38ha9443f7_2
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] numpy-base                1.20.1           py38h7d8b39e_0
[conda] torch                     1.9.1+cu111              pypi_0    pypi
[conda] torch-scatter             2.0.8                    pypi_0    pypi
[conda] torchaudio                0.9.1                    pypi_0    pypi
[conda] torchvision               0.10.1+cu111             pypi_0    pypi
```

Thank you.

@ngimel 


cc @VitalyFedyunin @ezyang @gchanan @zou3519