Fix XPU AMP tests (#65387)

dynamicheart · web-flow · commit 8ddbcb72f04f · 2024-06-28T13:50:50.000+08:00
* Fix XPU AMP tests

* fix

* update xhpc

* fix

* fix
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ if(NOT DEFINED XPU_XRE_BASE_VERSION)
   set(XPU_XRE_BASE_VERSION "4.32.0.1")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "20240621")
+  set(XPU_XHPC_BASE_DATE "20240626")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.2.1.2")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -265,9 +265,13 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32})},
       {"dropout_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"dropout",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"einsum", XPUKernelSet({phi::DataType::FLOAT32})},
       {"einsum_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"elementwise_add_grad",
diff --git a/paddle/phi/kernels/xpu/concat_kernel.cc b/paddle/phi/kernels/xpu/concat_kernel.cc
@@ -49,6 +49,15 @@ void ConcatKernel(const Context& dev_ctx,
                         axis,
                         x[0]->dims().size()));
 
+  std::vector<phi::DDim> x_dims;
+  for (size_t i = 0; i < x.size(); ++i) {
+    x_dims.push_back(x[i]->dims());
+  }
+
+  phi::DDim out_dims = phi::funcs::ComputeAndCheckShape(true, x_dims, axis);
+  out->Resize(out_dims);
+  dev_ctx.template Alloc<T>(out);
+
   // If axis is 0, the lod of the output is not the same as inputs.
 
   if (axis == 0 && x[0]->lod().size() > 0) {
@@ -81,8 +90,6 @@ void ConcatKernel(const Context& dev_ctx,
     }
   }
 
-  dev_ctx.template Alloc<T>(out);
-
   std::vector<std::vector<int>> xdims_list;
   std::vector<const XPUType*> ptrs;
   for (unsigned int i = 0; i < x.size(); ++i) {
diff --git a/paddle/phi/kernels/xpu/dropout_grad_kernel.cc b/paddle/phi/kernels/xpu/dropout_grad_kernel.cc
@@ -108,4 +108,5 @@ PD_REGISTER_KERNEL(dropout_grad,
                    ALL_LAYOUT,
                    phi::DropoutGradRawKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/dropout_kernel.cc b/paddle/phi/kernels/xpu/dropout_kernel.cc
@@ -138,7 +138,8 @@ PD_REGISTER_KERNEL(dropout,
                    ALL_LAYOUT,
                    phi::DropoutRawKernel,
                    float,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->OutputAt(1).SetDataType(phi::DataType::UINT8);
 }
diff --git a/python/paddle/amp/auto_cast.py b/python/paddle/amp/auto_cast.py
@@ -525,18 +525,18 @@ def amp_guard(
             if tracer._expected_place.is_xpu_place():
                 if (dtype == 'float16') and not _is_xpu_float16_supported():
                     xpu_verion = core.get_xpu_device_version(
-                        tracer._expected_place.get_device_id()
+                        _current_expected_place().get_device_id()
                     )
                     warnings.warn(
-                        '%d does not support float16 amp.' % xpu_verion
+                        'XPU%d does not support float16 amp.' % xpu_verion
                     )
                     enable = False
                 elif (dtype == 'bfloat16') and not _is_xpu_bfloat16_supported():
                     xpu_verion = core.get_xpu_device_version(
-                        tracer._expected_place.get_device_id()
+                        _current_expected_place().get_device_id()
                     )
                     warnings.warn(
-                        '%d does not support bfloat16 amp.' % xpu_verion
+                        'XPU%d does not support bfloat16 amp.' % xpu_verion
                     )
                     enable = False
             # For custom device:
diff --git a/test/amp/test_amp_master_weight.py b/test/amp/test_amp_master_weight.py
@@ -53,11 +53,6 @@ def forward(self, x):
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3, disable temporarily.",
-)
 class TestMasterWeight(AmpTestBase):
     def run_dygraph(self, dtype, level, use_promote, max_iters, x_data):
         losses = []
diff --git a/test/amp/test_amp_promote.py b/test/amp/test_amp_promote.py
@@ -36,11 +36,6 @@
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3 ,disable it temporarily.",
-)
 class TestStaticAmpPromoteStats(AmpTestBase):
     def check_promote_results(
         self, use_amp, dtype, level, use_promote, expected_op_calls, debug_info
@@ -142,11 +137,6 @@ def test_static_amp_o2(self):
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3 ,disable it temporarily.",
-)
 class TestEagerAmpPromoteStats(AmpTestBase):
     def check_promote_results(
         self, dtype, level, use_promote, expected_op_calls, debug_info
@@ -230,11 +220,6 @@ def test_o2_promote_off(self):
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3 ,disable it temporarily.",
-)
 class TestPirAmpPromoteStats(AmpTestBase):
     def check_promote_results(
         self, dtype, level, use_promote, expected_op_calls, debug_info
@@ -343,11 +328,6 @@ def test_o2_promote_off(self):
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3 ,disable it temporarily.",
-)
 class TestEagerAmpPromoteSimple(AmpTestBase):
     def setUp(self):
         self._conv = paddle.nn.Conv2D(
@@ -394,11 +374,6 @@ def test_o2_use_promote_off(self):
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3 ,disable it temporarily.",
-)
 class TestPirAmpPromoteSimple(AmpTestBase):
     def init_net(self):
         self._conv = paddle.nn.Conv2D(
diff --git a/test/amp/test_model_cast_to_bf16.py b/test/amp/test_model_cast_to_bf16.py
@@ -202,11 +202,6 @@ def test_graph_cast(self):
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3, disable it temporarily.",
-)
 class TestProgramBF16(AmpTestBase):
     def _check_optimizer(self, program, expected_num_mp):
         optimizers = []
@@ -275,11 +270,6 @@ def test_amp_bf16_o2(self):
     and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
     "run test when xpu's compute capability >= xpu3.",
 )
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "Bugs on XPU3, disable it temporarily.",
-)
 class TestStaticBF16(AmpTestBase):
     def _generate_feed_x(self):
         x = np.random.random(size=[16, 16]).astype("float32")
