Add RMSNorm (part of parallel_attn_blocks)

tests/modules/layers/test_normalizations.py

@@ -5,7 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from torchmultimodal.modules.layers.normalizations import Fp32GroupNorm, Fp32LayerNorm
+from tests.test_utils import assert_expected
+
+from torchmultimodal.modules.layers.normalizations import (
+    Fp32GroupNorm,
+    Fp32LayerNorm,
+    RMSNorm,
+)
 
 
 def test_fp32layernorm():
@@ -20,3 +26,38 @@ def test_fp32groupnorm():
     norm = Fp32GroupNorm(2, 4)
     output = norm(x)
     assert output.dtype == torch.float16
+
+
+def test_rms_norm_core_algo():
+    """compare RMSNorm with RMSNorm using F.norm version"""
+    dims = 10
+    rms_norm = RMSNorm(dims)
+
+    input_ones = torch.ones(dims, dtype=torch.float)
+
+    input_fixed = torch.tensor(
+        [0.999, 1.1111, 2.222, 3.333, 4.444, 5.555, 6.678, 7.987, 8.123, 9.101010],
+        dtype=torch.float16,
+    )
+    fixed_expected = torch.tensor(
+        [
+            0.1749,
+            0.1946,
+            0.3892,
+            0.5835,
+            0.7783,
+            0.9727,
+            1.1699,
+            1.3984,
+            1.4229,
+            1.5938,
+        ],
+        dtype=torch.float,
+    )
+
+    output_fixed = rms_norm(input_fixed)
+    output_ones = rms_norm(input_ones)
+
+    assert_expected(output_ones, input_ones)
+    assert_expected(output_fixed, fixed_expected, atol=1e-04, rtol=1e-05)
+    assert output_fixed.dtype == torch.float32

torchmultimodal/modules/layers/normalizations.py

@@ -6,6 +6,7 @@
 
 from typing import Any
 
+import torch
 from torch import nn, Tensor
 
 
@@ -45,3 +46,29 @@ def forward(self, x: Tensor) -> Tensor:
             self.eps,
         )
         return output.type_as(x)
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization
+    as proposed in: https://arxiv.org/abs/1910.07467
+
+    Calcs are done in fp32.
+
+    original impl: https://github.com/facebookresearch/llama/blob/main/llama/model.py
+
+    Args:
+        dim(int) = model size
+        eps(float) = epsilon
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x: Tensor) -> Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_normed = self._norm(x.float()).type_as(x)
+        return x_normed * self.scale