[XPU] fix bug in sigmoid_cross_entropy_with_logits and its grad (#70974)

houj04 · web-flow · commit 073b7882c83e · 2025-02-05T15:52:28.000+08:00
* [XPU] fix bug in sigmoid_cross_entropy_with_logits and its grad

* fix ut.

* update xhpc version
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 set(XPU_XPUDNN_LIB_NAME "libxpu_dnn.so")
 
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "dev/20250114")
+  set(XPU_XHPC_BASE_DATE "dev/20250204")
 endif()
 set(XPU_XCCL_BASE_VERSION "3.0.1.6") # For XRE5
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/kernels/xpu/flash_attn_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_kernel.cc
@@ -150,7 +150,14 @@ void FlashAttnKernelBase(const Context& ctx,
                              {},         // alibi_slopes_shape
                              -1,         // window_size_left
                              -1,         // window_size_right
-                             -1          // v_head_dim
+                             -1,         // v_head_dim
+                             nullptr,    // downstart_row_indices_data
+                             nullptr,    // downend_row_indices_data
+                             nullptr,    // upstart_row_indices_data
+                             nullptr,    // upend_row_indices_data
+                             0,          // flash_mask_head_num
+                             nullptr,    // flashmask_maxmin
+                             nullptr     // side_stream
       );
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_fwd");
 }
