Switch to PyTorch's built-in RMSNorm

tests/torchtune/modules/test_rms_norm.py

@@ -5,12 +5,10 @@
 # LICENSE file in the root directory of this source tree.
 
 import pytest
-
 import torch
 
 from tests.test_utils import assert_expected
 from torch.nn.functional import normalize
-
 from torchtune.modules.rms_norm import RMSNorm
 from torchtune.training.seed import set_seed
 
@@ -66,6 +64,7 @@ def test_forward_fp16(self, rms_norm, input_random_fp16, dim) -> None:
 
         # convert input to float since rms_norm computes in fp32
         expected_fp16 = normalize(input_random_fp16.float(), p=2, dim=-1) * (dim**0.5)
+        expected_fp16 = expected_fp16.to(torch.float16)
 
         assert_expected(output_fp16, expected_fp16, atol=1e-7, rtol=1e-3)
-        assert output_fp16.dtype == torch.float32
+        assert output_fp16.dtype == torch.float16

torchtune/modules/rms_norm.py

@@ -5,18 +5,15 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-
+import torch.nn.functional as F
 from torch import nn
 
 
 class RMSNorm(nn.Module):
     """
-    Implements Root Mean Square Normalization introduced in
-    https://arxiv.org/abs/1910.07467.
+    Root Mean Square Normalization in fp32.
 
-    Reference implementation (used for correctness verification)
-    can be found here:
-    https://github.com/facebookresearch/llama/blob/main/llama/model.py
+    See: https://pytorch.org/docs/stable/generated/torch.nn.RMSNorm.html
 
     Args:
         dim (int): embedding size
@@ -25,6 +22,7 @@ class RMSNorm(nn.Module):
 
     def __init__(self, dim: int, eps: float = 1e-6) -> None:
         super().__init__()
+        self.normalized_shape = (dim,)
         self.eps = eps
         self.scale = nn.Parameter(torch.ones(dim))
 
@@ -37,8 +35,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             torch.Tensor: The normalized and scaled tensor having the same shape as ``x``.
         """
         # computation is in fp32
-        x_fp32 = x.float()
-        x_normed = (
-            x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)
-        ).type_as(x)
-        return x_normed * self.scale
+        return F.rms_norm(
+            x.float(),
+            normalized_shape=self.normalized_shape,
+            weight=self.scale,
+            eps=self.eps,
+        ).to(x.dtype)