Support Half/BFloat16 in native_layer_norm

ghstack-source-id: f8d98dddce9c5fc56d75eb3bb48f228686a4dc76 ghstack-comment-id: 2599385748 Pull Request resolved: #7752
pytorch · Jan 17, 2025 · bff3ad3 · bff3ad3
1 parent 57a09f4
commit bff3ad3
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 7 deletions.
diff --git a/kernels/optimized/cpu/op_native_layer_norm.cpp b/kernels/optimized/cpu/op_native_layer_norm.cpp
@@ -155,7 +155,7 @@ std::tuple<Tensor&, Tensor&, Tensor&> opt_native_layer_norm_out(
       InvalidArgument,
       ret_val);
 
-  ET_SWITCH_FLOAT_TYPES(
+  ET_SWITCH_FLOATHBF16_TYPES(
       input.scalar_type(), ctx, "native_layer_norm.out", CTYPE, [&]() {
         layer_norm<CTYPE>(
             input,

diff --git a/kernels/portable/cpu/op_native_layer_norm.cpp b/kernels/portable/cpu/op_native_layer_norm.cpp
@@ -66,15 +66,16 @@ void layer_norm(
     bias_data = nullptr;
   }
 
+  const CTYPE ct_normalized = static_cast<CTYPE>(normalized);
   for (int i = 0; i < leading; ++i) {
     const CTYPE* x = input_data + i * normalized;
     CTYPE* y = out_data + i * normalized;
 
     // compute E[X] and Var[x] = E[x^2] - E[x]^2
-    CTYPE sum = reduce_add(x, normalized);
-    CTYPE sq_sum = vec_powerf(x, normalized);
-    CTYPE mean_value = sum / normalized;
-    CTYPE variance = sq_sum / normalized - mean_value * mean_value;
+    CTYPE sum = reduce_add(x, ct_normalized);
+    CTYPE sq_sum = vec_powerf(x, ct_normalized);
+    CTYPE mean_value = sum / ct_normalized;
+    CTYPE variance = sq_sum / ct_normalized - mean_value * mean_value;
     CTYPE std = std::sqrt(variance + eps);
 
     // Calculate the elements of output
@@ -167,7 +168,7 @@ std::tuple<Tensor&, Tensor&, Tensor&> native_layer_norm_out(
       InvalidArgument,
       ret_val);
 
-  ET_SWITCH_FLOAT_TYPES(
+  ET_SWITCH_FLOATHBF16_TYPES(
       input.scalar_type(), ctx, "native_layer_norm.out", CTYPE, [&]() {
         layer_norm<CTYPE>(
             input,

diff --git a/kernels/test/op_native_layer_norm_test.cpp b/kernels/test/op_native_layer_norm_test.cpp
@@ -95,7 +95,15 @@ class OpNativeLayerNormTest : public OperatorTest {
       EXPECT_TENSOR_CLOSE(out0, std::get<0>(result));
 
       Tensor expected = tf.make(test_case.sizes, test_case.expected_data);
-      EXPECT_TENSOR_CLOSE(out0, expected);
+      if constexpr (DTYPE == ScalarType::BFloat16) {
+        EXPECT_TENSOR_CLOSE_WITH_TOL(
+            out0,
+            expected,
+            1e-2,
+            executorch::runtime::testing::internal::kDefaultBFloat16Atol);
+      } else {
+        EXPECT_TENSOR_CLOSE(out0, expected);
+      }
     }
   }
 
@@ -393,6 +401,8 @@ std::vector<int64_t> vector_32_to_64(std::vector<int32_t> vector_32) {
 TEST_F(OpNativeLayerNormTest, FloatTensors) {
   run_floating_point_test_cases<ScalarType::Float>();
   run_floating_point_test_cases<ScalarType::Double>();
+  run_floating_point_test_cases<ScalarType::Half>();
+  run_floating_point_test_cases<ScalarType::BFloat16>();
 }
 
 TEST_F(OpNativeLayerNormTest, IntTensorsDies) {