PaddlePaddle
diff --git a/‎paddle/fluid/distributed/collective/process_group_bkcl.cc
Copy file name to clipboard
+147-151Lines changed: 147 additions & 151 deletions b/‎paddle/fluid/distributed/collective/process_group_bkcl.cc
Copy file name to clipboard
+147-151Lines changed: 147 additions & 151 deletions
diff --git a/‎paddle/fluid/distributed/collective/process_group_kernel_utils.h
Copy file name to clipboard
+20-9Lines changed: 20 additions & 9 deletions b/‎paddle/fluid/distributed/collective/process_group_kernel_utils.h
Copy file name to clipboard
+20-9Lines changed: 20 additions & 9 deletions
@@ -466,74 +466,72 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::AllToAll(
 
           int64_t nranks = size_;
 
-          if (in_row_size > 0 && out_row_size > 0) {
-            std::vector<int64_t> in_numel_vec(nranks);
-            std::vector<int64_t> in_offset_vec(nranks);
-            std::vector<int64_t> out_numel_vec(nranks);
-            std::vector<int64_t> out_offset_vec(nranks);
-
-            int64_t in_offset = 0;
-            int64_t out_offset = 0;
-            for (int64_t i = 0; i < nranks; i++) {
-              int64_t in_numel = in_split_sizes[i] * in_row_size;
-              int64_t out_numel = out_split_sizes[i] * out_row_size;
-
-              in_numel_vec[i] = in_numel;
-              in_offset_vec[i] = in_offset;
-              in_offset += in_numel;
-
-              out_numel_vec[i] = out_numel;
-              out_offset_vec[i] = out_offset;
-              out_offset += out_numel;
-            }
-
-            PADDLE_ENFORCE_GE(
-                in_tensor.place().GetDeviceId(),
-                0,
-                common::errors::PreconditionNotMet(
-                    "The all_to_all device id must greater or equal than 0."));
-            phi::XPUPlace place = in_tensor.place();
-            auto allocator = std::unique_ptr<phi::Allocator>(
-                new paddle::experimental::DefaultAllocator(place));
-            phi::DenseTensorMeta meta(phi::DataType::INT64, phi::DDim{nranks});
-
-            phi::DenseTensor in_size_tensor = {allocator.get(), meta};
-            phi::DenseTensor in_offset_tensor = {allocator.get(), meta};
-            phi::DenseTensor out_size_tensor = {allocator.get(), meta};
-            phi::DenseTensor out_offset_tensor = {allocator.get(), meta};
-
-            memory::Copy(place,
-                         in_size_tensor.data(),
-                         phi::CPUPlace(),
-                         in_numel_vec.data(),
-                         in_size_tensor.numel() * sizeof(int64_t));
-
-            memory::Copy(place,
-                         in_offset_tensor.data(),
-                         phi::CPUPlace(),
-                         in_offset_vec.data(),
-                         in_offset_tensor.numel() * sizeof(int64_t));
-
-            memory::Copy(place,
-                         out_size_tensor.data(),
-                         phi::CPUPlace(),
-                         out_numel_vec.data(),
-                         out_size_tensor.numel() * sizeof(int64_t));
-
-            memory::Copy(place,
-                         out_offset_tensor.data(),
-                         phi::CPUPlace(),
-                         out_offset_vec.data(),
-                         out_offset_tensor.numel() * sizeof(int64_t));
-
-            comm_context->AllToAllUnequalSplit(out_tensor,
-                                               in_tensor,
-                                               out_size_tensor,
-                                               out_offset_tensor,
-                                               in_size_tensor,
-                                               in_offset_tensor,
-                                               stream);
+          std::vector<int64_t> in_numel_vec(nranks);
+          std::vector<int64_t> in_offset_vec(nranks);
+          std::vector<int64_t> out_numel_vec(nranks);
+          std::vector<int64_t> out_offset_vec(nranks);
+
+          int64_t in_offset = 0;
+          int64_t out_offset = 0;
+          for (int64_t i = 0; i < nranks; i++) {
+            int64_t in_numel = in_split_sizes[i] * in_row_size;
+            int64_t out_numel = out_split_sizes[i] * out_row_size;
+
+            in_numel_vec[i] = in_numel;
+            in_offset_vec[i] = in_offset;
+            in_offset += in_numel;
+
+            out_numel_vec[i] = out_numel;
+            out_offset_vec[i] = out_offset;
+            out_offset += out_numel;
           }
+
+          PADDLE_ENFORCE_GE(
+              in_tensor.place().GetDeviceId(),
+              0,
+              common::errors::PreconditionNotMet(
+                  "The all_to_all device id must greater or equal than 0."));
+          phi::XPUPlace place = in_tensor.place();
+          auto allocator = std::unique_ptr<phi::Allocator>(
+              new paddle::experimental::DefaultAllocator(place));
+          phi::DenseTensorMeta meta(phi::DataType::INT64, phi::DDim{nranks});
+
+          phi::DenseTensor in_size_tensor = {allocator.get(), meta};
+          phi::DenseTensor in_offset_tensor = {allocator.get(), meta};
+          phi::DenseTensor out_size_tensor = {allocator.get(), meta};
+          phi::DenseTensor out_offset_tensor = {allocator.get(), meta};
+
+          memory::Copy(place,
+                       in_size_tensor.data(),
+                       phi::CPUPlace(),
+                       in_numel_vec.data(),
+                       in_size_tensor.numel() * sizeof(int64_t));
+
+          memory::Copy(place,
+                       in_offset_tensor.data(),
+                       phi::CPUPlace(),
+                       in_offset_vec.data(),
+                       in_offset_tensor.numel() * sizeof(int64_t));
+
+          memory::Copy(place,
+                       out_size_tensor.data(),
+                       phi::CPUPlace(),
+                       out_numel_vec.data(),
+                       out_size_tensor.numel() * sizeof(int64_t));
+
+          memory::Copy(place,
+                       out_offset_tensor.data(),
+                       phi::CPUPlace(),
+                       out_offset_vec.data(),
+                       out_offset_tensor.numel() * sizeof(int64_t));
+
+          comm_context->AllToAllUnequalSplit(out_tensor,
+                                             in_tensor,
+                                             out_size_tensor,
+                                             out_offset_tensor,
+                                             in_size_tensor,
+                                             in_offset_tensor,
+                                             stream);
         }
       },
       in_tensor,
