Torch's `LayerNorm` and Adam optimizer vs those in tensorflow
Labels: needs reproduction, module: numerical-stability, module: nn, module: optimizer, triaged

### 🐛 Describe the bug

Hello pytorch team, I am recently converting [OpenAI's tensorflow RLHF code](https://github.com/openai/lm-human-preferences) to pytorch. Given the same data and model, I was able even to align the gradients, but somehow training with pytorch resulted in significant instability. Upon further investigation, I was able to pinpoint the difference to **subtle differences in the gradient calculation in Torch's `LayerNorm` and Adam optimizer's update**. 

**The reproduction snippet is https://github.com/vwxyzjn/minimal-adam-layer-norm-bug-repro/blob/master/main.py.** 

## More detail

I used the same data, model, and model weights to perform a Proximal Policy Optimization (PPO) update.  After the first backward pass, you can see that most gradients match perfectly with those obtained in OpenAI's codebase via `tf.train.AdamOptimizer(learning_rate=0.00001, epsilon=1e-5)`.

![image](https://github.com/pytorch/pytorch/assets/5555347/a89ab761-b77d-4e18-a4e1-04771d5da0c5)

### Observation 1

Notice how the gradients of the `LayerNorm` parameters are significantly different. 

As an example, we print some gradients below. Notice how the first gradient for `h0/ln_1/b:0` is `2.4899011e-05` in OAI's codebase, but `1.8152525e-05` in `main.py`. This is a difference of `0.6746486e-05`, which is quite significant.

In comparison, the gradients of the other layers are much more similar. For example, the first gradient for `h0/attn/c_attn/w:0` is `2.88992633e-05` in OAI's codebase, but `2.88992633e-05` in `main.py`. This is a difference of `0.0`, which is much smaller than the difference in the `LayerNorm` parameters.



```
(Pdb) oname, ograd[:10], name, param.grad.detach().numpy()[:10]
('policy/model/h0/ln_1/b:0', array([ 2.4899011e-05, -1.1588502e-03,  1.7985557e-03,  7.4343453e-03,
       -2.5840786e-03, -3.5906259e-03, -6.6465489e-04,  1.8007826e-03,
       -1.6414827e-03, -6.6386913e-03], dtype=float32), 'transformer.h.0.ln_1.bias', array([ 1.8152525e-05, -1.1576341e-03,  1.7961735e-03,  7.4219629e-03,
       -2.5832835e-03, -3.5855419e-03, -6.7265466e-04,  1.8039590e-03,
       -1.6386800e-03, -6.6277790e-03], dtype=float32))
(Pdb) oname, ograd[:10], name, param.grad.detach().numpy()[:10]
('policy/model/h0/attn/c_attn/w:0', array([[[ 2.88992633e-05, -6.70402551e-06, -1.57610848e-05, ...,
         -1.05873929e-04, -9.40704340e-05,  1.00523466e-04],
        [ 7.87996178e-05, -5.04239551e-07, -8.35032733e-06, ...,
         -4.07231477e-04,  4.93751504e-05, -2.81412737e-04],
        [ 8.21374197e-05, -1.94475469e-05, -1.36382323e-05, ...,
         -1.95847577e-04, -4.09606873e-04,  2.84076581e-04],
        ...,
        [-6.02674390e-06,  4.23970641e-06, -7.39748998e-07, ...,
          1.90844381e-04, -8.59782376e-05, -6.60822116e-05],
        [ 3.50006849e-05, -1.32066066e-06, -3.52823263e-05, ...,
         -1.33828435e-04,  1.01715421e-04,  3.40739585e-04],
        [ 1.05423496e-04, -2.66656862e-05, -4.54609835e-05, ...,
         -4.23200603e-04, -1.64171652e-04,  2.63288792e-04]]],
      dtype=float32), 'transformer.h.0.attn.c_attn.weight', array([[ 2.88245592e-05, -6.68141320e-06, -1.57281083e-05, ...,
        -1.05716754e-04, -9.39845631e-05,  1.00422243e-04],
       [ 7.86117525e-05, -4.96559778e-07, -8.27262829e-06, ...,
        -4.06837964e-04,  4.93464249e-05, -2.81286135e-04],
       [ 8.19143170e-05, -1.94303120e-05, -1.35097052e-05, ...,
        -1.95316272e-04, -4.09374770e-04,  2.83872039e-04],
       ...,
       [-1.33238527e-05, -6.14432452e-08,  7.30297143e-06, ...,
         8.94646073e-05, -1.24311875e-04,  1.05930310e-04],
       [-1.15070456e-04,  1.79788076e-05,  3.04212826e-05, ...,
         6.06048678e-04,  3.23058601e-04, -4.77053138e-04],
       [-5.75132690e-05,  2.93947778e-05,  3.10599062e-05, ...,
         2.26184493e-05,  1.36010476e-05,  9.29407452e-06]], dtype=float32))
```

Then after a gradient pass (i.e., `optimizer.step()`) in `main.py`, I plot the parameter difference between `main.py` and `lm-human-preferences`:


![image](https://github.com/pytorch/pytorch/assets/5555347/66e00fdf-d671-41b3-b937-a145e1c0d264)




### Observation 2

Even though the gradients of most layers are similar between pytorch and OAI's codebase, the Adam optimizer causes a significant difference in the parameters post-update. For example, `transformers.h.11.mlp.c_proj_bias` have nearly identical gradients as its counterparts in TensorFlow, as indicated in the last section, but their weights become quite different after a gradient pass.



Then I did the same setup but with the SGD optimizer. The gradient difference is of course the same, but the parameter difference is much smaller and only relevant to the `LayerNorm` parameters:


![image](https://github.com/pytorch/pytorch/assets/5555347/7d472e7a-280b-4ad2-ac9b-8506bd8d161b)


### End-to-end

I then did end-to-end testing with a toy RLHF codebase (https://gist.github.com/vwxyzjn/010e0f92e057ef1a779028d656ab4705) using SGD and Adam, respectively, with 10 random seeds.


```bash
# The only difference is the optimizer used
diff train_policy.py train_policy_adam.py 
292c292
<     optimizer = optim.SGD(policy.parameters(), lr=args.ppo.lr)
---
>     optimizer = optim.Adam(policy.parameters(), lr=args.ppo.lr)
```

The results are staggering, with SGD converging to a good policy and Adam experiencing significant instability:

![image](https://github.com/pytorch/pytorch/assets/5555347/e5d77c92-7ffa-4c51-9faf-df8a6ac971f6)



While I can probably use SGD for now, but this issue is so puzzling and it may warrant further investigation. If you could take a look at it, I would be greatly appreciated.



### Versions

```
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Pop!_OS 21.10 (x86_64)
GCC version: (Ubuntu 11.2.0-7ubuntu2) 11.2.0
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.34

Python version: 3.9.5 (default, Jul 19 2021, 13:27:26)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.17.5-76051705-generic-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: 11.3.109
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3060 Ti
GPU 1: NVIDIA GeForce RTX 3060 Ti

Nvidia driver version: 470.103.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   43 bits physical, 48 bits virtual
CPU(s):                          24
On-line CPU(s) list:             0-23
Thread(s) per core:              2
Core(s) per socket:              12
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           113
Model name:                      AMD Ryzen 9 3900X 12-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         3800.000
CPU max MHz:                     3800.0000
CPU min MHz:                     2200.0000
BogoMIPS:                        7600.37
Virtualization:                  AMD-V
L1d cache:                       384 KiB
L1i cache:                       384 KiB
L2 cache:                        6 MiB
L3 cache:                        64 MiB
NUMA node0 CPU(s):               0-23
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, STIBP always-on, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es

Versions of relevant libraries:
[pip3] numpy==1.25.0
[pip3] torch==2.0.0
[pip3] triton==2.0.0
[conda] Could not collect
```

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @vincentqb @janeyx99 @crcrpar