fix spmd rule of flatten_grad, add xshape dist_attr in result (#64985)

PaddlePaddle · Jun 7, 2024 · ed8168d · ed8168d
1 parent 60fe41b
commit ed8168d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/paddle/phi/infermeta/spmd_rules/flatten.cc b/paddle/phi/infermeta/spmd_rules/flatten.cc
@@ -191,7 +191,13 @@ SpmdInfo FlattenInferSpmdReverse(const DistMetaTensor& x,
 
 SpmdInfo FlattenGradInferSpmd(const DistMetaTensor& xshape,
                               const DistMetaTensor& out_grad) {
-  return ReshapeGradInferSpmd(xshape, out_grad);
+  // TODO(jeff41404): when ReshapeInferSpmd and ReshapeGradInferSpmd can deliver
+  // distributed attribute of xshape, we will use ReshapeGradInferSpmd directly
+  // in future return ReshapeGradInferSpmd(xshape, out_grad);
+  auto shape = phi::vectorize(xshape.dims());
+  shape = std::vector<int64_t>(shape.begin() + 1, shape.end());
+  const auto& spmd = ReshapeInferSpmd(out_grad, shape);
+  return {{xshape.dist_attr(), spmd.first[0]}, {spmd.second[0]}};
 }
 
 }  // namespace distributed

diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
@@ -1908,13 +1908,14 @@ TEST(Flatten, Ctor) {
   check_dim_mapping(spmd4.second[0], {-1, -1});
   check_dim_mapping(spmd4.second[1], {-1, -1, -1, -1, -1, -1});  // x_shape
 
-  auto out_grad = build_input({2, 1024, 1024}, {-1, -1, -1});
-  auto xshape = build_input({0, 2, 1024, 4, 1024 / 4}, {-1, 0, 1, -1, -1});
+  auto out_grad = build_input({2, 1024, 1024}, {0, -1, 1});
+  auto xshape = build_input({0, 2, 1024, 4, 1024 / 4}, {-1, 0, -1, 1, -1});
   auto spmd_grad = FlattenGradInferSpmd(xshape, out_grad);
-  EXPECT_EQ(spmd_grad.first.size(), static_cast<size_t>(1));
+  EXPECT_EQ(spmd_grad.first.size(), static_cast<size_t>(2));
   EXPECT_EQ(spmd_grad.second.size(), static_cast<size_t>(1));
-  check_dim_mapping(spmd_grad.first[0], {0, 1, -1});
-  check_dim_mapping(spmd_grad.second[0], {0, 1, -1, -1});
+  check_dim_mapping(spmd_grad.first[0], {-1, 0, -1, 1, -1});
+  check_dim_mapping(spmd_grad.first[1], {0, -1, 1});
+  check_dim_mapping(spmd_grad.second[0], {0, -1, 1, -1});
 }
 
 }  // namespace auto_parallel