[XPU unit test] add fp16 and bf16 support for c_embedding and c_embedding grad unit test; add c_allreduce_max unit test (#65478)

RuohengMa · web-flow · commit 8c542c89f41c · 2024-07-04T17:30:25.000+08:00
* [XPU unit test] add fp16 and bf16 support for c_embedding and c_embedding grad unit test; add c_allreduce_max unit test

* fix
diff --git a/test/legacy_test/c_embedding_op_base.py b/test/legacy_test/c_embedding_op_base.py
@@ -15,7 +15,11 @@
 import unittest
 
 import numpy as np
-from op_test import OpTest
+from op_test import (
+    OpTest,
+    convert_float_to_uint16,
+    convert_uint16_to_float,
+)
 
 import paddle
 from paddle.framework import core
@@ -152,6 +156,77 @@ def init_dtype(self):
         self.ids_dtype = "int32"
 
 
+class TestCEmbeddingOpFP16(TestCEmbeddingOpBase):
+    def setUp(self):
+        self.init_dtype()
+        self.initcase()
+
+    def initcase(self):
+        self.op_type = "c_embedding"
+        self.python_api = c_embedding_wrapper
+        table = np.random.random((17, 64)).astype(self.dtype)
+        ids = np.random.randint(low=0, high=17 * 2, size=(2, 4)).astype(
+            self.ids_dtype
+        )
+        self.start_index = 10
+        ids[0][1] = 12
+        ids[0][2] = 12
+        ids[1][2] = 12
+        ids[1][3] = 12
+        self.end_index = self.start_index + 17
+
+        self.inputs = {'W': table, 'Ids': ids}
+        np_out = get_c_embedding(self.start_index, self.end_index, table, ids)
+        self.outputs = {'Out': np_out.reshape((2, 4, 64))}
+        self.attrs = {'start_index': self.start_index}
+
+        if core.is_compiled_with_xpu():
+            self.__class__.use_xpu = True
+        elif core.is_compiled_with_cuda():
+            self.__class__.exist_fp64_check_grad = True
+
+    def init_dtype(self):
+        self.dtype = "float16"
+        self.ids_dtype = "int32"
+
+
+class TestCEmbeddingOpBF16(TestCEmbeddingOpBase):
+    def setUp(self):
+        self.init_dtype()
+        self.initcase()
+
+    def initcase(self):
+        self.op_type = "c_embedding"
+        self.python_api = c_embedding_wrapper
+        table = np.random.random((17, 64)).astype('float32')
+        table_bf16 = convert_float_to_uint16(table)
+        table = convert_uint16_to_float(table_bf16)
+        ids = np.random.randint(low=0, high=17 * 2, size=(2, 4)).astype(
+            self.ids_dtype
+        )
+        self.start_index = 10
+        ids[0][1] = 12
+        ids[0][2] = 12
+        ids[1][2] = 12
+        ids[1][3] = 12
+        self.end_index = self.start_index + 17
+
+        self.inputs = {'W': table_bf16, 'Ids': ids}
+        np_out = get_c_embedding(self.start_index, self.end_index, table, ids)
+        np_out = convert_float_to_uint16(np_out)
+        self.outputs = {'Out': np_out.reshape((2, 4, 64))}
+        self.attrs = {'start_index': self.start_index}
+
+        if core.is_compiled_with_xpu():
+            self.__class__.use_xpu = True
+        elif core.is_compiled_with_cuda():
+            self.__class__.exist_fp64_check_grad = True
+
+    def init_dtype(self):
+        self.dtype = np.uint16
+        self.ids_dtype = "int32"
+
+
 class TestCEmbeddingOpComplex64(TestCEmbeddingOpBase):
     def setUp(self):
         self.init_dtype()
diff --git a/test/xpu/test_c_embedding_op_xpu.py b/test/xpu/test_c_embedding_op_xpu.py
@@ -17,6 +17,8 @@
 from c_embedding_op_base import (
     TestCEmbeddingCPU,
     TestCEmbeddingOpBase,
+    TestCEmbeddingOpBF16,
+    TestCEmbeddingOpFP16,
     TestCEmbeddingOpFP32,
 )
 
@@ -30,5 +32,9 @@
 
 TestCEmbeddingOpFP32()
 
+TestCEmbeddingOpFP16()
+
+TestCEmbeddingOpBF16()
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/xpu/test_collective_allreduce_xpu.py b/test/xpu/test_collective_allreduce_xpu.py
@@ -18,6 +18,7 @@
 from xpu.test_collective_api_base import TestDistBase
 
 import paddle
+import paddle.distributed as dist
 from paddle import core
 
 paddle.enable_static()
@@ -31,7 +32,7 @@ def _setup_config(self):
         not core.is_compiled_with_xpu() or paddle.device.xpu.device_count() < 2,
         "run test when having at least 2 XPUs.",
     )
-    def test_allreduce(self):
+    def test_allreduce_sum(self):
         support_types = get_xpu_op_support_types('c_allreduce_sum')
         for dtype in support_types:
             self.check_with_place(
@@ -40,11 +41,21 @@ def test_allreduce(self):
                 dtype=dtype,
             )
 
+    def test_allreduce_max(self):
+        support_types = get_xpu_op_support_types('c_allreduce_max')
+        for dtype in support_types:
+            self.check_with_place(
+                "collective_allreduce_api.py",
+                "allreduce",
+                dtype=dtype,
+                reduce_type=dist.ReduceOp.MAX,
+            )
+
     @unittest.skipIf(
         not core.is_compiled_with_xpu() or paddle.device.xpu.device_count() < 2,
         "run test when having at least 2 XPUs.",
     )
-    def test_allreduce_dygraph(self):
+    def test_allreduce_sum_dygraph(self):
         support_types = get_xpu_op_support_types('c_allreduce_sum')
         for dtype in support_types:
             self.check_with_place(
@@ -54,6 +65,17 @@ def test_allreduce_dygraph(self):
                 dtype=dtype,
             )
 
+    def test_allreduce_max_dygraph(self):
+        support_types = get_xpu_op_support_types('c_allreduce_max')
+        for dtype in support_types:
+            self.check_with_place(
+                "collective_allreduce_api_dygraph.py",
+                "allreduce",
+                static_mode="0",
+                dtype=dtype,
+                reduce_type=dist.ReduceOp.MAX,
+            )
+
 
 if __name__ == '__main__':
     unittest.main()
