Repository: NVIDIA/TransformerEngine

Issue:
Does `transformer_engine.pytorch.RMSNorm` support `zero_centered_gamma`?
The [doc](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.RMSNorm) explains the option `zero_centered_gamma`  clearly, but I found the implementation at  [normalization.cu:L221](https://github.com/NVIDIA/TransformerEngine/blob/82555b3fd1ddece3a886fac3c7a4aa8d6f8c1955/transformer_engine/pytorch/csrc/extensions/normalization.cu#L221) assert `zero_centered_gamma == false`.

Is there any limitation?


Discussion:
Yes, that's a good catch. I will add the support for this option.

Resolving patch:
diff --git a/transformer_engine/pytorch/module/rmsnorm.py b/transformer_engine/pytorch/module/rmsnorm.py
--- a/transformer_engine/pytorch/module/rmsnorm.py
+++ b/transformer_engine/pytorch/module/rmsnorm.py
@@ -114,7 +114,7 @@ class RMSNorm(torch.nn.Module):
                          the RMSNorm formula changes to
 
                          .. math::
-                            y = \frac{x}{RMS(x) + \varepsilon} * (1 + \gamma)
+                            y = \frac{x}{RMS_\varepsilon(x)} * (1 + \gamma)
     device : Union[torch.device, str], default = "cuda"
           The device on which the parameters of the model will allocated. It is the user's
           responsibility to ensure all parameters are moved to the GPU before running the
@@ -155,7 +155,7 @@ def __init__(
     def reset_rms_norm_parameters(self) -> None:
         """Init RMSNorm params"""
         warnings.warn(
-            ("This method will be deprecated in an upcoming release. "
+            ("This method is deprecated and will be removed in an upcoming release. "
              "Update your code to use RMSNorm.reset_parameters() instead."),
             DeprecationWarning,
             stacklevel=2


Test patch:
diff --git a/tests/cpp/operator/test_rmsnorm.cu b/tests/cpp/operator/test_rmsnorm.cu
--- a/tests/cpp/operator/test_rmsnorm.cu
+++ b/tests/cpp/operator/test_rmsnorm.cu
@@ -43,13 +43,17 @@ void compute_ref_stats(const InputType *data, float *rsigma, const size_t N, con
 template <typename InputType, typename OutputType>
 void compute_ref_output(const InputType *data, const InputType *gamma, OutputType *output,
                         const float *rsigma, const size_t N, const size_t H, float *amax,
-                        float scale) {
+                        float scale, const bool zero_centered_gamma) {
   using compute_t = float;
   compute_t current_max = -1e100;
   for (size_t i = 0; i < N; ++i) {
     for (size_t j = 0; j < H; ++j) {
       compute_t current = static_cast<compute_t>(data[i * H + j]);
-      compute_t tmp = current * rsigma[i] * static_cast<compute_t>(gamma[j]);
+      compute_t g = static_cast<compute_t>(gamma[j]);
+      if (zero_centered_gamma) {
+        g += 1;
+      }
+      compute_t tmp = current * rsigma[i] * g;
       output[i * H + j] = static_cast<OutputType>(tmp * scale);
       current_max = fmaxf(current_max, fabsf(tmp));
     }
@@ -60,7 +64,7 @@ void compute_ref_output(const InputType *data, const InputType *gamma, OutputTyp
 template <typename InputType, typename OutputType>
 void compute_ref_backward(const OutputType *output_grad, const InputType *data, const float *rsigma,
                           const InputType *gamma, InputType *data_grad, InputType *gamma_grad,
-                          const size_t N, const size_t H) {
+                          const size_t N, const size_t H, const bool zero_centered_gamma) {
   using compute_t = float;
   std::vector<compute_t> dgamma(H, 0.f);
 
@@ -70,7 +74,10 @@ void compute_ref_backward(const OutputType *output_grad, const InputType *data,
     for (size_t j = 0; j < H; ++j) {
       const compute_t x = static_cast<compute_t>(data[i * H + j]);
       const compute_t y = x * rsigma[i];
-      const compute_t g = static_cast<compute_t>(gamma[j]);
+      compute_t g = static_cast<compute_t>(gamma[j]);
+      if (zero_centered_gamma) {
+        g += 1;
+      }
       const compute_t dz = static_cast<compute_t>(output_grad[i * H + j]);
       const compute_t dy = g * dz;
       dgamma[j] += y * dz;
@@ -82,7 +89,10 @@ void compute_ref_backward(const OutputType *output_grad, const InputType *data,
     for (size_t j = 0; j < H; ++j) {
       const compute_t x = static_cast<compute_t>(data[i * H + j]);
       const compute_t y = x * rsigma[i];
-      const compute_t g = static_cast<compute_t>(gamma[j]);
+      compute_t g = static_cast<compute_t>(gamma[j]);
+      if (zero_centered_gamma) {
+        g += 1;
+      }
       const compute_t dz = static_cast<compute_t>(output_grad[i * H + j]);
       const compute_t dy = g * dz;
       const compute_t dx = rsigma[i] * (dy - mdyy * y);
@@ -97,7 +107,7 @@ void compute_ref_backward(const OutputType *output_grad, const InputType *data,
 }
 
 template <typename InputType, typename OutputType>
-void performTest(const size_t N, const size_t H) {
+void performTest(const size_t N, const size_t H, const bool zero_centered_gamma) {
   if (sizeof(InputType) < sizeof(OutputType)) {
     GTEST_SKIP() << "RMSNorm kernel does not support OutputType > InputType";
     return;
@@ -137,23 +147,25 @@ void performTest(const size_t N, const size_t H) {
 
   // Forward kernel
   float epsilon = 1e-5;
-  nvte_rmsnorm_fwd(input.data(), gamma.data(), epsilon, z.data(), rsigma.data(), 0,
-                   prop.multiProcessorCount, workspace.data(), barrier.data());
+  auto fwd_function = zero_centered_gamma ? nvte_rmsnorm1p_fwd : nvte_rmsnorm_fwd;
+  fwd_function(input.data(), gamma.data(), epsilon, z.data(), rsigma.data(), 0,
+               prop.multiProcessorCount, workspace.data(), barrier.data());
   workspace = Tensor(workspace.shape(), workspace.dtype());
   barrier = Tensor(barrier.shape(), barrier.dtype());
-  nvte_rmsnorm_fwd(input.data(), gamma.data(), epsilon, z.data(), rsigma.data(), 0,
-                   prop.multiProcessorCount, workspace.data(), barrier.data());
+  fwd_function(input.data(), gamma.data(), epsilon, z.data(), rsigma.data(), 0,
+               prop.multiProcessorCount, workspace.data(), barrier.data());
 
   // Backward kernel
-  nvte_rmsnorm_bwd(dz.data(), input.data(), rsigma.data(), gamma.data(), dx.data(), dgamma.data(),
-                   dgamma_part.data(), 0, prop.multiProcessorCount, workspace.data(),
-                   barrier.data());
+  auto bwd_function = zero_centered_gamma ? nvte_rmsnorm1p_bwd : nvte_rmsnorm_bwd;
+  bwd_function(dz.data(), input.data(), rsigma.data(), gamma.data(), dx.data(), dgamma.data(),
+               dgamma_part.data(), 0, prop.multiProcessorCount, workspace.data(),
+               barrier.data());
   workspace = Tensor(workspace.shape(), workspace.dtype());
   barrier = Tensor(barrier.shape(), barrier.dtype());
   dgamma_part = Tensor(dgamma_part.shape(), dgamma_part.dtype());
-  nvte_rmsnorm_bwd(dz.data(), input.data(), rsigma.data(), gamma.data(), dx.data(), dgamma.data(),
-                   dgamma_part.data(), 0, prop.multiProcessorCount, workspace.data(),
-                   barrier.data());
+  bwd_function(dz.data(), input.data(), rsigma.data(), gamma.data(), dx.data(), dgamma.data(),
+               dgamma_part.data(), 0, prop.multiProcessorCount, workspace.data(),
+               barrier.data());
 
   // Reference implementations
   // use the GPU stats to tighten the tolerances
@@ -162,10 +174,11 @@ void performTest(const size_t N, const size_t H) {
   compute_ref_stats(input.cpu_dptr<InputType>(), ref_rsigma.get(), N, H, epsilon);
   float ref_scale = isFp8Type(otype) ? z.scale() : 1.f;
   compute_ref_output(input.cpu_dptr<InputType>(), gamma.cpu_dptr<WeightType>(), ref_output.get(),
-                     rsigma.cpu_dptr<float>(), N, H, &ref_amax, ref_scale);
+                     rsigma.cpu_dptr<float>(), N, H, &ref_amax, ref_scale,
+                     zero_centered_gamma);
   compute_ref_backward(dz.cpu_dptr<WeightType>(), input.cpu_dptr<InputType>(),
                        rsigma.cpu_dptr<float>(), gamma.cpu_dptr<WeightType>(), ref_dx.get(),
-                       ref_dgamma.get(), N, H);
+                       ref_dgamma.get(), N, H, zero_centered_gamma);
 
   cudaDeviceSynchronize();
   auto err = cudaGetLastError();
@@ -197,9 +210,10 @@ std::vector<std::pair<size_t, size_t>> test_cases = {
 
 }  // namespace
 
-class RMSNormTestSuite
-    : public ::testing::TestWithParam<std::tuple<
-          transformer_engine::DType, transformer_engine::DType, std::pair<size_t, size_t>>> {};
+class RMSNormTestSuite : public ::testing::TestWithParam<std::tuple<transformer_engine::DType,
+                                                                    transformer_engine::DType,
+                                                                    std::pair<size_t, size_t>,
+                                                                    bool>> {};
 
 TEST_P(RMSNormTestSuite, TestRMSNorm) {
   using namespace transformer_engine;
@@ -208,11 +222,11 @@ TEST_P(RMSNormTestSuite, TestRMSNorm) {
   const DType input_type = std::get<0>(GetParam());
   const DType output_type = std::get<1>(GetParam());
   const auto size = std::get<2>(GetParam());
+  const bool zero_centered_gamma = std::get<3>(GetParam());
 
-  TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(
-      input_type, InputType,
-      TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(
-          output_type, OutputType, performTest<InputType, OutputType>(size.first, size.second);););
+  TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(input_type, InputType,
+    TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(output_type, OutputType,
+      performTest<InputType, OutputType>(size.first, size.second, zero_centered_gamma);););
 }
 
 INSTANTIATE_TEST_SUITE_P(OperatorTest, RMSNormTestSuite,
@@ -220,11 +234,14 @@ INSTANTIATE_TEST_SUITE_P(OperatorTest, RMSNormTestSuite,
                                                               DType::kFloat16),
                                             ::testing::Values(DType::kFloat32, DType::kBFloat16,
                                                               DType::kFloat16, DType::kFloat8E4M3),
-                                            ::testing::ValuesIn(test_cases)),
+                                            ::testing::ValuesIn(test_cases),
+                                            ::testing::Values(false, true)),
                          [](const testing::TestParamInfo<RMSNormTestSuite::ParamType> &info) {
-                           std::string name = test::typeName(std::get<0>(info.param)) + "X" +
-                                              test::typeName(std::get<1>(info.param)) + "X" +
-                                              std::to_string(std::get<2>(info.param).first) + "X" +
-                                              std::to_string(std::get<2>(info.param).second);
+                           std::string name =
+                             test::typeName(std::get<0>(info.param)) + "X" +
+                             test::typeName(std::get<1>(info.param)) + "X" +
+                             std::to_string(std::get<2>(info.param).first) + "X" +
+                             std::to_string(std::get<2>(info.param).second) + "X" +
+                             std::to_string(std::get<3>(info.param));
                            return name;
                          });
diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -21,7 +21,7 @@
 )
 from transformer_engine.pytorch import (
     DotProductAttention, LayerNormLinear, LayerNormMLP, Linear,
-    MultiheadAttention, RMSNorm, TransformerLayer
+    MultiheadAttention, RMSNorm, TransformerLayer, LayerNorm
 )
 from transformer_engine.pytorch.distributed import checkpoint as te_checkpoint
 from transformer_engine.pytorch.distributed import _set_cuda_rng_state, CudaRNGStatesTracker
@@ -215,15 +215,41 @@ def forward(
         return context_layer
 
 
+class TorchLayerNorm(nn.Module):
+    def __init__(self, in_features: int,
+                 eps: float,
+                 zero_centered_gamma: bool):
+        super().__init__()
+        self.eps = eps
+        self.in_features = in_features
+        self.zero_centered_gamma = zero_centered_gamma
+
+        initial_value = torch.ones(in_features) if zero_centered_gamma else torch.zeros(in_features)
+        self.weight = nn.Parameter(initial_value)
+        self.bias = nn.Parameter(torch.zeros(in_features))
+        self.register_parameter("weight", self.weight)
+        self.register_parameter("bias", self.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        w = self.weight if not self.zero_centered_gamma else 1 + self.weight
+        w = w.to(torch.float32)
+        b = self.bias.to(torch.float32)
+        inp = x.to(torch.float32)
+        out = torch.nn.functional.layer_norm(inp, (self.in_features,), weight=w,
+                                             bias=b, eps=self.eps)
+        return out.to(x.dtype)
+
 # Adapted from https://github.com/bzhangGo/rmsnorm/blob/c6691f20ec0af4128c8159c903071f7575404295/rmsnorm_torch.py
 class TorchRMSNorm(nn.Module):
-    def __init__(self, in_features, eps=1e-5):
+    def __init__(self, in_features, zero_centered_gamma, eps=1e-5):
         super().__init__()
 
         self.eps = eps
         self.in_features = in_features
+        self.zero_centered_gamma = zero_centered_gamma
 
-        self.weight = nn.Parameter(torch.ones(in_features))
+        initial_value = torch.ones(in_features) if zero_centered_gamma else torch.zeros(in_features)
+        self.weight = nn.Parameter(initial_value)
         self.register_parameter("weight", self.weight)
 
     def forward(self, x):
@@ -234,18 +260,24 @@ def forward(self, x):
         r_rms_x = rms_x2 ** (-1. / 2)
         x_normed = x * r_rms_x
 
-        return (self.weight.float() * x_normed).to(x.dtype)
+        w = self.weight.float()
+        if self.zero_centered_gamma:
+            w = 1 + w
+        return (w * x_normed).to(x.dtype)
 
 
 class TorchLayerNormLinear(nn.Module):
     def __init__(self, in_features: int, out_features: int,
                  eps: float, bias: bool = True,
-                 normalization: str = "LayerNorm"):
+                 normalization: str = "LayerNorm",
+                 zero_centered_gamma: bool = False):
         super().__init__()
         if normalization == "LayerNorm":
-            self.layernorm = nn.LayerNorm(in_features, eps=eps)
+            self.layernorm = TorchLayerNorm(in_features, eps=eps,
+                                            zero_centered_gamma=zero_centered_gamma)
         elif normalization == "RMSNorm":
-            self.layernorm = TorchRMSNorm(in_features, eps=eps)
+            self.layernorm = TorchRMSNorm(in_features, eps=eps,
+                                          zero_centered_gamma=zero_centered_gamma)
         else:
             raise RuntimeError("Unsupported normalization")
 
@@ -299,9 +331,11 @@ def __init__(self, hidden_size: int, ffn_hidden_size: int,
                  normalization: str = "LayerNorm"):
         super().__init__()
         if normalization == "LayerNorm":
-            self.ln = nn.LayerNorm(hidden_size, eps=eps)
+            self.ln = TorchLayerNorm(hidden_size, eps=eps,
+                                     zero_centered_gamma=False)
         elif normalization == "RMSNorm":
-            self.ln = TorchRMSNorm(hidden_size, eps=eps)
+            self.ln = TorchRMSNorm(hidden_size, eps=eps,
+                                   zero_centered_gamma=False)
         else:
             raise RuntimeError("Unsupported normalization")
         if 'glu' in activation:
@@ -893,13 +927,15 @@ def test_linear_accuracy(dtype, bs, model):
 @pytest.mark.parametrize("bs", batch_sizes)
 @pytest.mark.parametrize("model", model_configs.keys())
 @pytest.mark.parametrize("eps", [1e-1, 1e-3, 1e-5, 1e-7])
-def test_rmsnorm_accuracy(dtype, bs, model, eps):
+@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
+def test_rmsnorm_accuracy(dtype, bs, model, eps, zero_centered_gamma):
     config = model_configs[model]
 
     te_rmsnorm = (
         RMSNorm(
             config.hidden_size,
             eps=eps,
+            zero_centered_gamma=zero_centered_gamma
         )
         .to(dtype=dtype)
         .cuda()
@@ -910,6 +946,7 @@ def test_rmsnorm_accuracy(dtype, bs, model, eps):
         TorchRMSNorm(
             config.hidden_size,
             eps=eps,
+            zero_centered_gamma=zero_centered_gamma
         )
         .to(dtype=dtype)
         .cuda()
@@ -924,17 +961,64 @@ def test_rmsnorm_accuracy(dtype, bs, model, eps):
     torch_outputs = _test_granular_accuracy(torch_rmsnorm, bs, dtype, config)
 
     # Check output.
-    if dtype == torch.float32:
-        assert_allclose(te_outputs[0], torch_outputs[0], 1e-7)
-    else:
-        assert_allclose(te_outputs[0], torch_outputs[0], 2e-2)
+    atol = {torch.float32 : 1e-7,
+            torch.half    : 2e-3,
+            torch.bfloat16: 2e-2,
+    }
+    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype])
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("eps", [1e-1, 1e-3, 1e-5, 1e-7])
+@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
+def test_layernorm_accuracy(dtype, bs, model, eps, zero_centered_gamma):
+    config = model_configs[model]
+
+    te_layernorm = (
+        LayerNorm(
+            config.hidden_size,
+            eps=eps,
+            zero_centered_gamma=zero_centered_gamma
+        )
+        .to(dtype=dtype)
+        .cuda()
+        .eval()
+    )
+
+    torch_layernorm = (
+        TorchLayerNorm(
+            config.hidden_size,
+            eps=eps,
+            zero_centered_gamma=zero_centered_gamma
+        )
+        .to(dtype=dtype)
+        .cuda()
+        .eval()
+    )
+
+    # Share params
+    with torch.no_grad():
+        torch_layernorm.weight = Parameter(te_layernorm.weight.clone())
+        torch_layernorm.bias = Parameter(te_layernorm.bias.clone())
+
+    te_outputs = _test_granular_accuracy(te_layernorm, bs, dtype, config)
+    torch_outputs = _test_granular_accuracy(torch_layernorm, bs, dtype, config)
+
+    # Check output.
+    atol = {torch.float32 : 1e-7,
+            torch.half    : 2e-3,
+            torch.bfloat16: 2e-2,
+    }
+    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype])
 
 
 @pytest.mark.parametrize("dtype", param_types)
 @pytest.mark.parametrize("bs", batch_sizes)
 @pytest.mark.parametrize("model", model_configs.keys())
 @pytest.mark.parametrize("normalization", all_normalizations)
-def test_layernorm_linear_accuracy(dtype, bs, model, normalization):
+@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
+def test_layernorm_linear_accuracy(dtype, bs, model, normalization, zero_centered_gamma):
     config = model_configs[model]
 
     te_ln_linear = (
@@ -944,6 +1028,7 @@ def test_layernorm_linear_accuracy(dtype, bs, model, normalization):
             config.eps,
             bias=True,
             normalization=normalization,
+            zero_centered_gamma=zero_centered_gamma,
         )
         .to(dtype=dtype)
         .cuda()
@@ -957,6 +1042,7 @@ def test_layernorm_linear_accuracy(dtype, bs, model, normalization):
             config.eps,
             bias=True,
             normalization=normalization,
+            zero_centered_gamma=zero_centered_gamma,
         )
         .to(dtype=dtype)
         .cuda()
@@ -975,10 +1061,11 @@ def test_layernorm_linear_accuracy(dtype, bs, model, normalization):
     torch_outputs = _test_granular_accuracy(torch_ln_linear, bs, dtype, config)
 
     # Check output.
-    if dtype == torch.float32:
-        assert_allclose(te_outputs[0], torch_outputs[0], 5e-3)
-    else:
-        assert_allclose(te_outputs[0], torch_outputs[0], 5e-2)
+    atol = {torch.float32 : 2e-4,
+            torch.half    : 2e-3,
+            torch.bfloat16: 2e-2,
+    }
+    assert_allclose(te_outputs[0], torch_outputs[0], atol[dtype])
 
 
 @pytest.mark.parametrize("dtype", param_types)
diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py
--- a/tests/pytorch/test_onnx_export.py
+++ b/tests/pytorch/test_onnx_export.py
@@ -625,7 +625,7 @@ def __init__(self) -> None:
             eps = 1e-6 # An arbitrary small value
             dtype = torch.float if fake_bf16_io else precision
             self.ln = te.LayerNorm(inp_shape[1], eps, params_dtype=dtype,
-                zero_centered_gamma=False).eval().cuda()
+                zero_centered_gamma=zero_centered_gamma).eval().cuda()
 
         def forward(self, inp):
             ret = self.ln(inp)
@@ -679,6 +679,7 @@ def forward(self, inp):
             fname, inp, model, atol=atol, is_fp8=use_fp8, allow_cnt_errors=3, te_outputs=te_outputs)
 
 @pytest.mark.parametrize("scale_factor", [448, 112])
+@pytest.mark.parametrize("zero_centered_gamma", [False, True])
 @pytest.mark.parametrize(
     "use_fp8, precision,             atol", [
     [False,   torch.float32,         1e-7],
@@ -695,6 +696,7 @@ def test_export_rmsnorm(
     use_fp8: bool,
     scale_factor: float,
     precision: torch.dtype,
+    zero_centered_gamma: bool,
     atol: float
 ):
     fake_bf16_io = precision == "fake-torch.bfloat16"
@@ -713,7 +715,8 @@ def __init__(self) -> None:
             super().__init__()
             eps = 1e-6 # An arbitrary small value
             dtype = torch.float if fake_bf16_io else precision
-            self.ln = te.RMSNorm(inp_shape[1], eps, params_dtype=dtype).eval().cuda()
+            self.ln = te.RMSNorm(inp_shape[1], eps, params_dtype=dtype,
+                                 zero_centered_gamma=zero_centered_gamma).eval().cuda()
 
         def forward(self, inp):
             ret = self.ln(inp)
@@ -739,7 +742,7 @@ def forward(self, inp):
                 self.meta,
                 self.fp8_tensor,
                 self.fp8_type,
-                False)
+                zero_centered_gamma)
 
             ret = cast_from_fp8(
                 ret,
@@ -875,9 +878,6 @@ def test_export_layernorm_linear(
     if use_fp8 and not fp8_available:
         pytest.skip(reason_for_no_fp8)
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
-
     # Set dimensions (these are arbitrary).
     in_features = 64
     out_features = 256
@@ -946,8 +946,6 @@ def test_export_layernorm_mlp(
     if use_fp8 and not fp8_available:
         pytest.skip(reason_for_no_fp8)
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
 
     # Set dimensions (these are arbitrary).
     in_features = 64
diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py
--- a/tests/pytorch/test_sanity.py
+++ b/tests/pytorch/test_sanity.py
@@ -361,9 +361,6 @@ def test_sanity_layernorm_linear(dtype, fp8_recipe, model, skip_wgrad,
         if not config.is_fp8_supported():
             pytest.skip("Model config does not support FP8")
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
-
     sigma = 0.023
     init_method = init_method_normal(sigma)
 
@@ -427,9 +424,6 @@ def test_sanity_layernorm_mlp(dtype, fp8_recipe, model, skip_wgrad,
         if not config.is_fp8_supported():
             pytest.skip("Model config does not support FP8")
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
-
     sigma = 0.023
     init_method = init_method_normal(sigma)
     output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
@@ -472,9 +466,6 @@ def test_sanity_gpt(dtype, fp8_recipe, model, skip_wgrad,
         if not config.is_fp8_supported():
             pytest.skip("Model config does not support FP8")
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
-
     sigma = 0.023
     init_method = init_method_normal(sigma)
     output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
@@ -544,9 +535,6 @@ def test_sanity_bert(dtype, fp8_recipe, model, skip_wgrad, zero_centered_gamma,
         if not config.is_fp8_supported():
             pytest.skip("Model config does not support FP8")
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
-
     sigma = 0.023
     init_method = init_method_normal(sigma)
     output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
@@ -608,9 +596,6 @@ def test_sanity_T5(dtype, fp8_recipe, model, skip_wgrad, zero_centered_gamma,
         if not config.is_fp8_supported():
             pytest.skip("Model config does not support FP8")
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
-
     sigma = 0.023
     init_method = init_method_normal(sigma)
     output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
@@ -824,9 +809,6 @@ def test_gpt_cuda_graph(dtype, fp8_recipe, model, skip_wgrad, zero_centered_gamm
         if not config.is_fp8_supported():
             pytest.skip("Model config does not support FP8")
 
-    if normalization == "RMSNorm" and zero_centered_gamma:
-        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")
-
     sigma = 0.023
     init_method = init_method_normal(sigma)
     output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)