[auto parallel] add embedding backward spmd rules (#59003)

cxxly · web-flow · commit 7331bf95ef13 · 2023-11-29T10:45:43.000+08:00
* add embedding backward spmd rules

* [auto parallel] add dense embedding_grad spmd rule.

* fix embedding out_grad is reshard bugs
diff --git a/paddle/phi/api/lib/api_custom_impl.cc b/paddle/phi/api/lib/api_custom_impl.cc
@@ -264,12 +264,17 @@ void embedding_grad_impl(const Tensor& x,
       auto meta_dist_input_x = MakeDistMetaTensor(*x.impl());
       auto meta_dist_input_weight = MakeDistMetaTensor(*weight.impl());
       auto meta_dist_input_out_grad = MakeDistMetaTensor(*out_grad.impl());
-      auto spmd_info = phi::distributed::VariadicReplicatedInferSpmdDynamic(
-          meta_dist_input_weight, meta_dist_input_x, meta_dist_input_out_grad);
+      auto spmd_info =
+          phi::distributed::EmbeddingGradInferSpmd(meta_dist_input_x,
+                                                   meta_dist_input_weight,
+                                                   meta_dist_input_out_grad,
+                                                   padding_idx,
+                                                   sparse);
 
       // 2. Create Temporary Output & Prepare Dist and Dense Output
       std::shared_ptr<phi::distributed::DistTensor> shared_dist_out =
-          CreateKernelDistOutput(weight_grad, !rank_is_in_current_mesh);
+          CreateKernelDistOutput(
+              weight_grad, !rank_is_in_current_mesh, spmd_info.second[0]);
       phi::distributed::DistTensor* dist_out = shared_dist_out.get();
       phi::DenseTensor* dense_out = dist_out->unsafe_mutable_value();
       if (dense_out && !rank_is_in_current_mesh && !dist_out->defined()) {
@@ -284,37 +289,29 @@ void embedding_grad_impl(const Tensor& x,
       UnchangedInferMeta(MakeMetaTensor(*weight.impl()), &meta_dist_out);
 
       // 4. Set Output Dist Attr For Default Impl
-      auto current_process_mesh =
-          paddle::holds_alternative<phi::distributed::TensorDistAttr>(
-              spmd_info.first[0])
-              ? paddle::get<0>(spmd_info.first[0]).process_mesh()
-              : paddle::get<1>(spmd_info.first[0]).at(0).process_mesh();
-      SetReplicatedDistAttrForOutput(dist_out, current_process_mesh);
 
       if (rank_is_in_current_mesh) {
         // 5. Reshard Input
-        auto dist_input_weight =
-            ReshardApiInputToKernelInput(dev_ctx, weight, spmd_info.first[0]);
         auto dist_input_x =
-            ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[1]);
+            ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[0]);
+        auto dist_input_weight =
+            ReshardApiInputToKernelInput(dev_ctx, weight, spmd_info.first[1]);
         auto dist_input_out_grad =
             ReshardApiInputToKernelInput(dev_ctx, out_grad, spmd_info.first[2]);
 
         // 6. PrepareData (DataTransform & Prepare Dense Input)
-        dist_input_weight = PrepareDataForDistTensor(
-            dist_input_weight,
+        dist_input_x = PrepareDataForDistTensor(
+            dist_input_x,
             GetKernelInputArgDef(kernel.InputAt(0), kernel_key.backend()),
             {},
             kernel_result.is_stride_kernel);
-        auto input_weight = &dist_input_weight->value();
-
-        dist_input_x = PrepareDataForDistTensor(
-            dist_input_x,
+        auto input_x = &dist_input_x->value();
+        dist_input_weight = PrepareDataForDistTensor(
+            dist_input_weight,
             GetKernelInputArgDef(kernel.InputAt(1), kernel_key.backend()),
             {},
             kernel_result.is_stride_kernel);
-        auto input_x = &dist_input_x->value();
-
+        auto input_weight = &dist_input_weight->value();
         dist_input_out_grad = PrepareDataForDistTensor(
             dist_input_out_grad,
             GetKernelInputArgDef(kernel.InputAt(2), kernel_key.backend()),
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -362,6 +362,7 @@
   infer_meta :
     func : EmbeddingInferMeta
     param : [x, weight, padding_idx]
+    spmd_rule: EmbeddingInferSpmd
   kernel :
     func : embedding {dense, dense -> dense}
            sparse_weight_embedding {dense, selected_rows -> dense}
diff --git a/paddle/phi/infermeta/spmd_rules/embedding.cc b/paddle/phi/infermeta/spmd_rules/embedding.cc
@@ -19,6 +19,8 @@ limitations under the License. */
 #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/matmul.h"
+#include "paddle/phi/infermeta/spmd_rules/reshape.h"
 #include "paddle/phi/infermeta/spmd_rules/utils.h"
 
 namespace phi {
@@ -208,5 +210,115 @@ SpmdInfo EmbeddingInferSpmdReverse(const DistMetaTensor& x,
   return {{x_dist_attr, weight_dist_attr}, {out_dist_attr_src}};
 }
 
+SpmdInfo EmbeddingGradInferSpmd(const DistMetaTensor& x,
+                                const DistMetaTensor& weight,
+                                const DistMetaTensor& out_grad,
+                                int64_t padding_idx,
+                                bool sparse) {
+  PADDLE_ENFORCE_EQ(out_grad.dims().size(),
+                    out_grad.dist_attr().dims_mapping().size(),
+                    phi::errors::InvalidArgument(
+                        "The Tensor out_grad's rank [%d] and out_grad's "
+                        "dims_mapping size [%d] are not matched.",
+                        out_grad.dims(),
+                        out_grad.dist_attr().dims_mapping().size()));
+
+  if (sparse) {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "EmbeddingGradInferSpmd does't support sparse currently."));
+  }
+
+  // Propagate sharding info using composite operators.
+  // The whole mathematical expression of EmbeddingGrad is:
+  // w_grad = einsum('...j, ...k->jk', onehot(x, j), out_grad)
+
+  // TODO(cxxly): Simplifies the code logic of sharding propagation using
+  // primitive operators.
+  DistMetaTensor x_dst(x.dims(), x.dist_attr());
+  DistMetaTensor w_dst(weight.dims(), weight.dist_attr());
+  DistMetaTensor out_grad_dst(out_grad.dims(), out_grad.dist_attr());
+  DistMetaTensor w_grad(weight.dims(), weight.dist_attr());
+
+  // Step1: t0 = onehot(x_dst, w_dst.shape[0]) = eye(num_classes)[x_dst]
+  auto t0_dims_mapping = x_dst.dist_attr().dims_mapping();
+  t0_dims_mapping.emplace_back(-1);
+  TensorDistAttr t0_dist_attr(x.dist_attr());
+  t0_dist_attr.set_dims_mapping(t0_dims_mapping);
+  auto t0_shape = phi::vectorize(x.dims());
+  t0_shape.emplace_back(w_dst.dims()[0]);
+  DistMetaTensor t0(phi::make_ddim(t0_shape), t0_dist_attr);
+
+  // Step2: w_grad = einsum('...j, ...k -> jk', t0, out_grad_dst)
+  // Step 2.1: Build Einsum Notation
+  std::string alphabet = "abcdefghijlmnopqrstuvwxyz";
+  std::string t0_axes =
+      GetBroadcastAxes(t0.dims().size(), t0.dims().size(), alphabet);
+  std::string out_grad_dst_axes = t0_axes.substr(0, t0_axes.length() - 1) + "k";
+  std::string w_grad_axes = t0_axes.substr(t0_axes.length() - 1, 1) + "k";
+
+  // Step2.2: Sharding Propogation
+  // Step2.2.1: merge input shardings
+  auto axis_to_dim_map = ShardingMergeForTensors(
+      {{t0_axes, t0.dist_attr().dims_mapping()},
+       {out_grad_dst_axes, out_grad_dst.dist_attr().dims_mapping()}},
+      false);
+
+  // Step2.2.2: infer output's dims mapping.
+  auto w_grad_dist_attr = w_grad.dist_attr();
+  std::vector<int64_t> w_grad_dims_mapping =
+      GetDimsMappingForAxes(w_grad_axes, axis_to_dim_map);
+  w_grad_dist_attr.set_dims_mapping(w_grad_dims_mapping);
+
+  // Step2.2.3: merge potential conflict in inputs,
+  // update input dims mapping with merged shardings.
+  t0_dist_attr.set_dims_mapping(
+      GetDimsMappingForAxes(t0_axes, axis_to_dim_map));
+  auto out_grad_dst_dist_attr = out_grad_dst.dist_attr();
+  out_grad_dst_dist_attr.set_dims_mapping(
+      GetDimsMappingForAxes(out_grad_dst_axes, axis_to_dim_map));
+
+  // Step2.2.4: Handle Partial
+  std::vector<int64_t> partial_on_dims =
+      ResoluteOutputPartialDimension(axis_to_dim_map, w_grad_axes);
+  w_grad_dist_attr.set_partial_status(partial_on_dims);
+
+  // Step2.3: Update inputs info.
+  // NOTE: Reshard happend on intemediate operators must be ensure propagated
+  // back to first inputs.
+  t0 = DistMetaTensor(t0.dims(), t0_dist_attr);
+  const auto& t0_dims = t0.dist_attr().dims_mapping();
+  if (x_dst.dist_attr().dims_mapping() !=
+      std::vector<int64_t>(t0_dims.begin(), t0_dims.end() - 1)) {
+    TensorDistAttr t0_new(t0.dist_attr());
+    t0_new.set_dims_mapping(
+        std::vector<int64_t>(t0_dims.begin(), t0_dims.end() - 1));
+    x_dst = DistMetaTensor(x_dst.dims(), t0_new);
+  }
+  out_grad_dst = DistMetaTensor(out_grad_dst.dims(), out_grad_dst_dist_attr);
+  w_grad = DistMetaTensor(w_grad.dims(), w_grad_dist_attr);
+
+  VLOG(6) << "EmbeddingGradInferSpmd:\n"
+          << "Input x shape: [" << str_join(phi::vectorize(x.dims()))
+          << "], src_dims_mapping: [" << str_join(x.dist_attr().dims_mapping())
+          << "], dst_dims_mapping: ["
+          << str_join(x_dst.dist_attr().dims_mapping()) << "]\n"
+          << "Input weight shape: [" << str_join(phi::vectorize(weight.dims()))
+          << "], src_dims_mapping: ["
+          << str_join(weight.dist_attr().dims_mapping())
+          << "], dst_dims_mapping: ["
+          << str_join(w_dst.dist_attr().dims_mapping()) << "]\n"
+          << "Input out_grad shape: ["
+          << str_join(phi::vectorize(out_grad.dims()))
+          << "], src_dims_mapping: ["
+          << str_join(out_grad.dist_attr().dims_mapping())
+          << "], dst_dims_mapping: ["
+          << str_join(out_grad_dst.dist_attr().dims_mapping()) << "]\n"
+          << "Output w_grad shape: [" << str_join(phi::vectorize(w_grad.dims()))
+          << "], dims_mapping: [" << str_join(w_grad.dist_attr().dims_mapping())
+          << "]\n\n";
+
+  return {{x_dst.dist_attr(), w_dst.dist_attr(), out_grad_dst.dist_attr()},
+          {w_grad.dist_attr()}};
+}
 }  // namespace distributed
 }  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/embedding.h b/paddle/phi/infermeta/spmd_rules/embedding.h
@@ -29,13 +29,18 @@ namespace distributed {
 SpmdInfo EmbeddingInferSpmd(const DistMetaTensor& x,
                             const DistMetaTensor& weight,
                             int padding_idx,
-                            bool sparse);
+                            bool sparse = false);
 
 SpmdInfo EmbeddingInferSpmdReverse(const DistMetaTensor& x,
                                    const DistMetaTensor& weight,
                                    const DistMetaTensor& out,
                                    int padding_idx,
-                                   bool sparse);
+                                   bool sparse = false);
 
+SpmdInfo EmbeddingGradInferSpmd(const DistMetaTensor& x,
+                                const DistMetaTensor& weight,
+                                const DistMetaTensor& out_grad,
+                                int64_t padding_idx,
+                                bool sparse = false);
 }  // namespace distributed
 }  // namespace phi
