[Prim] Fix silu backward prim bug (#56280)

cyber-pioneer · web-flow · commit 08dbb33c7f2b · 2023-08-16T16:54:01.000+08:00
* fix silu backward prim bug

* fix silu double grad prim
diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -66,9 +66,21 @@ void silu_grad(const Tensor& x,
                const Tensor& out_grad,
                Tensor* x_grad) {
   if (x_grad) {
-    auto sigmoid = out / x;
-    auto res = out_grad * sigmoid * (1.0 + x * (1.0 - sigmoid));
-    set_output<T>(res, x_grad);
+    auto org_dtype = x.dtype();
+    bool need_cast = org_dtype == phi::DataType::FLOAT16 ||
+                     org_dtype == phi::DataType::BFLOAT16;
+    if (need_cast) {
+      auto x_cast = cast<T>(x, phi::DataType::FLOAT32);
+      auto out_cast = cast<T>(out, phi::DataType::FLOAT32);
+      auto out_grad_cast = cast<T>(out_grad, phi::DataType::FLOAT32);
+      auto sigmoid = 1.0 / (1.0 + exp<T>(-x_cast));
+      auto res = out_grad_cast * sigmoid * (1.0 + x_cast - out_cast);
+      set_output<T>(cast<T>(res, org_dtype), x_grad);
+    } else {
+      auto sigmoid = 1.0 / (1.0 + exp<T>(-x));
+      auto res = out_grad * sigmoid * (1.0 + x - out);
+      set_output<T>(res, x_grad);
+    }
   }
 }
 
diff --git a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h
@@ -439,16 +439,15 @@ void silu_double_grad(const Tensor& x,
                       const Tensor& grad_x_grad,
                       Tensor* grad_x,
                       Tensor* grad_out_grad) {
-  auto sigmoid = out / x;
+  auto sigmoid = 1 / (1 + exp<T>(-x));
   auto tmp1 = 1 - sigmoid;
   auto tmp2 = 1 + tmp1 * x;
   if (grad_out_grad) {
     auto ddout = grad_x_grad * sigmoid * tmp2;
     set_output<T>(ddout, grad_out_grad);
   }
   if (grad_x) {
-    auto dx =
-        sigmoid * grad_x_grad * out_grad * (1 + (tmp2 - x * sigmoid)) * tmp1;
+    auto dx = sigmoid * grad_x_grad * out_grad * (1 + (tmp2 - out)) * tmp1;
     set_output<T>(dx, grad_x);
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -439,16 +439,15 @@ void silu_double_grad(const Tensor& x,`
`439`	`439`	`const Tensor& grad_x_grad,`
`440`	`440`	`Tensor* grad_x,`
`441`	`441`	`Tensor* grad_out_grad) {`
`442`		`- auto sigmoid = out / x;`
	`442`	`+ auto sigmoid = 1 / (1 + exp<T>(-x));`
`443`	`443`	`auto tmp1 = 1 - sigmoid;`
`444`	`444`	`auto tmp2 = 1 + tmp1 * x;`
`445`	`445`	`if (grad_out_grad) {`
`446`	`446`	`auto ddout = grad_x_grad * sigmoid * tmp2;`
`447`	`447`	`set_output<T>(ddout, grad_out_grad);`
`448`	`448`	`}`
`449`	`449`	`if (grad_x) {`
`450`		`- auto dx =`
`451`		`- sigmoid * grad_x_grad * out_grad * (1 + (tmp2 - x * sigmoid)) * tmp1;`
	`450`	`+ auto dx = sigmoid * grad_x_grad * out_grad * (1 + (tmp2 - out)) * tmp1;`
`452`	`451`	`set_output<T>(dx, grad_x);`
`453`	`452`	`}`
`454`	`453`	`}`