[XPU] Support bf16 clip_grad and redirect xpu kernel to clamp_grad (#69723)

cqulilujia · web-flow · commit 28153dcaa6a1 · 2024-11-27T17:47:05.000+08:00
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -29,7 +29,7 @@ set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 set(XPU_XPUDNN_LIB_NAME "libxpu_dnn.so")
 
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "dev/20241118")
+  set(XPU_XHPC_BASE_DATE "dev/20241127")
 endif()
 set(XPU_XCCL_BASE_VERSION "3.0.0.5") # For XRE5
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -229,6 +229,7 @@ XPUOpMap& get_kl3_ops() {
       {"clip_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"coalesce_tensor",
diff --git a/paddle/phi/kernels/fusion/xpu/cross_attention_xpu_kernel.cc b/paddle/phi/kernels/fusion/xpu/cross_attention_xpu_kernel.cc
@@ -159,7 +159,7 @@ void CrossAttentionXPUKernelImpl(
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "qkv_attention_xpu");
 
   if (input_q.dtype() == DataType::FLOAT32) {
-    int r_cast_out = xpu::cast_v2<XPUTypeFP16, XPUTypeOut>(
+    int r_cast_out = xpu::cast<XPUTypeFP16, XPUTypeOut>(
         ctx.x_context(), qkv_temp_data, qkv_data, qkv->numel());
     PADDLE_ENFORCE_XDNN_SUCCESS(
         r_cast_out, "cross_attention_xpu(cast out from fp16 to fp32)");
diff --git a/paddle/phi/kernels/fusion/xpu/fused_layernorm_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_layernorm_kernel.cc
@@ -75,12 +75,12 @@ void FusedLayerNormKernel(const Context& dev_ctx,
                                 residual_alpha);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
 
-  r = baidu::xpu::api::cast_v2(
+  r = baidu::xpu::api::cast(
       xpu_ctx->x_context(),
       residual_alpha_tmp.data<float>(),
       reinterpret_cast<XPUType*>(residual_alpha_ptr.data<T>()),
       1);
-  PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast_v2");
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
 
   if (residual) {
     dev_ctx.template Alloc<T>(residual_out);
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h b/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h
@@ -274,69 +274,135 @@ void XPUFusedRotaryEveryTwo(const Context& dev_ctx,
                             DenseTensor* out_q,
                             DenseTensor* out_k,
                             DenseTensor* out_v) {
-  auto single_func = &xpu::rotary_embedding_v3_single<XPUType, XPUSCType>;
-  auto fusion_func = &xpu::rotary_embedding_v3<XPUType, XPUSCType>;
+  auto single_func_fwd = &xpu::rotary_embedding_v3_single<XPUType, XPUSCType>;
+  auto fusion_func_fwd = &xpu::rotary_embedding_v3<XPUType, XPUSCType>;
+  auto single_func_bwd =
+      &xpu::rotary_embedding_v3_single_grad<XPUType, XPUSCType>;
+  auto fusion_func_bwd = &xpu::rotary_embedding_v3_grad<XPUType, XPUSCType>;
   const char* single_func_name = "rotary_embedding_v3_single";
   const char* fusion_func_name = "rotary_embedding_v3";
   if (is_bwd) {
-    single_func = &xpu::rotary_embedding_v3_single_grad<XPUType, XPUSCType>;
-    fusion_func = &xpu::rotary_embedding_v3_grad<XPUType, XPUSCType>;
     single_func_name = "rotary_embedding_v3_single_grad";
     fusion_func_name = "rotary_embedding_v3_grad";
   }
-  if (!in_k) {
-    int ret = single_func(
-        dev_ctx.x_context(),
-        reinterpret_cast<const XPUType*>(in_q.data()),
-        cos_data,
-        sin_data,
-        reinterpret_cast<XPUType*>(out_q->data()),
-        batch_size,
-        seq_len,
-        num_heads,
-        head_dim,
-        {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
-        "BLHD",
-        true);
-    PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
+  if (is_bwd) {
+    if (!in_k) {
+      int ret = single_func_bwd(
+          dev_ctx.x_context(),
+          reinterpret_cast<const XPUType*>(in_q.data()),
+          cos_data,
+          sin_data,
+          reinterpret_cast<XPUType*>(out_q->data()),
+          batch_size,
+          seq_len,
+          num_heads,
+          head_dim,
+          {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
+          std::string("BLHD").c_str(),
+          true);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
+    } else {
+      int64_t num_heads_k = in_k->dims()[2];
+      int ret = fusion_func_bwd(
+          dev_ctx.x_context(),
+          reinterpret_cast<const XPUType*>(in_q.data()),
+          reinterpret_cast<const XPUType*>(in_k->data()),
+          cos_data,
+          sin_data,
+          reinterpret_cast<XPUType*>(out_q->data()),
+          reinterpret_cast<XPUType*>(out_k->data()),
+          batch_size,
+          seq_len,
+          num_heads,
+          head_dim,
+          {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
+          {seq_len * num_heads_k * head_dim,
+           num_heads_k * head_dim,
+           head_dim,
+           1},
+          num_heads_k,
+          std::string("BLHD").c_str(),
+          true);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, fusion_func_name);
+    }
+    if (in_v) {
+      int64_t num_heads_v = in_v->dims()[2];
+      int ret = single_func_bwd(dev_ctx.x_context(),
+                                reinterpret_cast<const XPUType*>(in_v->data()),
+                                cos_data,
+                                sin_data,
+                                reinterpret_cast<XPUType*>(out_v->data()),
+                                batch_size,
+                                seq_len,
+                                num_heads_v,
+                                head_dim,
+                                {seq_len * num_heads_v * head_dim,
+                                 num_heads_v * head_dim,
+                                 head_dim,
+                                 1},
+                                std::string("BLHD").c_str(),
+                                true);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
+    }
   } else {
-    int64_t num_heads_k = in_k->dims()[2];
-    int ret = fusion_func(
-        dev_ctx.x_context(),
-        reinterpret_cast<const XPUType*>(in_q.data()),
-        reinterpret_cast<const XPUType*>(in_k->data()),
-        cos_data,
-        sin_data,
-        reinterpret_cast<XPUType*>(out_q->data()),
-        reinterpret_cast<XPUType*>(out_k->data()),
-        batch_size,
-        seq_len,
-        num_heads,
-        head_dim,
-        {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
-        {seq_len * num_heads_k * head_dim, num_heads_k * head_dim, head_dim, 1},
-        num_heads_k,
-        "BLHD",
-        true);
-    PADDLE_ENFORCE_XDNN_SUCCESS(ret, fusion_func_name);
-  }
-
-  if (in_v) {
-    int64_t num_heads_v = in_v->dims()[2];
-    int ret = single_func(
-        dev_ctx.x_context(),
-        reinterpret_cast<const XPUType*>(in_v->data()),
-        cos_data,
-        sin_data,
-        reinterpret_cast<XPUType*>(out_v->data()),
-        batch_size,
-        seq_len,
-        num_heads_v,
-        head_dim,
-        {seq_len * num_heads_v * head_dim, num_heads_v * head_dim, head_dim, 1},
-        "BLHD",
-        true);
-    PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
+    if (!in_k) {
+      int ret = single_func_fwd(
+          dev_ctx.x_context(),
+          reinterpret_cast<const XPUType*>(in_q.data()),
+          cos_data,
+          sin_data,
+          reinterpret_cast<XPUType*>(out_q->data()),
+          batch_size,
+          seq_len,
+          num_heads,
+          head_dim,
+          {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
+          "BLHD",
+          true);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
+    } else {
+      int64_t num_heads_k = in_k->dims()[2];
+      int ret = fusion_func_fwd(
+          dev_ctx.x_context(),
+          reinterpret_cast<const XPUType*>(in_q.data()),
+          reinterpret_cast<const XPUType*>(in_k->data()),
+          cos_data,
+          sin_data,
+          reinterpret_cast<XPUType*>(out_q->data()),
+          reinterpret_cast<XPUType*>(out_k->data()),
+          batch_size,
+          seq_len,
+          num_heads,
+          head_dim,
+          {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
+          {seq_len * num_heads_k * head_dim,
+           num_heads_k * head_dim,
+           head_dim,
+           1},
+          num_heads_k,
+          "BLHD",
+          true);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, fusion_func_name);
+    }
+    if (in_v) {
+      int64_t num_heads_v = in_v->dims()[2];
+      int ret = single_func_fwd(dev_ctx.x_context(),
+                                reinterpret_cast<const XPUType*>(in_v->data()),
+                                cos_data,
+                                sin_data,
+                                reinterpret_cast<XPUType*>(out_v->data()),
+                                batch_size,
+                                seq_len,
+                                num_heads_v,
+                                head_dim,
+                                {seq_len * num_heads_v * head_dim,
+                                 num_heads_v * head_dim,
+                                 head_dim,
+                                 1},
+                                "BLHD",
+                                true);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
+    }
   }
 }
 
diff --git a/paddle/phi/kernels/fusion/xpu/multi_encoder_xpu_kernel.cc b/paddle/phi/kernels/fusion/xpu/multi_encoder_xpu_kernel.cc
@@ -104,7 +104,7 @@ void MultiEncoderXPUKernel(
   if (x_dtype == phi::DataType::FLOAT32) {
     auto* x_fp16_data_t = reinterpret_cast<XPUTypeFP16*>(
         ctx.template Alloc<phi::dtype::float16>(x_fp16));
-    int r_cast_x = xpu::cast_v2<float, XPUTypeFP16>(
+    int r_cast_x = xpu::cast<float, XPUTypeFP16>(
         ctx.x_context(), x.data<float>(), x_fp16_data_t, x.numel());
     PADDLE_ENFORCE_XDNN_SUCCESS(r_cast_x,
                                 "multi_encoder_xpu(cast x from fp32 to fp16)");
@@ -331,10 +331,10 @@ void MultiEncoderXPUKernel(
 
   if (x_dtype == phi::DataType::FLOAT32) {
     int r_cast_out =
-        xpu::cast_v2<XPUTypeFP16, float>(ctx.x_context(),
-                                         out_fp16_data,
-                                         ctx.template Alloc<float>(out),
-                                         out->numel());
+        xpu::cast<XPUTypeFP16, float>(ctx.x_context(),
+                                      out_fp16_data,
+                                      ctx.template Alloc<float>(out),
+                                      out->numel());
     PADDLE_ENFORCE_XDNN_SUCCESS(
         r_cast_out, "multi_encoder_xpu(cast out from fp16 to fp32)");
   }
diff --git a/paddle/phi/kernels/fusion/xpu/qkv_attention_xpu_kernel.cc b/paddle/phi/kernels/fusion/xpu/qkv_attention_xpu_kernel.cc
@@ -107,23 +107,27 @@ void QKVAttentionXPUKernelImpl(const Context& ctx,
         XPUTypeFP16* k_data_fp16 = nullptr;
         XPUTypeFP16* v_data_fp16 = nullptr;
         if (qkv_fc_fusion) {
-          r_cast_x = xpu::cast_v2<float, XPUTypeFP16>(
+          r_cast_x = xpu::cast<float, XPUTypeFP16>(
               ctx.x_context(), q.data<float>(), x_fp16_data_t, q.numel());
+          PADDLE_ENFORCE_XDNN_SUCCESS(r_cast_x, "cast");
           q_data_fp16 = x_fp16_data_t;
           k_data_fp16 = x_fp16_data_t + head_num * head_dim;
           v_data_fp16 = x_fp16_data_t + 2 * head_num * head_dim;
         } else {
-          r_cast_x = xpu::cast_v2<float, XPUTypeFP16>(
+          r_cast_x = xpu::cast<float, XPUTypeFP16>(
               ctx.x_context(), q.data<float>(), x_fp16_data_t, q.numel());
-          r_cast_x = xpu::cast_v2<float, XPUTypeFP16>(ctx.x_context(),
-                                                      k.data<float>(),
-                                                      x_fp16_data_t + q.numel(),
-                                                      k.numel());
-          r_cast_x = xpu::cast_v2<float, XPUTypeFP16>(
+          PADDLE_ENFORCE_XDNN_SUCCESS(r_cast_x, "cast");
+          r_cast_x = xpu::cast<float, XPUTypeFP16>(ctx.x_context(),
+                                                   k.data<float>(),
+                                                   x_fp16_data_t + q.numel(),
+                                                   k.numel());
+          PADDLE_ENFORCE_XDNN_SUCCESS(r_cast_x, "cast");
+          r_cast_x = xpu::cast<float, XPUTypeFP16>(
               ctx.x_context(),
               v.data<float>(),
               x_fp16_data_t + q.numel() + k.numel(),
               v.numel());
+          PADDLE_ENFORCE_XDNN_SUCCESS(r_cast_x, "cast");
           q_data_fp16 = x_fp16_data_t;
           k_data_fp16 = x_fp16_data_t + q.numel();
           v_data_fp16 = x_fp16_data_t + q.numel() + k.numel();
@@ -153,7 +157,7 @@ void QKVAttentionXPUKernelImpl(const Context& ctx,
                                             tmp_mask,
                                             qk_max_data);
         PADDLE_ENFORCE_XDNN_SUCCESS(r, "qkv_attention_xpu");
-        int r_cast_out = xpu::cast_v2<XPUTypeFP16, float>(
+        int r_cast_out = xpu::cast<XPUTypeFP16, float>(
             ctx.x_context(), out_fp16_data, qkv->data<float>(), qkv->numel());
         PADDLE_ENFORCE_XDNN_SUCCESS(
             r_cast_out, "multi_encoder_xpu(cast out from fp16 to fp32)");
diff --git a/paddle/phi/kernels/fusion/xpu/weight_only_linear_kernel_xpu.cc b/paddle/phi/kernels/fusion/xpu/weight_only_linear_kernel_xpu.cc
@@ -63,7 +63,7 @@ void WeightOnlyLinearKernel(const Context& dev_ctx,
             0,
             common::errors::Fatal(
                 "scale failed, scale related variable `r` is %d", r));
-        r = baidu::xpu::api::cast_v2<XPUType, float>(
+        r = baidu::xpu::api::cast<XPUType, float>(
             xpu_ctx->x_context(),
             reinterpret_cast<const XPUType*>(
                 max_value_fp16.data<phi::dtype::float16>()),
@@ -72,7 +72,7 @@ void WeightOnlyLinearKernel(const Context& dev_ctx,
         PADDLE_ENFORCE_EQ(r,
                           0,
                           common::errors::Fatal(
-                              "cast_v2 failed, related variable `r` is %d", r));
+                              "cast failed, related variable `r` is %d", r));
       } else if (weight_scale.dtype() == phi::DataType::FLOAT32) {
         r = baidu::xpu::api::scale(xpu_ctx->x_context(),
                                    weight_scale.data<float>(),
@@ -95,12 +95,13 @@ void WeightOnlyLinearKernel(const Context& dev_ctx,
           bias.get().dtype() == phi::DataType::FLOAT16) {
         bias_fp32.Resize(bias.get().dims());
         dev_ctx.template Alloc<float>(&bias_fp32);
-        r = baidu::xpu::api::cast_v2<XPUType, float>(
+        r = baidu::xpu::api::cast<XPUType, float>(
             xpu_ctx->x_context(),
             reinterpret_cast<const XPUType*>(
                 bias.get().data<phi::dtype::float16>()),
             bias_fp32.data<float>(),
             n);
+        PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
       }
       if (weight_dtype == "int8") {
         r = baidu::xpu::api::gpt_fc_fusion<XPUType, int8_t, XPUType, int8_wo_t>(
diff --git a/paddle/phi/kernels/xpu/clip_grad_kernel.cc b/paddle/phi/kernels/xpu/clip_grad_kernel.cc
@@ -29,13 +29,13 @@ void ClipGradKernel(const Context& ctx,
   ctx.template Alloc<T>(x_grad);
   using XPUDataType = typename XPUTypeTrait<T>::Type;
   int r =
-      xpu::clip_grad(ctx.x_context(),
-                     reinterpret_cast<const XPUDataType*>(x.data<T>()),
-                     reinterpret_cast<const XPUDataType*>(out_grad.data<T>()),
-                     reinterpret_cast<XPUDataType*>(x_grad->data<T>()),
-                     x.numel(),
-                     static_cast<XPUDataType>(min.to<T>()),
-                     static_cast<XPUDataType>(max.to<T>()));
+      xpu::clamp_grad(ctx.x_context(),
+                      reinterpret_cast<const XPUDataType*>(x.data<T>()),
+                      reinterpret_cast<const XPUDataType*>(out_grad.data<T>()),
+                      reinterpret_cast<XPUDataType*>(x_grad->data<T>()),
+                      x.numel(),
+                      static_cast<XPUDataType>(min.to<T>()),
+                      static_cast<XPUDataType>(max.to<T>()));
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "clip_grad");
 }
 }  // namespace phi
@@ -46,5 +46,6 @@ PD_REGISTER_KERNEL(clip_grad,
                    phi::ClipGradKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int64_t,
                    int) {}
diff --git a/test/legacy_test/test_clip_op.py b/test/legacy_test/test_clip_op.py
@@ -166,7 +166,7 @@ def initTestCase(self):
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
     or not core.is_bfloat16_supported(core.CUDAPlace(0)),
-    "core is not compiled with CUDA and not support the bfloat16",
+    "core is not compiled with CUDA or not support the bfloat16",
 )
 class TestClipBF16Op(OpTest):
     def setUp(self):
diff --git a/test/xpu/test_clip_op_xpu.py b/test/xpu/test_clip_op_xpu.py

Original file line number	Diff line number	Diff line change
`@@ -166,7 +166,7 @@ def initTestCase(self):`
`166`	`166`	`@unittest.skipIf(`
`167`	`167`	`not core.is_compiled_with_cuda()`
`168`	`168`	`or not core.is_bfloat16_supported(core.CUDAPlace(0)),`
`169`		`- "core is not compiled with CUDA and not support the bfloat16",`
	`169`	`+ "core is not compiled with CUDA or not support the bfloat16",`
`170`	`170`	`)`
`171`	`171`	`class TestClipBF16Op(OpTest):`
`172`	`172`	`def setUp(self):`