diff --git a/test/xpu/test_dropout_op_xpu.py b/test/xpu/test_dropout_op_xpu.py
@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+from op_test import convert_float_to_uint16, convert_uint16_to_float
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -41,21 +42,39 @@ def setUp(self):
             self.init_attrs()
             self.dtype = self.in_type
             self.op_type = 'dropout'
-            self.inputs = {'X': np.random.random(self.shape).astype(self.dtype)}
+            if self.dtype == np.uint16:
+                x = np.random.random(self.shape).astype('float32')
+                x = convert_float_to_uint16(x)
+                self.inputs = {'X': x}
+            else:
+                self.inputs = {
+                    'X': np.random.random(self.shape).astype(self.dtype)
+                }
             self.attrs = {
                 'dropout_prob': self.dropout_prob,
                 'fix_seed': self.fix_seed,
                 'is_test': self.is_test,
                 'dropout_implementation': self.dropout_implementation,
             }
 
-            out = self.inputs['X'] * (1.0 - self.dropout_prob)
+            if self.dtype == np.uint16:
+                x_float32 = convert_uint16_to_float(self.inputs['X'])
+                out = x_float32 * (1.0 - self.dropout_prob)
+                out = convert_float_to_uint16(out)
+            else:
+                out = self.inputs['X'] * (1.0 - self.dropout_prob)
+
             if not self.is_test:
                 mask = None
                 if self.dropout_prob == 0.0:
-                    mask = np.ones(self.shape).astype(self.dtype)
+                    mask = np.ones(self.shape)
                 elif self.dropout_prob == 1.0:
-                    mask = np.zeros(self.shape).astype(self.dtype)
+                    mask = np.zeros(self.shape)
+
+                if self.dtype == np.uint16:
+                    mask = convert_float_to_uint16(mask.astype('float32'))
+                else:
+                    mask = mask.astype(self.dtype)
                 self.outputs = {'Out': out, 'Mask': mask}
             else:
                 self.outputs = {'Out': out}
@@ -168,16 +187,27 @@ def setUp(self):
             self.places.append(base.XPUPlace(0))
 
         def cal_grad_upscale_train(self, mask, prob):
-            return mask.astype(self.in_type) / (1 - prob)
+            if self.in_type == np.uint16:
+                return convert_float_to_uint16(
+                    mask.astype("float32") / (1 - prob)
+                )
+            else:
+                return mask.astype(self.in_type) / (1 - prob)
 
         def cal_grad_downscale_in_infer(self, mask):
