【PaddlePaddle Hackathon 4】No.63 : add embedding fp16 test

paddle/phi/kernels/gpu/broadcast_tensors_grad_kernel.cu

@@ -108,4 +108,5 @@ PD_REGISTER_KERNEL(broadcast_tensors_grad,
                    int64_t,
                    float,
                    double,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

paddle/phi/kernels/gpu/broadcast_tensors_kernel.cu

@@ -27,4 +27,5 @@ PD_REGISTER_KERNEL(broadcast_tensors,
                    int64_t,
                    float,
                    double,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

paddle/phi/kernels/gpu/lerp_grad_kernel.cu

@@ -279,5 +279,6 @@ PD_REGISTER_KERNEL(lerp_grad,
                    ALL_LAYOUT,
                    phi::LerpGradKernel,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float,
                    double) {}

paddle/phi/kernels/gpu/lerp_kernel.cu

@@ -23,5 +23,6 @@ PD_REGISTER_KERNEL(lerp,
                    ALL_LAYOUT,
                    phi::LerpKernel,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float,
                    double) {}

python/paddle/fluid/tests/unittests/test_lerp_op.py

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from eager_op_test import OpTest
+from eager_op_test import OpTest, convert_float_to_uint16
 
 import paddle
 from paddle.fluid import core
@@ -220,5 +220,42 @@ def test_x_y_broadcast_w(self):
         paddle.enable_static()
 
 
+@unittest.skipIf(
+    not core.is_compiled_with_cuda()
+    or not core.is_bfloat16_supported(core.CUDAPlace(0)),
+    "core is not complied with CUDA and not support the bfloat16",
+)
+class TestLerpBF16(OpTest):
+    def setUp(self):
+        self.op_type = "lerp"
+        self.python_api = paddle.lerp
+        self.init_dtype()
+        self.init_shape()
+        self.__class__.op_type = self.op_type
+        x = np.arange(1.0, 101.0).astype(np.float32).reshape(self.shape)
+        y = np.full(100, 10.0).astype(np.float32).reshape(self.shape)
+        w = np.asarray([0.5]).astype(np.float32)
+        self.inputs = {
+            'X': convert_float_to_uint16(x),
+            'Y': convert_float_to_uint16(y),
+            'Weight': convert_float_to_uint16(w),
+        }
+        self.outputs = {'Out': convert_float_to_uint16(x + w * (y - x))}
+
+    def init_dtype(self):
+        self.dtype = np.uint16
+
+    def init_shape(self):
+        self.shape = [100]
+
+    def test_check_output(self):
+        place = core.CUDAPlace(0)
+        self.check_output_with_place(place)
+
+    def test_check_grad(self):
+        place = core.CUDAPlace(0)
+        self.check_grad_with_place(place, ['X', 'Y'], 'Out')
+
+
 if __name__ == "__main__":
     unittest.main()

python/paddle/fluid/tests/unittests/test_lookup_table_v2_op.py

@@ -47,11 +47,17 @@ class TestLookupTableOp(OpTest):
     def setUp(self):
         self.op_type = "lookup_table_v2"
         self.python_api = paddle.nn.functional.embedding
-        table = np.random.random((17, 31)).astype("float64")
+        self.init_dtype()
+
+        table = np.random.random((17, 32)).astype(self.dtype)
         ids = np.random.randint(0, 17, 4).astype(self.id_dtype())
+
         self.inputs = {'W': table, 'Ids': ids}
         self.outputs = {'Out': table[ids]}
 
+    def init_dtype(self):
+        self.dtype = "float64"
+
     def id_dtype(self):
         return "int64"
 
@@ -102,7 +108,7 @@ class TestLookupTableOpWithPadding(TestLookupTableOp):
     def test_check_output(self):
         ids = np.squeeze(self.inputs['Ids'])
         padding_idx = np.random.choice(ids, 1)[0]
-        self.outputs['Out'][ids == padding_idx] = np.zeros(31)
+        self.outputs['Out'][ids == padding_idx] = np.zeros(32)
         self.attrs = {'padding_idx': int(padding_idx)}
         self.check_output()
 
@@ -297,6 +303,11 @@ def test_param_dtype():
             )
 
 
+class TestEmbeddingFP16OP(TestLookupTableOp):
+    def init_dtype(self):
+        self.dtype = np.float16
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()