Add zero-centered RMSNorm support

GitOrigin-RevId: dd52991

Author: ds-hwang

axlearn/common/layers.py

@@ -352,10 +352,15 @@ def forward(self, x: Tensor, *, paddings: Optional[Tensor] = None) -> Tensor:
 
 
 class RMSNormStateless(BaseNormalizationLayer):
-    """Stateless version of https://github.com/bzhangGo/rmsnorm."""
+    """Stateless version of https://github.com/bzhangGo/rmsnorm.
+
+    Paper: https://arxiv.org/abs/1910.07467
+    """
 
     @config_class
     class Config(BaseNormalizationLayer.Config):
+        """RMSNormStateless.Config."""
+
         # The epsilon.
         eps: float = 1e-8
         # Cast input to this dtype for the 'forward' call. If None, do not cast.
@@ -364,14 +369,27 @@ class Config(BaseNormalizationLayer.Config):
         input_partition_spec: Optional[tuple[Optional[str]]] = None
         # If not None, how to partition output activation values.
         output_partition_spec: Optional[tuple[Optional[str]]] = None
+        # If True, center the input by subtracting the mean before RMS normalization.
+        # This variant is used in Gemma and Qwen3-Next and provides better numerical stability.
+        zero_centered: Optional[bool] = None
 
     def _forward(self, x: Tensor) -> Tensor:
+        """Applies RMS normalization to the input.
+
+        Args:
+            x: Input tensor to normalize.
+
+        Returns:
+            Normalized tensor with the same shape as input.
+        """
         cfg = self.config
         x = maybe_shard(x, cfg.input_partition_spec)
         x_dtype = x.dtype
         if cfg.forward_dtype is not None:
             x = x.astype(cfg.forward_dtype)
         moment2 = (x * x).mean(axis=-1, keepdims=True)
+        if cfg.zero_centered:
+            x = x - x.mean(axis=-1, keepdims=True)
         x = x * jax.lax.rsqrt(moment2 + cfg.eps)
         x = x.astype(x_dtype)
         return x

axlearn/common/layers_test.py

@@ -561,6 +561,35 @@ def test_rms_norm_partition_specs_constraint(self, mock_with_sharding_constraint
         self.assertEqual(calls[1].args[0].dtype, jnp.float32)
         np.testing.assert_array_equal(calls[1].args[0], outputs)
 
+    def test_rms_norm_zero_centered(self):
+        dim = 6
+        cfg = RMSNorm.default_config().set(name="norm", input_dim=dim, zero_centered=True)
+        layer: RMSNorm = cfg.instantiate(parent=None)
+        prng_key = jax.random.PRNGKey(123)
+        prng_key, init_key = jax.random.split(prng_key)
+        layer_params = layer.initialize_parameters_recursively(init_key)
+
+        # Random inputs with non-zero mean to test zero-centering.
+        prng_key, input_key = jax.random.split(prng_key)
+        inputs = jax.random.normal(input_key, [2, 3, dim]) + 5.0  # Add offset
+        outputs, _ = F(
+            layer,
+            inputs=(inputs,),
+            is_training=True,
+            state=layer_params,
+            prng_key=prng_key,
+        )
+
+        # Manually compute expected output based on zero_centered setting.
+        moment2 = (inputs * inputs).mean(axis=-1, keepdims=True)
+        inputs = inputs - inputs.mean(axis=-1, keepdims=True)
+        expected_normalized = inputs * jax.lax.rsqrt(moment2 + cfg.eps)
+
+        expected_outputs = expected_normalized * layer_params["scale"]
+        self.assertNestedAllClose(outputs, expected_outputs)
+        output_mean = outputs.mean(axis=-1, keepdims=True)
+        self.assertNestedAllClose(output_mean, np.zeros_like(output_mean))
+
     def test_l2_norm(self):
         cfg = L2Norm.default_config().set(name="norm")
         layer: L2Norm = cfg.instantiate(parent=None)