keep rms computation in full precision

riccardofelluga · pytorchmergebot · commit 7fd4753d7672 · 2025-03-07T19:04:18.000Z
diff --git a/aten/src/ATen/native/layer_norm.cpp b/aten/src/ATen/native/layer_norm.cpp
@@ -315,16 +315,16 @@ Tensor rms_norm_symint(
       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);
+    Tensor upcasted_result = upcasted_input.mul(rqrst_input);
 
     if (weight_opt.has_value()) {
-      result = result.mul(weight_opt.value());
+      upcasted_result = upcasted_result.mul(weight_opt.value());
     }
 
-    return result;
+    return upcasted_result;
   });
 
-  return result;
+  return result.type_as(input);
 
 }
 } // namespace at::native
diff --git a/torch/testing/_internal/common_modules.py b/torch/testing/_internal/common_modules.py
@@ -1945,10 +1945,9 @@ def rms_norm_reference_fn(m, p, i):
         dims = [ndim - i - 1 for i in range(len(normalized_shape))]
         upcasted_i = i.float()
         result = upcasted_i * torch.rsqrt(upcasted_i.pow(2).mean(dim=dims, keepdim=True) + m.eps)
-        result = result.type_as(i)
         if weight is not None:
             result *= weight
-        return result
+        return result.type_as(i)
 
     return [
         ModuleInput(

Original file line number	Diff line number	Diff line change
`@@ -315,16 +315,16 @@ Tensor rms_norm_symint(`
`315`	`315`	`rqrst_input = rsqrt(at::pow(upcasted_input, 2).mean(dims_to_reduce_ref, /keepdim=/true).add_(eps_val));`
`316`	`316`	`}`
`317`	`317`
`318`		`- Tensor result = upcasted_input.mul(rqrst_input).type_as(input);`
	`318`	`+ Tensor upcasted_result = upcasted_input.mul(rqrst_input);`
`319`	`319`
`320`	`320`	`if (weight_opt.has_value()) {`
`321`		`- result = result.mul(weight_opt.value());`
	`321`	`+ upcasted_result = upcasted_result.mul(weight_opt.value());`
`322`	`322`	`}`
`323`	`323`
`324`		`- return result;`
	`324`	`+ return upcasted_result;`
`325`	`325`	`});`
`326`	`326`
`327`		`- return result;`
	`327`	`+ return result.type_as(input);`
`328`	`328`
`329`	`329`	`}`
`330`	`330`	`} // namespace at::native`