[Auto Parallel] add dropout spmd rules (#70216)

deepllz · web-flow · commit 9fcf8c8b0f21 · 2024-12-16T20:09:57.000+08:00
diff --git a/paddle/phi/infermeta/spmd_rules/dropout.cc b/paddle/phi/infermeta/spmd_rules/dropout.cc
@@ -0,0 +1,60 @@
+/* Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    https://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/infermeta/spmd_rules/dropout.h"
+
+#include <vector>
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+#include "paddle/phi/infermeta/spmd_rules/elementwise.h"
+#include "paddle/phi/infermeta/spmd_rules/spmd_rule_macro_define.h"
+#include "paddle/phi/infermeta/spmd_rules/utils.h"
+
+namespace phi {
+namespace distributed {
+
+// args : (Tensor x, Tensor seed_tensor, Scalar p, bool is_test, str mode, int
+// seed, bool fix_seed) output : Tensor(out), Tensor(mask)
+SpmdInfo DropoutFwdInferSpmd(const DistMetaTensor& x,
+                             const DistMetaTensor& seed_tensor,
+                             Scalar p,
+                             bool is_test,
+                             const std::string& mode,
+                             int seed,
+                             bool fix_seed) {
+  std::vector<int64_t> seed_tensor_shape =
+      common::vectorize(seed_tensor.dims());
+  // seed_tensor is None currently
+  PADDLE_ENFORCE_EQ(
+      IsEmpty(seed_tensor_shape),
+      true,
+      common::errors::InvalidArgument("seed_tensor should be None"));
+  SpmdInfo info = ElementwiseUnaryInferSpmd(x);
+  info.first.push_back(seed_tensor.dist_attr());
+  info.second.push_back(info.second[0]);
+  return info;
+}
+
+// args : (Tensor mask, Tensor out_grad, Scalar p, bool is_test, str mode)
+// output : Tensor(x_grad)
+SpmdInfo DropoutBwdInferSpmd(const DistMetaTensor& mask,
+                             const DistMetaTensor& out_grad,
+                             Scalar p,
+                             bool is_test,
+                             const std::string& mode) {
+  return ElementwiseBinaryInferSpmd(mask, out_grad);
+}
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/dropout.h b/paddle/phi/infermeta/spmd_rules/dropout.h
@@ -0,0 +1,43 @@
+/* Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    https://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <string>
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+
+// args : (Tensor x, Tensor seed_tensor, Scalar p, bool is_test, str mode, int
+// seed, bool fix_seed) output : Tensor(out), Tensor(mask)
+SpmdInfo DropoutFwdInferSpmd(const DistMetaTensor& x,
+                             const DistMetaTensor& seed_tensor,
+                             Scalar p,
+                             bool is_test,
+                             const std::string& mode,
+                             int seed,
+                             bool fix_seed);
+
+// args : (Tensor mask, Tensor out_grad, Scalar p, bool is_test, str mode)
+// output : Tensor(x_grad)
+SpmdInfo DropoutBwdInferSpmd(const DistMetaTensor& mask,
+                             const DistMetaTensor& out_grad,
+                             Scalar p,
+                             bool is_test,
+                             const std::string& mode);
+// out: "x_grad", "gate_logits_grad"
+
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -26,6 +26,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/cross_entropy_with_softmax.h"
 #include "paddle/phi/infermeta/spmd_rules/cumsum.h"
 #include "paddle/phi/infermeta/spmd_rules/default_data_parallel.h"
+#include "paddle/phi/infermeta/spmd_rules/dropout.h"
 #include "paddle/phi/infermeta/spmd_rules/elementwise.h"
 #include "paddle/phi/infermeta/spmd_rules/embedding.h"
 #include "paddle/phi/infermeta/spmd_rules/expand_as.h"
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -810,6 +810,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [out_grad]
+    spmd_rule: DropoutBwdInferSpmd
   kernel :
     func : dropout_grad
   composite : dropout_grad(mask, out_grad, p, is_test, mode, x_grad)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1543,6 +1543,7 @@
   output : Tensor(out), Tensor(mask)
   infer_meta :
     func : DropoutInferMeta
+    spmd_rule: DropoutFwdInferSpmd
   kernel :
     func : dropout
     data_type : x
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
@@ -2260,6 +2260,38 @@ TEST(Conv2dGradSPMDRule, Ctor) {
   EXPECT_EQ(is_partial(infered_dist_attrs.second[1]), true);
 }
 
+TEST(Dropout, Ctor) {
+  std::vector<int64_t> mesh_shape = {2, 2};
+  std::vector<int64_t> process_ids = {0, 1, 2, 3};
+  std::vector<std::string> dim_names = {"dp", "mp"};
+  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
+
+  // test forward
+  auto t_dist_attr = TensorDistAttr();
+  t_dist_attr.set_process_mesh(process_mesh);
+  t_dist_attr.set_dims_mapping({0, -1, -1});
+  t_dist_attr.set_dynamic_dims({false, false, false});
+  phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
+      common::make_ddim({4, 6, 8}), t_dist_attr);
+  phi::distributed::DistMetaTensor seed_tensor;
+  phi::distributed::SpmdInfo forward_info =
+      phi::distributed::DropoutFwdInferSpmd(
+          x, seed_tensor, 0.5, false, "", 0, true);
+  check_dim_mapping(forward_info.first[0], {0, -1, -1});
+  check_dim_mapping(forward_info.first[1], {});
+  check_dim_mapping(forward_info.second[0], {0, -1, -1});
+  check_dim_mapping(forward_info.second[1], {0, -1, -1});
+  // test backward
+  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
+      common::make_ddim({4, 6, 8}), t_dist_attr);
+  phi::distributed::DistMetaTensor mask = out_grad;
+  phi::distributed::SpmdInfo backward_info =
+      phi::distributed::DropoutBwdInferSpmd(mask, out_grad, 0.5, false, "");
+  check_dim_mapping(backward_info.first[0], {0, -1, -1});
+  check_dim_mapping(backward_info.first[1], {0, -1, -1});
+  check_dim_mapping(backward_info.second[0], {0, -1, -1});
+}
+
 }  // namespace auto_parallel
 }  // namespace distributed
 }  // namespace paddle
