[Accuracy diff No.20、61、9-10] Fix accuracy diff for expand_grad API (#72992)

ooooo-create · web-flow · commit 461166c6b1b7 · 2025-06-12T11:24:22.000+08:00
* fix

* fix
diff --git a/paddle/phi/kernels/impl/expand_grad_kernel_impl.h b/paddle/phi/kernels/impl/expand_grad_kernel_impl.h
@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/cast_kernel.h"
 #include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
@@ -32,19 +33,50 @@ void ExpandBackward(const Context& dev_ctx,
   dev_ctx.template Alloc<T>(in_grad);
   in_grad->data<T>();
 
-  auto x_grad = EigenVector<T>::Flatten(*in_grad);
-  Eigen::DSizes<Eigen::DenseIndex, Dims * 2> reshape_dims;
-  for (size_t i = 0; i < reshape_size; ++i) {
-    reshape_dims[i] = reshape_dims_vec[i];
-  }
-  Eigen::DSizes<Eigen::DenseIndex, Dims> reduce_dims;
-  for (size_t i = 0; i < reduce_size; ++i) {
-    reduce_dims[i] = reduce_dims_vec[i];
+  if constexpr (std::is_same_v<T, dtype::float16> ||
+                std::is_same_v<T, dtype::bfloat16>) {
+    const DenseTensor out_grad_fp32 =
+        phi::Cast<T, Context>(dev_ctx, out_grad, DataType::FLOAT32);
+    DenseTensor in_grad_fp32;
+    in_grad_fp32.Resize(in_grad->dims());
+    dev_ctx.template Alloc<float>(&in_grad_fp32);
+
+    auto x_grad = EigenVector<float>::Flatten(in_grad_fp32);
+    Eigen::DSizes<Eigen::DenseIndex, Dims * 2> reshape_dims;
+    for (size_t i = 0; i < reshape_size; ++i) {
+      reshape_dims[i] = reshape_dims_vec[i];
+    }
+    Eigen::DSizes<Eigen::DenseIndex, Dims> reduce_dims;
+    for (size_t i = 0; i < reduce_size; ++i) {
+      reduce_dims[i] = reduce_dims_vec[i];
+    }
+    const auto out_grad0 = EigenVector<float>::Flatten(out_grad_fp32);
+    auto& place = *dev_ctx.eigen_device();
+    phi::funcs::EigenBroadcastGrad<std::decay_t<decltype(place)>, float, Dims>::
+        Eval(place, x_grad, out_grad0, reduce_dims, reshape_dims);
+
+    if constexpr (std::is_same_v<T, dtype::float16>) {
+      phi::CastKernel<float, Context>(
+          dev_ctx, in_grad_fp32, DataType::FLOAT16, in_grad);
+    } else {
+      phi::CastKernel<float, Context>(
+          dev_ctx, in_grad_fp32, DataType::BFLOAT16, in_grad);
+    }
+  } else {
+    auto x_grad = EigenVector<T>::Flatten(*in_grad);
+    Eigen::DSizes<Eigen::DenseIndex, Dims * 2> reshape_dims;
+    for (size_t i = 0; i < reshape_size; ++i) {
+      reshape_dims[i] = reshape_dims_vec[i];
+    }
+    Eigen::DSizes<Eigen::DenseIndex, Dims> reduce_dims;
+    for (size_t i = 0; i < reduce_size; ++i) {
+      reduce_dims[i] = reduce_dims_vec[i];
+    }
+    auto out_grad0 = EigenVector<T>::Flatten(out_grad);
+    auto& place = *dev_ctx.eigen_device();
+    phi::funcs::EigenBroadcastGrad<std::decay_t<decltype(place)>, T, Dims>::
+        Eval(place, x_grad, out_grad0, reduce_dims, reshape_dims);
   }
-  auto out_grad0 = EigenVector<T>::Flatten(out_grad);
-  auto& place = *dev_ctx.eigen_device();
-  phi::funcs::EigenBroadcastGrad<std::decay_t<decltype(place)>, T, Dims>::Eval(
-      place, x_grad, out_grad0, reduce_dims, reshape_dims);
 }
 
 template <typename T, typename Context>
