[XPU][PHI Kernels] Support bf16 for fused_rope (#60064)

lj970926 · web-flow · commit 2a424a848cd7 · 2023-12-19T16:22:43.000+08:00
* support fused_rope fp16

* bug fix

* bug fix
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -130,13 +130,18 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"c_embedding",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"c_embedding_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"c_identity",
        XPUKernelSet({phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT64,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"gather_part",
@@ -152,7 +157,8 @@ XPUOpMap& get_kl2_ops() {
       {"c_split",
        XPUKernelSet({phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
-                     phi::DataType::INT32})},
+                     phi::DataType::INT32,
+                     phi::DataType::BFLOAT16})},
       {"c_sync_calc_stream",
        XPUKernelSet({phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
@@ -442,9 +448,13 @@ XPUOpMap& get_kl2_ops() {
       {"fused_multi_transformer_int8_xpu",
        XPUKernelSet({phi::DataType::FLOAT16})},
       {"fused_rotary_position_embedding",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"fused_rotary_position_embedding_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"unfold",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"unfold_grad",
@@ -820,7 +830,8 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"sigmoid",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"sigmoid_grad", XPUKernelSet({phi::DataType::FLOAT32})},
@@ -1104,9 +1115,9 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
-      {"sin", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"sin", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"sin_grad", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"cos", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"cos", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"cos_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"linspace",
        XPUKernelSet({phi::DataType::FLOAT32,
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -129,13 +129,18 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"c_embedding",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"c_embedding_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"c_identity",
        XPUKernelSet({phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT64,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"c_softmax_with_cross_entropy", XPUKernelSet({phi::DataType::FLOAT32})},
@@ -145,6 +150,7 @@ XPUOpMap& get_kl3_ops() {
       {"c_split",
        XPUKernelSet({phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT32})},
       {"c_sync_calc_stream",
        XPUKernelSet({phi::DataType::FLOAT16,
@@ -438,9 +444,13 @@ XPUOpMap& get_kl3_ops() {
       {"fused_multi_transformer_xpu",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"fused_rotary_position_embedding",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"fused_rotary_position_embedding_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"unfold",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"unfold_grad",
@@ -807,7 +817,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"sigmoid",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"sigmoid_grad", XPUKernelSet({phi::DataType::FLOAT32})},
@@ -1119,9 +1130,9 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
-      {"sin", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"sin", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"sin_grad", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"cos", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"cos", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"cos_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"linspace",
        XPUKernelSet({phi::DataType::FLOAT32,
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_grad_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_rope_grad_kernel.cc
@@ -123,4 +123,5 @@ PD_REGISTER_KERNEL(fused_rotary_position_embedding_grad,
                    ALL_LAYOUT,
                    phi::fusion::FusedRopeGradKernel,
                    float,
-                   phi::dtype::float16){};
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16){};
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_rope_kernel.cc
@@ -123,4 +123,5 @@ PD_REGISTER_KERNEL(fused_rotary_position_embedding,
                    ALL_LAYOUT,
                    phi::fusion::FusedRopeKernel,
                    float,
-                   phi::dtype::float16){};
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16){};
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h b/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h
@@ -68,15 +68,29 @@ void XPUGetSinCosData(const Context& dev_ctx,
           phi::errors::InvalidArgument(
               "The batch_size and seq_len of position_ids must be the same as "
               "those of q."));
-      int ret = xpu::gather<XPUT, int64_t>(
-          dev_ctx.x_context(),
-          reinterpret_cast<const XPUT*>(sin_cos->data()),
-          position_ids->data<int64_t>(),
-          sin_cos_data,
-          {seq_len, head_dim},
-          batch_size * seq_len,
-          0);
-      PADDLE_ENFORCE_XDNN_SUCCESS(ret, "gather");
+      using XPUTFp16 = typename XPUTypeTrait<phi::dtype::float16>::Type;
+      using XPUTBf16 = typename XPUTypeTrait<phi::dtype::bfloat16>::Type;
+      if (std::is_same<XPUT, XPUTBf16>::value) {
+        int ret = xpu::gather<XPUTFp16, int64_t>(
+            dev_ctx.x_context(),
+            reinterpret_cast<const XPUTFp16*>(sin_cos->data()),
+            position_ids->data<int64_t>(),
+            reinterpret_cast<XPUTFp16*>(sin_cos_data),
+            {seq_len, head_dim},
+            batch_size * seq_len,
+            0);
+        PADDLE_ENFORCE_XDNN_SUCCESS(ret, "gather");
+      } else {
+        int ret = xpu::gather<XPUT, int64_t>(
+            dev_ctx.x_context(),
+            reinterpret_cast<const XPUT*>(sin_cos->data()),
+            position_ids->data<int64_t>(),
+            sin_cos_data,
+            {seq_len, head_dim},
+            batch_size * seq_len,
+            0);
+        PADDLE_ENFORCE_XDNN_SUCCESS(ret, "gather");
+      }
     } else {
       int ret =
           xpu::broadcast<XPUT>(dev_ctx.x_context(),
@@ -122,8 +136,8 @@ void XPUFusedRotaryHalf(const Context& dev_ctx,
            {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
            {seq_len * head_dim, head_dim, head_dim, 1});
   PADDLE_ENFORCE_XDNN_SUCCESS(ret,
-                              is_bwd ? "rotary_no_freqs_embedding_v2"
-                                     : "rotary_no_freqs_embedding_v2_grad");
+                              is_bwd ? "rotary_no_freqs_embedding_v2_grad"
+                                     : "rotary_no_freqs_embedding_v2");
 }
 }  // namespace fusion
 }  // namespace phi
diff --git a/paddle/phi/kernels/selected_rows/shape_kernel.cc b/paddle/phi/kernels/selected_rows/shape_kernel.cc
@@ -82,7 +82,8 @@ PD_REGISTER_KERNEL(shape_sr,
                    int64_t,
                    float,
                    double,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
   kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
diff --git a/paddle/phi/kernels/shape_kernel.cc b/paddle/phi/kernels/shape_kernel.cc
@@ -83,7 +83,8 @@ PD_REGISTER_KERNEL(shape,
                    int64_t,
                    float,
                    double,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
   kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
diff --git a/paddle/phi/kernels/xpu/activation_kernel.cc b/paddle/phi/kernels/xpu/activation_kernel.cc
@@ -633,6 +633,11 @@ PD_REGISTER_KERNEL(
 PD_REGISTER_KERNEL(
     relu6, XPU, ALL_LAYOUT, phi::Relu6Kernel, float, phi::dtype::float16) {}
 
+PD_REGISTER_KERNEL(
+    sin, XPU, ALL_LAYOUT, phi::SinKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(
+    cos, XPU, ALL_LAYOUT, phi::CosKernel, float, phi::dtype::float16) {}
+
 #define PD_REGISTER_ACTIVATION_KERNEL(name, func) \
   PD_REGISTER_KERNEL(name, XPU, ALL_LAYOUT, phi::func, float) {}
 
@@ -643,6 +648,4 @@ PD_REGISTER_ACTIVATION_KERNEL(mish, MishKernel)
 PD_REGISTER_ACTIVATION_KERNEL(pow, PowKernel)
 PD_REGISTER_ACTIVATION_KERNEL(reciprocal, ReciprocalKernel)
 PD_REGISTER_ACTIVATION_KERNEL(softplus, SoftplusKernel)
-PD_REGISTER_ACTIVATION_KERNEL(sin, SinKernel)
-PD_REGISTER_ACTIVATION_KERNEL(cos, CosKernel)
 PD_REGISTER_ACTIVATION_KERNEL(rsqrt, RsqrtKernel)
diff --git a/paddle/phi/kernels/xpu/c_embedding_kernel.cc b/paddle/phi/kernels/xpu/c_embedding_kernel.cc
@@ -75,4 +75,5 @@ PD_REGISTER_KERNEL(c_embedding,
                    ALL_LAYOUT,
                    phi::CEmbeddingKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/c_embedding_kernel_grad.cc b/paddle/phi/kernels/xpu/c_embedding_kernel_grad.cc
@@ -86,4 +86,5 @@ PD_REGISTER_KERNEL(c_embedding_grad,
                    ALL_LAYOUT,
                    phi::CEmbeddingGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/c_identity_kernel.cc b/paddle/phi/kernels/xpu/c_identity_kernel.cc
@@ -26,4 +26,5 @@ PD_REGISTER_KERNEL(c_identity,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/c_split_kernel.cc b/paddle/phi/kernels/xpu/c_split_kernel.cc
@@ -84,4 +84,5 @@ PD_REGISTER_KERNEL(c_split,
                    phi::CSplitKernel,
                    float,
                    int,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/test/xpu/test_collective_base_xpu.py b/test/xpu/test_collective_base_xpu.py
@@ -45,6 +45,8 @@ def DataTypeCast(date_type):
         np_dtype = np.int32
     elif date_type == "int64":
         np_dtype = np.int64
+    elif date_type == "bfloat16":
+        np_dtype = np.uint16
     else:
         raise ValueError("This data type is not support!")
 
diff --git a/test/xpu/test_collective_identity_xpu.py b/test/xpu/test_collective_identity_xpu.py
@@ -32,6 +32,7 @@ def test_identity(self):
             "float64",
             "int32",
             "int64",
+            "bfloat16",
         ]
         for dtype in dtypes_to_test:
             self.check_with_place(
diff --git a/test/xpu/test_fused_rotary_position_embedding_op_xpu.py b/test/xpu/test_fused_rotary_position_embedding_op_xpu.py
diff --git a/test/xpu/test_shape_op_xpu.py b/test/xpu/test_shape_op_xpu.py
