[XPU] Modify the layer_norm and layer_norm_grad operators. (#63827)

ZibinGuo · web-flow · commit 9c735a883864 · 2024-04-25T11:11:42.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -578,9 +578,13 @@ XPUOpMap& get_kl3_ops() {
       {"lars_momentum",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"layer_norm_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"layer_norm",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"leaky_relu_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"leaky_relu",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
diff --git a/paddle/phi/kernels/xpu/layer_norm_grad_kernel.cc b/paddle/phi/kernels/xpu/layer_norm_grad_kernel.cc
@@ -32,6 +32,27 @@ void LayerNormGradKernel(const Context& ctx,
                          DenseTensor* x_grad,
                          DenseTensor* scale_grad,
                          DenseTensor* bias_grad) {
+  auto x_dtype = x.dtype();
+  const auto* scale_ptr = scale.get_ptr();
+  const auto* bias_ptr = bias.get_ptr();
+  phi::DataType scale_bias_dtype;
+  if (scale_ptr != nullptr) {
+    scale_bias_dtype = scale_ptr->dtype();
+  } else {
+    if (bias_ptr != nullptr) {
+      scale_bias_dtype = bias_ptr->dtype();
+    } else {
+      scale_bias_dtype = x_dtype;
+    }
+  }
+
+  bool is_scale_bias_same_dtype_with_x = x_dtype == scale_bias_dtype;
+  if (!is_scale_bias_same_dtype_with_x) {
+    PADDLE_ENFORCE_EQ(scale_bias_dtype,
+                      phi::CppTypeToDataType<float>::Type(),
+                      phi::errors::InvalidArgument(
+                          "Unsupported data type of Scale and Bias"));
+  }
   using XPUType = typename XPUTypeTrait<T>::Type;
   const auto& x_dims = x.dims();
   auto matrix_dim = common::flatten_to_2d(x_dims, begin_norm_axis);
@@ -47,7 +68,8 @@ void LayerNormGradKernel(const Context& ctx,
   // scale
   const float* scale_data_fp32 = nullptr;
   float* scale_grad_data_fp32 = nullptr;
-  const auto* scale_ptr = scale.get_ptr();
+  const T* scale_data_T = nullptr;
+  T* scale_grad_data_T = nullptr;
   bool need_cast_scale = false;
   if (scale_ptr == nullptr) {
     // no scale, do nothing
@@ -69,14 +91,21 @@ void LayerNormGradKernel(const Context& ctx,
             : RAII_GUARD.alloc_l3_or_gm<float>(scale_ptr->numel());
   } else {
     // no need to cast
-    scale_data_fp32 = scale_ptr->data<float>();
-    scale_grad_data_fp32 =
-        scale_grad == nullptr ? nullptr : ctx.template Alloc<float>(scale_grad);
+    if (is_scale_bias_same_dtype_with_x) {
+      scale_data_T = scale_ptr->data<T>();
+      scale_grad_data_T =
+          scale_grad == nullptr ? nullptr : ctx.template Alloc<T>(scale_grad);
+    } else {
+      scale_data_fp32 = scale_ptr->data<float>();
+      scale_grad_data_fp32 = scale_grad == nullptr
+                                 ? nullptr
+                                 : ctx.template Alloc<float>(scale_grad);
+    }
   }
 
   // bias
   float* bias_grad_data_fp32 = nullptr;
-  const auto* bias_ptr = bias.get_ptr();
+  T* bias_grad_data_T = nullptr;
   bool need_cast_bias = false;
   if (bias_ptr == nullptr) {
     // no bias, do nothing
@@ -89,29 +118,49 @@ void LayerNormGradKernel(const Context& ctx,
             : RAII_GUARD.alloc_l3_or_gm<float>(bias_ptr->numel());
   } else {
     // no need to cast
-    bias_grad_data_fp32 =
-        bias_grad == nullptr ? nullptr : ctx.template Alloc<float>(bias_grad);
+    if (is_scale_bias_same_dtype_with_x) {
+      bias_grad_data_T =
+          bias_grad == nullptr ? nullptr : ctx.template Alloc<T>(bias_grad);
+    } else {
+      bias_grad_data_fp32 =
+          bias_grad == nullptr ? nullptr : ctx.template Alloc<float>(bias_grad);
+    }
   }
 
   auto* x_grad_data =
       (x_grad == nullptr ? nullptr : ctx.template Alloc<T>(x_grad));
 
-  // int layer_norm_grad(Context* ctx, const T* x, const T* dy, T* dx, int64_t
-  // m, int64_t n, float eps, const float* scale, const float* mean, const
-  // float* var, float* dscale, float* dbias, bool is_rstd = false);
-  int r = xpu::layer_norm_grad(ctx.x_context(),
-                               reinterpret_cast<const XPUType*>(x_data),
-                               reinterpret_cast<const XPUType*>(out_grad_data),
-                               reinterpret_cast<XPUType*>(x_grad_data),
-                               left,
-                               right,
-                               epsilon,
-                               scale_data_fp32,
-                               mean_data,
-                               variance_data,
-                               scale_grad_data_fp32,
-                               bias_grad_data_fp32);
-  PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm_grad");
+  if (!is_scale_bias_same_dtype_with_x) {
+    int r =
+        xpu::layer_norm_grad(ctx.x_context(),
+                             reinterpret_cast<const XPUType*>(x_data),
+                             reinterpret_cast<const XPUType*>(out_grad_data),
+                             reinterpret_cast<XPUType*>(x_grad_data),
+                             left,
+                             right,
+                             epsilon,
+                             scale_data_fp32,
+                             mean_data,
+                             variance_data,
+                             scale_grad_data_fp32,
+                             bias_grad_data_fp32);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm_grad");
+  } else {
+    int r =
+        xpu::layer_norm_grad(ctx.x_context(),
+                             reinterpret_cast<const XPUType*>(x_data),
+                             reinterpret_cast<const XPUType*>(out_grad_data),
+                             reinterpret_cast<XPUType*>(x_grad_data),
+                             left,
+                             right,
+                             epsilon,
+                             reinterpret_cast<const XPUType*>(scale_data_T),
+                             mean_data,
+                             variance_data,
+                             reinterpret_cast<XPUType*>(scale_grad_data_T),
+                             reinterpret_cast<XPUType*>(bias_grad_data_T));
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm_grad");
+  }
 
   if (need_cast_scale) {
     int r = xpu::cast<float, XPUType>(
@@ -137,7 +186,8 @@ PD_REGISTER_KERNEL(layer_norm_grad,
                    ALL_LAYOUT,
                    phi::LayerNormGradKernel,
                    float,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   if (kernel_key.dtype() == phi::DataType::FLOAT16) {
     kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
     kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
diff --git a/paddle/phi/kernels/xpu/layer_norm_kernel.cc b/paddle/phi/kernels/xpu/layer_norm_kernel.cc
@@ -29,6 +29,34 @@ void LayerNormKernel(const Context& ctx,
                      DenseTensor* out,
                      DenseTensor* mean,
                      DenseTensor* variance) {
+  bool valid_scale = (scale.get_ptr() != nullptr);
+  bool valid_bias = (bias.get_ptr() != nullptr);
+  auto* void_scale_data = valid_scale ? scale->data() : nullptr;
+  auto* void_bias_data = valid_bias ? bias->data() : nullptr;
+
+  auto x_dtype = x.dtype();
+  phi::DataType scale_bias_dtype;
+  if (valid_scale) {
+    scale_bias_dtype = scale->dtype();
+    if (valid_bias) {
+      PADDLE_ENFORCE_EQ(
+          scale->dtype(),
+          bias->dtype(),
+          phi::errors::InvalidArgument("This Scale and Bias of layer_norm op "
+                                       "should have the same data type."));
+    }
+  } else {
+    scale_bias_dtype = valid_bias ? bias->dtype() : x_dtype;
+  }
+
+  bool is_scale_bias_same_dtype_with_x = x_dtype == scale_bias_dtype;
+  if (!is_scale_bias_same_dtype_with_x) {
+    PADDLE_ENFORCE_EQ(scale_bias_dtype,
+                      phi::CppTypeToDataType<float>::Type(),
+                      phi::errors::InvalidArgument(
+                          "Unsupported data type of Scale and Bias"));
+  }
+
   using XPUType = typename XPUTypeTrait<T>::Type;
   const auto& x_dims = x.dims();
   auto matrix_dim = common::flatten_to_2d(x_dims, begin_norm_axis);
@@ -56,7 +84,9 @@ void LayerNormKernel(const Context& ctx,
     scale_data_fp32 = scale_data_temp;
   } else {
     // no need to cast
-    scale_data_fp32 = scale_ptr->data<float>();
+    if (!is_scale_bias_same_dtype_with_x) {
+      scale_data_fp32 = scale_ptr->data<float>();
+    }
   }
 
   // bias
@@ -76,27 +106,40 @@ void LayerNormKernel(const Context& ctx,
     bias_data_fp32 = bias_data_temp;
   } else {
     // no need to cast
-    bias_data_fp32 = bias_ptr->data<float>();
+    if (!is_scale_bias_same_dtype_with_x) {
+      bias_data_fp32 = bias_ptr->data<float>();
+    }
   }
 
   auto* out_data = ctx.template Alloc<T>(out);
   auto* mean_data = ctx.template Alloc<float>(mean);
   auto* variance_data = ctx.template Alloc<float>(variance);
 
-  // int layer_norm(Context* ctx, const T* x, T* y, int64_t m, int64_t n, float
-  // eps, const float* scale, const float* bias, float* mean, float* var, bool
-  // is_rstd = false);
-  int r = xpu::layer_norm(ctx.x_context(),
-                          reinterpret_cast<const XPUType*>(x_data),
-                          reinterpret_cast<XPUType*>(out_data),
-                          left,
-                          right,
-                          epsilon,
-                          scale_data_fp32,
-                          bias_data_fp32,
-                          mean_data,
-                          variance_data);
-  PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm");
+  if (!is_scale_bias_same_dtype_with_x) {
+    int r = xpu::layer_norm(ctx.x_context(),
+                            reinterpret_cast<const XPUType*>(x_data),
+                            reinterpret_cast<XPUType*>(out_data),
+                            left,
+                            right,
+                            epsilon,
+                            scale_data_fp32,
+                            bias_data_fp32,
+                            mean_data,
+                            variance_data);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm");
+  } else {
+    int r = xpu::layer_norm(ctx.x_context(),
+                            reinterpret_cast<const XPUType*>(x_data),
+                            reinterpret_cast<XPUType*>(out_data),
+                            left,
+                            right,
+                            epsilon,
+                            reinterpret_cast<const XPUType*>(void_scale_data),
+                            reinterpret_cast<const XPUType*>(void_bias_data),
+                            mean_data,
+                            variance_data);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm");
+  }
 }
 }  // namespace phi
 
@@ -105,4 +148,8 @@ PD_REGISTER_KERNEL(layer_norm,
                    ALL_LAYOUT,
                    phi::LayerNormKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+  kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
+}
diff --git a/test/xpu/test_layer_norm_op_xpu.py b/test/xpu/test_layer_norm_op_xpu.py
@@ -22,9 +22,11 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
+from paddle.framework import core
 
 paddle.enable_static()
 
@@ -57,15 +59,19 @@ def __init__(self):
     class TestXPULayerNormOp(XPUOpTest):
         def setUp(self):
             self.op_type = "layer_norm"
-            self.dtype = self.in_type
+            if self.in_type == np.uint16:
+                self.dtype = np.float32
+            else:
+                self.dtype = self.in_type
             self.shape = [2, 3, 4, 5]
             self.epsilon = 1e-05
             self.begin_norm_axis = 1
             self.use_fp16_scale_bias = False
+            self.use_bf16_scale_bias = False
             self.set_attrs()
 
             self.atol = 1e-4
-            if self.dtype == np.float16:
+            if self.dtype == np.float16 or self.in_type == np.uint16:
                 self.atol = 1e-2
 
             right = reduce(
@@ -78,6 +84,11 @@ def setUp(self):
             if self.dtype == np.float16 and self.use_fp16_scale_bias:
                 scale_np = scale_np.astype('float16')
                 bias_np = scale_np.astype('float16')
+            if (
+                self.dtype == np.uint16 and self.use_fp16_scale_bias
+            ):  # bfloat16 actually
+                scale_np = convert_float_to_uint16(scale_np)
+                bias_np = convert_float_to_uint16(bias_np)
             ref_y_np, ref_mean_np, ref_variance_np = ref_layer_norm(
                 x_np, scale_np, bias_np, self.epsilon, self.begin_norm_axis
             )
@@ -123,17 +134,43 @@ def set_attrs(self):
 
     class TestXPULayerNormOpFP16(TestXPULayerNormOp):
         def set_attrs(self):
-            self.use_fp16_scale_bias = True
+            self.use_fp16_scale_bias = False
 
     class TestXPULayerNormOpFP16_2D(TestXPULayerNormOp):
         def set_attrs(self):
             self.shape = [10, 12]
-            self.use_fp16_scale_bias = True
+            self.use_fp16_scale_bias = False
 
     class TestXPULayerNormOpFP16_3D(TestXPULayerNormOp):
         def set_attrs(self):
             self.shape = [4, 5, 6]
-            self.use_fp16_scale_bias = True
+            self.use_fp16_scale_bias = False
+
+    class TestXPULayerNormOpBF16(TestXPULayerNormOp):
+        def set_attrs(self):
+            self.use_bf16_scale_bias = True
+            if core.get_xpu_device_version(0) == core.XPUVersion.XPU3:
+                self.dtype = np.uint16
+            else:
+                self.dtype = np.float32
+
+    class TestXPULayerNormOpBF16_2D(TestXPULayerNormOp):
+        def set_attrs(self):
+            self.shape = [10, 12]
+            self.use_bf16_scale_bias = True
+            if core.get_xpu_device_version(0) == core.XPUVersion.XPU3:
+                self.dtype = np.uint16
+            else:
+                self.dtype = np.float32
+
+    class TestXPULayerNormOpBF16_3D(TestXPULayerNormOp):
+        def set_attrs(self):
+            self.shape = [4, 5, 6]
+            self.use_bf16_scale_bias = False
+            if core.get_xpu_device_version(0) == core.XPUVersion.XPU3:
+                self.dtype = np.uint16
+            else:
+                self.dtype = np.float32
 
 
 support_types = get_xpu_op_support_types('layer_norm')