diff --git a/paddle/phi/kernels/xpu/sigmoid_cross_entropy_with_logits_grad_kernel.cc b/paddle/phi/kernels/xpu/sigmoid_cross_entropy_with_logits_grad_kernel.cc
@@ -51,20 +51,19 @@ void SigmoidCrossEntropyWithLogitsGradKernel(
   auto pos_weight_data =
       (pos_weight.get_ptr() == nullptr ? nullptr
                                        : pos_weight.get_ptr()->data<T>());
-  // int sigmoid_cross_entropy_with_logits_grad(Context* ctx, const T* x, const
-  // T* label, const T* dy, T* dx, int64_t m, int64_t n, TH* hit = nullptr,
-  // int64_t ignore_index = -100, const T* pos_weight = nullptr);
-  int r = xpu::sigmoid_cross_entropy_with_logits_grad(
+  // int paddle_sigmoid_cross_entropy_with_logits_grad(Context* ctx, const T* x,
+  // const T* label, const T* pos_weight, const T* dy, T* dx, int* hit, int
+  // ignore_index, int64_t n);
+  int r = xpu::paddle_sigmoid_cross_entropy_with_logits_grad(
       dev_ctx.x_context(),
       reinterpret_cast<const XPUType*>(x.data<T>()),
       reinterpret_cast<const XPUType*>(label.data<T>()),
+      reinterpret_cast<const XPUType*>(pos_weight_data),
       reinterpret_cast<const XPUType*>(out_grad.data<T>()),
       reinterpret_cast<XPUType*>(in_grad->data<T>()),
-      1,
-      x.numel(),
       hit,
       ignore_index,
-      reinterpret_cast<const XPUType*>(pos_weight_data));
+      x.numel());
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "sigmoid_cross_entropy_with_logits");
   if (normalize) {
     int* non_zero = RAII_GUARD.alloc_l3_or_gm<int>(1);
diff --git a/paddle/phi/kernels/xpu/sigmoid_cross_entropy_with_logits_kernel.cc b/paddle/phi/kernels/xpu/sigmoid_cross_entropy_with_logits_kernel.cc
@@ -48,19 +48,18 @@ void SigmoidCrossEntropyWithLogitsKernel(
   auto pos_weight_data =
       (pos_weight.get_ptr() == nullptr ? nullptr
                                        : pos_weight.get_ptr()->data<T>());
-  // int sigmoid_cross_entropy_with_logits(Context* ctx, const T* x, const T*
-  // label, T* y, int64_t m, int64_t n, TH* hit = nullptr, int64_t ignore_index
-  // = -100, const T* pos_weight = nullptr);
-  int r = xpu::sigmoid_cross_entropy_with_logits(
+  // int paddle_sigmoid_cross_entropy_with_logits(Context* ctx, const T* x,
+  // const T* label, const T* pos_weight, T* y, int* hit, int ignore_index,
+  // int64_t n);
+  int r = xpu::paddle_sigmoid_cross_entropy_with_logits(
       dev_ctx.x_context(),
       reinterpret_cast<const XPUType*>(x.data<T>()),
       reinterpret_cast<const XPUType*>(label.data<T>()),
+      reinterpret_cast<const XPUType*>(pos_weight_data),
       reinterpret_cast<XPUType*>(out->data<T>()),
-      1,
-      x.numel(),
       hit,
       ignore_index,
-      reinterpret_cast<const XPUType*>(pos_weight_data));
+      x.numel());
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "sigmoid_cross_entropy_with_logits");
   if (normalize) {
     int* non_zero = RAII_GUARD.alloc_l3_or_gm<int>(1);
diff --git a/test/legacy_test/test_sigmoid_cross_entropy_with_logits_grad_with_auto_grad.py b/test/legacy_test/test_sigmoid_cross_entropy_with_logits_grad_with_auto_grad.py
@@ -27,14 +27,19 @@ def setUp(self):
         np.random.seed(2023)
         paddle.seed(2023)
         self.places = []
-        if (
-            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
-            in ['1', 'true', 'on']
-            or not base.core.is_compiled_with_cuda()
+        if os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower() in [
+            '1',
+            'true',
+            'on',
+        ] or (
+            not base.core.is_compiled_with_cuda()
+            and not base.core.is_compiled_with_xpu()
         ):
             self.places.append(base.CPUPlace())
         if base.core.is_compiled_with_cuda():
             self.places.append(base.CUDAPlace(0))
+        if base.core.is_compiled_with_xpu():
+            self.places.append(base.XPUPlace(0))
         self.batch_size = 64
         self.num_classes = 20
 
@@ -85,7 +90,10 @@ def cal(fn, place):
             if idx == 0:
                 paddle.set_device('cpu')
             else:
-                paddle.set_device('gpu')
+                if base.core.is_compiled_with_cuda():
+                    paddle.set_device('gpu')
+                if base.core.is_compiled_with_xpu():
+                    paddle.set_device('xpu')
 
             ref = cal(fn_ref, p)
             actual = cal(fn_comp, p)
diff --git a/test/xpu/test_sigmoid_cross_entropy_with_logits_op_xpu.py b/test/xpu/test_sigmoid_cross_entropy_with_logits_op_xpu.py
@@ -266,14 +266,13 @@ def set_inputs(self):
         def set_output(self):
             # Fw Pass is implemented as elementwise sigmoid followed by
             # elementwise logistic loss
-            # Label * -log(sigmoid(X)) + (1 - label) * -log(1 - sigmoid(X))
-            term1 = np.maximum(self.inputs['X'], 0)
-            term2 = self.inputs['X'] * self.inputs['Label']
-            term3 = (
-                np.log(1 + np.exp(-1 * np.abs(self.inputs['X'])))
-                * self.inputs['pos_weight']
+            max_val = np.clip(-self.inputs['X'], 0, np.finfo(np.float64).max)
+            term1 = (1 - self.inputs['Label']) * self.inputs['X']
+            term2 = np.log(
+                np.exp(-max_val) + np.exp(-self.inputs['X'] - max_val)
             )
-            self.outputs = {'Out': term1 - term2 + term3}
+            out = term1 + self.inputs['pos_weight'] * (term2 + max_val)
+            self.outputs = {'Out': out}
 
     class TestSigmoidCrossEntropyWithLogitsNorm(
         TestSigmoidCrossEntropyWithLogitsOp
diff --git a/tools/xpu/disable_ut_xpu_kl3.local b/tools/xpu/disable_ut_xpu_kl3.local
@@ -43,7 +43,6 @@ test_auto_checkpoint_multiple
 test_auto_recompute_dy2static
 test_autograd_functional_dynamic
 test_backward_without_params
-test_binary_cross_entropy_with_logits_op
 test_c_comm_init_op
 test_c_concat
 test_c_split
