Fixing Pytorch RMS norm implementation

kkontny · kkontny · commit c79577d3d554 · 2024-08-09T11:39:50.000+02:00
diff --git a/aten/src/ATen/native/layer_norm.cpp b/aten/src/ATen/native/layer_norm.cpp
@@ -294,8 +294,13 @@ Tensor rms_norm(
     } else {
       eps_val = eps.value();
     }
+    at::Tensor result;
 
-    auto result = input.mul(at::rsqrt(at::pow(input, 2).mean(dims_to_reduce_ref, /*keep_dim=*/true).add_(eps_val)));
+    if constexpr (std::is_same_v<scalar_t, c10::Half>) {
+      result = input.mul(at::rsqrt(at::pow(input.to(at::ScalarType::Float), 2).mean(dims_to_reduce_ref, /*keep_dim=*/true).add_(eps_val)).to(at::ScalarType::Half));
+    } else {
+      result = input.mul(at::rsqrt(at::pow(input, 2).mean(dims_to_reduce_ref, /*keep_dim=*/true).add_(eps_val)));
+    }
 
     if (weight_opt.has_value()) {
       result = result.mul(weight_opt.value());