@@ -614,95 +612,93 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::AllToAll(
           out_numel_sum += (*out_tensors)[i].numel();
         }
 
-        if (in_numel_sum > 0 || out_numel_sum > 0) {
-          std::vector<int64_t> in_numel_vec(nranks);
-          std::vector<int64_t> in_offset_vec(nranks);
-          std::vector<int64_t> out_numel_vec(nranks);
-          std::vector<int64_t> out_offset_vec(nranks);
-
-          int64_t in_offset = 0;
-          int64_t out_offset = 0;
-          for (int64_t i = 0; i < nranks; i++) {
-            int64_t in_numel = in_tensors[i].numel();
-            int64_t out_numel = (*out_tensors)[i].numel();
-
-            in_numel_vec[i] = in_numel;
-            in_offset_vec[i] = in_offset;
-            in_offset += in_numel;
-
-            out_numel_vec[i] = out_numel;
-            out_offset_vec[i] = out_offset;
-            out_offset += out_numel;
-          }
-
-          PADDLE_ENFORCE_GE(
-              in_tensors[0].place().GetDeviceId(),
-              0,
-              common::errors::PreconditionNotMet(
-                  "The all_to_all device id must greater or equal than 0."));
-          phi::XPUPlace place = in_tensors[0].place();
-          auto allocator = std::unique_ptr<phi::Allocator>(
-              new paddle::experimental::DefaultAllocator(place));
-
-          phi::DenseTensorMeta concated_in_tensor_meta(in_tensors[0].dtype(),
-                                                       phi::DDim{in_numel_sum});
-          phi::DenseTensorMeta concated_out_tensor_meta(
-              (*out_tensors)[0].dtype(), phi::DDim{out_numel_sum});
-          phi::DenseTensorMeta split_meta(phi::DataType::INT64,
-                                          phi::DDim{nranks});
-
-          phi::DenseTensor concated_in_tensor = {allocator.get(),
-                                                 concated_in_tensor_meta};
-          phi::DenseTensor concated_out_tensor = {allocator.get(),
-                                                  concated_out_tensor_meta};
-          phi::DenseTensor in_size_tensor = {allocator.get(), split_meta};
-          phi::DenseTensor in_offset_tensor = {allocator.get(), split_meta};
-          phi::DenseTensor out_size_tensor = {allocator.get(), split_meta};
-          phi::DenseTensor out_offset_tensor = {allocator.get(), split_meta};
-
-          if (in_numel_sum > 0) {
-            ConcatTensorByNumel(*GetDeviceContext(place, use_calc_stream),
-                                in_tensors,
-                                &concated_in_tensor);
-          }
+        std::vector<int64_t> in_numel_vec(nranks);
+        std::vector<int64_t> in_offset_vec(nranks);
+        std::vector<int64_t> out_numel_vec(nranks);
+        std::vector<int64_t> out_offset_vec(nranks);
 
-          memory::Copy(place,
-                       in_size_tensor.data(),
-                       phi::CPUPlace(),
-                       in_numel_vec.data(),
-                       in_size_tensor.numel() * sizeof(int64_t));
-
-          memory::Copy(place,
-                       in_offset_tensor.data(),
-                       phi::CPUPlace(),
-                       in_offset_vec.data(),
-                       in_offset_tensor.numel() * sizeof(int64_t));
+        int64_t in_offset = 0;
+        int64_t out_offset = 0;
+        for (int64_t i = 0; i < nranks; i++) {
+          int64_t in_numel = in_tensors[i].numel();
+          int64_t out_numel = (*out_tensors)[i].numel();
 
-          memory::Copy(place,
-                       out_size_tensor.data(),
-                       phi::CPUPlace(),
-                       out_numel_vec.data(),
-                       out_size_tensor.numel() * sizeof(int64_t));
+          in_numel_vec[i] = in_numel;
+          in_offset_vec[i] = in_offset;
+          in_offset += in_numel;
 
-          memory::Copy(place,
-                       out_offset_tensor.data(),
-                       phi::CPUPlace(),
-                       out_offset_vec.data(),
-                       out_offset_tensor.numel() * sizeof(int64_t));
+          out_numel_vec[i] = out_numel;
+          out_offset_vec[i] = out_offset;
+          out_offset += out_numel;
+        }
 
-          comm_context->AllToAllUnequalSplit(&concated_out_tensor,
-                                             concated_in_tensor,
-                                             out_size_tensor,
-                                             out_offset_tensor,
-                                             in_size_tensor,
-                                             in_offset_tensor,
-                                             stream);
+        PADDLE_ENFORCE_GE(
+            in_tensors[0].place().GetDeviceId(),
+            0,
+            common::errors::PreconditionNotMet(
+                "The all_to_all device id must greater or equal than 0."));
+        phi::XPUPlace place = in_tensors[0].place();
+        auto allocator = std::unique_ptr<phi::Allocator>(
+            new paddle::experimental::DefaultAllocator(place));
+
+        phi::DenseTensorMeta concated_in_tensor_meta(in_tensors[0].dtype(),
+                                                     phi::DDim{in_numel_sum});
+        phi::DenseTensorMeta concated_out_tensor_meta((*out_tensors)[0].dtype(),
+                                                      phi::DDim{out_numel_sum});
+        phi::DenseTensorMeta split_meta(phi::DataType::INT64,
+                                        phi::DDim{nranks});
+
+        phi::DenseTensor concated_in_tensor = {allocator.get(),
+                                               concated_in_tensor_meta};
+        phi::DenseTensor concated_out_tensor = {allocator.get(),
+                                                concated_out_tensor_meta};
+        phi::DenseTensor in_size_tensor = {allocator.get(), split_meta};
+        phi::DenseTensor in_offset_tensor = {allocator.get(), split_meta};
+        phi::DenseTensor out_size_tensor = {allocator.get(), split_meta};
+        phi::DenseTensor out_offset_tensor = {allocator.get(), split_meta};
+
+        if (in_numel_sum > 0) {
+          ConcatTensorByNumel(*GetDeviceContext(place, use_calc_stream),
+                              in_tensors,
+                              &concated_in_tensor);
+        }
 
-          if (out_numel_sum > 0) {
-            SplitTensorByNumel(*GetDeviceContext(place, use_calc_stream),
-                               concated_out_tensor,
-                               out_tensors);
-          }
+        memory::Copy(place,
+                     in_size_tensor.data(),
+                     phi::CPUPlace(),
+                     in_numel_vec.data(),
+                     in_size_tensor.numel() * sizeof(int64_t));
+
+        memory::Copy(place,
+                     in_offset_tensor.data(),
+                     phi::CPUPlace(),
+                     in_offset_vec.data(),
+                     in_offset_tensor.numel() * sizeof(int64_t));
+
+        memory::Copy(place,
+                     out_size_tensor.data(),
+                     phi::CPUPlace(),
+                     out_numel_vec.data(),
+                     out_size_tensor.numel() * sizeof(int64_t));
+
+        memory::Copy(place,
+                     out_offset_tensor.data(),
+                     phi::CPUPlace(),
+                     out_offset_vec.data(),
+                     out_offset_tensor.numel() * sizeof(int64_t));
+
+        comm_context->AllToAllUnequalSplit(&concated_out_tensor,
+                                           concated_in_tensor,
+                                           out_size_tensor,
+                                           out_offset_tensor,
+                                           in_size_tensor,
+                                           in_offset_tensor,
+                                           stream);
+
+        if (out_numel_sum > 0) {
+          SplitTensorByNumel(*GetDeviceContext(place, use_calc_stream),
+                             concated_out_tensor,
+                             out_tensors);
         }
       },
       in_tensors,
 
@@ -29,6 +29,10 @@ struct ConcatDenseTensorByNumel {
   void operator()(const DeviceContext &context,
                   const std::vector<phi::DenseTensor> &in,
                   phi::DenseTensor *out) {
+    if (out->numel() == 0) {
+      return;
+    }
+
     auto out_dims = common::vectorize(out->dims());
     auto flattened_out_dims = {out->numel()};
     std::vector<phi::DenseTensor> in_flatten;
@@ -39,11 +43,12 @@ struct ConcatDenseTensorByNumel {
 
     int64_t in_numel_sum = 0;
     for (auto &tensor : in) {
-      phi::DenseTensor tensor_flatten(tensor.Holder(), tensor.meta());
-      tensor_flatten.Resize({tensor.numel()});
-      in_flatten.push_back(tensor_flatten);
-
-      in_numel_sum += tensor.numel();
+      if (tensor.numel() > 0) {
+        phi::DenseTensor tensor_flatten(tensor.Holder(), tensor.meta());
+        tensor_flatten.Resize({tensor.numel()});
+        in_flatten.push_back(tensor_flatten);
+        in_numel_sum += tensor.numel();
+      }
     }
     PADDLE_ENFORCE_EQ(
         out->numel(),
@@ -105,6 +110,10 @@ struct SplitDenseTensorByNumel {
   void operator()(const DeviceContext &context,
                   const phi::DenseTensor &in,
                   std::vector<phi::DenseTensor> *out) {
+    if (in.numel() == 0) {
+      return;
+    }
+
     phi::DenseTensor in_flatten(in.Holder(), in.meta());
     in_flatten.Resize({in.numel()});
 
@@ -115,10 +124,12 @@ struct SplitDenseTensorByNumel {
     int64_t out_numel_sum = 0;
 
     for (auto &tensor : *out) {
-      phi::DenseTensor tensor_flatten(tensor.Holder(), tensor.meta());
-      tensor_flatten.Resize({tensor.numel()});
-      out_flatten.push_back(tensor_flatten);
-      out_numel_sum += tensor.numel();
+      if (tensor.numel() > 0) {
+        phi::DenseTensor tensor_flatten(tensor.Holder(), tensor.meta());
+        tensor_flatten.Resize({tensor.numel()});
+        out_flatten.push_back(tensor_flatten);
+        out_numel_sum += tensor.numel();
+      }
     }
     for (auto &tensor : out_flatten) {
       shape_refer.push_back(&tensor);