fix maximum_grad minimum_grad bug (#71716)

wanghuancoder · web-flow · commit 6e225b3f04c5 · 2025-03-20T14:19:30.000+08:00
* fix maximum_grad minimum_grad bug
diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -1413,9 +1413,28 @@ void maximum_grad(const Tensor& x,
                   const Tensor& out_grad,
                   Tensor* x_grad,
                   Tensor* y_grad) {
+  Tensor half_tensor;
+  Tensor out_grad_copy = out_grad;
+  if (x_grad || y_grad) {
+    // cast, because divide and add kernel is not support bf16 and fp16 on CPU
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      out_grad_copy = cast<T>(out_grad, phi::DataType::FLOAT32);
+    }
+    auto equal_tensor = cast<T>(equal<T>(x, y), out_grad_copy.dtype());
+    auto tmp_tensor =
+        full<T>({1}, 2.0, out_grad_copy.dtype(), out_grad_copy.place());
+    half_tensor = (out_grad_copy / tmp_tensor) * equal_tensor;
+  }
+
   if (x_grad) {
-    auto x_tmp = cast<T>(greater_than<T>(x, y), out_grad.dtype());
-    auto dx_res = out_grad * x_tmp;
+    auto x_tmp = cast<T>(greater_than<T>(x, y), out_grad_copy.dtype());
+    auto dx_res = out_grad_copy * x_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dx_res = cast<T>(dx_res, out_grad.dtype());
+    }
+
     if (out_grad.dims() != x.dims()) {
       // Maybe need reduce here
       auto reduce_dim = get_reduce_dims(x.dims(), out_grad.dims());
@@ -1438,8 +1457,12 @@ void maximum_grad(const Tensor& x,
   }
 
   if (y_grad) {
-    auto y_tmp = cast<T>(less_equal<T>(x, y), out_grad.dtype());
-    auto dy_res = out_grad * y_tmp;
+    auto y_tmp = cast<T>(less_than<T>(x, y), out_grad_copy.dtype());
+    auto dy_res = out_grad_copy * y_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dy_res = cast<T>(dy_res, out_grad.dtype());
+    }
     if (out_grad.dims() != y.dims()) {
       // Maybe need reduce here
       phi::DDim reduce_dim = get_reduce_dims(y.dims(), out_grad.dims());
@@ -1897,9 +1920,27 @@ void minimum_grad(const Tensor& x,
                   const Tensor& out_grad,
                   Tensor* x_grad,
                   Tensor* y_grad) {
+  Tensor half_tensor;
+  Tensor out_grad_copy = out_grad;
+  if (x_grad || y_grad) {
+    // cast, because divide and add kernel is not support bf16 and fp16 on CPU
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      out_grad_copy = cast<T>(out_grad, phi::DataType::FLOAT32);
+    }
+    auto equal_tensor = cast<T>(equal<T>(x, y), out_grad_copy.dtype());
+    auto tmp_tensor =
+        full<T>({1}, 2.0, out_grad_copy.dtype(), out_grad_copy.place());
+    half_tensor = (out_grad_copy / tmp_tensor) * equal_tensor;
+  }
+
   if (x_grad) {
-    auto x_tmp = cast<T>(less_than<T>(x, y), out_grad.dtype());
-    auto dx_res = out_grad * x_tmp;
+    auto x_tmp = cast<T>(less_than<T>(x, y), out_grad_copy.dtype());
+    auto dx_res = out_grad_copy * x_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dx_res = cast<T>(dx_res, out_grad.dtype());
+    }
     if (out_grad.dims() != x.dims()) {
       // Maybe need reduce here
       auto reduce_dim = get_reduce_dims(x.dims(), out_grad.dims());
@@ -1922,8 +1963,12 @@ void minimum_grad(const Tensor& x,
   }
 
   if (y_grad) {
-    auto y_tmp = cast<T>(greater_equal<T>(x, y), out_grad.dtype());
-    auto dy_res = out_grad * y_tmp;
+    auto y_tmp = cast<T>(greater_than<T>(x, y), out_grad_copy.dtype());
+    auto dy_res = out_grad_copy * y_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dy_res = cast<T>(dy_res, out_grad.dtype());
+    }
     if (out_grad.dims() != y.dims()) {
       // Maybe need reduce here
       phi::DDim reduce_dim = get_reduce_dims(y.dims(), out_grad.dims());
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -1217,9 +1217,27 @@ void maximum_grad(const Tensor& x,
                   const Tensor& out_grad,
                   Tensor* x_grad,
                   Tensor* y_grad) {
+  Tensor half_tensor;
+  Tensor out_grad_copy = out_grad;
+  if (x_grad || y_grad) {
+    // cast, because divide and add kernel is not support bf16 and fp16 on CPU
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      out_grad_copy = cast<T>(out_grad, phi::DataType::FLOAT32);
+    }
+    auto equal_tensor = cast<T>(equal<T>(x, y), out_grad_copy.dtype());
+    auto tmp_tensor =
+        full<T>({1}, 2.0, out_grad_copy.dtype(), out_grad_copy.place());
+    half_tensor = (out_grad_copy / tmp_tensor) * equal_tensor;
+  }
+
   if (x_grad) {
-    auto x_tmp = cast<T>(greater_than<T>(x, y), out_grad.dtype());
-    auto dx_res = out_grad * x_tmp;
+    auto x_tmp = cast<T>(greater_than<T>(x, y), out_grad_copy.dtype());
+    auto dx_res = out_grad_copy * x_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dx_res = cast<T>(dx_res, out_grad.dtype());
+    }
     if (has_dynamic_shape(x.shape()) || has_dynamic_shape(out_grad.shape()) ||
         out_grad.dims() != x.dims()) {
       auto dx_reduce_res = reduce_as<T>(dx_res, x);
@@ -1230,8 +1248,12 @@ void maximum_grad(const Tensor& x,
   }
 
   if (y_grad) {
-    auto y_tmp = cast<T>(less_equal<T>(x, y), out_grad.dtype());
-    auto dy_res = out_grad * y_tmp;
+    auto y_tmp = cast<T>(less_than<T>(x, y), out_grad_copy.dtype());
+    auto dy_res = out_grad_copy * y_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dy_res = cast<T>(dy_res, out_grad.dtype());
+    }
     if (has_dynamic_shape(y.shape()) || has_dynamic_shape(out_grad.shape()) ||
         out_grad.dims() != y.dims()) {
       auto dy_reduce_res = reduce_as<T>(dy_res, y);
@@ -2243,9 +2265,27 @@ void minimum_grad(const Tensor& x,
                   const Tensor& out_grad,
                   Tensor* x_grad,
                   Tensor* y_grad) {
+  Tensor half_tensor;
+  Tensor out_grad_copy = out_grad;
+  if (x_grad || y_grad) {
+    // cast, because divide and add kernel is not support bf16 and fp16 on CPU
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      out_grad_copy = cast<T>(out_grad, phi::DataType::FLOAT32);
+    }
+    auto equal_tensor = cast<T>(equal<T>(x, y), out_grad_copy.dtype());
+    auto tmp_tensor =
+        full<T>({1}, 2.0, out_grad_copy.dtype(), out_grad_copy.place());
+    half_tensor = (out_grad_copy / tmp_tensor) * equal_tensor;
+  }
+
   if (x_grad) {
-    auto x_tmp = cast<T>(less_than<T>(x, y), out_grad.dtype());
-    auto dx_res = out_grad * x_tmp;
+    auto x_tmp = cast<T>(less_than<T>(x, y), out_grad_copy.dtype());
+    auto dx_res = out_grad_copy * x_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dx_res = cast<T>(dx_res, out_grad.dtype());
+    }
     if (has_dynamic_shape(x.shape()) || has_dynamic_shape(out_grad.shape()) ||
         out_grad.dims() != x.dims()) {
       auto dx_reduce_res = reduce_as<T>(dx_res, x);
@@ -2256,8 +2296,12 @@ void minimum_grad(const Tensor& x,
   }
 
   if (y_grad) {
-    auto y_tmp = cast<T>(greater_equal<T>(x, y), out_grad.dtype());
-    auto dy_res = out_grad * y_tmp;
+    auto y_tmp = cast<T>(greater_than<T>(x, y), out_grad_copy.dtype());
+    auto dy_res = out_grad_copy * y_tmp + half_tensor;
+    if (out_grad.dtype() == phi::DataType::BFLOAT16 ||
+        out_grad.dtype() == phi::DataType::FLOAT16) {
+      dy_res = cast<T>(dy_res, out_grad.dtype());
+    }
     if (has_dynamic_shape(y.shape()) || has_dynamic_shape(out_grad.shape()) ||
         out_grad.dims() != y.dims()) {
       auto dy_reduce_res = reduce_as<T>(dy_res, y);
diff --git a/paddle/phi/kernels/funcs/elementwise_functor.h b/paddle/phi/kernels/funcs/elementwise_functor.h
@@ -477,14 +477,16 @@ struct MaximumFunctor {
 template <typename T>
 struct MaxGradXFunctor {
   inline HOSTDEVICE T operator()(const T x, const T y, const T dout) const {
-    return dout * static_cast<T>(x > y);
+    return dout * static_cast<T>(x > y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 
 template <typename T>
 struct MaxGradYFunctor {
   inline HOSTDEVICE T operator()(const T x, const T y, const T dout) const {
-    return dout * static_cast<T>(x <= y);
+    return dout * static_cast<T>(x < y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 
@@ -494,10 +496,14 @@ struct MaxGradXYFunctor {
                                                    const InT y,
                                                    const InT dout) {
     phi::Array<OutT, 2> outs;
-    // dx = dout * (x > y)
-    outs[0] = static_cast<OutT>(dout * static_cast<InT>(x > y));
-    // dy = dout * (x <= y)
-    outs[1] = static_cast<OutT>(dout * static_cast<InT>(x <= y));
+    // dx = dout * (x > y) + dout / 2 * (x == y)
+    outs[0] = static_cast<OutT>(dout * static_cast<InT>(x > y) +
+                                (dout / static_cast<InT>(2)) *
+                                    static_cast<InT>(x == y));
+    // dy = dout * (x < y) + dout / 2 * (x == y)
+    outs[1] = static_cast<OutT>(dout * static_cast<InT>(x < y) +
+                                (dout / static_cast<InT>(2)) *
+                                    static_cast<InT>(x == y));
     return outs;
   }
 };
@@ -512,13 +518,15 @@ struct MinimumFunctor {
 template <typename T>
 struct MinGradXFunctor {
   inline HOSTDEVICE T operator()(const T x, const T y, const T dout) const {
-    return dout * static_cast<T>(x < y);
+    return dout * static_cast<T>(x < y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 template <typename T>
 struct MinGradYFunctor {
   inline HOSTDEVICE T operator()(const T x, const T y, const T dout) const {
-    return dout * static_cast<T>(x >= y);
+    return dout * static_cast<T>(x > y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 
@@ -528,10 +536,14 @@ struct MinGradXYFunctor {
                                                    const InT y,
                                                    const InT dout) {
     phi::Array<OutT, 2> outs;
-    // dx = dout * (x < y)
-    outs[0] = static_cast<OutT>(dout * static_cast<InT>(x < y));
-    // dy = dout * (x >= y)
-    outs[1] = static_cast<OutT>(dout * static_cast<InT>(x >= y));
+    // dx = dout * (x < y) + dout / 2 * (x == y)
+    outs[0] = static_cast<OutT>(dout * static_cast<InT>(x < y) +
+                                (dout / static_cast<InT>(2)) *
+                                    static_cast<InT>(x == y));
+    // dy = dout * (x > y) + dout / 2 * (x == y)
+    outs[1] = static_cast<OutT>(dout * static_cast<InT>(x > y) +
+                                (dout / static_cast<InT>(2)) *
+                                    static_cast<InT>(x == y));
     return outs;
   }
 };
diff --git a/paddle/phi/kernels/impl/elementwise_grad_kernel_impl.h b/paddle/phi/kernels/impl/elementwise_grad_kernel_impl.h
@@ -1305,14 +1305,16 @@ void MultiplyTripleGradKernel(const Context& dev_ctx,
 template <typename T>
 struct MaxGradDx {
   HOSTDEVICE T operator()(T x, T y, T out UNUSED, T dout) const {
-    return dout * static_cast<T>(x > y);
+    return dout * static_cast<T>(x > y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 
 template <typename T>
 struct MaxGradDy {
   HOSTDEVICE T operator()(T x, T y, T out UNUSED, T dout) const {
-    return dout * static_cast<T>(x <= y);
+    return dout * static_cast<T>(x < y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 
@@ -1324,14 +1326,16 @@ struct MaxGradDy {
 template <typename T>
 struct MinGradDx {
   HOSTDEVICE T operator()(T x, T y, T out UNUSED, T dout) const {
-    return dout * static_cast<T>(x < y);
+    return dout * static_cast<T>(x < y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 
 template <typename T>
 struct MinGradDy {
   HOSTDEVICE T operator()(T x, T y, T out UNUSED, T dout) const {
-    return dout * static_cast<T>(x >= y);
+    return dout * static_cast<T>(x > y) +
+           (dout / static_cast<T>(2)) * static_cast<T>(x == y);
   }
 };
 
diff --git a/test/legacy_test/test_elementwise_max_op.py b/test/legacy_test/test_elementwise_max_op.py
@@ -370,5 +370,22 @@ def setUp(self):
         self.outputs = {'Out': np.maximum(self.inputs['X'], self.inputs['Y'])}
 
 
+class TestElementwiseOpEqualInput(TestElementwiseOp):
+    def init_data(self):
+        self.x = np.ones([13, 17]).astype(np.float32)
+        self.y = np.ones([13, 17]).astype(np.float32)
+
+    def setUp(self):
+        self.init_data()
+        self.op_type = "elementwise_max"
+        self.prim_op_type = "prim"
+        self.if_enable_cinn()
+        self.python_api = paddle.maximum
+        self.dtype = np.float32
+        self.public_python_api = paddle.maximum
+        self.inputs = {'X': self.x, 'Y': self.y}
+        self.outputs = {'Out': np.maximum(self.inputs['X'], self.inputs['Y'])}
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/legacy_test/test_elementwise_min_op.py b/test/legacy_test/test_elementwise_min_op.py
@@ -451,5 +451,22 @@ def init_data(self):
         )
 
 
+class TestElementwiseOpEqualInput(TestElementwiseOp):
+    def init_data(self):
+        self.x = np.ones([13, 17]).astype(np.float32)
+        self.y = np.ones([13, 17]).astype(np.float32)
+
+    def setUp(self):
+        self.init_data()
+        self.op_type = "elementwise_min"
+        self.prim_op_type = "prim"
+        self.if_enable_cinn()
+        self.python_api = paddle.minimum
+        self.dtype = np.float32
+        self.public_python_api = paddle.minimum
+        self.inputs = {'X': self.x, 'Y': self.y}
+        self.outputs = {'Out': np.minimum(self.inputs['X'], self.inputs['Y'])}
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/legacy_test/test_maximum_op.py b/test/legacy_test/test_maximum_op.py
@@ -120,6 +120,50 @@ def test_dynamic_api(self):
         res = res.numpy()
         np.testing.assert_allclose(res, self.np_expected4, rtol=1e-05)
 
+    @unittest.skipIf(
+        core.is_compiled_with_xpu(),
+        "XPU need fix the bug",
+    )
+    def test_equal_tensors(self):
+        numpy_tensor = np.ones([10000]).astype("float32")
+        paddle_x = paddle.to_tensor(numpy_tensor)
+        paddle_x.stop_gradient = False
+        numpy_tensor = np.ones([10000]).astype("float32")
+        paddle_x2 = paddle.to_tensor(numpy_tensor)
+        paddle_x2.stop_gradient = False
+
+        numpy_tensor = np.ones([10000]).astype("float32")
+        paddle_outgrad = paddle.to_tensor(numpy_tensor)
+
+        paddle_out = paddle.maximum(paddle_x, paddle_x2)
+        paddle_x_grad, paddle_x2_grad = paddle.grad(
+            [paddle_out],
+            [paddle_x, paddle_x2],
+            grad_outputs=[paddle_outgrad],
+            allow_unused=True,
+        )
+
+        np.testing.assert_allclose(
+            paddle_out.numpy(),
+            numpy_tensor,
+            1e-2,
+            1e-2,
+        )
+
+        np.testing.assert_allclose(
+            paddle_x_grad.numpy(),
+            numpy_tensor * 0.5,
+            1e-2,
+            1e-2,
+        )
+
+        np.testing.assert_allclose(
+            paddle_x2_grad.numpy(),
+            numpy_tensor * 0.5,
+            1e-2,
+            1e-2,
+        )
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/legacy_test/test_minimum_op.py b/test/legacy_test/test_minimum_op.py
