[Zero-Dim] distributed scatter/all_to_all support input 0D tensor (#53186)

paddle/fluid/operators/collective/c_allgather_op.cc

@@ -31,8 +31,13 @@ class CAllGatherOp : public framework::OperatorWithKernel {
                       platform::errors::InvalidArgument(
                           "The value of nranks should be >=2."));
     framework::DDim dim = ctx->GetInputDim("X");
-    dim[0] = dim[0] * nranks;
-    if (dim[0] < 0) dim[0] = -1;
+    // 0D use stack/unstack while others use concat/split
+    if (dim.size() == 0) {
+      dim = phi::make_ddim({nranks});
+    } else {
+      dim[0] = dim[0] * nranks;
+      if (dim[0] < 0) dim[0] = -1;
+    }
     ctx->SetOutputDim("Out", dim);
   }
 };

paddle/fluid/operators/collective/c_allgather_op.cu.cc

@@ -57,13 +57,9 @@ class CAllGatherOpCUDAKernel : public framework::OpKernel<T> {
         platform::errors::InvalidArgument(
             "nranks: %s should equal to %s", nranks, comm->nranks()));
 
-    framework::DDim out_dims = in->dims();
-    out_dims[0] *= nranks;
-    out->mutable_data<T>(out_dims, place);
-
     int64_t send_numel = in->numel();
     const T* send_buff = in->data<T>();
