[XPU] Add fp16 and bf16 support for XPU embedding_grad op (#59454)

cqulilujia · web-flow · commit 64ce72c5fef0 · 2023-11-29T13:54:31.000+08:00
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -24,7 +24,7 @@ set(XPU_XFT_LIB_NAME "libxft.so")
 set(XPU_XPTI_LIB_NAME "libxpti.so")
 
 if(NOT DEFINED XPU_BASE_DATE)
-  set(XPU_BASE_DATE "20231114")
+  set(XPU_BASE_DATE "20231128")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.1.6.1")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -324,7 +324,10 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT64})},
-      {"embedding_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"embedding_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"embedding_sparse_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"equal",
        XPUKernelSet({phi::DataType::INT64,
@@ -567,9 +570,14 @@ XPUOpMap& get_kl2_ops() {
       {"logical_not", XPUKernelSet({phi::DataType::BOOL})},
       {"logical_or", XPUKernelSet({phi::DataType::BOOL})},
       {"logical_xor", XPUKernelSet({phi::DataType::BOOL})},
-      {"lookup_table_v2_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"lookup_table_v2_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"lookup_table_v2",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"masked_select",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -308,7 +308,10 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT64})},
-      {"embedding_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"embedding_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"embedding_sparse_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"equal",
        XPUKernelSet({phi::DataType::INT64,
@@ -541,7 +544,10 @@ XPUOpMap& get_kl3_ops() {
       {"logical_not", XPUKernelSet({phi::DataType::BOOL})},
       {"logical_or", XPUKernelSet({phi::DataType::BOOL})},
       {"logical_xor", XPUKernelSet({phi::DataType::BOOL})},
-      {"lookup_table_v2_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"lookup_table_v2_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"lookup_table_v2",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/kernels/cpu/embedding_grad_kernel.cc b/paddle/phi/kernels/cpu/embedding_grad_kernel.cc
@@ -208,6 +208,7 @@ PD_REGISTER_KERNEL(embedding_grad,
                    phi::EmbeddingGradKernel,
                    float,
                    double,
+                   phi::dtype::float16,
                    phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(embedding_sparse_grad,
diff --git a/paddle/phi/kernels/cpu/embedding_kernel.cc b/paddle/phi/kernels/cpu/embedding_kernel.cc
@@ -122,4 +122,5 @@ PD_REGISTER_KERNEL(embedding,
                    float,
                    double,
                    int8_t,
+                   phi::dtype::float16,
                    phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/embedding_grad_kernel.cc b/paddle/phi/kernels/xpu/embedding_grad_kernel.cc
@@ -150,8 +150,13 @@ void EmbeddingSparseGradKernel(const Context& ctx,
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    embedding_grad, XPU, ALL_LAYOUT, phi::EmbeddingGradKernel, float) {}
+PD_REGISTER_KERNEL(embedding_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::EmbeddingGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 PD_REGISTER_KERNEL(embedding_sparse_grad,
                    XPU,
                    ALL_LAYOUT,
diff --git a/test/xpu/test_lookup_table_v2_op_xpu.py b/test/xpu/test_lookup_table_v2_op_xpu.py
@@ -55,7 +55,9 @@ def test_check_output(self):
             self.check_output_with_place(self.place)
 
         def test_check_grad(self):
-            self.check_grad_with_place(self.place, ['X'], 'Out')
+            self.check_grad_with_place(
+                self.place, ['W'], 'Out', no_grad_set=set('Ids')
+            )
 
         def init_config(self):
             self.input_shape = (17, 31)
