# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import numpy as np from op_test import get_cuda_version, get_device, is_custom_device import paddle from paddle.base import core # define the e4m3/e5m2 constants E4M3_MAX_POS = 448.0 E5M2_MAX_POS = 57344.0 def check_fp8_support() -> bool: """Return if fp8 support is available""" gpu_arch = ( paddle.device.cuda.get_device_capability()[0] * 10 + paddle.device.cuda.get_device_capability()[1] ) if gpu_arch >= 90: # hopper and above return True # Device compute capability 8.9 or higher required for FP8 execution. if gpu_arch < 89: # pre-ada return False if get_cuda_version() < 12010: return False return True class TestFP8CastOp(unittest.TestCase): def setUp(self): 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): if core.is_compiled_with_cuda() or is_custom_device(): for self.device in ["cpu", get_device()]: 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) 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)) class TestFP8FullOp(unittest.TestCase): def setUp(self): 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): if core.is_compiled_with_cuda() or is_custom_device(): for self.device in ["cpu", get_device()]: 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") self.assertTrue(paddle.equal_all(expect, input_fp32)) def test_zeros(self): if core.is_compiled_with_cuda() or is_custom_device(): for self.device in ["cpu", get_device()]: 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) 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)) @unittest.skipIf( not (core.is_compiled_with_cuda() or is_custom_device()) or not check_fp8_support(), "Fp8 matmul requires CUDA >= 12.1 on Ada arch or hopper arch", ) class TestFP8MatmulOp(unittest.TestCase): def gelu(self, x): return ( 0.5 * x * (1.0 + np.tanh(0.7978845608 * (x + 0.044715 * np.power(x, 3)))) ) def setUp(self): self.dtype_dict = { "float8_e4m3fn": core.VarDesc.VarType.FP8_E4M3FN, "float8_e5m2": core.VarDesc.VarType.FP8_E5M2, } def test_matmul(self): for self.device in [get_device()]: paddle.device.set_device(self.device) for self.dtype in ["float8_e4m3fn"]: input1 = paddle.ones([4, 16, 32], dtype=self.dtype) input2 = paddle.ones([4, 64, 32], dtype=self.dtype) bias_fp16 = paddle.ones([64], dtype="float16") bias_bf16 = paddle.ones([64], dtype="bfloat16") input3 = np.ones((4, 64, 32)).astype("float32") input4 = np.ones((4, 32, 64)).astype("float32") bias_float32 = paddle.ones([64], dtype="float32") output_fp16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, output_dtype="float16", ) output_bf16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, output_dtype="bfloat16", ) output_bias_fp16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, bias=bias_fp16, scale=1.0, output_dtype="float16", ) output_bias_bf16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, bias=bias_bf16, scale=1.0, output_dtype="bfloat16", ) output_gelu_fp16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, scale=1.0, act="gelu", output_dtype="float16", ) output_gelu_bf16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, scale=1.0, act="gelu", output_dtype="bfloat16", ) output_bias_gelu_fp16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, bias=bias_fp16, scale=1.0, act="gelu", output_dtype="float16", ) output_bias_gelu_bf16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, bias=bias_bf16, scale=1.0, act="gelu", output_dtype="bfloat16", ) output_bias_relu_fp16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, bias=bias_fp16, scale=1.0, act="relu", output_dtype="float16", ) output_bias_relu_bf16 = paddle.linalg.fp8_fp8_half_gemm_fused( input1, input2, transpose_x=False, transpose_y=True, bias=bias_bf16, scale=1.0, act="relu", output_dtype="bfloat16", ) expect_result = np.matmul(input3, input4) if self.device == "gpu": self.assertTrue( paddle.equal_all( paddle.cast(output_fp16, "float32"), paddle.cast(output_bf16, "float32"), paddle.to_tensor(expect_result), ) ) self.assertTrue( paddle.equal_all( paddle.cast(output_gelu_fp16, "float32"), paddle.cast(output_gelu_bf16, "float32"), paddle.to_tensor(self.gelu(expect_result)), ) ) self.assertTrue( paddle.equal_all( paddle.cast(output_bias_fp16, "float32"), paddle.cast(output_bias_bf16, "float32"), paddle.to_tensor(expect_result + bias_float32), ) ) self.assertTrue( paddle.equal_all( paddle.cast(output_bias_gelu_fp16, "float32"), paddle.cast(output_bias_gelu_bf16, "float32"), paddle.to_tensor( self.gelu(expect_result) + bias_float32 ), ) ) self.assertTrue( paddle.equal_all( paddle.cast(output_bias_relu_fp16, "float32"), paddle.cast(output_bias_relu_bf16, "float32"), paddle.to_tensor( np.maximum(expect_result, 0) + bias_float32 ), ) ) if __name__ == "__main__": unittest.main()