From 5df1296d3e03a4b54545e0d7deaec9b2997bc2f9 Mon Sep 17 00:00:00 2001
From: Chitsing KUI <kuizhiqing@msn.com>
Date: Sun, 9 Apr 2023 12:59:54 +0800
Subject: [PATCH] fix fused_dropout_add bug (#52644)

---
 paddle/phi/api/yaml/fused_backward.yaml             |  5 +++--
 paddle/phi/infermeta/backward.cc                    | 13 +++++++++++++
 paddle/phi/infermeta/backward.h                     |  5 +++++
 .../kernels/fusion/gpu/fused_dropout_add_utils.h    |  2 +-
 python/paddle/distributed/auto_parallel/utils.py    |  2 ++
 5 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/paddle/phi/api/yaml/fused_backward.yaml b/paddle/phi/api/yaml/fused_backward.yaml
index 95914828461..bdea08a1dd6 100644
--- a/paddle/phi/api/yaml/fused_backward.yaml
+++ b/paddle/phi/api/yaml/fused_backward.yaml
@@ -9,8 +9,9 @@
   args : (Tensor seed_offset, Tensor out_grad, Scalar p, bool is_test, str mode, bool fix_seed)
   output : Tensor(x_grad), Tensor(y_grad)
   infer_meta :
-    func : GeneralBinaryGradInferMeta
-    param : [out_grad, out_grad]
+    func : FusedDropoutAddGradInferMeta
+    param : [seed_offset, out_grad]
   kernel :
     func : fused_dropout_add_grad
+    data_type : out_grad
   support_dygraph_mode : true
diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
index 971591f0384..185cc09260d 100644
--- a/paddle/phi/infermeta/backward.cc
+++ b/paddle/phi/infermeta/backward.cc
@@ -215,6 +215,19 @@ void FlashAttnGradInferMeta(const MetaTensor& q,
   }
 }
 
+void FusedDropoutAddGradInferMeta(const MetaTensor& seed_offset,
+                                  const MetaTensor& out_grad,
+                                  MetaTensor* x_grad,
+                                  MetaTensor* y_grad) {
+  if (x_grad != nullptr) {
+    x_grad->share_meta(out_grad);
+  }
+
+  if (y_grad != nullptr) {
+    y_grad->share_meta(out_grad);
+  }
+}
+
 void CrossEntropyWithSoftmaxGradInferMeta(const MetaTensor& label,
                                           const MetaTensor& softmax,
                                           const MetaTensor& loss_grad,
diff --git a/paddle/phi/infermeta/backward.h b/paddle/phi/infermeta/backward.h
index 6abb8d46c33..b95d164e748 100644
--- a/paddle/phi/infermeta/backward.h
+++ b/paddle/phi/infermeta/backward.h
@@ -179,6 +179,11 @@ void FlashAttnGradInferMeta(const MetaTensor& q,
                             MetaTensor* dk,
                             MetaTensor* dv);
 
+void FusedDropoutAddGradInferMeta(const MetaTensor& seed_offset,
+                                  const MetaTensor& out_grad,
+                                  MetaTensor* x_grad,
+                                  MetaTensor* y_grad);
+
 void GatherNdGradInferMeta(const MetaTensor& x,
                            const MetaTensor& index,
                            const MetaTensor& out_grad,
diff --git a/paddle/phi/kernels/fusion/gpu/fused_dropout_add_utils.h b/paddle/phi/kernels/fusion/gpu/fused_dropout_add_utils.h
index 615bb9d868e..83d180b42ae 100644
--- a/paddle/phi/kernels/fusion/gpu/fused_dropout_add_utils.h
+++ b/paddle/phi/kernels/fusion/gpu/fused_dropout_add_utils.h
@@ -20,7 +20,7 @@ namespace phi {
 namespace fusion {
 
 template <typename Context>
-static inline std::vector<size_t> GetRandomCudaProp(int numel,
+static inline std::vector<size_t> GetRandomCudaProp(int64_t numel,
                                                     const Context& dev_ctx) {
   constexpr int kVecSize = funcs::uniform_distribution<float>::kReturnsCount;
   auto gpu_config =
diff --git a/python/paddle/distributed/auto_parallel/utils.py b/python/paddle/distributed/auto_parallel/utils.py
index c248b9957a4..153ac060e59 100644
--- a/python/paddle/distributed/auto_parallel/utils.py
+++ b/python/paddle/distributed/auto_parallel/utils.py
@@ -1260,6 +1260,7 @@ def set_grad_var_shape(program, dist_context):
                 "exp_grad",
                 "sigmoid_grad",
                 "unsqueeze2_grad",
+                "fused_dropout_add_grad",
             ]
             forward_list = [
                 "reshape2",
@@ -1281,6 +1282,7 @@ def set_grad_var_shape(program, dist_context):
                 "exp",
                 "sigmoid",
                 "unsqueeze2",
+                "fused_dropout_add",
             ]
             if op.type in need_set_shape_list:
                 for forward_op in block.ops:
-- 
GitLab