diff --git a/paddle/phi/kernels/cpu/embedding_grad_kernel.cc b/paddle/phi/kernels/cpu/embedding_grad_kernel.cc
@@ -58,7 +58,6 @@ struct EmbeddingGradCPUFunctor {
       auto* d_table_data = weight_grad_->data<T>();
 
       memset(d_table_data, 0, weight_grad_->numel() * sizeof(T));
-
       for (int64_t i = 0; i < ids_num; ++i) {
         if (padding_idx_ != kNoPadding && ids_data[i] == padding_idx_) {
           // the gradient of padding_idx should be 0, already done by memset, so
diff --git a/test/auto_parallel/semi_auto_parallel_for_embedding.py b/test/auto_parallel/semi_auto_parallel_for_embedding.py
@@ -21,7 +21,7 @@
 from paddle.distributed import Replicate, Shard
 
 
-class TestCustomEmbeddingGradApiForSemiAutoParallel:
+class TestEmbeddingApiForSemiAutoParallel:
     def __init__(self):
         self._dtype = os.getenv("dtype")
         self._backend = os.getenv("backend")
@@ -100,28 +100,37 @@ def test_w_col_shard(self):
         )
 
     def test_x_row_w_col_shard(self):
-        self.test_body(
-            x_shape=[12, 16],
-            w_shape=[10, 4],
-            x_placements=[Shard(0)],
-            w_placements=[Shard(1)],
-        )
+        try:
+            self.test_body(
+                x_shape=[12, 16],
+                w_shape=[10, 4],
+                x_placements=[Shard(0)],
+                w_placements=[Shard(1)],
+            )
+        except RuntimeError as e:
+            assert 'sharded by same mesh dimension ' in str(e)
 
     def test_x_col_w_row_shard(self):
-        self.test_body(
-            x_shape=[12, 16],
-            w_shape=[10, 4],
-            x_placements=[Shard(1)],
-            w_placements=[Shard(0)],
-        )
+        try:
+            self.test_body(
+                x_shape=[12, 16],
+                w_shape=[10, 4],
+                x_placements=[Shard(1)],
+                w_placements=[Shard(0)],
+            )
+        except RuntimeError as e:
+            assert 'sharded by same mesh dimension ' in str(e)
 
     def test_both_col_shard(self):
-        self.test_body(
-            x_shape=[12, 16],
-            w_shape=[10, 4],
-            x_placements=[Shard(1)],
-            w_placements=[Shard(1)],
-        )
+        try:
+            self.test_body(
+                x_shape=[12, 16],
+                w_shape=[10, 4],
+                x_placements=[Shard(1)],
+                w_placements=[Shard(1)],
+            )
+        except RuntimeError as e:
+            assert 'sharded by same mesh dimension', str(e)
 
     def run_test_case(self):
         if self._backend == "cpu":
@@ -134,12 +143,13 @@ def run_test_case(self):
         self.test_non_shard()
         self.test_x_row_shard()
         self.test_x_col_shard()
-        self.test_w_row_shard()
+        # Sharding along weight's row axis is not supported by raw embdding kernel.
+        # self.test_w_row_shard()
         self.test_w_col_shard()
         self.test_x_row_w_col_shard()
         self.test_x_col_w_row_shard()
         self.test_both_col_shard()
 
 
 if __name__ == '__main__':
-    TestCustomEmbeddingGradApiForSemiAutoParallel().run_test_case()
+    TestEmbeddingApiForSemiAutoParallel().run_test_case()
diff --git a/test/auto_parallel/test_semi_auto_parallel_basic.py b/test/auto_parallel/test_semi_auto_parallel_basic.py
@@ -138,13 +138,13 @@ def test_flash_attention_api(self):
                 user_defined_envs=envs,
             )
 
-    def test_custom_embedding_grad_api(self):
+    def test_embedding_api(self):
         envs_list = test_base.gen_product_envs_list(
             self._default_envs, self._changeable_envs
         )
         for envs in envs_list:
             self.run_test_case(
-                "semi_auto_parallel_for_embedding_grad.py",
+                "semi_auto_parallel_for_embedding.py",
                 user_defined_envs=envs,
             )
 
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc

Original file line number	Diff line number	Diff line change
`@@ -138,13 +138,13 @@ def test_flash_attention_api(self):`
`138`	`138`	`user_defined_envs=envs,`
`139`	`139`	`)`
`140`	`140`
`141`		`- def test_custom_embedding_grad_api(self):`
	`141`	`+ def test_embedding_api(self):`
`142`	`142`	`envs_list = test_base.gen_product_envs_list(`
`143`	`143`	`self._default_envs, self._changeable_envs`
`144`	`144`	`)`
`145`	`145`	`for envs in envs_list:`
`146`	`146`	`self.run_test_case(`
`147`		`- "semi_auto_parallel_for_embedding_grad.py",`
	`147`	`+ "semi_auto_parallel_for_embedding.py",`
`148`	`148`	`user_defined_envs=envs,`
`149`	`149`	`)`
`150`	`150`