[BUG] Post Layer Normalization model gradients not matching

Describe the bug Test inputs in test_cuda_backward.py that have is_preln = False will often fail but the same input with is_preln = True will pass. I can see that in the file all the test cases with post layer normalization are commented out. Is post layer normalization no longer supported or is the test case not able to properly test models with post layer normalization?

To Reproduce Steps to reproduce the behavior:

1. Go to ‘test_cuda_backward.py’
2. Add a test cases on line 260

(64,1600,128,2,24,True,True, 0.2),
(64,1600,128,2,24,False,True, 0.2),

1. Run the test_cuda_backward.py.

Expected behavior Both of the test cases should pass however only the test case that uses pre-layer normalization works.

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/peter/anaconda3/lib/python3.8/site-packages/torch']
torch version .................... 1.9.1+cu111
torch cuda version ............... 11.1
nvcc version ..................... 11.1
deepspeed install path ........... ['/home/peter/anaconda3/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.5.4, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.9, cuda 11.1

System info (please complete the following information):

• OS: Ubuntu 20.04
• GPU count: 1 X RTX 6000
• Python version: Python 3.8

Launcher context Just running the command locally

[PeterDykas]

Awesome thank you!

[RezaYazdaniAminabadi]

Hi @PeterDykas ,

Thanks for pointing this issue. I will look into this and send a PR to fix this soon.

Thanks, Reza