fixed rmsnorm bug. Changed division to multiply since using torch.rsqrt

mosaicml · Jun 26, 2023 · df0667a · df0667a
1 parent b1ac9da
commit df0667a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/models/layers/norm.py b/llmfoundry/models/layers/norm.py
@@ -52,7 +52,7 @@ def forward(self, x):
 
 
 def rms_norm(x, weight=None, eps=1e-5):
-    output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+    output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
     if weight is not None:
         return output * weight
     return output