[XPU] bfloat16 for fused_gemm_epilogue_grad (#71954)

houj04 · web-flow · commit d12bf33dc667 · 2025-03-28T11:09:00.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -1751,7 +1751,9 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
       {"fused_gemm_epilogue_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"fused_linear_param_grad_add",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"fused_attention",
diff --git a/paddle/phi/kernels/funcs/fused_gemm_epilogue_xpu.h b/paddle/phi/kernels/funcs/fused_gemm_epilogue_xpu.h
@@ -71,12 +71,14 @@ void ComputeFusedGemmEpilogueBackwardXPU(const phi::XPUContext& dev_ctx,
                        dout->numel());
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "relu_grad");
   } else if (activation_grad == "gelu") {
+    // int gelu_grad(Context* ctx, const T* x, const T* dy, T* dx, int64_t len,
+    // bool approximate);
     r = xpu::gelu_grad(xpu_ctx,
-                       reserve_space_ptr,
                        reserve_space_ptr,
                        dout_ptr,
                        d_act_input_ptr,
-                       dout->numel());
+                       dout->numel(),
+                       false);
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "gelu_grad");
   } else if (activation_grad == "none") {
     // pass
diff --git a/paddle/phi/kernels/fusion/xpu/fused_gemm_epilogue_grad_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_gemm_epilogue_grad_kernel.cc
@@ -74,4 +74,5 @@ PD_REGISTER_KERNEL(fused_gemm_epilogue_grad,
                    ALL_LAYOUT,
                    phi::fusion::FusedGemmEpilogueXPUGradKernel,
                    float,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16) {}
diff --git a/test/xpu/test_fused_gemm_epilogue_grad_op_xpu.py b/test/xpu/test_fused_gemm_epilogue_grad_op_xpu.py
@@ -21,6 +21,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -40,80 +41,82 @@ def __init__(self):
         self.op_name = 'fused_gemm_epilogue_grad'
         self.use_dynamic_create_class = False
 
-    class TestFuseGemmEpilogueGradOpDXYBias1(XPUOpTest):
+    class TestFuseGemmEpilogueGradOp(XPUOpTest):
         def setUp(self):
             paddle.enable_static()
             self.op_type = "fused_gemm_epilogue_grad"
             self.__class__.no_need_check_grad = True
 
             self.dtype = self.in_type
+            self.init_shape()
             self.init_data()
+            self.init_output()
+
+        def init_shape(self):
+            self.x_shape = (8, 4)
+            self.y_shape = (4, 128)
+            self.out_shape = (8, 128)
 
         def init_data(self):
+            self.dout = np.random.uniform(-0.5, 0.5, self.out_shape)
+            self.x = np.random.uniform(-0.5, 0.5, self.x_shape)
+            self.y = np.random.uniform(-0.5, 0.5, self.y_shape)
+
+            if self.dtype != np.uint16:
+                self.dout = self.dout.astype(self.dtype)
+                self.x = self.x.astype(self.dtype)
+                self.y = self.y.astype(self.dtype)
+
+            # numpy outputs
+            self.dx, self.dy, self.dbias = get_outputs(
+                self.dout, self.x, self.y
+            )
+
+            # special for bfloat16
+            if self.dtype == np.uint16:
+                self.dout = convert_float_to_uint16(self.dout)
+                self.x = convert_float_to_uint16(self.x)
+                self.y = convert_float_to_uint16(self.y)
+                self.dx = convert_float_to_uint16(self.dx)
+                self.dy = convert_float_to_uint16(self.dy)
+                self.dbias = convert_float_to_uint16(self.dbias)
+
+            # inputs
             self.inputs = {
-                'DOut': np.random.random((8, 128)).astype(self.dtype) - 0.5,
-                'X': np.random.random((8, 4)).astype(self.dtype) - 0.5,
-                'Y': np.random.random((4, 128)).astype(self.dtype) - 0.5,
+                'DOut': self.dout,
+                'X': self.x,
+                'Y': self.y,
             }
 
+            # attrs
             self.attrs = {"activation_grad": 'none'}
 
-            DX, DY, DBias = get_outputs(
-                self.inputs['DOut'], self.inputs['X'], self.inputs['Y']
-            )
-            self.outputs = {'DX': DX, 'DY': DY, 'DBias': DBias}
+        def init_output(self):
+            self.outputs = {'DX': self.dx, 'DY': self.dy, 'DBias': self.dbias}
 
         def test_check_output(self):
             self.atol = 5e-4
             if self.dtype == np.float16:
                 self.atol = 1e-3
             self.check_output_with_place(core.XPUPlace(0), atol=self.atol)
 
-    class TestFuseGemmEpilogueGradOpDXYBias2(XPUOpTest):
-        def init_data(self):
-            self.inputs = {
-                'DOut': np.random.random((8, 128)).astype(self.dtype) - 0.5,
-                'X': np.random.random((8, 4)).astype(self.dtype) - 0.5,
-                'Y': np.random.random((4, 128)).astype(self.dtype) - 0.5,
-            }
-
-            self.attrs = {"activation_grad": 'none'}
-
-            _, DY, DBias = get_outputs(
-                self.inputs['DOut'], self.inputs['X'], self.inputs['Y']
-            )
-            self.outputs = {'DY': DY, 'DBias': DBias}
+    class TestFuseGemmEpilogueGradOp2(TestFuseGemmEpilogueGradOp):
+        def init_output(self):
+            self.outputs = {'DY': self.dy, 'DBias': self.dbias}
 
-    class TestFuseGemmEpilogueGradOpDXYBias3(XPUOpTest):
-        def init_data(self):
-            self.inputs = {
-                'DOut': np.random.random((8, 128)).astype(self.dtype) - 0.5,
-                'X': np.random.random((8, 4)).astype(self.dtype) - 0.5,
-                'Y': np.random.random((4, 128)).astype(self.dtype) - 0.5,
+    class TestFuseGemmEpilogueGradOp3(XPUOpTest):
+        def init_output(self):
+            self.outputs = {
+                'DY': self.dy,
             }
 
-            self.attrs = {"activation_grad": 'none'}
-
-            _, DY, _ = get_outputs(
-                self.inputs['DOut'], self.inputs['X'], self.inputs['Y']
-            )
-            self.outputs = {'DY': DY}
-
-    class TestFuseGemmEpilogueGradOpDXYBias4(XPUOpTest):
-        def init_data(self):
-            self.inputs = {
-                'DOut': np.random.random((8, 128)).astype(self.dtype) - 0.5,
-                'X': np.random.random((8, 4)).astype(self.dtype) - 0.5,
-                'Y': np.random.random((4, 128)).astype(self.dtype) - 0.5,
+    class TestFuseGemmEpilogueGradOp4(XPUOpTest):
+        def init_output(self):
+            self.outputs = {
+                'DX': self.dx,
+                'DY': self.dy,
             }
 
-            self.attrs = {"activation_grad": 'none'}
-
-            DX, DY, _ = get_outputs(
-                self.inputs['DOut'], self.inputs['X'], self.inputs['Y']
-            )
-            self.outputs = {'DX': DX, 'DY': DY}
-
 
 support_types = get_xpu_op_support_types('fused_gemm_epilogue_grad')
 for stype in support_types:
