fix reshard dist_attr (#60535)

PaddlePaddle · Jan 3, 2024 · 2ad9e24 · 2ad9e24
1 parent 5d01382
commit 2ad9e24
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/paddle/phi/core/distributed/auto_parallel/dist_tensor.cc b/paddle/phi/core/distributed/auto_parallel/dist_tensor.cc
@@ -117,6 +117,9 @@ DistTensor::DistTensor() : value_(std::make_shared<DenseTensor>()) {}
 DistTensor::DistTensor(const std::shared_ptr<phi::DenseTensor>& global_value,
                        const TensorDistAttr& dist_attr)
     : global_dims_(global_value->dims()), dist_attr_(dist_attr) {
+  process_mesh_ = dist_attr_.process_mesh();
+  placements_ = ToPlacements(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())) {

diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc
@@ -52,6 +52,8 @@ void ReshardFunction::SetDistProps(DistTensor* tensor,
 
   tensor->global_dims_ = dims;
   tensor->dist_attr_ = dist_attr;
+  tensor->process_mesh_ = dist_attr.process_mesh();
+  tensor->placements_ = ToPlacements(dist_attr);
 }
 
 void ReshardFunction::SetDistProps(DistTensor* tensor,
@@ -64,6 +66,8 @@ void ReshardFunction::SetDistProps(DistTensor* tensor,
                         str_join(vectorize(tensor->dims()))));
 
   tensor->dist_attr_ = dist_attr;
+  tensor->process_mesh_ = dist_attr.process_mesh();
+  tensor->placements_ = ToPlacements(dist_attr);
 }
 
 DenseTensor* ReshardFunction::GetMutableTensor(DistTensor* tensor) {