chore: import upstream snapshot with attribution

python/paddle/distributed/communication/__init__.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .all_gather import all_gather, all_gather_object  # noqa: F401
+from .all_reduce import all_reduce  # noqa: F401
+from .all_to_all import alltoall, alltoall_single  # noqa: F401
+from .batch_isend_irecv import P2POp, batch_isend_irecv  # noqa: F401
+from .broadcast import broadcast, broadcast_object_list  # noqa: F401
+from .gather import gather  # noqa: F401
+from .group import (  # noqa: F401
+    barrier,
+    destroy_process_group,
+    get_backend,
+    get_group,
+    is_initialized,
+    wait,
+)
+from .recv import irecv, recv, recv_object_list  # noqa: F401
+from .reduce import ReduceOp, reduce  # noqa: F401
+from .reduce_scatter import reduce_scatter  # noqa: F401
+from .scatter import scatter, scatter_object_list  # noqa: F401
+from .send import isend, send, send_object_list  # noqa: F401

python/paddle/distributed/communication/all_gather.py

@@ -0,0 +1,156 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+import numpy as np
+
+import paddle
+from paddle import framework
+from paddle.distributed.communication import stream
+
+from .serialization_utils import (
+    convert_object_to_tensor,
+    convert_tensor_to_object,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+    _T = TypeVar("_T")
+
+
+def all_gather(
+    tensor_list: list[Tensor],
+    tensor: Tensor,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task | None:
+    """
+
+    Gather tensors from all participators and all get the result. As shown
+    below, one process is started with a GPU and the data of this process is represented
+    by its group rank. Through the all_gather operator, each GPU will have data
+    from all GPUs.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/allgather.png
+        :width: 800
+        :alt: all_gather
+        :align: center
+
+    Args:
+        tensor_list (list): A list of output Tensors. Every element in the list must be a Tensor whose data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool, bfloat16, complex64 or complex128.
+        tensor (Tensor): The Tensor to send. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool, bfloat16, complex64 or complex128.
+        group (Group|None, optional): The group instance return by new_group or None for global default group.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> tensor_list = []  # type: ignore
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> dist.all_gather(tensor_list, data)
+            >>> print(tensor_list)
+            >>> # [[[4, 5, 6], [4, 5, 6]], [[1, 2, 3], [1, 2, 3]]] (2 GPUs)
+    """
+    return stream.all_gather(tensor_list, tensor, group, sync_op)
+
+
+def all_gather_object(
+    object_list: list[_T] | list[None], obj: _T, group: Group = None
+) -> None:
+    """
+
+    Gather picklable objects from all participators and all get the result. Similar to all_gather(), but python object can be passed in.
+
+    After the call, ``object_list[i]`` holds the object gathered from rank ``i``. Both
+    initialization styles below are supported and produce the same result, which is
+    consistent with :func:`torch.distributed.all_gather_object`:
+
+    - Pre-allocated list of length ``world_size`` (PyTorch style):
+      ``object_list = [None for _ in range(dist.get_world_size())]``
+    - Empty list (Paddle legacy style): ``object_list = []`` - the list is extended in
+      place to hold ``world_size`` items.
+
+    Args:
+        object_list (list): A list of output object. The datatype of every element in the list is same as the input obj.
+        obj (Any): The picklable object to send.
+        group (Group): The group instance return by new_group or None for global default group.
+
+    Returns:
+        None.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> object_list = [None for _ in range(dist.get_world_size())]
+            >>> if dist.get_rank() == 0:
+            ...     obj = {"foo": [1, 2, 3]}
+            >>> else:
+            ...     obj = {"bar": [4, 5, 6]}
+            >>> dist.all_gather_object(object_list, obj)
+            >>> print(object_list)
+            >>> # [{'foo': [1, 2, 3]}, {'bar': [4, 5, 6]}] (2 GPUs)
+    """
+    assert framework.in_dynamic_mode(), (
+        "all_gather_object doesn't support static graph mode."
+    )
+
+    tensor, len_of_tensor = convert_object_to_tensor(obj)
+
+    # gather len_of_tensor from all ranks
+    list_len_of_tensor = []
+    all_gather(list_len_of_tensor, len_of_tensor, group)
+    # get the max length from list
+    max_len_of_tensor = int(max(list_len_of_tensor).item())
+    # resize the input tensor to max length avoid hang in all gather
+    # Note(liyurui): Maybe we should support various length all_gather?
+    # Now this operation is efficient for we don't support resize in python.
+    numpy_data = tensor.numpy()
+    numpy_data = np.resize(numpy_data, [max_len_of_tensor])
+    input_tensor = paddle.to_tensor(numpy_data)
+
+    tensor_list = []
+    all_gather(tensor_list, input_tensor, group)
+    # Ensure object_list has enough slots for all gathered objects
+    while len(object_list) < len(tensor_list):
+        object_list.append(None)
+
+    for i, tensor in enumerate(tensor_list):
+        object_list[i] = convert_tensor_to_object(tensor, list_len_of_tensor[i])

python/paddle/distributed/communication/all_reduce.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.distributed.communication import stream
+from paddle.distributed.communication.reduce import ReduceOp
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+    from paddle.distributed.communication.reduce import _ReduceOp
+
+
+def all_reduce(
+    tensor: Tensor,
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task:
+    """
+
+    Reduce a tensor over all ranks so that all get the result.
+    As shown below, one process is started with a GPU and the data of this process is represented
+    by its group rank. The reduce operator is sum. Through all_reduce operator,
+    each GPU will have the sum of the data from all GPUs.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/allreduce.png
+        :width: 800
+        :alt: all_reduce
+        :align: center
+
+    Args:
+        tensor (Tensor): The input Tensor. It also works as the output Tensor. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8 or bool.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD|ReduceOp.AVG, optional): The operation used. Default value is ReduceOp.SUM.
+        group (Group|None, optional): The group instance return by new_group or None for global default group.
+        sync_op (bool, optional): Whether this op is a sync op. Default value is True.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> dist.all_reduce(data)
+            >>> print(data)
+            >>> # [[5, 7, 9], [5, 7, 9]] (2 GPUs)
+    """
+    # AVG is only supported when nccl >= 2.10
+    if op == ReduceOp.AVG and paddle.base.core.nccl_version() < 21000:
+        group = (
+            paddle.distributed.collective._get_global_group()
+            if group is None
+            else group
+        )
+        tensor.scale_(1.0 / group.nranks)
+        return stream.all_reduce(
+            tensor,
+            op=ReduceOp.SUM,
+            group=group,
+            sync_op=sync_op,
+            use_calc_stream=False,
+        )
+
+    return stream.all_reduce(
+        tensor, op=op, group=group, sync_op=sync_op, use_calc_stream=False
+    )

python/paddle/distributed/communication/all_to_all.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle.distributed.communication import stream
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def alltoall(
+    out_tensor_list: list[Tensor],
+    in_tensor_list: list[Tensor],
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task:
+    """
+    Scatter tensors in in_tensor_list to all participators averagely and gather the result tensors in out_tensor_list.
+    As shown below, the in_tensor_list in GPU0 includes 0_0 and 0_1, and GPU1 includes 1_0 and 1_1.
+    Through alltoall operator, the 0_0 in GPU0 will be sent to GPU0 and 0_1 to GPU1, 1_0 in GPU1 sent to GPU0 and 1_1 to GPU1.
+    Finally the out_tensor_list in GPU0 includes 0_0 and 1_0, and GPU1 includes 0_1 and 1_1.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/alltoall.png
+        :width: 800
+        :alt: alltoall
+        :align: center
+
+    Args:
+        out_tensor_list (List[Tensor]): List of tensors to be gathered one per rank. The data type of each tensor should be the same as the input tensors.
+        in_tensor_list (List[Tensor]): List of tensors to scatter one per rank. The data type of each tensor
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+
+            >>> # all_to_all with equal split sizes
+            >>> out_tensor_list = []  # type: ignore
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([[1, 2, 3], [4, 5, 6]])
+            ...     data2 = paddle.to_tensor([[7, 8, 9], [10, 11, 12]])
+            >>> else:
+            ...     data1 = paddle.to_tensor([[13, 14, 15], [16, 17, 18]])
+            ...     data2 = paddle.to_tensor([[19, 20, 21], [22, 23, 24]])
+            >>> dist.alltoall(out_tensor_list, [data1, data2])
+            >>> print(out_tensor_list)
+            >>> # [[[1, 2, 3], [4, 5, 6]], [[13, 14, 15], [16, 17, 18]]] (2 GPUs, out for rank 0)
+            >>> # [[[7, 8, 9], [10, 11, 12]], [[19, 20, 21], [22, 23, 24]]] (2 GPUs, out for rank 1)
+
+            >>> # all_to_all with unequal split sizes
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([[1, 2, 3], [4, 5, 6]])  # shape: (2, 3)
+            ...     data2 = paddle.to_tensor([7])  # shape: (1, )
+            ...     out_data1 = paddle.empty((2, 3), dtype=data1.dtype)
+            ...     out_data2 = paddle.empty((3, 2), dtype=data1.dtype)
+            >>> else:
+            ...     data1 = paddle.to_tensor([[8, 9], [10, 11], [12, 13]]) # shape: (3, 2)
+            ...     data2 = paddle.to_tensor([[14, 15, 16, 17]])           # shape: (1, 4)
+            ...     out_data1 = paddle.empty((1,), dtype=data1.dtype)
+            ...     out_data2 = paddle.empty((1, 4), dtype=data1.dtype)
+            >>> dist.alltoall([out_data1, out_data2], [data1, data2])
+            >>> print([out_data1, out_data2])
+            >>> # [[[1, 2, 3], [4, 5, 6]], [[8, 9], [10, 11], [12, 13]]]  (2 GPUs, out for rank 0)
+            >>> # [[7], [[14, 15, 16, 17]]]                               (2 GPUs, out for rank 1)
+    """
+    return stream.alltoall(
+        out_tensor_list, in_tensor_list, group, sync_op, False
+    )
+
+
+def alltoall_single(
+    out_tensor: Tensor,
+    in_tensor: Tensor,
+    in_split_sizes: list[int] | None = None,
+    out_split_sizes: list[int] | None = None,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task:
+    """
+    Scatter a single input tensor to all participators and gather the received tensors in out_tensor.
+
+    Note:
+        ``alltoall_single`` is only supported in eager mode.
+
+    Args:
+        out_tensor (Tensor): Output Tensor. The data type should be the same as the data type of the input Tensor.
+        in_tensor (Tensor): Input tensor. The data type should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        in_split_sizes (list[int]|None, optional): Split sizes of ``in_tensor`` for dim[0]. If not given, dim[0] of ``in_tensor``
+            must be divisible by group size and ``in_tensor`` will be scattered averagely to all participators. Default: None.
+        out_split_sizes (list[int]|None, optional): Split sizes of ``out_tensor`` for dim[0]. If not given, dim[0] of ``out_tensor``
+            must be divisible by group size and ``out_tensor`` will be gathered averagely from all participators. Default: None.
+        group (Group|None, optional): The group instance return by ``new_group`` or None for global default group. Default: None.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> rank = dist.get_rank()
+            >>> size = dist.get_world_size()
+
+            >>> # case 1 (2 GPUs)
+            >>> data = paddle.arange(2, dtype='int64') + rank * 2
+            >>> # data for rank 0: [0, 1]
+            >>> # data for rank 1: [2, 3]
+            >>> output = paddle.empty([2], dtype='int64')
+            >>> dist.alltoall_single(output, data)
+            >>> print(output)
+            >>> # output for rank 0: [0, 2]
+            >>> # output for rank 1: [1, 3]
+
+            >>> # case 2 (2 GPUs)
+            >>> in_split_sizes = [i + 1 for i in range(size)]
+            >>> # in_split_sizes for rank 0: [1, 2]
+            >>> # in_split_sizes for rank 1: [1, 2]
+            >>> out_split_sizes = [rank + 1 for i in range(size)]
+            >>> # out_split_sizes for rank 0: [1, 1]
+            >>> # out_split_sizes for rank 1: [2, 2]
+            >>> data = paddle.ones([sum(in_split_sizes), size], dtype='float32') * rank
+            >>> # data for rank 0: [[0., 0.], [0., 0.], [0., 0.]]
+            >>> # data for rank 1: [[1., 1.], [1., 1.], [1., 1.]]
+            >>> output = paddle.empty([(rank + 1) * size, size], dtype='float32')
+            >>> group = dist.new_group([0, 1])
+            >>> task = dist.alltoall_single(
+            ...     data,
+            ...     output,
+            ...     in_split_sizes,
+            ...     out_split_sizes,
+            ...     sync_op=False,
+            ...     group=group,
+            ... )
+            >>> task.wait()
+            >>> print(output)
+            >>> # output for rank 0: [[0., 0.], [1., 1.]]
+            >>> # output for rank 1: [[0., 0.], [0., 0.], [1., 1.], [1., 1.]]
+
+    """
+    return stream.alltoall_single(
+        out_tensor,
+        in_tensor,
+        out_split_sizes,
+        in_split_sizes,
+        group,
+        sync_op,
+        False,
+    )

python/paddle/distributed/communication/batch_isend_irecv.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import contextlib
+from typing import TYPE_CHECKING
+
+import paddle.distributed as dist
+from paddle import framework
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator, Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed import Group
+
+    _P2POpType = Callable[[Tensor, int, Group], task]
+
+
+class P2POp:
+    """
+    A class that makes point-to-point operations for "batch_isend_irecv".
+
+    This class creates the type of P2P operation, communication buffer, peer rank,
+    Group. Instances of this class will be passed to
+    ``paddle.distributed.batch_isend_irecv`` for point-to-point communication.
+
+    Args:
+        op (callable): A function to send data to or receive data from a peer process.
+            The type of ``op`` is either ``paddle.distributed.isend`` or ``paddle.distributed.irecv``.
+        tensor (Tensor): Tensor to send or receive.
+        peer (int): The destination or source rank.
+        group (Group, optional): The group instance return by new_group or None for global
+            default group. Default: None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> rank = dist.get_rank()
+            >>> world_size = dist.get_world_size()
+
+            >>> send_t = paddle.arange(2) + rank
+            >>> # paddle.tensor([0, 1])  # Rank-0
+            >>> # paddle.tensor([1, 2])  # Rank-1
+
+            >>> recv_t = paddle.empty(shape=[2], dtype=send_t.dtype)
+
+            >>> send_op = dist.P2POp(dist.isend, send_t, (rank + 1) % world_size)
+            >>> recv_op = dist.P2POp(dist.irecv, recv_t, (rank - 1 + world_size) % world_size)
+
+    """
+
+    op: _P2POpType
+    tensor: Tensor
+    peer: int
+    group: Group | None
+
+    def __init__(
+        self,
+        op: _P2POpType,
+        tensor: Tensor,
+        peer: int,
+        group: Group | None = None,
+    ) -> None:
+        if op not in [dist.isend, dist.irecv]:
+            raise RuntimeError(
+                "Invalid ``op`` function. Expected ``op`` "
+                "to be of type ``paddle.distributed.isend`` or "
+                "``paddle.distributed.irecv``."
+            )
+
+        self.op = op
+        self.tensor = tensor
+        self.peer = peer
+        self.group = _get_global_group() if group is None else group
+
+
+@contextlib.contextmanager
+def _coalescing_manager(
+    group: Group, tasks: task | None = None
+) -> Generator[None, None, None]:
+    group = _get_global_group() if group is None else group
+    pg = group.process_group
+    pg._start_coalescing()
+    try:
+        yield
+    finally:
+        if tasks is None or len(tasks) == 0:
+            pg._end_coalescing()
+        else:
+            pg._end_coalescing(tasks)
+
+
+def _check_p2p_op_list(p2p_op_list: Sequence[P2POp]) -> None:
+    """
+    Helper to check that the ``p2p_op_list`` is a list of P2POp instances and
+    all ops use the same backend.
+    """
+    if not isinstance(p2p_op_list, list) or not all(
+        isinstance(p2p_op, P2POp) for p2p_op in p2p_op_list
+    ):
+        raise RuntimeError(
+            "Invalid ``p2p_op_list``. Each op is expected to "
+            "to be of type ``paddle.distributed.P2POp``."
+        )
+
+    backend = p2p_op_list[0].group.backend
+    if not all(backend == p2p_op.group.backend for p2p_op in p2p_op_list):
+        raise RuntimeError("All groups need to use the same backend.")
+
+
+def batch_isend_irecv(p2p_op_list: list[P2POp]) -> list[task]:
+    """
+    Send or Receive a batch of tensors asynchronously and return a list of requests.
+
+    Process each of the point-to-point operations in ``p2p_op_list`` and return the
+    corresponding tasks. NCCL are currently supported.
+
+    Args:
+        p2p_op_list (List[P2POp]): A list of point-to-point operations(type of each operator is
+            ``paddle.distributed.P2POp``). The order of the isend/irecv in the list
+            matters and it needs to match with corresponding isend/irecv on the
+            remote end.
+
+    Returns:
+        A list of distributed tasks returned by calling the corresponding
+        op in the op_list.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> rank = dist.get_rank()
+            >>> world_size = dist.get_world_size()
+
+            >>> send_t = paddle.arange(2) + rank
+            >>> # paddle.tensor([0, 1])  # Rank-0
+            >>> # paddle.tensor([1, 2])  # Rank-1
+
+            >>> recv_t = paddle.empty(shape=[2], dtype=send_t.dtype)
+
+            >>> send_op = dist.P2POp(dist.isend, send_t, (rank + 1) % world_size)
+            >>> recv_op = dist.P2POp(dist.irecv, recv_t, (rank - 1 + world_size) % world_size)
+
+            >>> tasks = dist.batch_isend_irecv([send_op, recv_op])
+
+            >>> for task in tasks:
+            ...     task.wait()
+
+            >>> print(recv_t)
+            >>> # paddle.tensor([1, 2])     # Rank-0
+            >>> # paddle.tensor([0, 1])     # Rank-1
+    """
+    _check_p2p_op_list(p2p_op_list)
+    group = p2p_op_list[0].group
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        backend = group.backend
+        tasks = []
+        with _coalescing_manager(group, tasks):
+            for p2p_op in p2p_op_list:
+                op = p2p_op.op
+                tensor = p2p_op.tensor
+                peer = p2p_op.peer
+                comm_group = p2p_op.group
+                task = op(tensor, peer, comm_group)
+                if task is not None:
+                    tasks.append(task)
+        return tasks
+    else:
+        raise RuntimeError("Don't support static graph mode currently.")

python/paddle/distributed/communication/broadcast.py

@@ -0,0 +1,149 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+from paddle.distributed.communication import stream
+
+from .serialization_utils import (
+    convert_object_to_tensor,
+    convert_tensor_to_object,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def broadcast(
+    tensor: Tensor, src: int, group: Group | None = None, sync_op: bool = True
+) -> task:
+    """
+
+    Broadcast a tensor from the source to all others.
+    As shown below, one process is started with a GPU and GPU0 owns data 0. Through broadcast operator,
+    data 0 will be sent to all GPUs from GPU0.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/broadcast.png
+        :width: 800
+        :alt: broadcast
+        :align: center
+
+    Args:
+        tensor (Tensor): The tensor to send if current rank is the source, or the tensor to receive otherwise. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool, bfloat16, complex64 or complex128.
+        src (int): The source rank in global view.
+        group (Group, optional): The group instance return by new_group or None for global default group.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> dist.broadcast(data, src=1)
+            >>> print(data)
+            >>> # [[1, 2, 3], [1, 2, 3]] (2 GPUs)
+    """
+    return stream.broadcast(
+        tensor,
+        src,
+        group=group,
+        sync_op=sync_op,
+        use_calc_stream=False,
+    )
+
+
+def broadcast_object_list(
+    object_list: list[Any], src: int, group: Group | None = None
+) -> None:
+    """
+
+    Broadcast picklable objects from the source to all others. Similar to broadcast(), but python object can be passed in.
+
+    Args:
+        object_list (list): The list of objects to send if current rank is the source, or the list of objects to receive otherwise.
+        src (int): The source rank in global view.
+        group (Group): The group instance return by new_group or None for global default group.
+
+    Returns:
+        None.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     object_list = [{"foo": [1, 2, 3]}]
+            >>> else:
+            ...     object_list = [{"bar": [4, 5, 6]}]
+            >>> dist.broadcast_object_list(object_list, src=1)
+            >>> print(object_list)
+            >>> # [{"bar": [4, 5, 6]}] (2 GPUs)
+    """
+    assert framework.in_dynamic_mode(), (
+        "broadcast_object_list doesn't support static graph mode."
+    )
+
+    rank = dist.get_rank()
+    obj_tensors = []
+    obj_nums = len(object_list)
+
+    if rank == src:
+        obj_sizes = []
+        for obj in object_list:
+            obj_tensor, obj_size = convert_object_to_tensor(obj)
+            obj_tensors.append(obj_tensor)
+            obj_sizes.append(obj_size)
+        obj_size_tensor = paddle.stack(obj_sizes)
+    else:
+        obj_size_tensor = paddle.empty([obj_nums], dtype="int64")
+    broadcast(obj_size_tensor, src, group)
+
+    if rank == src:
+        # cast to uint8 to keep the same dtype
+        obj_data_tensor = paddle.concat(obj_tensors).cast("uint8")
+    else:
+        data_len = paddle.sum(obj_size_tensor).item()
+        obj_data_tensor = paddle.empty([data_len], dtype="uint8")
+    broadcast(obj_data_tensor, src, group)
+
+    offset = 0
+    for i in range(obj_nums):
+        data_len = obj_size_tensor[i]
+        object_list[i] = convert_tensor_to_object(
+            obj_data_tensor[offset : offset + data_len], data_len
+        )
+        offset += data_len

python/paddle/distributed/communication/deep_ep/__init__.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from .buffer import Buffer, M2NBuffer
+from .utils import (
+    EventOverlap,
+    get_event_from_calc_stream,
+    get_event_from_comm_stream,
+    get_event_from_custom_stream,
+)
+
+__all__ = [
+    "Buffer",
+    "M2NBuffer",
+    "EventOverlap",
+    "get_event_from_calc_stream",
+    "get_event_from_comm_stream",
+    "get_event_from_custom_stream",
+]

python/paddle/distributed/communication/deep_ep/utils.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# The file has been adapted from DeepSeek DeepEP project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import paddle
+    from paddle.base.core import EventHandle
+
+import paddle
+
+
+class EventOverlap:
+    """
+    A wrapper class to manage CUDA events, also for better overlapping convenience.
+
+    Attributes:
+        event: the CUDA event captured.
+        extra_tensors: an easier way to simulate tensor `record_stream`, may be useful with CUDA graph.
+    """
+
+    def __init__(
+        self,
+        event: EventHandle | None = None,
+        extra_tensors: tuple[paddle.Tensor] | None = None,
+    ) -> None:
+        """
+        Initialize the class.
+
+        Arguments:
+            event: the CUDA event captured.
+            extra_tensors: an easier way to simulate tensor `record_stream`, may be useful with CUDA graph.
+        """
+        self.event = event
+
+        # NOTES: we use extra tensors to achieve stream recording, otherwise,
+        # stream recording will be incompatible with CUDA graph.
+        self.extra_tensors = extra_tensors
+
+    def current_stream_wait(self) -> None:
+        """
+        The current stream waits for the event to be finished.
+        """
+        assert self.event is not None
+        self.event.current_stream_wait()
+
+    def calc_stream_wait(self, group_idx) -> None:
+        self.event.calc_stream_wait(group_idx)
+
+    def comm_stream_wait(self, group_idx) -> None:
+        self.event.comm_stream_wait(group_idx)
+
+    def __enter__(self) -> Any:
+        """
+        Utility for overlapping and Python `with` syntax.
+
+        You can overlap the kernels on the current stream with the following example:
+        ```python
+        event_overlap = event_after_all_to_all_kernels()
+        with event_overlap():
+            do_something_on_current_stream()
+        # After exiting the `with` scope, the current stream with wait the event to be finished.
+        ```
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """
+        Utility for overlapping and Python `with` syntax.
+
+        Please follow the example in the `__enter__` function.
+        """
+        if self.event is not None:
+            self.event.current_stream_wait()
+
+
+def get_event_from_calc_stream(group_id: int) -> EventOverlap:
+    return EventOverlap(
+        event=paddle.base.core.get_event_handle_from_calc_stream(group_id)
+    )
+
+
+def get_event_from_comm_stream(group_id: int) -> EventOverlap:
+    return EventOverlap(
+        event=paddle.base.core.get_event_handle_from_comm_stream(group_id)
+    )
+
+
+def get_event_from_custom_stream(stream) -> EventOverlap:
+    return EventOverlap(
+        event=paddle.base.core.get_event_handle_from_custom_stream(stream)
+    )

python/paddle/distributed/communication/gather.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import framework
+from paddle.distributed.communication import stream
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def gather(
+    tensor: Tensor,
+    gather_list: list[Tensor] | None = None,
+    dst: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task | None:
+    """
+
+    Gather tensors from all participators.
+
+    Args:
+        tensor (Tensor): The input Tensor. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        gather_list (list): A list of Tensors to hold the gathered tensors. Every element in the list must be a Tensor whose data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16. Default value is None.
+        dst (int): The dst rank id. Default value is 0.
+        group (Group, optional): The group instance return by new_group or None for global default group.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Async work handle,which can be wait on, if async_op is set to True.
+        None, if not async_op
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> gather_list = []  # type: ignore
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([1, 2, 3])
+            ...     dist.gather(data, gather_list, dst=0)
+            >>> else:
+            ...     data = paddle.to_tensor([4, 5, 6])
+            ...     dist.gather(data, gather_list, dst=0)
+            >>> print(gather_list)
+            >>> # [[1, 2, 3], [4, 5, 6]] (2 GPUs, out for rank 0)
+            >>> # [] (2 GPUs, out for rank 1)
+    """
+    assert framework.in_dynamic_mode(), (
+        "gather doesn't support static graph mode yet."
+    )
+    return stream.gather(tensor, gather_list, dst, group, sync_op)

python/paddle/distributed/communication/group.py

@@ -0,0 +1,469 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Literal
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import ProcessGroup
+
+
+class Group:
+    """
+    The abstract representation of group.
+    """
+
+    def __init__(
+        self,
+        rank_in_group: int,
+        id: int,
+        ranks: list[int],
+        pg: ProcessGroup | None = None,
+        name: str | None = None,
+    ) -> None:
+        self._rank_in_group = rank_in_group
+        self._world_size = len(ranks) if rank_in_group >= 0 else -1
+        self._id = id
+        self._ranks = ranks
+        self._pg = pg
+        self._name = name
+
+    @property
+    def rank(self) -> int:
+        return self._rank_in_group
+
+    @property
+    def ranks(self) -> list[int]:
+        return self._ranks
+
+    @property
+    def nranks(self) -> int:
+        return len(self._ranks)
+
+    @property
+    def name(self) -> str | None:
+        return self._name
+
+    @property
+    def process_group(self) -> ProcessGroup:
+        return self._pg
+
+    @property
+    def world_size(self) -> int:
+        return self._world_size
+
+    @property
+    def backend(self) -> str:
+        return self._pg.name()
+
+    @property
+    def id(self) -> int:
+        return self._id
+
+    def is_member(self) -> bool:
+        if self.rank < 0:
+            return False
+        if self.nranks < 2:
+            return False
+        return True
+
+    def get_group_rank(self, rank: int) -> int | Literal[-1]:
+        if self.is_member():
+            return self.ranks.index(rank)
+        else:
+            return -1
+
+    def get_global_rank(self, rank: int) -> int | Literal[-1]:
+        """
+        Get the global rank of a process within a group.
+
+        Args:
+            rank (int): The local rank within the group.
+
+        Returns:
+            If the current process is a member of the group, returns the corresponding global rank;
+            otherwise returns -1.
+
+        """
+        if self.is_member():
+            return self.ranks[rank]
+        else:
+            return -1
+
+    def __repr__(self) -> str:
+        debug_str = (
+            f"rank: {self.rank}, nranks: {self.nranks}, id: {self.id}, ranks: "
+        )
+        debug_str += ", ".join(map(str, self.ranks))
+        debug_str += "; name: "
+        debug_str += self.name if self.name else "None"
+        return debug_str
+
+
+class _GroupManager:
+    global_group_id = 0
+    group_map_by_id = {}
+
+
+class _DistGroupMeta(type):
+    """Metaclass exposing :attr:`group.WORLD` as a dynamic class property."""
+
+    @property
+    def WORLD(cls) -> Group | None:
+        try:
+            return _get_global_group()
+        except RuntimeError:
+            return None
+
+    @WORLD.setter
+    def WORLD(cls, value: Group | None) -> None:
+        # Validate before mutating any registry so a rejected assignment
+        # leaves the existing default group intact.
+        if value is not None:
+            if not isinstance(value, Group):
+                raise TypeError(
+                    "group.WORLD must be a Group instance or None, got "
+                    f"{type(value).__name__}"
+                )
+            if value.id != _GroupManager.global_group_id:
+                raise ValueError(
+                    f"group.WORLD expects a Group with id="
+                    f"{_GroupManager.global_group_id}, got id={value.id}"
+                )
+
+        # Lazy import: ``collective`` imports from this module at its top.
+        from paddle.distributed import collective as _coll
+
+        prev = _GroupManager.group_map_by_id.pop(
+            _GroupManager.global_group_id, None
+        )
+        _coll._group_map.pop(_coll._global_env_gid, None)
+        _coll._group_map_by_name.pop(_coll._default_group_name, None)
+        if prev is not None:
+            _coll._group_map_backend.pop(prev, None)
+
+        if value is None:
+            return
+
+        _GroupManager.group_map_by_id[_GroupManager.global_group_id] = value
+        _coll._group_map[_coll._global_env_gid] = value
+        _coll._group_map_by_name[_coll._default_group_name] = value
+        if value._pg is not None:
+            # ``ProcessGroup.name()`` returns the C++ backend name in upper
+            # case (e.g. ``NCCL``); the registry is keyed by the lower-case
+            # Python form used in ``_valid_backend_list``.
+            _coll._group_map_backend[value] = value._pg.name().lower()
+
+
+class _DistGroupNamespace(metaclass=_DistGroupMeta):
+    """Namespace exposing :attr:`WORLD`, re-exported as
+    :data:`paddle.distributed.group`.
+    """
+
+
+def _get_global_group():
+    if _GroupManager.global_group_id not in _GroupManager.group_map_by_id:
+        raise RuntimeError("The global group is not initialized.")
+    return _GroupManager.group_map_by_id[_GroupManager.global_group_id]
+
+
+def _add_new_group(group):
+    if group.id in _GroupManager.group_map_by_id:
+        raise RuntimeError(f"The group with id {group.id} already exist.")
+    _GroupManager.group_map_by_id[group.id] = group
+
+
+def _is_global_group(group):
+    return group.id == _GroupManager.global_group_id
+
+
+def _warn_cur_rank_not_in_group(group):
+    global_rank = dist.get_rank()
+    if group and not group.is_member():
+        return True
+    return False
+
+
+def _get_or_throw_group_rank(global_rank, group):
+    group_rank = group.get_group_rank(global_rank)
+    assert group_rank >= 0, (
+        f"The input rank {global_rank} can not be found inside the group {group.name}"
+    )
+    return group_rank
+
+
+def is_initialized() -> bool:
+    """
+
+    Check whether the distributed environment has been initialized
+
+    Returns:
+        `True` if distributed environment has been initialized, otherwise `False`.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+
+            >>> print(paddle.distributed.is_initialized())
+            False
+
+            >>> paddle.distributed.init_parallel_env()
+            >>> print(paddle.distributed.is_initialized())
+            True
+
+    """
+    return _GroupManager.global_group_id in _GroupManager.group_map_by_id
+
+
+def destroy_process_group(group: Group | None = None) -> None:
+    """
+    Destroy a given group for communication
+
+    Args:
+        group (Group, optional): The group to be destroyed. All of process groups, including
+                                        the default group, will be destroyed and the distributed
+                                        environment will be deinitialized.
+
+    Returns : None
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> group = dist.new_group([0, 1])
+
+            >>> dist.destroy_process_group(group)
+            >>> print(dist.is_initialized())
+            True
+            >>> dist.destroy_process_group()
+            >>> print(dist.is_initialized())
+            False
+
+    """
+    group = _get_global_group() if group is None else group
+    assert group.id in _GroupManager.group_map_by_id, (
+        f"Destroy group with id {group.id} is invalid."
+    )
+    if _is_global_group(group):
+        _GroupManager.group_map_by_id.clear()
+        # The default group is also registered in the collective-layer
+        # registries by ``init_parallel_env``; clear those slots too so a
+        # follow-up ``init_process_group`` re-creates the default group
+        # rather than hitting ``init_parallel_env``'s early-return path.
+        from paddle.distributed import collective as _coll
+
+        _coll._group_map.pop(_coll._global_env_gid, None)
+        _coll._group_map_by_name.pop(_coll._default_group_name, None)
+        _coll._group_map_backend.pop(group, None)
+    else:
+        del _GroupManager.group_map_by_id[group.id]
+
+
+def get_group(id: int = 0) -> Group:
+    """
+
+    Get group instance by group id.
+
+    Args:
+        id (int): the group id. Default value is 0.
+
+    Returns:
+        Group: the group instance.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> gid = paddle.distributed.new_group([2, 4, 6])
+            >>> paddle.distributed.get_group(gid.id)
+
+    """
+
+    if id in _GroupManager.group_map_by_id:
+        return _GroupManager.group_map_by_id[id]
+    warnings.warn(f"Group {id} is not initialized.")
+    return None
+
+
+def _sync_calc_stream(tensor):
+    if framework.in_dynamic_mode():
+        return paddle._C_ops.sync_calc_stream(tensor)
+    else:
+        op_type = 'c_sync_calc_stream'
+        helper = framework.LayerHelper(op_type, **locals())
+        helper.append_op(
+            type=op_type,
+            inputs={'X': [tensor]},
+            outputs={'Out': [tensor]},
+        )
+
+
+def _sync_comm_stream(tensor, ring_id=0):
+    if framework.in_dynamic_mode():
+        return paddle._C_ops.sync_comm_stream([tensor], ring_id)
+    else:
+        op_type = 'c_sync_comm_stream'
+        helper = framework.LayerHelper(op_type, **locals())
+        helper.append_op(
+            type=op_type,
+            inputs={'X': [tensor]},
+            outputs={'Out': [tensor]},
+            attrs={'ring_id': ring_id},
+        )
+
+
+def wait(
+    tensor: Tensor, group: Group | None = None, use_calc_stream: bool = True
+) -> None:
+    """
+
+    wait to sync stream for group.
+
+    Args:
+        tensor (Tensor): The Tensor used before sync.
+        group (Group): The Group instance to perform sync.
+        use_calc_stream (bool): Whether to use calculation stream (True) or communication stream (False).
+            Default to True.
+
+    Returns:
+        None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+
+            >>> paddle.distributed.init_parallel_env()
+            >>> tindata = paddle.randn(shape=[2, 3])
+            >>> paddle.distributed.all_reduce(tindata, sync_op=True)
+            >>> paddle.distributed.wait(tindata)
+
+    """
+    if group is not None and not group.is_member():
+        return
+
+    if use_calc_stream:
+        _sync_calc_stream(tensor)
+    else:
+        ring_id = 0 if group is None else group.id
+        _sync_comm_stream(tensor, ring_id)
+
+
+def barrier(group: Group | None = None) -> None:
+    """
+
+    Barrier among all participators in the group.
+
+    Args:
+        group (Group): The group instance return by new_group or None for global default group.
+
+    Returns:
+        None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> from paddle.distributed import init_parallel_env
+
+            >>> paddle.set_device(f'gpu:{paddle.distributed.ParallelEnv().dev_id}')
+            >>> init_parallel_env()
+            >>> paddle.distributed.barrier()
+    """
+    if group is not None and not group.is_member():
+        return
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        place = framework._current_expected_place()
+        if isinstance(place, framework.CPUPlace):
+            task = group.process_group.barrier()
+        else:
+            device_id = place.get_device_id()
+            task = group.process_group.barrier(device_id)
+        task.wait()
+        return
+
+    ring_id = 0 if group is None else group.id
+
+    barrier_tensor = paddle.full([1], 1, dtype="int32")
+    if framework.in_dynamic_mode():
+        # barrier is not available in xpu for now
+        if not paddle.framework.core.is_compiled_with_xpu():
+            return paddle._legacy_C_ops.barrier(
+                barrier_tensor, barrier_tensor, 'ring_id', ring_id
+            )
+    else:
+        op_type = 'barrier'
+        if not isinstance(ring_id, int):
+            raise ValueError("The type of 'group' for barrier must be int.")
+        helper = framework.LayerHelper(op_type, **locals())
+        helper.append_op(
+            type=op_type,
+            inputs={'X': [barrier_tensor]},
+            outputs={'Out': [barrier_tensor]},
+            attrs={'ring_id': ring_id},
+        )
+
+
+def get_backend(group: Group | None = None) -> str:
+    """
+    Get the backend of given group.
+
+    Args:
+        group (Group): The group to work on. Use the global group as default.
+
+    Returns:
+        Returns the name of the given group backend.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+
+            >>> paddle.distributed.init_parallel_env()
+            >>> paddle.distributed.get_backend()
+            NCCL
+    """
+    if _warn_cur_rank_not_in_group(group):
+        raise RuntimeError("Invalid group specified")
+
+    group = _get_global_group() if group is None else group
+    return group.backend

python/paddle/distributed/communication/recv.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import paddle
+from paddle.distributed.communication import stream
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.serialization_utils import (
+    convert_tensor_to_object,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def recv(
+    tensor: Tensor,
+    src: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task:
+    """
+    Receive a tensor to the sender.
+
+    Args:
+        tensor (Tensor): The tensor to receive. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        src (int): The source rank id.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([7, 8, 9])
+            ...     dist.send(data, dst=1)
+            >>> else:
+            ...     data = paddle.to_tensor([1, 2, 3])
+            ...     dist.recv(data, src=0)
+            >>> print(data)
+            >>> # [7, 8, 9] (2 GPUs)
+    """
+    return stream.recv(
+        tensor, src=src, group=group, sync_op=sync_op, use_calc_stream=False
+    )
+
+
+def irecv(
+    tensor: Tensor, src: int | None = None, group: Group | None = None
+) -> task:
+    """
+    Receive a tensor to the sender.
+
+    Args:
+        tensor (Tensor): The Tensor to receive. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        src (int): The source rank id.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([7, 8, 9])
+            ...     task = dist.isend(data, dst=1)
+            >>> else:
+            ...     data = paddle.to_tensor([1, 2, 3])
+            ...     task = dist.irecv(data, src=0)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> print(data)
+            >>> # [7, 8, 9] (2 GPUs)
+    """
+    return recv(tensor, src, group, sync_op=False)
+
+
+def recv_object_list(
+    object_list: list[Any],
+    src: int | None = None,
+    group: Group | None = None,
+    src_in_group: int | None = None,
+):
+    """
+    Receive a list of Python objects from the sender.
+
+    Args:
+        object_list (list): The list to store received objects. Must be pre-allocated with correct size.
+        src (int, optional): The source rank id. Default: 0.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        src_in_group (int, optional): The source rank within the group. Cannot be specified together with src. Default: None.
+
+    Returns:
+        This function does not return any value.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = ["hello", {"key": 100}, [1, 2, 3]]
+            ...     dist.send_object_list(data, dst=1)
+            >>> else:
+            ...     data = [None] * 3
+            ...     dist.recv_object_list(data, src=0)
+            >>> print(data)
+            >>> # ["hello", {"key": 100}, [1, 2, 3]] (2 GPUs)
+    """
+    if object_list is None or len(object_list) == 0:
+        raise ValueError("object_list cannot be None or empty")
+
+    group = _get_global_group() if group is None else group
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if src_in_group is not None:
+        if src is not None:
+            raise ValueError(
+                "Cannot specify both 'src' and 'src_in_group' arguments."
+            )
+        src = group.get_global_rank(src_in_group)
+    else:
+        src = 0 if src is None else src
+
+    object_sizes_tensor = paddle.empty((len(object_list),), dtype='int64')
+    recv(object_sizes_tensor, src=src, group=group)
+
+    total_size = paddle.sum(object_sizes_tensor).item()
+    object_tensor = paddle.empty((total_size,), dtype=paddle.uint8)
+    recv(object_tensor, src=src, group=group)
+
+    offset = 0
+    for i, obj_size in enumerate(object_sizes_tensor):
+        obj_size = obj_size.item()
+        obj_view = object_tensor[offset : offset + obj_size]
+        object_list[i] = convert_tensor_to_object(obj_view, obj_size)
+        offset += obj_size

python/paddle/distributed/communication/reduce.py

@@ -0,0 +1,194 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar, Literal
+
+import paddle
+from paddle import framework
+from paddle.distributed.communication import stream
+
+if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+    _ReduceOp: TypeAlias = Literal[0, 1, 2, 3, 4]
+
+
+class ReduceOp:
+    """
+
+    Specify the type of operation used for element-wise reductions.
+    It should be one of the following values:
+
+        ReduceOp.SUM
+
+        ReduceOp.MAX
+
+        ReduceOp.MIN
+
+        ReduceOp.PROD
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> dist.all_reduce(data, op=dist.ReduceOp.SUM)
+            >>> print(data)
+            >>> # [[5, 7, 9], [5, 7, 9]] (2 GPUs)
+    """
+
+    SUM: ClassVar[Literal[0]] = 0
+    MAX: ClassVar[Literal[1]] = 1
+    MIN: ClassVar[Literal[2]] = 2
+    PROD: ClassVar[Literal[3]] = 3
+    AVG: ClassVar[Literal[4]] = 4
+
+
+def _get_reduce_op(reduce_op):
+    if reduce_op == ReduceOp.SUM:
+        return framework.core.ReduceOp.SUM
+    elif reduce_op == ReduceOp.MAX:
+        return framework.core.ReduceOp.MAX
+    elif reduce_op == ReduceOp.MIN:
+        return framework.core.ReduceOp.MIN
+    elif reduce_op == ReduceOp.PROD:
+        return framework.core.ReduceOp.PRODUCT
+    elif reduce_op == ReduceOp.AVG:
+        return framework.core.ReduceOp.AVG
+
+    raise ValueError(f"Unknown reduce_op type for {reduce_op}.")
+
+
+def _to_inplace_op(op_name):
+    return f"{op_name}_"
+
+
+def reduce(
+    tensor: Tensor,
+    dst: int,
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task:
+    """
+
+    Reduce a tensor to the destination from all others. As shown below, one process is started with a GPU and the data of this process is represented
+    by its group rank. The destination of the reduce operator is GPU0 and the process is sum. Through reduce operator,
+    the GPU0 will owns the sum of all data from all GPUs.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/reduce.png
+        :width: 800
+        :alt: reduce
+        :align: center
+
+    Args:
+        tensor (Tensor): The output Tensor for the destination and the input Tensor otherwise. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        dst (int): The destination rank id.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD|ReduceOp.AVG, optional): The operation used. Default value is ReduceOp.SUM.
+        group (Group|None, optional): The group instance return by new_group or None for global default group.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> dist.reduce(data, dst=0)
+            >>> print(data)
+            >>> # [[5, 7, 9], [5, 7, 9]] (2 GPUs, out for rank 0)
+            >>> # [[1, 2, 3], [1, 2, 3]] (2 GPUs, out for rank 1)
+    """
+    # AVG is only supported when nccl >= 2.10
+    if op == ReduceOp.AVG and (not is_avg_reduce_op_supported()):
+        group = (
+            paddle.distributed.collective._get_global_group()
+            if group is None
+            else group
+        )
+        tensor.scale_(1.0 / group.nranks)
+        return stream.reduce(
+            tensor,
+            dst=dst,
+            op=ReduceOp.SUM,
+            group=group,
+            sync_op=sync_op,
+            use_calc_stream=False,
+        )
+    return stream.reduce(
+        tensor,
+        dst=dst,
+        op=op,
+        group=group,
+        sync_op=sync_op,
+        use_calc_stream=False,
+    )
+
+    # code below will be removed after we remove the old dygraph
+    if group is not None and not group.is_member():
+        return
+    use_calc_stream = sync_op
+    ring_id = 0 if group is None else group.id
+    gdst = dst if group is None else group.get_group_rank(dst)
+    assert gdst >= 0, "dst rank out of group, need global rank"
+
+    if (
+        op == ReduceOp.SUM
+        or op == ReduceOp.MAX
+        or op == ReduceOp.MIN
+        or op == ReduceOp.PROD
+        or op == ReduceOp.AVG
+    ):
+        return paddle._C_ops.reduce(
+            tensor,
+            tensor,
+            'ring_id',
+            ring_id,
+            'root_id',
+            gdst,
+            'reduce_type',
+            op,
+        )
+    else:
+        raise ValueError(f"Unknown parameter: {op}.")
+
+
+def is_avg_reduce_op_supported() -> bool:
+    if paddle.is_compiled_with_cuda():
+        return paddle.base.core.nccl_version() >= 21000
+    else:
+        return False

python/paddle/distributed/communication/reduce_scatter.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.distributed.communication import stream
+from paddle.distributed.communication.reduce import ReduceOp
+from paddle.distributed.communication.stream.reduce_scatter import (
+    _reduce_scatter_base as _reduce_scatter_base_stream,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+    from paddle.distributed.communication.reduce import _ReduceOp
+
+
+def reduce_scatter(
+    tensor: Tensor,
+    tensor_list: list[Tensor],
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task:
+    """
+    Reduces, then scatters a list of tensors to all processes in a group
+
+    Args:
+        tensor (Tensor): The output tensor on each rank. The result will overwrite this tenor after communication. Support
+            float16, float32, float64, int32, int64, int8, uint8 or bool as the input data type.
+        tensor_list (List[Tensor]]): List of tensors to reduce and scatter. Every element in the list must be a Tensor whose data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD|ReduceOp.AVG, optional): The reduction used. If none is given, use ReduceOp.SUM as default.
+        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([0, 1])
+            ...     data2 = paddle.to_tensor([2, 3])
+            >>> else:
+            ...     data1 = paddle.to_tensor([4, 5])
+            ...     data2 = paddle.to_tensor([6, 7])
+            >>> dist.reduce_scatter(data1, [data1, data2])
+            >>> print(data1)
+            >>> # [4, 6] (2 GPUs, out for rank 0)
+            >>> # [8, 10] (2 GPUs, out for rank 1)
+
+    """
+    if op not in [
+        ReduceOp.AVG,
+        ReduceOp.MAX,
+        ReduceOp.MIN,
+        ReduceOp.PROD,
+        ReduceOp.SUM,
+    ]:
+        raise RuntimeError(
+            "Invalid ``op`` function. Expected ``op`` to be of type ``ReduceOp.SUM``, ``ReduceOp.Max``, ``ReduceOp.MIN``, ``ReduceOp.PROD`` or ``ReduceOp.AVG``."
+        )
+    # AVG is only supported when nccl >= 2.10
+    if op == ReduceOp.AVG and paddle.base.core.nccl_version() < 21000:
+        group = (
+            paddle.distributed.collective._get_global_group()
+            if group is None
+            else group
+        )
+        tensor.scale_(1.0 / group.nranks)
+        return stream.reduce_scatter(
+            tensor,
+            tensor_list,
+            op=ReduceOp.SUM,
+            group=group,
+            sync_op=sync_op,
+            use_calc_stream=False,
+        )
+    return stream.reduce_scatter(
+        tensor,
+        tensor_list,
+        op=op,
+        group=group,
+        sync_op=sync_op,
+        use_calc_stream=False,
+    )
+
+
+def _reduce_scatter_base(
+    output: Tensor,
+    input: Tensor,
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task | None:
+    """
+    Reduces, then scatters a flattened tensor to all processes in a group.
+
+    Args:
+        output (Tensor): Output tensor. Its data type should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        input (Tensor): Input tensor that is of size output tensor size times world size. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD): Optional. The operation used. Default: ReduceOp.SUM.
+        group (ProcessGroup, optional): The process group to work on. If None,
+            the default process group will be used.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Async task handle, if sync_op is set to False.
+        None, if sync_op or if not part of the group.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> rank = dist.get_rank()
+            >>> data = paddle.arange(4) + rank
+            >>> # [0, 1, 2, 3] (2 GPUs, for rank 0)
+            >>> # [1, 2, 3, 4] (2 GPUs, for rank 1)
+            >>> output = paddle.empty(shape=[2], dtype=data.dtype)
+            >>> dist.collective._reduce_scatter_base(output, data)
+            >>> print(output)
+            >>> # [1, 3] (2 GPUs, out for rank 0)
+            >>> # [5, 7] (2 GPUs, out for rank 1)
+
+    """
+    if op not in [ReduceOp.MAX, ReduceOp.MIN, ReduceOp.PROD, ReduceOp.SUM]:
+        raise RuntimeError(
+            "Invalid ``op`` function. Expected ``op`` to be of type ``ReduceOp.SUM``, ``ReduceOp.Max``, ``ReduceOp.MIN`` or ``ReduceOp.PROD``."
+        )
+    return _reduce_scatter_base_stream(
+        output,
+        input,
+        op=op,
+        group=group,
+        sync_op=sync_op,
+        use_calc_stream=False,
+    )

python/paddle/distributed/communication/scatter.py

@@ -0,0 +1,165 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+from paddle.distributed.communication import stream
+
+from .serialization_utils import (
+    convert_object_to_tensor,
+    convert_tensor_to_object,
+)
+
+
+def scatter(
+    tensor: Tensor,
+    tensor_list: Sequence[Tensor] | None = None,
+    src: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task | None:
+    """
+
+    Scatter a tensor to all participators. As shown below, one process is started with a GPU and the source of the scatter
+    is GPU0. Through scatter operator, the data in GPU0 will be sent to all GPUs averagely.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/scatter.png
+        :width: 800
+        :alt: scatter
+        :align: center
+
+    Args:
+        tensor (Tensor): The output Tensor. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        tensor_list (list|tuple): A list/tuple of Tensors to scatter. Every element in the list must be a Tensor whose data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16. Default value is None.
+        src (int): The source rank id. Default value is 0.
+        group (Group, optional): The group instance return by new_group or None for global default group.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([7, 8, 9])
+            ...     data2 = paddle.to_tensor([10, 11, 12])
+            ...     dist.scatter(data1, src=1)
+            >>> else:
+            ...     data1 = paddle.to_tensor([1, 2, 3])
+            ...     data2 = paddle.to_tensor([4, 5, 6])
+            ...     dist.scatter(data1, tensor_list=[data1, data2], src=1)
+            >>> print(data1, data2)
+            >>> # [1, 2, 3] [10, 11, 12] (2 GPUs, out for rank 0)
+            >>> # [4, 5, 6] [4, 5, 6] (2 GPUs, out for rank 1)
+    """
+    return stream.scatter(tensor, tensor_list, src, group, sync_op)
+
+
+def scatter_object_list(
+    out_object_list: list[Any],
+    in_object_list: list[Any] | None = None,
+    src: int = 0,
+    group: Group | None = None,
+) -> None:
+    """
+
+    Scatter picklable objects from the source to all others. Similar to scatter(), but python object can be passed in.
+
+    Args:
+        out_object_list (list): The list of objects to store the scattered objects.
+        in_object_list (list): The list of objects to scatter. Only objects on the src rank will be scattered.
+        src (int): The source rank in global view.
+        group (Group): The group instance return by new_group or None for global default group.
+
+    Returns:
+        None.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> out_object_list = []  # type: ignore
+            >>> if dist.get_rank() == 0:
+            ...     in_object_list = [{'foo': [1, 2, 3]}, {'foo': [4, 5, 6]}]
+            >>> else:
+            ...     in_object_list = [{'bar': [1, 2, 3]}, {'bar': [4, 5, 6]}]
+            >>> dist.scatter_object_list(out_object_list, in_object_list, src=1)
+            >>> print(out_object_list)
+            >>> # [{'bar': [1, 2, 3]}] (2 GPUs, out for rank 0)
+            >>> # [{'bar': [4, 5, 6]}] (2 GPUs, out for rank 1)
+    """
+    assert framework.in_dynamic_mode(), (
+        "scatter_object_list doesn't support static graph mode."
+    )
+
+    rank = dist.get_rank()
+    in_obj_tensors = []
+    in_obj_sizes = []
+
+    if rank == src:
+        for obj in in_object_list:
+            obj_tensor, obj_size = convert_object_to_tensor(obj)
+            in_obj_tensors.append(obj_tensor)
+            in_obj_sizes.append(obj_size)
+        max_obj_size_tensor = max(in_obj_sizes)
+    else:
+        max_obj_size_tensor = paddle.empty([], dtype="int64")
+    stream.broadcast(max_obj_size_tensor, src)
+    max_obj_size = int(max_obj_size_tensor.item())
+
+    # resize to the same size
+    in_tensor_list = []
+    for tensor in in_obj_tensors:
+        numpy_data = tensor.numpy()
+        numpy_data = np.resize(numpy_data, [max_obj_size])
+        in_tensor = paddle.to_tensor(numpy_data)
+        in_tensor_list.append(in_tensor)
+    out_tensor = paddle.empty([max_obj_size], dtype="uint8")
+    scatter(out_tensor, in_tensor_list if rank == src else None, src, group)
+
+    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()
+    out_object_list.append(
+        convert_tensor_to_object(out_tensor, out_tensor_size.item())
+    )

python/paddle/distributed/communication/send.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import paddle
+from paddle.distributed.communication import stream
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.serialization_utils import (
+    convert_object_to_tensor,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def send(
+    tensor: Tensor,
+    dst: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+) -> task | None:
+    """
+    Send a tensor to the receiver.
+
+    Args:
+        tensor (Tensor): The Tensor to send. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        dst (int): The destination rank id.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([7, 8, 9])
+            ...     dist.send(data, dst=1)
+            >>> else:
+            ...     data = paddle.to_tensor([1, 2, 3])
+            ...     dist.recv(data, src=0)
+            >>> print(data)
+            >>> # [7, 8, 9] (2 GPUs)
+    """
+    return stream.send(
+        tensor, dst=dst, group=group, sync_op=sync_op, use_calc_stream=False
+    )
+
+
+def isend(tensor: Tensor, dst: int, group: Group | None = None) -> task | None:
+    """
+    Send tensor asynchronously
+
+    Args:
+        tensor (Tensor): The Tensor to send. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        dst (int): The destination rank.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([7, 8, 9])
+            ...     task = dist.isend(data, dst=1)
+            >>> else:
+            ...     data = paddle.to_tensor([1, 2, 3])
+            ...     task = dist.irecv(data, src=0)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> print(data)
+            >>> # [7, 8, 9] (2 GPUs)
+
+    """
+    return send(tensor, dst, group, sync_op=False)
+
+
+def send_object_list(
+    object_list: list[Any],
+    dst: int | None = None,
+    group: Group | None = None,
+    dst_in_group: int | None = None,
+):
+    """
+    Send a list of Python objects to the receiver.
+
+    Args:
+        object_list (list): The list of Python objects to send.
+        dst (int, optional): The destination rank id. Default: 0.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        dst_in_group (int, optional): The destination rank within the group. Cannot be specified together with dst. Default: None.
+
+    Returns:
+        This function does not return any value.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = ["hello", {"key": 100}, [1, 2, 3]]
+            ...     dist.send_object_list(data, dst=1)
+            >>> else:
+            ...     data = [None] * 3
+            ...     dist.recv_object_list(data, src=0)
+            >>> print(data)
+            >>> # ["hello", {"key": 100}, [1, 2, 3]] (2 GPUs)
+    """
+    if object_list is None or len(object_list) == 0:
+        raise ValueError("object_list cannot be None or empty")
+
+    group = _get_global_group() if group is None else group
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if dst_in_group is not None:
+        if dst is not None:
+            raise ValueError(
+                "Cannot specify both 'dst' and 'dst_in_group' arguments."
+            )
+        dst = group.get_global_rank(dst_in_group)
+    else:
+        dst = 0 if dst is None else dst
+
+    # Convert objects to tensors and get their sizes
+    tensor_list, size_list = zip(
+        *[convert_object_to_tensor(obj) for obj in object_list]
+    )
+    size_list_values = [size.item() for size in size_list]
+
+    # Send sizes first
+    object_sizes_tensor = paddle.to_tensor(size_list_values, dtype='int64')
+    send(object_sizes_tensor, dst=dst, group=group)
+
+    # Send object data
+    if len(tensor_list) == 1:
+        object_tensor = tensor_list[0]
+    else:
+        object_tensor = paddle.concat(tensor_list)
+    send(object_tensor, dst=dst, group=group)

python/paddle/distributed/communication/serialization_utils.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import pickle
+
+import numpy as np
+
+import paddle
+
+
+def convert_object_to_tensor(obj):
+    _pickler = pickle.Pickler
+    f = io.BytesIO()
+    _pickler(f).dump(obj)
+    data = np.frombuffer(f.getvalue(), dtype=np.uint8)
+    tensor = paddle.to_tensor(data)
+    return tensor, tensor.numel()
+
+
+def convert_tensor_to_object(tensor, len_of_tensor):
+    _unpickler = pickle.Unpickler
+    return _unpickler(io.BytesIO(tensor.numpy()[:len_of_tensor])).load()

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

@@ -0,0 +1,38 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .all_gather import all_gather
+from .all_reduce import all_reduce
+from .all_to_all import alltoall, alltoall_single
+from .broadcast import broadcast
+from .gather import gather
+from .recv import recv
+from .reduce import reduce
+from .reduce_scatter import reduce_scatter
+from .scatter import scatter
+from .send import send
+
+__all__ = [
+    "all_gather",
+    "all_reduce",
+    "alltoall",
+    "alltoall_single",
+    "broadcast",
+    "reduce",
+    "reduce_scatter",
+    "recv",
+    "scatter",
+    "send",
+    "gather",
+]

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

@@ -0,0 +1,220 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import _get_global_group
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+from paddle.distributed.utils.stream_utils import ExecutionStreamType
+
+
+def _all_gather_into_tensor_in_dygraph(
+    out_tensor: Tensor,
+    in_tensor: Tensor,
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> task:
+    group = _get_global_group() if group is None else group
+
+    if use_calc_stream:
+        return group.process_group.all_gather_into_tensor_on_calc_stream(
+            out_tensor,
+            in_tensor,
+        )
+
+    task = group.process_group.all_gather_into_tensor(
+        out_tensor, in_tensor, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _all_gather_in_dygraph(
+    tensor_list: list[Tensor],
+    tensor: Tensor,
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> task:
+    group = _get_global_group() if group is None else group
+
+    if len(tensor_list) == 0:
+        tensor_list += [paddle.empty_like(tensor) for _ in range(group.nranks)]
+
+    if use_calc_stream:
+        return group.process_group.all_gather_on_calc_stream(
+            tensor_list, tensor
+        )
+
+    task = group.process_group.all_gather(tensor_list, tensor, sync_op)
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _all_gather_in_static_mode(
+    tensor_list: list[Tensor], tensor: Tensor, group: Group, sync_op: bool
+) -> None:
+    op_type = 'all_gather'
+    helper = framework.LayerHelper(op_type, **locals())
+    out = helper.create_variable_for_type_inference(dtype=tensor.dtype)
+    for elem in tensor_list:
+        data_feeder.check_variable_and_dtype(
+            elem,
+            'tensor_list',
+            [
+                'float16',
+                'float32',
+                'float64',
+                'int32',
+                'int64',
+                'bool',
+                'int8',
+                'uint8',
+                'complex64',
+                'complex128',
+            ],
+            'all_gather',
+        )
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        [
+            'float16',
+            'float32',
+            'float64',
+            'int32',
+            'int64',
+            'bool',
+            'int8',
+            'uint8',
+            'complex64',
+            'complex128',
+        ],
+        'all_gather',
+    )
+
+    ring_id = 0 if group is None else group.id
+    nranks = dist.get_world_size()
+    op = helper.append_op(
+        type=op_type,
+        inputs={'x': [tensor]},
+        outputs={'out': [out]},
+        attrs={
+            'ring_id': ring_id,
+            'nranks': nranks,
+        },
+    )
+    if sync_op:
+        op.dist_attr.execution_stream = ExecutionStreamType.DefaultStream.value
+    tensor_list.clear()
+    # 0-D 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(
+    tensor_or_tensor_list: Tensor | list[Tensor],
+    tensor: Tensor,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Gather tensors across devices to a correctly-sized tensor or a tensor list.
+
+    Args:
+        tensor_or_tensor_list (Union[Tensor, List[Tensor]]): The output. If it is a tensor, it should be correctly-sized. If it is a list, it
+            should be empty or contain correctly-sized tensors.
+        tensor (Tensor): The input tensor on each rank. The result will overwrite this tenor after communication. Support
+            float16, float32, float64, int32, int64, int8, uint, bool, complex64 or complex128 as the input data type.
+        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> local_rank = dist.get_rank()
+            >>> tensor_list = []  # type: ignore
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> task = dist.stream.all_gather(tensor_list, data, sync_op=False)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> print(tensor_list)
+            [[[4, 5, 6], [4, 5, 6]], [[1, 2, 3], [1, 2, 3]]] (2 GPUs)
+    """
+    if group is not None and not group.is_member():
+        raise RuntimeError(
+            "The group should not be None and all ranks which invoke this operation should be the member of this group."
+        )
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        if paddle.is_tensor(tensor_or_tensor_list):
+            return _all_gather_into_tensor_in_dygraph(
+                tensor_or_tensor_list, tensor, group, sync_op, use_calc_stream
+            )
+        else:
+            return _all_gather_in_dygraph(
+                tensor_or_tensor_list, tensor, group, sync_op, use_calc_stream
+            )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        if paddle.is_tensor(tensor_or_tensor_list):
+            raise RuntimeError(
+                "Only support passing a tensor list to `all_gather` in static graph mode now."
+            )
+        else:
+            return _all_gather_in_static_mode(
+                tensor_or_tensor_list, tensor, group, sync_op
+            )

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

@@ -0,0 +1,166 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops, framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.reduce import (
+    ReduceOp,
+    _get_reduce_op,
+    _to_inplace_op,
+)
+from paddle.framework import in_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+    from ..all_reduce import _ReduceOp
+
+
+def _all_reduce_in_dygraph(
+    tensor: Tensor,
+    op: _ReduceOp,
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> task:
+    op_type = _get_reduce_op(op)
+
+    if use_calc_stream:
+        return group.process_group.all_reduce_on_calc_stream(tensor, op_type)
+
+    task = group.process_group.all_reduce(tensor, op_type, sync_op)
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _all_reduce_in_static_mode(
+    tensor: Tensor,
+    op: _ReduceOp,
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> None:
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        [
+            'float16',
+            'float32',
+            'float64',
+            'int32',
+            'int64',
+            'int8',
+            'uint8',
+            'bool',
+            'uint16',
+        ],
+        'all_reduce',
+    )
+
+    ring_id = 0 if group is None else group.id
+
+    if not isinstance(ring_id, int):
+        raise ValueError("The type of 'ring_id' for all_reduce should be int.")
+
+    if in_pir_mode():
+        op_type: str = _to_inplace_op(op)
+        _C_ops.all_reduce_(tensor, ring_id, op)
+        return
+
+    # TODO: Support task and use task.wait in static graph mode
+    #       Use use_calc_stream rather than sync_op
+    op_type = _get_reduce_op(op)
+    helper = framework.LayerHelper(op_type, **locals())
+    helper.append_op(
+        type=op_type,
+        inputs={'X': [tensor]},
+        outputs={'Out': [tensor]},
+        attrs={'ring_id': ring_id, 'use_calc_stream': sync_op},
+    )
+
+
+def all_reduce(
+    tensor: Tensor,
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Perform specific reduction (for example, sum, max) on inputs across devices.
+
+    Args:
+        tensor (Tensor): The input tensor on each rank. The result will overwrite this tenor after communication. Support
+            float16, float32, float64, int32, int64, int8, uint8 or bool as the input data type.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD, optional): The reduction used. If none is given, use ReduceOp.SUM as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> local_rank = dist.get_rank()
+            >>> data = None
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> task = dist.stream.all_reduce(data, sync_op=False)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> out = data
+            >>> print(out)
+            [[5, 7, 9], [5, 7, 9]]
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        return _all_reduce_in_dygraph(
+            tensor, op, group, sync_op, use_calc_stream
+        )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        return _all_reduce_in_static_mode(
+            tensor, op, group, sync_op, use_calc_stream
+        )

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

@@ -0,0 +1,386 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def _all_to_all_tensor_in_dygraph(
+    out_tensor: Tensor,
+    in_tensor: Tensor,
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> task:
+    if use_calc_stream:
+        return group.process_group.all_to_all_tensor_on_calc_stream(
+            out_tensor, in_tensor
+        )
+
+    task = group.process_group.all_to_all_tensor(out_tensor, in_tensor, sync_op)
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _all_to_all_in_dygraph(
+    out_tensor_list: Sequence[Tensor],
+    in_tensor_list: Sequence[Tensor],
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> task:
+    if len(in_tensor_list) == 0:
+        raise RuntimeError("The input tensor_list should not be empty.")
+
+    if len(out_tensor_list) == 0:
+        out_tensor_list += [
+            paddle.empty_like(tensor) for tensor in in_tensor_list
+        ]
+
+    if use_calc_stream:
+        return group.process_group.all_to_all_on_calc_stream(
+            out_tensor_list, in_tensor_list
+        )
+
+    task = group.process_group.all_to_all(
+        out_tensor_list, in_tensor_list, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _all_to_all_in_static_mode(
+    out_tensor_or_tensor_list: Tensor | Sequence[Tensor],
+    in_tensor_or_tensor_list: Tensor | Sequence[Tensor],
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> None:
+    op_type = 'all_to_all'
+    ring_id = 0 if group is None else group.id
+    nranks = dist.get_world_size()
+    helper = framework.LayerHelper(op_type, **locals())
+
+    in_tensor = in_tensor_or_tensor_list
+    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.")
+        # 0-D 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:
+            raise ValueError(
+                "The 'out_tensor_list' for all_to_all must be an empty list."
+            )
+        out_tensor = helper.create_variable_for_type_inference(
+            dtype=in_tensor.dtype
+        )
+
+    data_feeder.check_variable_and_dtype(
+        in_tensor,
+        'in_tensor',
+        ['float16', 'float32', 'float64', 'int32', 'int64', 'uint16'],
+        'all_to_all',
+    )
+    op = helper.append_op(
+        type=op_type,
+        inputs={'x': [in_tensor]},
+        outputs={'out': [out_tensor]},
+        attrs={
+            'ring_id': ring_id,
+        },
+    )
+    if sync_op:
+        op.dist_attr.execution_stream = "default"
+    # NOTE(liyurui): If the argument `out_tensor_or_tensor_list` is a tensor_list,
+    # we need to split the result. So we should wait the result of all_to_all
+    # before split if the communication is not on calc stream.
+    if isinstance(out_tensor_or_tensor_list, list):
+        if not sync_op:
+            dist.wait(out_tensor, use_calc_stream=False)
+        # 0-D 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)
+            )
+
+
+def alltoall(
+    out_tensor_or_tensor_list: Tensor | Sequence[Tensor],
+    in_tensor_or_tensor_list: Tensor | Sequence[Tensor],
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Scatter a tensor (or a tensor list) across devices and gather outputs to another tensor (or a tensor list, respectively).
+
+    Args:
+        out_tensor_or_tensor_list (Union[Tensor, List[Tensor]]): The output. If it is a tensor, it should be correctly-sized.
+        If it is a list, it should be empty or contain correctly-sized tensors. Its data type should be the same as the input.
+        in_tensor_or_tensor_list (Union[Tensor, List[Tensor]]): The input to scatter (must be specified on the source rank).
+            If it is a tensor, it should be correctly-sized. If it is a list, it should contain correctly-sized tensors. Support
+            float16, float32, float64, int32, int64, int8, uint8 or bool as the input data type.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+
+            >>> # all_to_all with equal split sizes
+            >>> out_tensor_list = []  # type: ignore[var-annotated]
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([[1, 2, 3], [4, 5, 6]])
+            ...     data2 = paddle.to_tensor([[7, 8, 9], [10, 11, 12]])
+            >>> else:
+            ...     data1 = paddle.to_tensor([[13, 14, 15], [16, 17, 18]])
+            ...     data2 = paddle.to_tensor([[19, 20, 21], [22, 23, 24]])
+            >>> task = dist.stream.alltoall(out_tensor_list, [data1, data2], sync_op=False)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> print(out_tensor_list)
+            >>> # [[[1, 2, 3], [4, 5, 6]], [[13, 14, 15], [16, 17, 18]]]    (2 GPUs, out for rank 0)
+            >>> # [[[7, 8, 9], [10, 11, 12]], [[19, 20, 21], [22, 23, 24]]] (2 GPUs, out for rank 1)
+
+            >>> # all_to_all with unequal split sizes
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([[1, 2, 3], [4, 5, 6]])  # shape: (2, 3)
+            ...     data2 = paddle.to_tensor([7])  # shape: (1, )
+            ...     out_data1 = paddle.empty((2, 3), dtype=data1.dtype)
+            ...     out_data2 = paddle.empty((3, 2), dtype=data1.dtype)
+            >>> else:
+            ...     data1 = paddle.to_tensor([[8, 9], [10, 11], [12, 13]]) # shape: (3, 2)
+            ...     data2 = paddle.to_tensor([[14, 15, 16, 17]])           # shape: (1, 4)
+            ...     out_data1 = paddle.empty((1,), dtype=data1.dtype)
+            ...     out_data2 = paddle.empty((1, 4), dtype=data1.dtype)
+            >>> dist.alltoall([out_data1, out_data2], [data1, data2])
+            >>> print([out_data1, out_data2])
+            >>> # [[[1, 2, 3], [4, 5, 6]], [[8, 9], [10, 11], [12, 13]]]  (2 GPUs, out for rank 0)
+            >>> # [[7], [[14, 15, 16, 17]]]                               (2 GPUs, out for rank 1)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    if out_tensor_or_tensor_list is None:
+        raise RuntimeError("The output should be specified.")
+    if in_tensor_or_tensor_list is None:
+        raise RuntimeError("The input should be specified.")
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        out_is_tensor = paddle.is_tensor(out_tensor_or_tensor_list)
+        in_is_tensor = paddle.is_tensor(in_tensor_or_tensor_list)
+        if out_is_tensor and in_is_tensor:
+            return _all_to_all_tensor_in_dygraph(
+                out_tensor_or_tensor_list,
+                in_tensor_or_tensor_list,
+                group,
+                sync_op,
+                use_calc_stream,
+            )
+        elif not out_is_tensor and not in_is_tensor:
+            return _all_to_all_in_dygraph(
+                out_tensor_or_tensor_list,
+                in_tensor_or_tensor_list,
+                group,
+                sync_op,
+                use_calc_stream,
+            )
+        else:
+            raise RuntimeError(
+                "The output and input should be both tensor or tensor list."
+            )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        return _all_to_all_in_static_mode(
+            out_tensor_or_tensor_list,
+            in_tensor_or_tensor_list,
+            group,
+            sync_op,
+            use_calc_stream,
+        )
+
+
+def _alltoall_single_in_dygraph(
+    out_tensor: Tensor,
+    in_tensor: Tensor,
+    out_split_sizes: list[int],
+    in_split_sizes: list[int],
+    group: Group,
+    sync_op: bool,
+    use_calc_stream: bool,
+) -> task:
+    if out_split_sizes is None:
+        out_split_sizes = []
+    if in_split_sizes is None:
+        in_split_sizes = []
+
+    if use_calc_stream:
+        return group.process_group.all_to_all_single_on_calc_stream(
+            out_tensor, in_tensor, out_split_sizes, in_split_sizes
+        )
+
+    task = group.process_group.all_to_all_single(
+        out_tensor, in_tensor, out_split_sizes, in_split_sizes, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def alltoall_single(
+    out_tensor: Tensor,
+    in_tensor: Tensor,
+    out_split_sizes: list[int] | None = None,
+    in_split_sizes: list[int] | None = None,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task:
+    """
+
+    Split and Scatter the split input tensor to the out tensor across devices.
+
+    Args:
+        out_tensor(Tensor): The output tensor. Its data type should be the same as the input.
+        in_tensor (Tensor): The input tensor. Its data type should be float16, float32, float64, int32, int64, int8, uint8 or bool.
+        out_split_sizes (List[int]|None, optional): Split sizes of out_tensor for dim[0]. If not given, dim[0] of out_tensor must be divisible
+            by group size and out_tensor will be gathered averagely from all participators. If none is given, use a empty list as default.
+        in_split_sizes (List[int]|None, optional): Split sizes of in_tensor for dim[0]. If not given, dim[0] of in_tensor must be divisible
+        by group size and in_tensor will be scattered averagely to all participators. If none is given, use a empty list as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> local_rank = dist.get_rank()
+
+            >>> # case 1
+            >>> output = paddle.empty([2], dtype="int64")
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([0, 1])
+            >>> else:
+            ...     data = paddle.to_tensor([2, 3])
+            >>> task = dist.stream.alltoall_single(output, data, sync_op=False)
+            >>> task.wait()
+            >>> out = output.numpy()
+            >>> print(out)
+            >>> # [0, 2] (2 GPUs, out for rank 0)
+            >>> # [1, 3] (2 GPUs, out for rank 1)
+
+            >>> # case 2
+            >>> size = dist.get_world_size()
+            >>> output = paddle.empty([(local_rank + 1) * size, size], dtype='float32')
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1., 1.], [1., 1.], [1., 1.]])
+            >>> out_split_sizes = [local_rank + 1 for i in range(size)]
+            >>> in_split_sizes = [i + 1 for i in range(size)]
+            >>> task = dist.stream.alltoall_single(
+            ...     output,
+            ...     data,
+            ...     out_split_sizes,
+            ...     in_split_sizes,
+            ...     sync_op=False,
+            ... )
+            >>> task.wait()
+            >>> out = output.numpy()
+            >>> print(out)
+            >>> # [[0., 0.], [1., 1.]]                     (2 GPUs, out for rank 0)
+            >>> # [[0., 0.], [0., 0.], [1., 1.], [1., 1.]] (2 GPUs, out for rank 1)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        return _alltoall_single_in_dygraph(
+            out_tensor,
+            in_tensor,
+            out_split_sizes,
+            in_split_sizes,
+            group,
+            sync_op,
+            use_calc_stream,
+        )
+
+    raise RuntimeError(
+        "paddle.distributed.stream.alltoall_single is only supported in dygraph mode now."
+    )

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

@@ -0,0 +1,156 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops, framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _get_or_throw_group_rank,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.reduce import _to_inplace_op
+from paddle.framework import in_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def _broadcast_in_dygraph(
+    tensor, src_rank_in_group, group, sync_op, use_calc_stream
+):
+    if use_calc_stream:
+        return group.process_group.broadcast_on_calc_stream(
+            tensor, src_rank_in_group
+        )
+
+    task = group.process_group.broadcast(tensor, src_rank_in_group, sync_op)
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _broadcast_in_static_mode(
+    tensor, src_rank_in_group, group, sync_op, use_calc_stream
+):
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        [
+            'float16',
+            'float32',
+            'float64',
+            'int32',
+            'int64',
+            'int8',
+            'uint8',
+            'bool',
+        ],
+        'broadcast',
+    )
+
+    op_type = 'broadcast'
+    helper = framework.LayerHelper(op_type, **locals())
+    ring_id = 0 if group is None else group.id
+
+    if in_pir_mode():
+        op_type = _to_inplace_op(op_type)
+        getattr(_C_ops, op_type)(tensor, ring_id, src_rank_in_group, sync_op)
+        return
+
+    op = helper.append_op(
+        type=op_type,
+        inputs={'x': [tensor]},
+        outputs={'out': [tensor]},
+        attrs={
+            'root': src_rank_in_group,
+            'ring_id': ring_id,
+        },
+    )
+    if sync_op:
+        op.dist_attr.execution_stream = "default"
+
+
+def broadcast(
+    tensor: Tensor,
+    src: int,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Broadcast a tensor to all devices.
+
+    Args:
+        tensor (Tensor): The tensor to broadcast. Support float16, float32, float64, int32, int64, int8, uint8 or bool as its data type.
+        src (int, optional): Rank of the source device.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> local_rank = dist.get_rank()
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> task = dist.stream.broadcast(data, src=1, sync_op=False)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> out = data.numpy()
+            >>> print(out)
+            >>> # [[1, 2, 3], [1, 2, 3]] (2 GPUs)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be True in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        src_rank_in_group = _get_or_throw_group_rank(src, group)
+
+        return _broadcast_in_dygraph(
+            tensor, src_rank_in_group, group, sync_op, use_calc_stream
+        )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        return _broadcast_in_static_mode(
+            tensor, src, group, sync_op, use_calc_stream
+        )

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

@@ -0,0 +1,138 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _get_or_throw_group_rank,
+    _warn_cur_rank_not_in_group,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def _gather_in_dygraph(
+    tensor, gather_list, dst_rank_in_group, group, sync_op, use_calc_stream
+):
+    nranks = group.nranks
+    if group.rank == dst_rank_in_group:
+        if len(gather_list) == 0:
+            gather_list += [paddle.empty_like(tensor) for _ in range(nranks)]
+    else:
+        gather_list = [tensor for _ in range(nranks)]
+
+    assert len(gather_list) == nranks, (
+        f" gather_list length {len(gather_list)} and nrankd {nranks} not equal"
+    )
+
+    task = group.process_group.gather(
+        tensor, gather_list, dst_rank_in_group, sync_op, use_calc_stream
+    )
+
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def gather(
+    tensor: Tensor,
+    gather_list: Sequence[Tensor] | None = None,
+    dst: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Gather tensors from all participators.
+
+    Args:
+        tensor (Tensor): The input Tensor. Its data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
+        gather_list (list|None): A list of Tensors to hold the gathered tensors. Every element in the list must be a Tensor whose data type
+            should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16. Default value is None.
+        dst (int): The dst rank id. Default value is 0.
+        group (Group|None, optional): The group instance return by new_group or None for global default group.
+        sync_op (bool, optional): Whether this op is a sync op. The default value is True.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Async work handle,which can be wait on, if async_op is set to True.
+        None, if not async_op
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> gather_list = []  # type: ignore[var-annotated]
+            >>> if dist.get_rank() == 0:
+            ...     data = paddle.to_tensor([1, 2, 3])
+            ...     dist.stream.gather(data, gather_list, dst=0)
+            >>> else:
+            ...     data = paddle.to_tensor([4, 5, 6])
+            ...     dist.stream.gather(data, gather_list, dst=0)
+            >>> print(gather_list)
+            >>> # [[1, 2, 3], [4, 5, 6]] (2 GPUs, out for rank 0)
+            >>> # [] (2 GPUs, out for rank 1)
+    """
+
+    assert framework.in_dynamic_mode(), (
+        "gather doesn't support static graph mode yet."
+    )
+
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    # NOTE(liuzhenhai): Only the dst rank needs to specific the gather_list argument.
+    # Other ranks which pass this argument in will be ignored with a warning.
+    # The passed in type for non-dst rank is meaningless, for it will be ignored.
+    if dst != dist.get_rank():
+        if gather_list is not None:
+            warnings.warn(
+                "Specific `gather_list` is meaningless for rank which is not dst."
+            )
+        gather_list = []
+    else:
+        assert gather_list is not None, (
+            "gather_list must not be none for dst rank"
+        )
+
+    group = _get_global_group() if group is None else group
+    dst_rank_in_group = _get_or_throw_group_rank(dst, group)
+    return _gather_in_dygraph(
+        tensor, gather_list, dst_rank_in_group, group, sync_op, use_calc_stream
+    )

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

@@ -0,0 +1,136 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _get_or_throw_group_rank,
+    _warn_cur_rank_not_in_group,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def _recv_in_dygraph(
+    tensor, src_rank_in_group, group, sync_op, use_calc_stream
+):
+    if use_calc_stream:
+        return group.process_group.recv_on_calc_stream(
+            tensor, src_rank_in_group
+        )
+
+    task = group.process_group.recv(tensor, src_rank_in_group, sync_op)
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _recv_in_static_mode(
+    tensor, src_rank_in_group, group, sync_op, use_calc_stream
+):
+    op_type = 'recv_v2'
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        ['float16', 'float32', 'float64', 'int32', 'int64', 'uint16'],
+        'recv',
+    )
+    ring_id = 0 if group is None else group.id
+    helper = framework.LayerHelper(op_type, **locals())
+    helper.append_op(
+        type=op_type,
+        outputs={'Out': [tensor]},
+        attrs={
+            'ring_id': ring_id,
+            'peer': src_rank_in_group,
+            'out_shape': tensor.shape,
+            'dtype': tensor.dtype,
+            'use_calc_stream': sync_op,
+        },
+    )
+
+
+def recv(
+    tensor: Tensor,
+    src: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Receive a tensor from the source device.
+
+    Args:
+        tensor (Tensor): The tensor to receive. Support float16, float32, float64, int32, int64, int8, uint8 or bool as its data type.
+        src (int, optional): Rank of the source device. If none is given, use `0` as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> local_rank = dist.get_rank()
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            ...     task = dist.stream.send(data, dst=1, sync_op=False)
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            ...     task = dist.stream.recv(data, src=0, sync_op=False)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> out = data.numpy()
+            >>> print(out)
+            >>> # [[4, 5, 6], [4, 5, 6]] (2 GPUs)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be True in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        src_rank_in_group = _get_or_throw_group_rank(src, group)
+
+        return _recv_in_dygraph(
+            tensor, src_rank_in_group, group, sync_op, use_calc_stream
+        )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        return _recv_in_static_mode(
+            tensor, src, group, sync_op, use_calc_stream
+        )

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

@@ -0,0 +1,156 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _get_or_throw_group_rank,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.reduce import ReduceOp, _get_reduce_op
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+    from paddle.distributed.communication.reduce import _ReduceOp
+
+
+def _reduce_in_dygraph(
+    tensor, dst_rank_in_group, op, group, sync_op, use_calc_stream
+):
+    op_type = _get_reduce_op(op)
+    if use_calc_stream:
+        return group.process_group.reduce_on_calc_stream(
+            tensor, dst_rank_in_group, op_type
+        )
+
+    task = group.process_group.reduce(
+        tensor, dst_rank_in_group, op_type, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _reduce_in_static_mode(
+    tensor, dst_rank_in_group, reduce_type, group, sync_op, use_calc_stream
+):
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        [
+            'float16',
+            'float32',
+            'float64',
+            'int32',
+            'int64',
+            'int8',
+            'uint8',
+            'bool',
+        ],
+        'reduce',
+    )
+
+    op_type = "reduce"
+    ring_id = 0 if group is None else group.id
+
+    helper = framework.LayerHelper(op_type, **locals())
+    helper.append_op(
+        type=op_type,
+        inputs={'x': [tensor]},
+        outputs={'out': [tensor]},
+        attrs={
+            'ring_id': ring_id,
+            'root_id': dst_rank_in_group,
+            'reduce_type': int(reduce_type),
+        },
+    )
+
+
+def reduce(
+    tensor: Tensor,
+    dst: int = 0,
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Perform specific reduction (for example, sum, max) on a tensor across devices and send to the destination device.
+
+    Args:
+        tensor (Tensor): The input tensor on each rank. The result will overwrite this tenor after communication. Support
+            float16, float32, float64, int32, int64, int8, uint8 or bool as the input data type.
+        dst (int, optional): Rank of the destination device. If none is given, use `0` as default.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD, optional): The reduction used. If none is given, use ReduceOp.SUM as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> local_rank = dist.get_rank()
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            >>> task = dist.stream.reduce(data, dst=0, sync_op=False)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> out = data.numpy()
+            >>> print(out)
+            >>> # [[5, 7, 9], [5, 7, 9]] (2 GPUs, out for rank 0)
+            >>> # [[1, 2, 3], [1, 2, 3]] (2 GPUs, out for rank 1)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        dst_rank_in_group = _get_or_throw_group_rank(dst, group)
+        return _reduce_in_dygraph(
+            tensor, dst_rank_in_group, op, group, sync_op, use_calc_stream
+        )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        return _reduce_in_static_mode(
+            tensor, dst, op, group, sync_op, use_calc_stream
+        )

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

@@ -0,0 +1,273 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.reduce import ReduceOp, _get_reduce_op
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+    from paddle.distributed.communication.reduce import _ReduceOp
+
+
+def _reduce_scatter_tensor_in_dygraph(
+    out_tensor,
+    in_tensor,
+    op,
+    group,
+    sync_op,
+    use_calc_stream,
+    caller="reduce_scatter",
+):
+    op_type = _get_reduce_op(op)
+
+    if use_calc_stream:
+        return group.process_group.reduce_scatter_tensor_on_calc_stream(
+            out_tensor, in_tensor, op_type
+        )
+
+    task = group.process_group.reduce_scatter_tensor(
+        out_tensor, in_tensor, op_type, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _reduce_scatter_in_dygraph(
+    tensor, tensor_list, op, group, sync_op, use_calc_stream
+):
+    op_type = _get_reduce_op(op)
+
+    if use_calc_stream:
+        return group.process_group.reduce_scatter_on_calc_stream(
+            tensor, tensor_list, op_type
+        )
+
+    task = group.process_group.reduce_scatter(
+        tensor, tensor_list, op_type, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _reduce_scatter_in_static_mode(tensor, tensor_or_tensor_list, group):
+    op_type = 'reduce_scatter'
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        [
+            'float16',
+            'float32',
+            'float64',
+            'int32',
+            'int64',
+            'int8',
+            'uint8',
+            'bool',
+            'uint16',
+        ],
+        op_type,
+    )
+
+    helper = framework.LayerHelper(op_type, **locals())
+    ring_id = 0 if group is None else group.id
+    nranks = dist.get_world_size()
+
+    helper.append_op(
+        type=op_type,
+        inputs={'x': [tensor_or_tensor_list]},
+        outputs={'out': [tensor]},
+        attrs={
+            'ring_id': ring_id,
+            'nranks': nranks,
+        },
+    )
+
+
+def reduce_scatter(
+    tensor: Tensor,
+    tensor_or_tensor_list: Tensor | Sequence[Tensor],
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Reduce, then scatter a tensor (or a tensor list) across devices.
+
+    Args:
+        tensor (Tensor): The output tensor on each rank. The result will overwrite this tenor after communication. Support
+            float16, float32, float64, int32, int64, int8, uint8 or bool as the input data type.
+        tensor_or_tensor_list (Union[Tensor, List[Tensor]]): The input to scatter.
+            If it is a tensor, it should be correctly-sized. If it is a list, it should contain correctly-sized tensors.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD, optional): The reduction used. If none is given, use ReduceOp.SUM as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([0, 1])
+            ...     data2 = paddle.to_tensor([2, 3])
+            >>> else:
+            ...     data1 = paddle.to_tensor([4, 5])
+            ...     data2 = paddle.to_tensor([6, 7])
+            >>> dist.stream.reduce_scatter(data1, [data1, data2])
+            >>> out = data1.numpy()
+            >>> print(out)
+            >>> # [4, 6]  (2 GPUs, out for rank 0)
+            >>> # [8, 10] (2 GPUs, out for rank 1)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        if paddle.is_tensor(tensor_or_tensor_list):
+            return _reduce_scatter_tensor_in_dygraph(
+                tensor,
+                tensor_or_tensor_list,
+                op,
+                group,
+                sync_op,
+                use_calc_stream,
+            )
+        else:
+            return _reduce_scatter_in_dygraph(
+                tensor,
+                tensor_or_tensor_list,
+                op,
+                group,
+                sync_op,
+                use_calc_stream,
+            )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        return _reduce_scatter_in_static_mode(
+            tensor, tensor_or_tensor_list, group
+        )
+
+
+def _reduce_scatter_base(
+    out_tensor,
+    in_tensor,
+    op=ReduceOp.SUM,
+    group=None,
+    sync_op=True,
+    use_calc_stream=False,
+):
+    """
+
+    Reduce, then scatter a flattened tensor across devices.
+
+    Args:
+        out_tensor (Tensor): The output tensor on each rank. The result will overwrite this tenor after communication. Support
+            float16, float32, float64, int32 or int64 as the input data type.
+        in_tensor (Tensor): The input tensor to reduce and scatter.
+        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD, optional): The reduction used. If none is given, use ReduceOp.SUM as default.
+        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API will be deprecated in the future, and only supports the dygraph mode now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([7, 8, 9])
+            ...     data2 = paddle.to_tensor([10, 11, 12])
+            ...     dist.stream.scatter(data1, src=1)
+            >>> else:
+            ...     data1 = paddle.to_tensor([1, 2, 3])
+            ...     data2 = paddle.to_tensor([4, 5, 6])
+            ...     dist.stream.scatter(data1, [data1, data2], src=1)
+            >>> out = data1.numpy()
+            >>> print(out)
+            >>> # [1, 2, 3] (2 GPUs, out for rank 0)
+            >>> # [4, 5, 6] (2 GPUs, out for rank 1)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        return _reduce_scatter_tensor_in_dygraph(
+            out_tensor,
+            in_tensor,
+            op,
+            group,
+            sync_op,
+            use_calc_stream,
+            "_reduce_scatter_base",
+        )
+
+    raise RuntimeError(
+        "paddle.distributed.stream._reduce_scatter_base is only supported in dygraph mode now."
+    )

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

@@ -0,0 +1,246 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import paddle
+import paddle.distributed as dist
+from paddle import framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _get_or_throw_group_rank,
+    _warn_cur_rank_not_in_group,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def _scatter_tensor_in_dygraph(
+    out_tensor, in_tensor, src_rank_in_group, group, sync_op, use_calc_stream
+):
+    nranks = group.nranks
+
+    if use_calc_stream:
+        return group.process_group.scatter_tensor_on_calc_stream(
+            out_tensor, in_tensor, src_rank_in_group
+        )
+
+    task = group.process_group.scatter_tensor(
+        out_tensor, in_tensor, src_rank_in_group, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _scatter_in_dygraph(
+    tensor, tensor_list, src_rank_in_group, group, sync_op, use_calc_stream
+):
+    nranks = group.nranks
+    if group.rank == src_rank_in_group:
+        if len(tensor_list) == 0:
+            raise RuntimeError(
+                "The tensor_list should not be empty on src rank."
+            )
+    else:
+        tensor_list = [tensor for _ in range(nranks)]
+
+    if use_calc_stream:
+        return group.process_group.scatter_on_calc_stream(
+            tensor, tensor_list, src_rank_in_group
+        )
+
+    task = group.process_group.scatter(
+        tensor, tensor_list, src_rank_in_group, sync_op
+    )
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _scatter_in_static_mode(
+    tensor,
+    tensor_or_tensor_list,
+    src_rank_in_group,
+    group,
+    sync_op,
+    use_calc_stream,
+):
+    nranks = dist.get_world_size() if group is None else group.nranks
+    rank = dist.get_rank()
+
+    input_tensor = tensor_or_tensor_list
+    if isinstance(tensor_or_tensor_list, list):
+        tensor_list = tensor_or_tensor_list
+        if rank == src_rank_in_group:
+            if len(tensor_list) == 0:
+                raise RuntimeError(
+                    "The tensor_list should not be empty on src rank."
+                )
+        else:
+            tensor_list = [tensor for _ in range(nranks)]
+        # 0-D 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
+
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        [
+            'float16',
+            'float32',
+            'float64',
+            'int32',
+            'int64',
+            'int8',
+            'uint8',
+            'bool',
+        ],
+        'scatter',
+    )
+
+    op_type = 'c_scatter'
+    helper = framework.LayerHelper(op_type, **locals())
+    helper.append_op(
+        type=op_type,
+        inputs={'X': [input_tensor]},
+        outputs={'Out': [tensor]},
+        attrs={
+            'ring_id': ring_id,
+            'root': src_rank_in_group,
+            'use_calc_stream': sync_op,
+            'nranks': nranks,
+        },
+    )
+
+
+def scatter(
+    tensor: Tensor,
+    tensor_or_tensor_list: Tensor | Sequence[Tensor] | None = None,
+    src: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Scatter a tensor (or a tensor list) across devices.
+
+    Args:
+        tensor (Tensor): The output tensor on each rank. The result will overwrite this tenor after communication. Support
+            float16, float32, float64, int32, int64, int8, uint8 or bool as the input data type.
+        tensor_or_tensor_list (Union[Tensor, List[Tensor]]): The input to scatter (default is `None`, must be specified on the source rank).
+            If it is a tensor, it should be correctly-sized. If it is a list, it should contain correctly-sized tensors.
+        src (int, optional): Rank of the source device. If none is given, use `0` as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Warning:
+        This API only supports the dygraph mode now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data1 = paddle.to_tensor([7, 8, 9])
+            ...     data2 = paddle.to_tensor([10, 11, 12])
+            ...     dist.stream.scatter(data1, src=1)
+            >>> else:
+            ...     data1 = paddle.to_tensor([1, 2, 3])
+            ...     data2 = paddle.to_tensor([4, 5, 6])
+            ...     dist.stream.scatter(data1, [data1, data2], src=1)
+            >>> out = data1.numpy()
+            >>> print(out)
+            >>> # [1, 2, 3] (2 GPUs, out for rank 0)
+            >>> # [4, 5, 6] (2 GPUs, out for rank 1)
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be true in sync op behavior."
+        )
+
+    # NOTE(liyurui): Only the source rank needs to specific the tensor_or_tensor_list argument.
+    # Other ranks which pass this argument in will be ignored with a warning.
+    # If a tensor_list passed in, we need to concat it to a tensor before invoke C++ API.
+    # If a tensor passed in, concat is not needed.
+    # The passed in type for non-src rank is meaningless, for it will be ignored.
+    if src != dist.get_rank():
+        if tensor_or_tensor_list is not None:
+            warnings.warn(
+                "Specific `tensor_or_tensor_list` is meaningless for rank which is not src."
+            )
+        tensor_or_tensor_list = []
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        src_rank_in_group = _get_or_throw_group_rank(src, group)
+        if paddle.is_tensor(tensor_or_tensor_list):
+            return _scatter_tensor_in_dygraph(
+                tensor,
+                tensor_or_tensor_list,
+                src_rank_in_group,
+                group,
+                sync_op,
+                use_calc_stream,
+            )
+        else:
+            return _scatter_in_dygraph(
+                tensor,
+                tensor_or_tensor_list,
+                src_rank_in_group,
+                group,
+                sync_op,
+                use_calc_stream,
+            )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+
+        return _scatter_in_static_mode(
+            tensor,
+            tensor_or_tensor_list,
+            src,
+            group,
+            sync_op,
+            use_calc_stream,
+        )

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

@@ -0,0 +1,135 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import framework
+from paddle.base import data_feeder
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _get_or_throw_group_rank,
+    _warn_cur_rank_not_in_group,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
+
+def _send_in_dygraph(
+    tensor, dst_rank_in_group, group, sync_op, use_calc_stream
+):
+    if use_calc_stream:
+        return group.process_group.send_on_calc_stream(
+            tensor, dst_rank_in_group
+        )
+
+    task = group.process_group.send(tensor, dst_rank_in_group, sync_op)
+    if sync_op:
+        task.wait()
+
+    return task
+
+
+def _send_in_static_mode(
+    tensor, dst_rank_in_group, group, sync_op, use_calc_stream
+):
+    op_type = 'send_v2'
+    data_feeder.check_variable_and_dtype(
+        tensor,
+        'tensor',
+        ['float16', 'float32', 'float64', 'int32', 'int64', 'uint16'],
+        'send',
+    )
+
+    ring_id = 0 if group is None else group.id
+    helper = framework.LayerHelper(op_type, **locals())
+    helper.append_op(
+        type=op_type,
+        inputs={'X': [tensor]},
+        attrs={
+            'ring_id': ring_id,
+            'peer': dst_rank_in_group,
+            'use_calc_stream': sync_op,
+        },
+    )
+
+
+def send(
+    tensor: Tensor,
+    dst: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
+    """
+
+    Send a tensor to the destination device.
+
+    Args:
+        tensor (Tensor): The tensor to send. Support float16, float32, float64, int32, int64, int8, uint8 or bool as its data type.
+        dst (int, optional): Rank of the destination device. If none is given, use `0` as default.
+        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
+        use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
+            option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
+
+    Returns:
+        Return a task object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> local_rank = dist.get_rank()
+            >>> if local_rank == 0:
+            ...     data = paddle.to_tensor([[4, 5, 6], [4, 5, 6]])
+            ...     task = dist.stream.send(data, dst=1, sync_op=False)
+            >>> else:
+            ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
+            ...     task = dist.stream.recv(data, src=0, sync_op=False)
+            >>> task.wait()  # type: ignore[union-attr]
+            >>> out = data.numpy()
+            >>> print(out)
+            [[4, 5, 6], [4, 5, 6]]
+    """
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if not sync_op and use_calc_stream:
+        raise RuntimeError(
+            "use_calc_stream can only be True in sync op behavior."
+        )
+
+    if framework.in_dynamic_mode():
+        group = _get_global_group() if group is None else group
+        dst_rank_in_group = _get_or_throw_group_rank(dst, group)
+
+        return _send_in_dygraph(
+            tensor, dst_rank_in_group, group, sync_op, use_calc_stream
+        )
+    else:
+        assert group is None, (
+            "Group can not be used in static graph mode for now."
+        )
+        return _send_in_static_mode(
+            tensor, dst, group, sync_op, use_calc_stream
+        )