chore: import upstream snapshot with attribution

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
+        )