-            return mask.astype(self.in_type)
+            if self.in_type == np.uint16:
+                return convert_float_to_uint16(mask.astype("float32"))
+            else:
+                return mask.astype(self.in_type)
 
         def test_backward_downscale_in_infer(self):
             for place in self.places:
                 with base.dygraph.guard(place):
                     prob = 0.1
-                    input = paddle.uniform([100, 40], dtype=self.in_type)
+                    if self.in_type == np.uint16:
+                        input = paddle.uniform([100, 40], dtype="bfloat16")
+                    else:
+                        input = paddle.uniform([100, 40], dtype=self.in_type)
                     input.stop_gradient = False
                     out, mask = _legacy_C_ops.dropout(
                         input, 'dropout_prob', prob
@@ -197,7 +227,10 @@ def test_backward_upscale_train(self):
             for place in self.places:
                 with base.dygraph.guard(place):
                     prob = 0.5
-                    input = paddle.uniform([100, 40], dtype=self.in_type)
+                    if self.in_type == np.uint16:
+                        input = paddle.uniform([100, 40], dtype="bfloat16")
+                    else:
+                        input = paddle.uniform([100, 40], dtype=self.in_type)
                     input.stop_gradient = False
                     out, mask = _legacy_C_ops.dropout(
                         input,
@@ -221,7 +254,10 @@ def test_backward_upscale_train_2(self):
             for place in self.places:
                 with base.dygraph.guard(place):
                     prob = 0.2
-                    input = paddle.uniform([100, 40], dtype=self.in_type)
+                    if self.in_type == np.uint16:
+                        input = paddle.uniform([100, 40], dtype="bfloat16")
+                    else:
+                        input = paddle.uniform([100, 40], dtype=self.in_type)
                     input.stop_gradient = False
                     out, mask = _legacy_C_ops.dropout(
                         input,
diff --git a/tools/xpu/pack_paddle_depence.sh b/tools/xpu/pack_paddle_depence.sh
@@ -73,9 +73,9 @@ function xhpc_prepare() {
 
   check_files ${XHPC_DIR_NAME}/xfa/include/flash_api.h ${XHPC_DIR_NAME}/xfa/so/libxpu_flash_attention.so
 
-  # remove '#include "xpu/refactor/core/quant.h"' in flash_api.h
+  # remove '#include "xpu/flash_impl.h"' in flash_api.h
   # TODO(houj04): remove this hack when compile issue is resolved in XHPC
-  sed -i '8d' ${XHPC_DIR_NAME}/xfa/include/flash_api.h
+  sed -i '3d' ${XHPC_DIR_NAME}/xfa/include/flash_api.h
 
   cp -r ${XHPC_DIR_NAME}/xfa/include/* xpu/include/xhpc/xfa
   cp -r ${XHPC_DIR_NAME}/xfa/so/libxpu_flash_attention.so xpu/lib/

Original file line number	Diff line number	Diff line change
`@@ -525,18 +525,18 @@ def amp_guard(`
`525`	`525`	`if tracer._expected_place.is_xpu_place():`
`526`	`526`	`if (dtype == 'float16') and not _is_xpu_float16_supported():`
`527`	`527`	`xpu_verion = core.get_xpu_device_version(`
`528`		`- tracer._expected_place.get_device_id()`
	`528`	`+ _current_expected_place().get_device_id()`
`529`	`529`	`)`
`530`	`530`	`warnings.warn(`
`531`		`- '%d does not support float16 amp.' % xpu_verion`
	`531`	`+ 'XPU%d does not support float16 amp.' % xpu_verion`
`532`	`532`	`)`
`533`	`533`	`enable = False`
`534`	`534`	`elif (dtype == 'bfloat16') and not _is_xpu_bfloat16_supported():`
`535`	`535`	`xpu_verion = core.get_xpu_device_version(`
`536`		`- tracer._expected_place.get_device_id()`
	`536`	`+ _current_expected_place().get_device_id()`
`537`	`537`	`)`
`538`	`538`	`warnings.warn(`
`539`		`- '%d does not support bfloat16 amp.' % xpu_verion`
	`539`	`+ 'XPU%d does not support bfloat16 amp.' % xpu_verion`
`540`	`540`	`)`
`541`	`541`	`enable = False`
`542`	`542`	`# For custom device:`