-    T* recv_buff = out->data<T>();
+    T* recv_buff = out->mutable_data<T>(place);
 
     gpuStream_t stream = nullptr;
     if (ctx.Attr<bool>("use_calc_stream")) {

paddle/fluid/operators/collective/c_allgather_op.h

@@ -39,15 +39,12 @@ class CAllGatherOpCPUKernel : public framework::OpKernel<T> {
 #if defined(PADDLE_WITH_GLOO)
     auto in = ctx.Input<phi::DenseTensor>("X");
     auto out = ctx.Output<phi::DenseTensor>("Out");
-    framework::DDim out_dims = in->dims();
     auto place = ctx.GetPlace();
-
     auto gloo = paddle::framework::GlooWrapper::GetInstance();
     auto nranks = gloo->Size();
-    out_dims[0] *= nranks;
     int64_t send_numel = in->numel();
     const T* send_buff = in->data<T>();
-    T* recv_buff = out->mutable_data<T>(out_dims, place);
+    T* recv_buff = out->mutable_data<T>(place);
 
     PADDLE_ENFORCE_EQ(
         gloo->IsInitialized(),

paddle/fluid/operators/collective/c_allgather_op_xpu.cc

@@ -43,12 +43,9 @@ class CAllGatherOpXPUKernel : public framework::OpKernel<T> {
         platform::errors::InvalidArgument(
             "nranks: %s should equal to %s", nranks, comm->nranks()));
 
-    framework::DDim out_dims = in->dims();
-    out_dims[0] *= nranks;
-
     size_t numel = in->numel();
     const void* sendbuff = in->data<T>();
-    void* recvbuff = out->mutable_data<T>(out_dims, place);
+    void* recvbuff = out->mutable_data<T>(place);
 
     XPUStream stream = nullptr;
     if (ctx.Attr<bool>("use_calc_stream")) {

paddle/phi/core/ddim.cc

@@ -154,7 +154,7 @@ DDim flatten_to_1d(const DDim& src) { return DDim({product(src)}); }
 DDim stride(const DDim& ddim) {
   DDim strides;
   strides.rank_ = ddim.size();
-  strides[ddim.size() - 1] = 1;
+  if (ddim.size() > 0) strides[ddim.size() - 1] = 1;
   for (int i = ddim.size() - 2; i >= 0; --i) {
     strides[i] = strides[i + 1] * ddim[i + 1];
   }
@@ -164,7 +164,7 @@ DDim stride(const DDim& ddim) {
 DDim stride_numel(const DDim& ddim) {
   DDim strides;
   strides.rank_ = ddim.size();
-  strides[ddim.size() - 1] = ddim[ddim.size() - 1];
+  if (ddim.size() > 0) strides[ddim.size() - 1] = ddim[ddim.size() - 1];
   for (int i = ddim.size() - 2; i >= 0; --i) {
     strides[i] = strides[i + 1] * ddim[i];
   }

python/paddle/distributed/communication/scatter.py

@@ -122,8 +122,7 @@ def scatter_object_list(
             in_obj_sizes.append(obj_size)
         max_obj_size_tensor = max(in_obj_sizes)
     else:
-        # NOTE: shape can be [] after 0D tensor support
-        max_obj_size_tensor = paddle.empty([1], dtype="int64")
+        max_obj_size_tensor = paddle.empty([], dtype="int64")
     stream.broadcast(max_obj_size_tensor, src)
     max_obj_size = int(max_obj_size_tensor.item())
 
@@ -137,8 +136,7 @@ def scatter_object_list(
     out_tensor = paddle.empty([max_obj_size], dtype="uint8")
     scatter(out_tensor, in_tensor_list if rank == src else None, src, group)
 
-    # NOTE: shape can be [] after 0D tensor support
-    out_tensor_size = paddle.empty([1], dtype="int64")
+    out_tensor_size = paddle.empty([], dtype="int64")
     scatter(out_tensor_size, in_obj_sizes if rank == src else None, src, group)
 
     out_object_list.clear()

python/paddle/distributed/communication/stream/all_gather.py

@@ -108,7 +108,11 @@ def _all_gather_in_static_mode(tensor_list, tensor, group, sync_op):
         },
     )
     tensor_list.clear()
-    tensor_list.extend(paddle.split(out, nranks, 0))
+    # 0D use stack/unstack while others use concat/split
+    if len(tensor.shape) == 0:
+        tensor_list.extend(paddle.unstack(out, 0))
+    else:
+        tensor_list.extend(paddle.split(out, nranks, 0))
 
 
 def all_gather(

python/paddle/distributed/communication/stream/all_to_all.py

@@ -78,7 +78,12 @@ def _all_to_all_in_static_mode(
     if isinstance(in_tensor_or_tensor_list, list):
         if len(in_tensor_or_tensor_list) == 0:
             raise RuntimeError("The input tensor_list should not be empty.")
-        in_tensor = paddle.concat(in_tensor_or_tensor_list, axis=0)
+        # 0D use stack/unstack while others use concat/split
+        if len(in_tensor_or_tensor_list[0].shape) == 0:
+            in_tensor = paddle.stack(in_tensor_or_tensor_list, axis=0)
+        else:
+            in_tensor = paddle.concat(in_tensor_or_tensor_list, axis=0)
+
     out_tensor = out_tensor_or_tensor_list
     if isinstance(out_tensor_or_tensor_list, list):
         if len(out_tensor_or_tensor_list) != 0:
@@ -110,7 +115,13 @@ def _all_to_all_in_static_mode(
     if isinstance(out_tensor_or_tensor_list, list):
         if not sync_op:
             dist.wait(out_tensor, use_calc_stream=False)
-        out_tensor_or_tensor_list.extend(paddle.split(out_tensor, nranks, 0))
+        # 0D use stack/unstack while others use concat/split
+        if len(in_tensor_or_tensor_list[0].shape) == 0:
+            out_tensor_or_tensor_list.extend(paddle.unstack(out_tensor, 0))
+        else:
+            out_tensor_or_tensor_list.extend(
+                paddle.split(out_tensor, nranks, 0)
+            )
 
     return None
 

python/paddle/distributed/communication/stream/scatter.py

@@ -91,7 +91,11 @@ def _scatter_in_static_mode(
                 )
         else:
             tensor_list = [tensor for _ in range(nranks)]
-        input_tensor = paddle.concat(tensor_list, axis=0)
+        # 0D use stack/unstack while others use concat/split
+        if len(tensor_list[0].shape) == 0:
+            input_tensor = paddle.stack(tensor_list, axis=0)
+        else:
+            input_tensor = paddle.concat(tensor_list, axis=0)
 
     ring_id = 0 if group is None else group.id