Fix fp8 data type bugs on cpu (#65539)

Sunny-bot1 · web-flow · commit d8d4d73a14d9 · 2024-07-08T16:13:47.000+08:00
* fix fp8 data type bug on cpu
diff --git a/paddle/phi/common/float8_e4m3fn.h b/paddle/phi/common/float8_e4m3fn.h
@@ -69,7 +69,7 @@ struct PADDLE_ALIGN(1) float8_e4m3fn {
     const uint32_t sign = fb.ui & UINT32_C(0x80000000);
     fb.ui ^= sign;
     if (fb.ui >= fp8_max) {
-      result = 0x7f;
+      result = 0x7e;
     } else {
       if (fb.ui < (UINT32_C(121) << 23)) {
         fb.f = fb.f + denorm_mask.f;
diff --git a/paddle/phi/common/float8_e5m2.h b/paddle/phi/common/float8_e5m2.h
@@ -77,7 +77,7 @@ struct PADDLE_ALIGN(1) float8_e5m2 {
     fb.ui ^= sign;
 
     if (fb.ui >= fp8_max) {
-      result = fb.ui > fp32_inf ? UINT8_C(0x7F) : UINT8_C(0x7C);
+      result = fb.ui > fp32_inf ? UINT8_C(0x7F) : UINT8_C(0x7B);
     } else {
       if (fb.ui < (UINT32_C(113) << 23)) {
         fb.f = fb.f + denorm_mask.f;
@@ -172,7 +172,8 @@ struct PADDLE_ALIGN(1) float8_e5m2 {
     // https://github.com/pytorch/pytorch/blob/main/c10/util/Float8_e5m2.h
     uint16_t half_representation = x;
     half_representation <<= 8;
-    return static_cast<float>(static_cast<float16>(half_representation));
+    return static_cast<float>(
+        phi::dtype::raw_uint16_to_float16(half_representation));
 #endif
   }
 
diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -44,6 +44,8 @@
     VarDesc.VarType.INT8: DataType.INT8,
     VarDesc.VarType.COMPLEX64: DataType.COMPLEX64,
     VarDesc.VarType.COMPLEX128: DataType.COMPLEX128,
+    VarDesc.VarType.FP8_E4M3FN: DataType.FLOAT8_E4M3FN,
+    VarDesc.VarType.FP8_E5M2: DataType.FLOAT8_E5M2,
 }
 
 datatype_to_vartype = {v: k for k, v in vartype_to_datatype.items()}
@@ -73,6 +75,8 @@
     np.int8: DataType.INT8,
     np.complex64: DataType.COMPLEX64,
     np.complex128: DataType.COMPLEX128,
+    "float8_e4m3fn": DataType.FLOAT8_E4M3FN,
+    "float8_e5m2": DataType.FLOAT8_E5M2,
 }
 
 _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE = {
@@ -88,6 +92,8 @@
     DataType.UINT8: 'uint8',
     DataType.COMPLEX64: 'complex64',
     DataType.COMPLEX128: 'complex128',
+    DataType.FLOAT8_E4M3FN: 'float8_e4m3fn',
+    DataType.FLOAT8_E5M2: 'float8_e5m2',
 }
 
 
@@ -108,6 +114,10 @@ def convert_np_dtype_to_dtype_(np_dtype):
         # since there is still no support for bfloat16 in NumPy,
         # uint16 is used for casting bfloat16
         dtype = np.dtype("uint16")
+    elif isinstance(np_dtype, str) and np_dtype == "float8_e4m3fn":
+        dtype = 'float8_e4m3fn'
+    elif isinstance(np_dtype, str) and np_dtype == "float8_e5m2":
+        dtype = 'float8_e5m2'
     else:
         dtype = np.dtype(np_dtype)
 
diff --git a/test/legacy_test/test_float8.py b/test/legacy_test/test_float8.py
@@ -54,20 +54,47 @@ def check_fp8_support() -> bool:
     return True
 
 
-@unittest.skipIf(
-    not core.is_compiled_with_cuda() or not get_cuda_version() >= 11080,
-    "FP8 date type has some bug in CPU",
-)
 class TestFP8CastOp(unittest.TestCase):
     def setUp(self):
-        self.dtype_dict = {
-            "float8_e4m3fn": core.VarDesc.VarType.FP8_E4M3FN,
-            "float8_e5m2": core.VarDesc.VarType.FP8_E5M2,
-        }
+        if paddle.framework.use_pir_api():
+            self.dtype_dict = {
+                "float8_e4m3fn": core.DataType.FLOAT8_E4M3FN,
+                "float8_e5m2": core.DataType.FLOAT8_E5M2,
+            }
+        else:
+            self.dtype_dict = {
+                "float8_e4m3fn": core.VarDesc.VarType.FP8_E4M3FN,
+                "float8_e5m2": core.VarDesc.VarType.FP8_E5M2,
+            }
         self.shape = (16, 16)
 
     def test_cast(self):
-        for self.device in ["cpu", "gpu"]:
+        if core.is_compiled_with_cuda():
+            for self.device in ["cpu", "gpu"]:
+                paddle.device.set_device(self.device)
+                for self.dtype in ["float8_e4m3fn", "float8_e5m2"]:
+                    # test fp32 to fp8 (dtype)
+                    input = paddle.full(self.shape, 100000.0)
+                    input1 = input.astype(self.dtype)
+                    self.assertTrue(input1.dtype == self.dtype_dict[self.dtype])
+                    # test fp8 to fp32 (dtype)
+                    input2 = input1.astype("float32")
+                    if paddle.framework.use_pir_api():
+                        self.assertTrue(input2.dtype == core.DataType.FLOAT32)
+                    else:
+                        self.assertTrue(
+                            input2.dtype == core.VarDesc.VarType.FP32
+                        )
+                    # test fp32 to fp8 (value clip)
+                    expect = paddle.full(
+                        self.shape,
+                        E4M3_MAX_POS
+                        if self.dtype == "float8_e4m3fn"
+                        else E5M2_MAX_POS,
+                    )
+                    self.assertTrue(paddle.equal_all(input2, expect))
+        else:
+            self.device = "cpu"
             paddle.device.set_device(self.device)
             for self.dtype in ["float8_e4m3fn", "float8_e5m2"]:
                 # test fp32 to fp8 (dtype)
@@ -76,44 +103,65 @@ def test_cast(self):
                 self.assertTrue(input1.dtype == self.dtype_dict[self.dtype])
                 # test fp8 to fp32 (dtype)
                 input2 = input1.astype("float32")
-                self.assertTrue(input2.dtype == core.VarDesc.VarType.FP32)
+                if paddle.framework.use_pir_api():
+                    self.assertTrue(input2.dtype == core.DataType.FLOAT32)
+                else:
+                    self.assertTrue(input2.dtype == core.VarDesc.VarType.FP32)
                 # test fp32 to fp8 (value clip)
                 expect = paddle.full(
                     self.shape,
                     E4M3_MAX_POS
                     if self.dtype == "float8_e4m3fn"
                     else E5M2_MAX_POS,
                 )
-                # there exists some problem in cpu fp8 cast
-                if self.device == "gpu":
-                    self.assertTrue(paddle.equal_all(input2, expect))
+                self.assertTrue(paddle.equal_all(input2, expect))
 
 
-@unittest.skipIf(
-    not core.is_compiled_with_cuda() or not get_cuda_version() >= 11080,
-    "FP8 date type has some bug in CPU",
-)
 class TestFP8FullOp(unittest.TestCase):
     def setUp(self):
-        self.dtype_dict = {
-            "float8_e4m3fn": core.VarDesc.VarType.FP8_E4M3FN,
-            "float8_e5m2": core.VarDesc.VarType.FP8_E5M2,
-        }
+        if paddle.framework.use_pir_api():
+            self.dtype_dict = {
+                "float8_e4m3fn": core.DataType.FLOAT8_E4M3FN,
+                "float8_e5m2": core.DataType.FLOAT8_E5M2,
+            }
+        else:
+            self.dtype_dict = {
+                "float8_e4m3fn": core.VarDesc.VarType.FP8_E4M3FN,
+                "float8_e5m2": core.VarDesc.VarType.FP8_E5M2,
+            }
 
     def test_ones(self):
-        for self.device in ["cpu", "gpu"]:
+        if core.is_compiled_with_cuda():
+            for self.device in ["cpu", "gpu"]:
+                paddle.device.set_device(self.device)
+                for self.dtype in ["float8_e4m3fn", "float8_e5m2"]:
+                    input = paddle.ones([1, 2], dtype=self.dtype)
+                    self.assertTrue(input.dtype == self.dtype_dict[self.dtype])
+                    input_fp32 = input.astype("float32")
+                    expect = paddle.to_tensor([[1, 1]]).astype("float32")
+                    self.assertTrue(paddle.equal_all(expect, input_fp32))
+        else:
+            self.device = "cpu"
             paddle.device.set_device(self.device)
             for self.dtype in ["float8_e4m3fn", "float8_e5m2"]:
                 input = paddle.ones([1, 2], dtype=self.dtype)
                 self.assertTrue(input.dtype == self.dtype_dict[self.dtype])
                 input_fp32 = input.astype("float32")
                 expect = paddle.to_tensor([[1, 1]]).astype("float32")
-                # there exists some problem in cpu fp8 full
-                if self.device == "gpu":
-                    self.assertTrue(paddle.equal_all(expect, input_fp32))
+                self.assertTrue(paddle.equal_all(expect, input_fp32))
 
     def test_zeros(self):
-        for self.device in ["cpu", "gpu"]:
+        if core.is_compiled_with_cuda():
+            for self.device in ["cpu", "gpu"]:
+                paddle.device.set_device(self.device)
+                for self.dtype in ["float8_e4m3fn", "float8_e5m2"]:
+                    input = paddle.zeros([1, 2], dtype=self.dtype)
+                    self.assertTrue(input.dtype == self.dtype_dict[self.dtype])
+                    input_fp32 = input.astype("float32")
+                    expect = paddle.to_tensor([[0, 0]]).astype("float32")
+                    self.assertTrue(paddle.equal_all(expect, input_fp32))
+        else:
+            self.device = "cpu"
             paddle.device.set_device(self.device)
             for self.dtype in ["float8_e4m3fn", "float8_e5m2"]:
                 input = paddle.zeros([1, 2], dtype=self.dtype)
