Fix rms_norm fp16/bf16 (Lightning-AI#1751)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: riccardofelluga

thunder/tests/opinfos.py

@@ -7967,16 +7967,11 @@ def rms_norm_error_generator(op, device, **kwargs):
                 dtypes=(datatypes.float16,),
                 devicetypes=(devices.DeviceType.CPU,),
             ),
-            # See issue - https://github.com/Lightning-AI/lightning-thunder/issues/1395
             DecorateInfo(
-                custom_comparator(partial(assert_close, atol=1e-2, rtol=1e-2)),
-                dtypes=(datatypes.float16,),
+                pytest.mark.xfail,
+                dtypes=(datatypes.float16, datatypes.bfloat16),
                 devicetypes=(devices.DeviceType.CUDA,),
-            ),
-            DecorateInfo(
-                pytest.mark.skip(reason="Flaky. See https://github.com/Lightning-AI/lightning-thunder/issues/1678"),
-                "test_core_vs_torch_consistency",
-                dtypes=(datatypes.bfloat16,),
+                active_if=LooseVersion(torch.__version__) < "2.7",
             ),
         ),
     )

thunder/torch/__init__.py

@@ -3800,13 +3800,22 @@ def rms_norm(
     weight: None | TensorLike = None,
     eps: None | float = None,
 ):
+    input_dtype = a.dtype
+
+    if a.dtype in (thunder.float16, thunder.bfloat16):
+        a = clang.maybe_convert_to_dtype(a, thunder.float32, enforce_safe_casting=True)
+
     if eps is None:
         eps = torch.finfo(to_torch_dtype(a.dtype)).eps
+
     reduction_dims = _check_normalized_shape_and_get_reduction_dims(a, normalized_shape, weight)
-    norm_a = mean(a * a, dim=reduction_dims, keepdim=True)
+    norm_a = mean(a * a, dim=reduction_dims, keepdim=True, dtype=None)
     a_normed = a * rsqrt(norm_a + eps)
+
     if weight is not None:
         a_normed = a_normed * weight
+
+    a_normed = clang.maybe_convert_to_dtype(a_normed, input_dtype, enforce_safe_casting=True)
     return a_normed