[Semi-auto]Add srp in dist_tensor (#59683)

ForFishes · web-flow · commit ceea9042f154 · 2023-12-06T11:41:12.000+08:00
* add srp in disttensor

* add srp in disttensor

* add srp in disttensor

* add srp in disttensor

* add srp in disttensor
diff --git a/paddle/fluid/pybind/eager_properties.cc b/paddle/fluid/pybind/eager_properties.cc
@@ -836,6 +836,42 @@ PyObject* tensor_properties_get_place_str(TensorObject* self, void* closure) {
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+PyObject* tensor_properties_get_placements_str(TensorObject* self,
+                                               void* closure) {
+  EAGER_TRY
+  if (self->tensor.is_dist_tensor()) {
+#ifdef PADDLE_WITH_DISTRIBUTE
+    phi::distributed::DistTensor* dist_tensor =
+        static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get());
+
+    std::stringstream ostr;
+    ostr << "[";
+
+    bool isFirst = true;
+    for (const auto& p : dist_tensor->placements()) {
+      if (p) {
+        if (!isFirst) {
+          ostr << ", ";
+        }
+        ostr << p->to_string();
+        isFirst = false;
+      }
+    }
+    ostr << "]";
+    return ToPyObject(ostr.str());
+#else
+    PADDLE_THROW(platform::errors::Unavailable(
+        "The `placements()` property of (Dist)Tensor is not supported in the "
+        "current PaddlePaddle, please recompile and installPaddlePaddle with "
+        "the "
+        "option of `WITH_DISTRIBUTE=ON`."));
+#endif
+  } else {
+    RETURN_PY_NONE
+  }
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 PyDoc_STRVAR(tensor_dtype__doc__,
              R"DOC(dtype
 
@@ -998,6 +1034,11 @@ struct PyGetSetDef variable_properties[] = {  // NOLINT
      nullptr,
      nullptr,
      nullptr},
+    {"_placements_str",
+     (getter)tensor_properties_get_placements_str,
+     nullptr,
+     nullptr,
+     nullptr},
     {"dtype",
      (getter)tensor_properties_get_dtype,
      nullptr,
diff --git a/paddle/phi/core/distributed/auto_parallel/dist_tensor.cc b/paddle/phi/core/distributed/auto_parallel/dist_tensor.cc
@@ -34,11 +34,89 @@ inline void check_defined(const DistTensor& dist_tensor,
           method_hint));
 }
 
+TensorDistAttr ToTensorDistAttr(const ProcessMesh& process_mesh,
+                                const Placements& placements,
+                                const DDim& dims) {
+  TensorDistAttr dist_attr(vectorize(dims));
+  // Step1: set process_mesh
+  dist_attr.set_process_mesh(process_mesh);
+
+  // Step2: set dim_mapping
+  int64_t ndim = dims.size();
+  std::vector<int64_t> dim_map(ndim, -1);
+  for (size_t i = 0; i < placements.size(); i++) {
+    auto& placement = placements[i];
+    if (placement->is_shard()) {
+      auto shard_dim = dynamic_cast<const Shard&>(*placement).get_dim();
+      PADDLE_ENFORCE_EQ(
+          dim_map[shard_dim],
+          -1,
+          phi::errors::InvalidArgument(
+              "Tensor dim %lld is already sharded on mesh dim %lld,"
+              " DistTensor operator implementation does not support things "
+              "like hybrid"
+              " sharding strategies yet (i.e. [Shard(0), Shard(0)])",
+              shard_dim,
+              dim_map[shard_dim]));
+      dim_map[shard_dim] = i;
+    }
+  }
+  dist_attr.set_dims_mapping(dim_map);
+
+  // Step3: set partial_status
+  paddle::flat_hash_map<int64_t, ReduceType> partial_status;
+  for (size_t i = 0; i < placements.size(); ++i) {
+    auto& p = placements[i];
+    if (p->is_partial()) {
+      partial_status.insert({i, dynamic_cast<Partial&>(*p).get_reduce_type()});
+    }
+  }
+  dist_attr.set_partial_status(partial_status);
+
+  // Step4: mark annotated
+  dist_attr.mark_annotated("process_mesh");
+  dist_attr.mark_annotated("dims_mapping");
+  return dist_attr;
+}
+
+Placements ToPlacements(const TensorDistAttr& dist_attr) {
+  auto& process_mesh = dist_attr.process_mesh();
+  Placements placements;
+  placements.resize(process_mesh.ndim(), std::make_shared<Replicate>());
+
+  auto& partial_status = dist_attr.partial_status();
+  for (const auto& pair : partial_status) {
+    placements[pair.first] = std::make_shared<Partial>(pair.second);
+  }
+
+  auto& dim_mapping = dist_attr.dims_mapping();
+  for (size_t i = 0; i < dim_mapping.size(); ++i) {
+    auto& mesh_id = dim_mapping[i];
+    if (mesh_id >= 0) {
+      auto& p = placements[mesh_id];
+
+      if (p->is_shard()) {
+        PADDLE_THROW(phi::errors::PreconditionNotMet(
+            "ProcessMesh dimension cann't be mapped to two  dimension of the "
+            "same tensor: {%d} and {%d}",
+            i,
+            dynamic_cast<Shard&>(*p).get_dim()));
+      } else if (p->is_partial()) {
+        PADDLE_THROW(phi::errors::PreconditionNotMet(
+            "ProcessMesh dimension {%d} cannot be both shard and partial!",
+            mesh_id));
+      }
+      placements[mesh_id] = std::make_shared<Shard>(i);
+    }
+  }
+  return placements;
+}
+
 DistTensor::DistTensor() : value_(std::make_shared<DenseTensor>()) {}
 
 DistTensor::DistTensor(const std::shared_ptr<phi::DenseTensor>& global_value,
                        const TensorDistAttr& dist_attr)
-    : dims_(global_value->dims()), dist_attr_(dist_attr) {
+    : global_dims_(global_value->dims()), dist_attr_(dist_attr) {
   // If the current rank doesn't in process_mesh, we should create an
   // uninitialized tensor only with tensor_meta.
   if (IsCurRankInMesh(dist_attr.process_mesh())) {
@@ -67,32 +145,15 @@ DistTensor::DistTensor(const std::shared_ptr<phi::DenseTensor>& global_value,
 DistTensor::DistTensor(const std::shared_ptr<phi::DenseTensor>& global_value,
                        const ProcessMesh& process_mesh,
                        const Placements& placements)
-    : dims_(global_value->dims()) {
-  dist_tensor_meta_ = DistTensorMeta(
-      process_mesh,
-      placements,
-      DenseTensorMeta(global_value->dtype(), global_value->dims()));
-
-  std::vector<int64_t> partial_dims;
-  size_t idx = 0;
-  for (auto p : placements) {
-    if (p->is_partial()) {
-      partial_dims.push_back(idx);
-    }
-    idx++;
-  }
-  TensorDistAttr dist_attr(common::vectorize(dist_tensor_meta_.dims()));
-  dist_attr.set_process_mesh(dist_tensor_meta_.process_mesh());
-  dist_attr.set_dims_mapping(dist_tensor_meta_.dim_mapping());
-  dist_attr.set_partial_status(partial_dims);
-  dist_attr.mark_annotated("process_mesh");
-  dist_attr.mark_annotated("dims_mapping");
-  dist_attr_ = dist_attr;
+    : global_dims_(global_value->dims()) {
+  process_mesh_ = process_mesh;
+  placements_ = placements;
+  dist_attr_ = ToTensorDistAttr(process_mesh_, placements_, global_dims_);
 
   // If the current rank doesn't in process_mesh, we should create an
   // uninitialized tensor only with dist_tensor_meta_.
   if (IsCurRankInMesh(process_mesh)) {
-    if (!dist_tensor_meta_.is_replicated()) {
+    if (!dist_attr_.is_replicated()) {
       value_ = std::make_shared<DenseTensor>();
       // 1. create replicated global tensor
       TensorDistAttr replicated_dist_attr(
@@ -101,9 +162,9 @@ DistTensor::DistTensor(const std::shared_ptr<phi::DenseTensor>& global_value,
       DistTensor replicated_tensor(global_value, replicated_dist_attr);
 
       // 2. reshard from replicated to other state
-      auto* func = ChooseProperReshardFunction(replicated_tensor, dist_attr);
+      auto* func = ChooseProperReshardFunction(replicated_tensor, dist_attr_);
       auto* dev_ctx = DeviceContextPool::Instance().Get(global_value->place());
-      func->Eval(dev_ctx, replicated_tensor, dist_attr, this);
+      func->Eval(dev_ctx, replicated_tensor, dist_attr_, this);
     } else {
       value_ = global_value;
     }
@@ -115,16 +176,19 @@ DistTensor::DistTensor(const std::shared_ptr<phi::DenseTensor>& global_value,
 }
 
 DistTensor::DistTensor(const DDim& dims, const TensorDistAttr& dist_attr)
-    : dims_(dims),
+    : global_dims_(dims),
       dist_attr_(dist_attr),
-      value_(std::make_shared<DenseTensor>()) {}
+      value_(std::make_shared<DenseTensor>()) {
+  process_mesh_ = dist_attr.process_mesh();
+  placements_ = ToPlacements(dist_attr);
+}
 
 void DistTensor::unsafe_set_dims(const DDim& dims) {
   if (this->initialized()) {
     VLOG(3) << "You try to set an initialized DistTensor's global dims. "
                "Make sure you are aware of where you change its dims.";
   }
-  dims_ = dims;
+  global_dims_ = dims;
 }
 
 void DistTensor::unsafe_set_dist_attr(const TensorDistAttr& dist_attr) {
@@ -133,12 +197,14 @@ void DistTensor::unsafe_set_dist_attr(const TensorDistAttr& dist_attr) {
                "Make sure you are aware of where you change its dist attr.";
   }
   dist_attr_ = dist_attr;
+  process_mesh_ = dist_attr.process_mesh();
+  placements_ = ToPlacements(dist_attr);
 }
 
 int64_t DistTensor::numel() const {
   // DistTensor with uninitialized local tensor can
   // also have numel.
-  return product(dims_);
+  return product(global_dims_);
 }
 
 const DDim& DistTensor::local_dims() const {
diff --git a/paddle/phi/core/distributed/auto_parallel/dist_tensor.h b/paddle/phi/core/distributed/auto_parallel/dist_tensor.h
@@ -16,14 +16,25 @@
 
 #include <memory>
 
+#include "paddle/phi/common/reduce_type.h"
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/placement_types.h"
 #include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
+#include "paddle/phi/core/enforce.h"
 
 namespace phi {
 namespace distributed {
 class ReshardFunction;
+class Shard;
+class Partial;
+class Replicate;
+
+TensorDistAttr ToTensorDistAttr(const ProcessMesh& process_mesh,
+                                const Placements& placements,
+                                const DDim& dims);
+
+Placements ToPlacements(const TensorDistAttr& dist_attr);
 
 class DistTensor final
     : public phi::TensorBase,
@@ -61,7 +72,7 @@ class DistTensor final
 
   /// \brief Returns the global dims of the dist tensor.
   /// \return The global dims of the dist tensor.
-  const DDim& dims() const override { return dims_; }
+  const DDim& dims() const override { return global_dims_; }
 
   /// \brief Set the global dims of the dist tensor.
   /// \return void
@@ -73,20 +84,24 @@ class DistTensor final
 
   /// \brief Returns the process_mesh of current dist tensor.
   /// \return The ProcessMesh's const reference
-  const ProcessMesh& process_mesh() const {
-    return dist_tensor_meta_.process_mesh();
-  }
+  const ProcessMesh& process_mesh() const { return process_mesh_; }
 
   /// \brief Returns the placements of current dist tensor.
   /// \return The Placements's const reference
-  const Placements& placements() const {
-    return dist_tensor_meta_.placements();
-  }
+  const Placements& placements() const { return placements_; }
 
   /// \brief Returns the num_shard of current dist tensor.
   /// \return int64_t
-  int64_t num_shard() const { return dist_tensor_meta_.num_shard(); }
-
+  int64_t num_shard() const {
+    int64_t num_shard = 1;
+    const auto& mesh_shape = process_mesh_.shape();
+    for (size_t i = 0; i < placements_.size(); i++) {
+      if (placements_[i]->is_shard()) {
+        num_shard *= mesh_shape[i];
+      }
+    }
+    return num_shard;
+  }
   /// \brief Set the dist attr of current dist tensor.
   /// \return void
   void unsafe_set_dist_attr(const TensorDistAttr& dist_attr);
@@ -147,13 +162,15 @@ class DistTensor final
   friend class ReshardFunction;
 
   // The global dimensions(shape), will move to DistTensorMeta
-  DDim dims_;
+  DDim global_dims_;
   // The distributed attributes, will remove in the future
   TensorDistAttr dist_attr_;
   // The local DenseTensor value
   std::shared_ptr<DenseTensor> value_;
 
-  DistTensorMeta dist_tensor_meta_;
+  ProcessMesh process_mesh_;
+
+  Placements placements_;
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/placement_types.h b/paddle/phi/core/distributed/auto_parallel/placement_types.h
@@ -133,6 +133,7 @@ class Partial : public Placement {
  public:
   explicit Partial(ReduceType reduce_type) : reduce_type_(reduce_type) {}
   bool is_partial() const override { return true; }
+  ReduceType get_reduce_type() const { return reduce_type_; }
 
   bool operator==(const Placement& other) const override {
     const Partial* other_partial = dynamic_cast<const Partial*>(&other);
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc
@@ -50,7 +50,7 @@ void ReshardFunction::SetDistProps(DistTensor* tensor,
                         dist_attr.to_string(),
                         str_join(vectorize(dims))));
 
-  tensor->dims_ = dims;
+  tensor->global_dims_ = dims;
   tensor->dist_attr_ = dist_attr;
 }
 
diff --git a/python/paddle/tensor/to_string.py b/python/paddle/tensor/to_string.py
@@ -364,26 +364,28 @@ def dist_tensor_to_string(tensor, prefix='Tensor'):
         dtype = 'bfloat16'
 
     if not tensor._is_dense_tensor_hold_allocation():
-        _template = "{prefix}(shape={shape}, dtype={dtype}, place={place}, stop_gradient={stop_gradient}, dist_attr={dist_attr}, GlobalDenseTensor Not initialized)"
+        _template = "{prefix}(shape={shape}, dtype={dtype}, place={place}, stop_gradient={stop_gradient}, process_mesh={process_mesh}, placements={placements}, GlobalDenseTensor Not initialized)"
         return _template.format(
             prefix=prefix,
             shape=tensor.shape,
             dtype=dtype,
             place=tensor._place_str,
             stop_gradient=tensor.stop_gradient,
-            dist_attr=tensor.dist_attr,
+            process_mesh=tensor.process_mesh,
+            placements=tensor._placements_str,
         )
     else:
         indent = len(prefix) + 1
         data = _format_dense_tensor(tensor, indent)
-        _template = "{prefix}(shape={shape}, dtype={dtype}, place={place}, stop_gradient={stop_gradient}, dist_attr={dist_attr}, GlobalDenseTensor=\n{indent}{data})"
+        _template = "{prefix}(shape={shape}, dtype={dtype}, place={place}, stop_gradient={stop_gradient}, process_mesh={process_mesh}, placements={placements}, GlobalDenseTensor=\n{indent}{data})"
         return _template.format(
             prefix=prefix,
             shape=tensor.shape,
             dtype=dtype,
             place=tensor._place_str,
             stop_gradient=tensor.stop_gradient,
-            dist_attr=tensor.dist_attr,
+            process_mesh=tensor.process_mesh,
+            placements=tensor._placements_str,
             indent=' ' * indent,
             data=data,
         )
diff --git a/test/auto_parallel/reshard_p_to_r.py b/test/auto_parallel/reshard_p_to_r.py
@@ -41,6 +41,8 @@ def run_test_case(self):
 
         input_tensor = dist.shard_tensor(a, self._mesh, [dist.Partial()])
         out = dist.reshard(input_tensor, self._mesh, [dist.Replicate()])
+        print(input_tensor)
+        print(out)
 
         assert np.equal(out.shape, input_tensor.shape).all()
         np.testing.assert_equal(out._local_value().numpy(), a.numpy())

Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ void ReshardFunction::SetDistProps(DistTensor* tensor,`
`50`	`50`	`dist_attr.to_string(),`
`51`	`51`	`str_join(vectorize(dims))));`
`52`	`52`
`53`		`- tensor->dims_ = dims;`
	`53`	`+ tensor->global_dims_ = dims;`
`54`	`54`	`tensor->dist_attr_ = dist_attr;`
`55`	`55`	`}`
`56`	`56`