Fix RMSNorm epsilon value type for BF16 or FP16 (pytorch#142848)

Fixes pytorch#140092

Here's what this PR does:

In before, we create a `scalar_t eps_val;` variable, and the `eps` is mostly a double scalar which passed from python frontend, like 1e-6.

While we do `eps_val = std::numeric_limits<at::scalar_value_type<scalar_t>::type>::epsilon();` or `eps_val = eps.value();`, we down cast this epsilon to match input tensor dtype (`scalar_t`), in case of BFloat16, the 1e-6 double would be cast to `1.00136e-05`.

However, while we act `auto rqrst_input = rsqrt(at::pow(upcasted_input, 2).mean(dims_to_reduce_ref, /*keepdim=*/true).add_(eps_val));`, we up cast `eps_val` to match the `opmath_t`, the conversion between these two dtypes is UNNECESSARY, so we could just make the `opmath_t eps_val` instead of `scalar_t`.

Pull Request resolved: pytorch#142848
Approved by: https://github.com/mikaylagawarecki

Author: fmo-mt

aten/src/ATen/native/layer_norm.cpp

@@ -278,18 +278,27 @@ Tensor rms_norm_symint(
         input.scalar_type(),
         "rms_norm",
         [&] {
-    scalar_t eps_val;
-    if (!eps.has_value()) {
-      eps_val = std::numeric_limits<at::scalar_value_type<scalar_t>::type>::epsilon();
-    } else {
-      eps_val = eps.value();
-    }
-
     // upcast is needed for fp16 and bf16
     c10::ScalarType opmath_t = toOpMathType(input.scalar_type());
     Tensor upcasted_input = input.to(opmath_t);
 
-    auto rqrst_input = rsqrt(at::pow(upcasted_input, 2).mean(dims_to_reduce_ref, /*keepdim=*/true).add_(eps_val));
+    Tensor rqrst_input;
+
+    // opmath_t would be one of [Double, Float, ComplexFloat, ComplexDouble]
+    if (opmath_t == at::ScalarType::Float || opmath_t == at::ScalarType::ComplexFloat) {
+      float eps_val = std::numeric_limits<float>::epsilon();
+      if (eps.has_value()) {
+        eps_val = eps.value();
+      }
+      rqrst_input = rsqrt(at::pow(upcasted_input, 2).mean(dims_to_reduce_ref, /*keepdim=*/true).add_(eps_val));
+    } else {
+      double eps_val = std::numeric_limits<double>::epsilon();
+      if (eps.has_value()) {
+        eps_val = eps.value();
+      }
+      rqrst_input = rsqrt(at::pow(upcasted_input, 2).mean(dims_to_reduce_ref, /*keepdim=*/true).add_(eps_val));
+    }
+
     Tensor result = upcasted_input.mul(rqrst_input).type_as(input);
 
     if (weight_opt.has_value()) {