Enable FusedRMSNorm

[hubertlu-tw]

There are two introduced unit test failures from tests/L0/run_fused_layer_norm/test_fused_layer_norm.py.

========================================================================= short test summary info ==========================================================================
FAILED test_fused_layer_norm.py::TestFusedLayerNormElemWiseHalf::test_layer_norm - RuntimeError: expected scalar type Float but found Half
FAILED test_fused_layer_norm.py::TestFusedLayerNormElemWiseBFloat16::test_layer_norm - RuntimeError: expected scalar type Float but found BFloat16
================================================================= 2 failed, 15 passed, 3 skipped in 42.61s =================================================================

[hubertlu-tw]

After updating tests/L0/run_fused_layer_norm/test_fused_layer_norm.py, the results are as follows:

========================================================== 17 passed, 3 skipped, 17 warnings in 68.83s (0:01:08) ===========================================================

[hubertlu-tw]

The results from CI check "rocm-pytorch-release" show:

======================================================================
ERROR: test_autocast (test_fused_layer_norm.TestAutocastFusedRMSNorm) [torch.float16-True]
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/apex/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py", line 295, in test_autocast
    self._run_test(dtype, elementwise_affine)
  File "/apex/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py", line 277, in _run_test
    expected = native(native_x.cpu())
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1129, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/apex-0.1-py3.7-linux-x86_64.egg/apex/normalization/fused_layer_norm.py", line 383, in forward
    return manual_rms_norm(input, self.normalized_shape, self.weight, self.eps)
  File "/opt/conda/lib/python3.7/site-packages/apex-0.1-py3.7-linux-x86_64.egg/apex/normalization/fused_layer_norm.py", line 29, in manual_rms_norm
    return weight * input
  File "/opt/conda/lib/python3.7/site-packages/apex-0.1-py3.7-linux-x86_64.egg/apex/amp/wrap.py", line 53, in wrapper
    return orig_fn(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/apex-0.1-py3.7-linux-x86_64.egg/apex/amp/wrap.py", line 53, in wrapper
    return orig_fn(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/apex-0.1-py3.7-linux-x86_64.egg/apex/amp/wrap.py", line 53, in wrapper
    return orig_fn(*args, **kwargs)
  [Previous line repeated 24 more times]
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

----------------------------------------------------------------------

The failed test cannot be reproduced locally in the rocm/pytorch:latest docker
(PyTorch: 1.12.0a0+git2a932eb, ROCm: 5.2.0)

[hubertlu-tw]

jenkins: retest this please

[hubertlu-tw]

One of the failed test in the CI check "rocm-pytorch-release" related to the code changes of this PR is:
run_fused_layer_norm/test_fused_layer_norm.py::TestFusedLayerNormElemWiseBFloat16::test_layer_norm

It is marked as flaky test in upstream: https://github.com/NVIDIA/apex/blob/684c4733b1d44f6edfe4b190bf0cc55eb1ae1940/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py#L164

[hubertlu-tw]

apex-rocm-pytorch-master (#30)

• test_adam_option (test_fused_optimizer.TestFusedAdam): due to a regression introduced from upstream PyTorch
• test_multi_device (test_fused_optimizer.TestFusedAdam): due to a regression introduced from upstream PyTorch
• test_float (test_fused_optimizer.TestFusedAdam): due to a regression introduced from upstream PyTorch
• test_multi_device (test_lamb.TestFusedMixedPrecisionLamb): the test failure is irrelevant to the code changes of this PR.
• test_autocast (test_fused_layer_norm.TestAutocastFusedRMSNorm) [torch.float16-True]: skipped since the test failure cannot be reproduced locally.

apex-rocm-pytorch-release (#176)

• test_autocast (test_fused_layer_norm.TestAutocastFusedRMSNorm) [torch.float16-True]: skipped since the test failure cannot be reproduced locally.

[jithunnair-amd]

Among the unit test failures, this one concerns me: test_autocast (test_fused_layer_norm.TestAutocastFusedRMSNorm) because it is related to FusedRMSNorm, which is what this PR is trying to enable. Can we run the CI again to see if this unit test failure still occurs on both master and release CI runs? If so, we should be able to reproduce this reliably. Alternately, if you wish to skip the unit test for now, please create a github issue for the unit test failures introduced by this PR.

[jithunnair-amd]

test_autocast (test_fused_layer_norm.TestAutocastFusedRMSNorm): #81

test_adam_option (test_fused_optimizer.TestFusedAdam): due to https://github.com/pytorch/pytorch/issues/80809#issuecomment-1175211598
test_multi_device (test_fused_optimizer.TestFusedAdam): due to https://github.com/pytorch/pytorch/issues/80809#issuecomment-1175211598
test_float (test_fused_optimizer.TestFusedAdam): due to https://github.com/pytorch/pytorch/issues/80809#issuecomment-1175211598

: #82
test_multi_device (test_lamb.TestFusedMixedPrecisionLamb): #83