support to fallback to sharding by one mesh dim from by more than one (#74396)

* support to fallback to shard by one mesh dim from by more than one

* add test case

Author: liufengwei0103

paddle/phi/core/distributed/auto_parallel/dist_attr.cc

@@ -560,15 +560,31 @@ TensorDistAttr::DimMapProxy::operator const std::vector<int64_t>&() const {
 void TensorDistAttr::DimMapProxy::sync_1d_map() const {
   dims_mapping_1d.resize(dims_mapping_2d->size());
   for (size_t i = 0; i < dims_mapping_2d->size(); ++i) {
-    PADDLE_ENFORCE_LE(dims_mapping_2d->at(i).size(),
-                      1,
-                      common::errors::InvalidArgument(
-                          "There are %d mesh dim sharded on tensor dim %d,"
-                          "you should call \"multi_dims_mapping()\"",
-                          dims_mapping_2d->at(i).size(),
-                          i));
-    dims_mapping_1d[i] =
-        (*dims_mapping_2d)[i].empty() ? -1 : (*dims_mapping_2d)[i][0];
+    size_t num_mesh_dim = dims_mapping_2d->at(i).size();
+    if (num_mesh_dim <= 1) {
+      dims_mapping_1d[i] =
+          (*dims_mapping_2d)[i].empty() ? -1 : (*dims_mapping_2d)[i][0];
+      continue;
+    }
+
+    int64_t max_mesh_dim = (*dims_mapping_2d)[i][0];
+    int64_t max_mesh_dim_size = process_mesh.shape()[max_mesh_dim];
+
+    for (size_t j = 1; j < num_mesh_dim; ++j) {
+      int64_t cur_mesh_dim = (*dims_mapping_2d)[i][j];
+      int64_t cur_mesh_dim_size = process_mesh.shape()[cur_mesh_dim];
+
+      if (cur_mesh_dim_size > max_mesh_dim_size) {
+        max_mesh_dim = cur_mesh_dim;
+        max_mesh_dim_size = cur_mesh_dim_size;
+      }
+    }
+
+    LOG(WARNING) << "There are " << num_mesh_dim << " shared on tensor dim "
+                 << i << ". Now fallback to sharding by mesh dim "
+                 << max_mesh_dim << ".";
+
+    dims_mapping_1d[i] = max_mesh_dim;
   }
 }
 

paddle/phi/core/distributed/auto_parallel/dist_attr.h

@@ -231,8 +231,9 @@ class TEST_API TensorDistAttr {
   // delete it after all 1d vector dims_mapping_ have been upgraded to 2d.
   class DimMapProxy final {
    public:
-    DimMapProxy(std::vector<std::vector<int64_t>>* dims_mapping_2d)
-        : dims_mapping_2d(dims_mapping_2d) {}
+    DimMapProxy(std::vector<std::vector<int64_t>>* dims_mapping_2d,
+                const ProcessMesh& process_mesh)
+        : dims_mapping_2d(dims_mapping_2d), process_mesh(process_mesh) {}
 
     DimMapProxy& operator=(
         const std::vector<std::vector<int64_t>>& dims_mapping);
@@ -248,6 +249,7 @@ class TEST_API TensorDistAttr {
     void sync_2d_map();
     mutable std::vector<int64_t> dims_mapping_1d;
     std::vector<std::vector<int64_t>>* dims_mapping_2d;
+    const ProcessMesh& process_mesh;
   };
 
   static std::vector<std::string> fields_;
@@ -266,7 +268,7 @@ class TEST_API TensorDistAttr {
 
   std::vector<std::vector<int64_t>> dims_mapping_;
   // for short time, backward compatible for existing spmd relus.
-  DimMapProxy dims_mapping_proxy{&dims_mapping_};
+  DimMapProxy dims_mapping_proxy{&dims_mapping_, process_mesh_};
 };
 
 inline std::ostream& operator<<(std::ostream& os, const TensorDistAttr& obj) {

test/auto_parallel/co_shard.py

@@ -189,6 +189,11 @@ def run_test_case_4(self):
             str(relu_out.placements[1]), "Shard(dim=0, shard_order=1)"
         )
 
+        # test fallback to shard by one dim.
+        add_out = paddle.add(relu_out, relu_out)
+        np.testing.assert_equal(str(add_out.placements[0]), "Shard(dim=0)")
+        np.testing.assert_equal(str(add_out.placements[1]), "Replicate()")
+
     def run_test_case_main(self):
         self.basic_interface_case()
         self.run_test_case_0()