collective.py 72.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
#   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.

import numpy as np
import os
17
from datetime import timedelta
18
from ..fluid.layer_helper import LayerHelper
19
from ..fluid.framework import Variable
20
from ..fluid.framework import in_dygraph_mode
21
from ..fluid.framework import OpProtoHolder
J
Jiabin Yang 已提交
22
from ..fluid.framework import _non_static_mode
23
from ..fluid.framework import _in_legacy_dygraph
24
from ..fluid.framework import convert_np_dtype_to_dtype_
J
Jiangxinz 已提交
25
from ..fluid.framework import _varbase_creator
26 27 28 29
from ..fluid.data_feeder import convert_dtype
from ..fluid.data_feeder import check_variable_and_dtype
from ..fluid.data_feeder import check_type
from ..fluid.data_feeder import check_dtype
30 31
from ..fluid.layers.tensor import fill_constant
from ..fluid.layers import utils
B
Baibaifan 已提交
32
from ..fluid.dygraph import layers
33 34 35 36
from ..fluid.dygraph.parallel import prepare_context
import paddle
import paddle.fluid as fluid
import paddle.fluid.core as core
W
wanghuancoder 已提交
37
from paddle import _C_ops
J
Jiangxinz 已提交
38
import paddle.fluid.dygraph_utils as dygraph_utils
39

40
__all__ = []
41 42 43


class ReduceOp:
L
lilong12 已提交
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
    """
    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:: python

            import numpy as np
            import paddle
            from paddle.distributed import ReduceOp
            from paddle.distributed import init_parallel_env

            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
            init_parallel_env()
            if paddle.distributed.ParallelEnv().local_rank == 0:
                np_data = np.array([[4, 5, 6], [4, 5, 6]])
            else:
                np_data = np.array([[1, 2, 3], [1, 2, 3]])
            data = paddle.to_tensor(np_data)
            paddle.distributed.all_reduce(data, op=ReduceOp.SUM)
            out = data.numpy()
            # [[5, 7, 9], [5, 7, 9]]
    """
75 76 77 78
    SUM = 0
    MAX = 1
    MIN = 2
    PROD = 3
79
    AVG = 4
80 81


K
kuizhiqing 已提交
82 83 84 85
class Group():
    """
    The abstract representation of group.
    """
86

87
    def __init__(self, rank, rank_num, id=0, ranks=[], pg=None, name=None):
88 89
        self.rank = rank
        self.nranks = rank_num
K
kuizhiqing 已提交
90 91
        self.id = id
        self.ranks = ranks
92 93
        self.pg = pg
        self.name = name
K
kuizhiqing 已提交
94 95 96 97 98 99 100 101 102 103 104 105 106 107

    def is_member(self):
        if self.rank < 0:
            return False
        if self.nranks < 2:
            return False
        return True

    def get_group_rank(self, rank):
        if self.is_member() and rank in self.ranks:
            return self.ranks.index(rank)
        else:
            return -1

108 109 110 111
    @property
    def process_group(self):
        return self.pg

112 113 114 115
    def __repr__(self):
        debug_str = "rank: {}, nranks: {}, id: {}, ranks: ".format(
            self.rank, self.nranks, self.id)
        debug_str += ", ".join(map(str, self.ranks))
116 117
        debug_str += "; name: "
        debug_str += self.name if self.name else "None"
118 119
        return debug_str

K
kuizhiqing 已提交
120 121 122 123 124 125 126 127 128 129 130 131 132 133 134

_global_env = None


def _get_global_env():
    global _global_env
    if not _global_env:
        _global_env = paddle.distributed.ParallelEnv()
    return _global_env


# group map : the map of all group, 0 for GlobalGroup
# Dict[int, Group]
_group_map = {}

135 136 137 138 139 140 141
# group map by name : the map of all groups from their names
# Dict[name, Group]
_group_map_by_name = {}

# Name of the default group for init_parallel_env
_default_group_name = "_default_pg"

142
_valid_backend_list = ['nccl', 'gloo', 'hccl', 'heter']
143 144 145
_default_store = None  # the default tcp store
_default_backend = None

K
kuizhiqing 已提交
146

L
lilong12 已提交
147 148 149 150 151 152 153 154 155 156
def _set_default_backend(backend):
    global _default_backend
    _default_backend = backend


def _set_default_store(store):
    global _default_store
    _default_store = store


K
kuizhiqing 已提交
157 158 159 160
def _get_group_map():
    global _group_map
    if not _group_map:
        genv = _get_global_env()
161 162 163
        _group_map[0] = Group(genv.rank,
                              genv.world_size,
                              ranks=list(range(genv.world_size)))
K
kuizhiqing 已提交
164 165 166 167 168 169 170
    return _group_map


def _get_global_group():
    return _get_group_map()[0]


171 172 173 174 175 176
def _get_group_map_by_name():
    global _group_map_by_name
    return _group_map_by_name


def _get_default_group():
L
lilong12 已提交
177
    global _group_map_by_name
178 179 180 181 182 183
    assert _default_group_name in _group_map_by_name, (
        "Call paddle.distributed.init_parallel_env first "
        "to initialize the distributed environment.")
    return _get_group_map_by_name()[_default_group_name]


L
lilong12 已提交
184 185 186 187 188 189 190 191 192 193 194 195
def _set_group_map(gid, group):
    global _group_map
    assert gid not in _group_map
    _group_map[gid] = group


def _set_group_map_by_name(name, group):
    global _group_map_by_name
    assert name not in _group_map_by_name
    _group_map_by_name[name] = group


K
kuizhiqing 已提交
196 197 198 199 200 201 202 203 204 205
def _new_ring_id():
    return len(_get_group_map()) + max(_get_global_env().nrings, 9)


def get_group(id=0):
    """

    Get group instance by group id.

    Args:
K
kuizhiqing 已提交
206
        id (int): the group id. Default value is 0.
K
kuizhiqing 已提交
207 208 209 210 211 212 213 214 215 216 217 218 219 220

    Returns:
        Group: the group instance.

    Examples:
        .. code-block:: python

            ...
            gid = paddle.distributed.new_group([2,4,6])
            paddle.distributed.get_group(gid.id)

    """

    gm = _get_group_map()
J
Jiangxinz 已提交
221
    return gm[id] if id in gm else None
K
kuizhiqing 已提交
222 223


224 225 226 227 228 229
def _new_process_group_impl(backend,
                            store,
                            rank,
                            world_size,
                            group_name,
                            pg_options,
L
lilong12 已提交
230 231 232
                            group_id=0,
                            src_rank=None,
                            dst_rank=None):
233
    pg = None
234
    genv = _get_global_env()
L
lilong12 已提交
235 236 237 238
    if backend != 'heter':
        assert src_rank is None and dst_rank is None, (
            "src_rank and dst_rank "
            "can only be set for heter backend.")
L
lilong12 已提交
239
    assert backend in _valid_backend_list, "Unsupported backend: %s." % backend
240
    if backend == "gloo":
241 242
        place = core.CPUPlace()
        pg = core.ProcessGroupGloo(store, rank, world_size, place, group_id)
243
    elif backend == "nccl":
244 245
        place = core.CUDAPlace(genv.device_id)
        pg = core.ProcessGroupNCCL(store, rank, world_size, place, group_id)
246
    elif backend == "hccl":
247 248
        place = core.NPUPlace(genv.device_id)
        pg = core.ProcessGroupHCCL(store, rank, world_size, place, group_id)
249
    elif backend == "heter":
250 251 252 253 254
        place = None
        if core.is_compiled_with_cuda():
            place = core.CUDAPlace(genv.device_id)
        elif core.is_compiled_with_npu():
            place = core.NPUPlace(genv.device_id)
255 256 257 258 259 260 261 262 263 264 265 266 267
        cluster_id = int(os.getenv("CLUSTER_ID", "-1"))
        assert cluster_id >= 0, "please set the CLUSTER_ID variable."
        cluster_size = os.getenv("CLUSTER_SIZE", None)
        assert cluster_size, "please set the CLUSTER_SIZE variable."
        cluster_size = cluster_size.split(",")
        cluster_size = [int(s) for s in cluster_size]
        switch_ep = os.getenv("CLUSTER_SWITCH", None)
        assert switch_ep, "please set the CLUSTER_SWITCH variable."
        cluster_size_cumsum = np.cumsum(cluster_size)
        cluster_offset = 0 if cluster_id == 0 else cluster_size_cumsum[
            cluster_id - 1]
        global_rank = cluster_offset + rank
        global_world_size = cluster_size_cumsum[-1]
268 269 270 271 272 273 274 275 276 277 278 279 280
        pg = core.ProcessGroupHeter(store,
                                    rank=global_rank,
                                    world_size=global_world_size,
                                    place=place,
                                    gid=group_id,
                                    local_rank=rank,
                                    local_size=world_size,
                                    gloo_rank=cluster_id,
                                    gloo_size=len(cluster_size),
                                    with_switch=True,
                                    switch_endpoint=switch_ep,
                                    src_rank=src_rank,
                                    dst_rank=dst_rank)
281 282 283 284

    return pg


S
ShenLiang 已提交
285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308
def barrier(group=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:: python

            import paddle
            from paddle.distributed import init_parallel_env

            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
            init_parallel_env()
            paddle.distributed.barrier()
    """
    if group is not None and not group.is_member():
        return

L
lilong12 已提交
309
    if in_dygraph_mode():
310 311 312 313 314
        group = _get_default_group() if group is None else group
        task = group.process_group.barrier()
        task.wait()
        return

S
ShenLiang 已提交
315 316 317
    ring_id = 0 if group is None else group.id

    temp = fill_constant([1], dtype="int32", value="1")
J
Jiabin Yang 已提交
318
    if _non_static_mode():
W
wanghuancoder 已提交
319
        return _C_ops.barrier(temp, temp, 'ring_id', ring_id)
W
wanghuancoder 已提交
320 321 322

    op_type = 'barrier'

S
ShenLiang 已提交
323 324 325
    if not isinstance(ring_id, int):
        raise ValueError("The type of 'group' for barrier must be int.")
    helper = LayerHelper(op_type, **locals())
326 327 328 329
    helper.append_op(type=op_type,
                     inputs={'X': [temp]},
                     outputs={'Out': [temp]},
                     attrs={'ring_id': ring_id})
S
ShenLiang 已提交
330 331


L
lilong12 已提交
332 333 334 335 336 337 338
# _custom_gid provides a way for users to
# set the group id, which is usually useful
# to be compatible with the static mode.
_custom_gid = None


def _set_custom_gid(gid):
339
    global _custom_gid
L
lilong12 已提交
340 341 342
    _custom_gid = gid


K
kuizhiqing 已提交
343 344 345
def new_group(ranks=None, backend=None):
    """

K
kuizhiqing 已提交
346
    Creates a new distributed communication group.
K
kuizhiqing 已提交
347 348

    Args:
K
kuizhiqing 已提交
349
        ranks (list): The global ranks of group members.
K
kuizhiqing 已提交
350 351 352
        backend (str): The backend used to create group, only nccl is supported now.

    Returns:
K
kuizhiqing 已提交
353
        Group: The group instance.
K
kuizhiqing 已提交
354 355 356 357 358 359 360

    Examples:
        .. code-block:: python

            import paddle

            paddle.distributed.init_parallel_env()
K
kuizhiqing 已提交
361 362 363
            tindata = paddle.randn(shape=[2, 3])
            gp = paddle.distributed.new_group([2,4,6])
            paddle.distributed.all_reduce(tindata, group=gp, use_calc_stream=False)
K
kuizhiqing 已提交
364 365

    """
366
    global _custom_gid
367
    global _group_map
L
lilong12 已提交
368
    if in_dygraph_mode():
369
        global _default_group_name
L
lilong12 已提交
370
        gid = _custom_gid if _custom_gid else _new_ring_id()
371
        group_name = _default_group_name + str(gid)
L
lilong12 已提交
372
        if backend != 'heter' and (ranks is None or len(ranks) > 1):
373 374 375 376 377 378 379 380 381
            global_group = _get_default_group()
            global_rank = global_group.rank
            global_ranks = global_group.ranks
            backend = _default_backend if backend is None else backend
            if ranks is None:
                ranks = global_ranks
            assert len(ranks) <= len(global_ranks), (
                "Size of new group must be less than or "
                "equal to that of the default global group.")
382 383
        size = len(ranks)
        ranks = sorted(ranks)
L
lilong12 已提交
384 385 386 387
        if backend == 'heter' or (size > 1 and global_rank in ranks):
            rank = 0 if backend == 'heter' else ranks.index(global_rank)
            src_rank = ranks[0] if backend == 'heter' else None
            dst_rank = ranks[1] if backend == 'heter' else None
388 389 390 391 392 393 394 395 396
            pg = _new_process_group_impl(backend,
                                         _default_store,
                                         rank,
                                         size,
                                         group_name,
                                         pg_options=None,
                                         group_id=gid,
                                         src_rank=src_rank,
                                         dst_rank=dst_rank)
397 398 399 400 401 402 403
        else:
            rank = -1
            pg = None
        group = Group(rank, size, id=gid, ranks=ranks, pg=pg, name=group_name)
        _group_map_by_name[group_name] = group
        _group_map[gid] = group

404
        # TODO(shenliang03): This is a temporary solution to solve the problem of
405
        # hang caused by tcp
406
        paddle.distributed.barrier(group=group)
407
        return group
K
kuizhiqing 已提交
408 409 410 411 412 413 414 415 416 417 418 419 420 421

    if not backend:
        backend = 'nccl'
    assert backend == 'nccl', ("backend other than nccl is not supported yet")

    genv = _get_global_env()
    global_rank = genv.rank

    ring_id = _new_ring_id()

    if global_rank not in ranks:
        gp = Group(-1, -1, ring_id, ranks)
        _group_map[ring_id] = gp
    else:
422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441
        ranks = sorted(ranks)
        group_rank = ranks.index(global_rank)
        group_size = len(ranks)
        gp = Group(group_rank, group_size, ring_id, ranks)
        _group_map[ring_id] = gp

        if group_size >= 2:
            strategy = core.ParallelStrategy()
            strategy.nranks = group_size
            strategy.local_rank = group_rank
            strategy.trainer_endpoints = [
                genv.trainer_endpoints[i] for i in ranks
            ]
            strategy.current_endpoint = genv.current_endpoint
            strategy.nrings = 1

            if core.is_compiled_with_cuda():
                place = core.CUDAPlace(genv.device_id)
                core.NCCLParallelContext(strategy,
                                         place).init_with_ring_id(ring_id)
442 443 444 445
            elif core.is_compiled_with_npu():
                place = core.NPUPlace(genv.device_id)
                core.HCCLParallelContext(strategy,
                                         place).init_with_ring_id(ring_id)
446 447 448 449
            elif core.is_compiled_with_mlu():
                place = core.MLUPlace(genv.device_id)
                core.CNCLParallelContext(strategy,
                                         place).init_with_ring_id(ring_id)
450 451 452 453 454
            else:
                assert False, ("no cuda device found")
        else:
            return gp

455
    # TODO(shenliang03): This is a temporary solution to solve the problem of
456
    # hang caused by cross-creation of new_group
457
    tmp = paddle.to_tensor(
J
Jiabin Yang 已提交
458
        [1], dtype="int32") if _non_static_mode() else fill_constant(
459
            [0], dtype="int32", value="1")
460 461
    paddle.distributed.all_reduce(tmp, use_calc_stream=True)
    paddle.distributed.wait(tmp)
K
kuizhiqing 已提交
462 463
    return gp

464

K
kuizhiqing 已提交
465 466 467 468 469 470 471 472
def wait(tensor, group=None, use_calc_stream=True):
    """

    wait to sync stream for group.

    Args:
        tensor (Tensor): The Tensor used before sync.
        group (Group): The Group instance to perform sync.
K
kuizhiqing 已提交
473 474
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
K
kuizhiqing 已提交
475 476 477 478 479 480 481 482 483 484

    Returns:
        None.

    Examples:
        .. code-block:: python

            import paddle

            paddle.distributed.init_parallel_env()
K
kuizhiqing 已提交
485
            tindata = paddle.randn(shape=[2, 3])
K
kuizhiqing 已提交
486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503
            paddle.distributed.all_reduce(tindata, use_calc_stream=True)
            paddle.distributed.wait(tindata)

    """

    if group is not None and not group.is_member():
        return

    ring_id = 0 if group is None else group.id

    if use_calc_stream:
        _sync_calc_stream(tensor)
    else:
        _sync_comm_stream(tensor, ring_id)


def _sync_calc_stream(tensor):

J
Jiabin Yang 已提交
504
    if _non_static_mode():
W
wanghuancoder 已提交
505
        return _C_ops.c_sync_calc_stream(tensor, tensor)
K
kuizhiqing 已提交
506 507 508 509 510 511 512

    op_type = 'c_sync_calc_stream'

    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [tensor]},
513 514
        outputs={'Out': [tensor]},
    )
515

516

K
kuizhiqing 已提交
517
def _sync_comm_stream(tensor, ring_id=0):
518

J
Jiabin Yang 已提交
519
    if _non_static_mode():
W
wanghuancoder 已提交
520
        return _C_ops.c_sync_comm_stream([tensor], [tensor], 'ring_id', ring_id)
521

K
kuizhiqing 已提交
522
    op_type = 'c_sync_comm_stream'
523

K
kuizhiqing 已提交
524 525 526 527 528
    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [tensor]},
        outputs={'Out': [tensor]},
529 530
        attrs={'ring_id': ring_id},
    )
K
kuizhiqing 已提交
531 532 533


def broadcast(tensor, src, group=None, use_calc_stream=True):
534 535 536
    """

    Broadcast a tensor from the source to all others.
537 538 539 540 541 542 543
    As shown below, 4 GPUs each start 4 processes and GPU0 owns data 0. Through broadcast operator,
    the 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
544 545 546 547 548

    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 or int64.
        src (int): The source rank.
K
kuizhiqing 已提交
549
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
550 551
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
552 553 554 555 556 557 558

    Returns:
        None.

    Examples:
        .. code-block:: python

559
            # required: distributed
560 561 562 563 564 565 566 567 568 569 570 571 572 573
            import numpy as np
            import paddle
            from paddle.distributed import init_parallel_env

            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
            init_parallel_env()
            if paddle.distributed.ParallelEnv().local_rank == 0:
                np_data = np.array([[4, 5, 6], [4, 5, 6]])
            else:
                np_data = np.array([[1, 2, 3], [1, 2, 3]])
            data = paddle.to_tensor(np_data)
            paddle.distributed.broadcast(data, 1)
            out = data.numpy()
            # [[1, 2, 3], [1, 2, 3]]
574
    """
K
kuizhiqing 已提交
575 576 577 578 579 580 581

    if group is not None and not group.is_member():
        return

    if not isinstance(src, int):
        raise ValueError("src should be int.")

L
lilong12 已提交
582
    if in_dygraph_mode():
583 584 585 586 587 588 589 590 591 592 593
        group = _get_default_group() if group is None else group
        gsrc = group.get_group_rank(src)
        assert gsrc >= 0, ("src rank out of group, need global rank")
        task = group.process_group.broadcast(tensor, gsrc)
        if use_calc_stream:
            task.wait()
            return None
        else:
            return task

    ring_id = ring_id = 0 if group is None else group.id
K
kuizhiqing 已提交
594
    gsrc = src if group is None else group.get_group_rank(src)
K
kuizhiqing 已提交
595
    assert gsrc >= 0, ("src rank out of group, need global rank")
K
kuizhiqing 已提交
596

J
Jiabin Yang 已提交
597
    if _non_static_mode():
W
wanghuancoder 已提交
598 599 600
        return _C_ops.c_broadcast(tensor, tensor, 'root', gsrc,
                                  'use_calc_stream', use_calc_stream, 'ring_id',
                                  ring_id)
601 602 603 604 605 606 607

    op_type = 'c_broadcast'
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'broadcast')

    helper = LayerHelper(op_type, **locals())
608 609 610 611 612 613 614 615
    helper.append_op(type=op_type,
                     inputs={'X': [tensor]},
                     outputs={'Out': [tensor]},
                     attrs={
                         'root': gsrc,
                         'use_calc_stream': use_calc_stream,
                         'ring_id': ring_id,
                     })
616 617


K
kuizhiqing 已提交
618
def all_reduce(tensor, op=ReduceOp.SUM, group=None, use_calc_stream=True):
619 620 621
    """

    Reduce a tensor over all ranks so that all get the result.
622 623 624 625 626 627 628 629
    As shown below, 4 GPUs each start 4 processes and the data on each GPU is represnted
    by the GPU number. 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
630 631 632 633

    Args:
        tensor (Tensor): The input Tensor. It also works as the output Tensor. Its data type
            should be float16, float32, float64, int32 or int64.
K
kuizhiqing 已提交
634
        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.Min|ReduceOp.PROD): Optional. The operation used. Default value is ReduceOp.SUM.
K
kuizhiqing 已提交
635
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
636 637
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
638 639 640 641 642 643 644

    Returns:
        None.

    Examples:
        .. code-block:: python

645
            # required: distributed
646 647 648 649 650 651 652 653 654 655 656 657 658 659 660
            import numpy as np
            import paddle
            from paddle.distributed import ReduceOp
            from paddle.distributed import init_parallel_env

            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
            init_parallel_env()
            if paddle.distributed.ParallelEnv().local_rank == 0:
                np_data = np.array([[4, 5, 6], [4, 5, 6]])
            else:
                np_data = np.array([[1, 2, 3], [1, 2, 3]])
            data = paddle.to_tensor(np_data)
            paddle.distributed.all_reduce(data)
            out = data.numpy()
            # [[5, 7, 9], [5, 7, 9]]
661
    """
K
kuizhiqing 已提交
662 663 664
    if group is not None and not group.is_member():
        return

L
lilong12 已提交
665
    if in_dygraph_mode():
666 667 668 669 670 671
        if op == ReduceOp.SUM:
            op_type = core.ReduceOp.SUM
        elif op == ReduceOp.MAX:
            op_type = core.ReduceOp.MAX
        elif op == ReduceOp.MIN:
            op_type = core.ReduceOp.MIN
672 673
        elif op == ReduceOp.PROD:
            op_type = core.ReduceOp.PRODUCT
674 675 676 677 678 679 680 681 682 683
        else:
            raise ValueError("Unknown reduce_op type for allreduce.")
        group = _get_default_group() if group is None else group
        task = group.process_group.allreduce(tensor, op_type)
        if use_calc_stream:
            task.wait()
            return None
        else:
            return task

K
kuizhiqing 已提交
684
    ring_id = 0 if group is None else group.id
J
Jiabin Yang 已提交
685
    if _non_static_mode():
686
        if op == ReduceOp.SUM:
W
wanghuancoder 已提交
687 688
            return _C_ops.c_allreduce_sum_(tensor, 'use_calc_stream',
                                           use_calc_stream, 'ring_id', ring_id)
689
        elif op == ReduceOp.MAX:
W
wanghuancoder 已提交
690 691
            return _C_ops.c_allreduce_max_(tensor, 'use_calc_stream',
                                           use_calc_stream, 'ring_id', ring_id)
692
        elif op == ReduceOp.MIN:
W
wanghuancoder 已提交
693 694
            return _C_ops.c_allreduce_min_(tensor, 'use_calc_stream',
                                           use_calc_stream, 'ring_id', ring_id)
695
        elif op == ReduceOp.PROD:
W
wanghuancoder 已提交
696 697
            return _C_ops.c_allreduce_prod_(tensor, 'use_calc_stream',
                                            use_calc_stream, 'ring_id', ring_id)
698 699 700 701 702 703 704 705 706 707 708 709 710 711
        else:
            raise ValueError("Unknown parameter: {}.".format(op))

    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'all_reduce')
    if op == ReduceOp.SUM:
        op_type = 'c_allreduce_sum'
    elif op == ReduceOp.MAX:
        op_type = 'c_allreduce_max'
    elif op == ReduceOp.MIN:
        op_type = 'c_allreduce_min'
    elif op == ReduceOp.PROD:
        op_type = 'c_allreduce_prod'
K
kuizhiqing 已提交
712 713
    if not isinstance(ring_id, int):
        raise ValueError("The type of 'ring_id' for all_reduce should be int.")
714
    helper = LayerHelper(op_type, **locals())
715 716 717 718 719 720 721
    helper.append_op(type=op_type,
                     inputs={'X': [tensor]},
                     outputs={'Out': [tensor]},
                     attrs={
                         'ring_id': ring_id,
                         'use_calc_stream': use_calc_stream
                     })
722 723


K
kuizhiqing 已提交
724
def reduce(tensor, dst, op=ReduceOp.SUM, group=None, use_calc_stream=True):
725 726
    """

727 728 729 730 731 732 733 734
    Reduce a tensor to the destination from all others. As shown below, 4 GPUs each start 4 processes and the data on each GPU is respresnted
    by the GPU number. 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
735 736 737 738 739

    Args:
        tensor (Tensor): The output Tensor for the destination and the input Tensor otherwise. Its data type
            should be float16, float32, float64, int32 or int64.
        dst (int): The destination rank id.
K
kuizhiqing 已提交
740
        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.Min|ReduceOp.PROD): Optional. The operation used. Default value is ReduceOp.SUM.
K
kuizhiqing 已提交
741
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
742 743
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
744 745 746 747 748 749 750

    Returns:
        None.

    Examples:
        .. code-block:: python

751
            # required: distributed
752 753 754 755 756 757 758 759 760 761 762 763 764 765
            import numpy as np
            import paddle
            from paddle.distributed import init_parallel_env

            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
            init_parallel_env()
            if paddle.distributed.ParallelEnv().local_rank == 0:
                np_data = np.array([[4, 5, 6], [4, 5, 6]])
            else:
                np_data = np.array([[1, 2, 3], [1, 2, 3]])
            data = paddle.to_tensor(np_data)
            paddle.distributed.reduce(data, 0)
            out = data.numpy()
            # [[5, 7, 9], [5, 7, 9]]
766
    """
K
kuizhiqing 已提交
767 768 769
    if group is not None and not group.is_member():
        return

L
lilong12 已提交
770
    if in_dygraph_mode():
771 772 773 774 775 776
        if op == ReduceOp.SUM:
            op_type = core.ReduceOp.SUM
        elif op == ReduceOp.MAX:
            op_type = core.ReduceOp.MAX
        elif op == ReduceOp.MIN:
            op_type = core.ReduceOp.MIN
777 778
        elif op == ReduceOp.PROD:
            op_type = core.ReduceOp.PRODUCT
779 780 781 782 783 784 785 786 787 788 789
        else:
            raise ValueError("Unknown reduce_op type for reduce.")
        group = _get_default_group() if group is None else group
        gdst = group.get_group_rank(dst)
        assert gdst >= 0, ("dst rank out of group, need global rank")
        task = group.process_group.reduce(tensor, gdst, op_type)
        if use_calc_stream:
            task.wait()
            return None
        else:
            return task
K
kuizhiqing 已提交
790 791 792

    ring_id = 0 if group is None else group.id
    gdst = dst if group is None else group.get_group_rank(dst)
K
kuizhiqing 已提交
793
    assert gdst >= 0, ("dst rank out of group, need global rank")
K
kuizhiqing 已提交
794

J
Jiabin Yang 已提交
795
    if _non_static_mode():
796
        if op == ReduceOp.SUM:
W
wanghuancoder 已提交
797 798 799
            return _C_ops.c_reduce_sum(tensor, tensor, 'use_calc_stream',
                                       use_calc_stream, 'ring_id', ring_id,
                                       'root_id', gdst)
800
        elif op == ReduceOp.MAX:
W
wanghuancoder 已提交
801 802 803
            return _C_ops.c_reduce_max(tensor, tensor, 'use_calc_stream',
                                       use_calc_stream, 'ring_id', ring_id,
                                       'root_id', gdst)
804
        elif op == ReduceOp.MIN:
W
wanghuancoder 已提交
805 806 807
            return _C_ops.c_reduce_min(tensor, tensor, 'use_calc_stream',
                                       use_calc_stream, 'ring_id', ring_id,
                                       'root_id', gdst)
808
        elif op == ReduceOp.PROD:
W
wanghuancoder 已提交
809 810 811
            return _C_ops.c_reduce_prod(tensor, tensor, 'use_calc_stream',
                                        use_calc_stream, 'ring_id', ring_id,
                                        'root_id', gdst)
812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829
        else:
            raise ValueError("Unknown parameter: {}.".format(op))

    op_type = 'c_reduce'
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'all_reduce')

    if op == ReduceOp.SUM:
        op_type = 'c_reduce_sum'
    elif op == ReduceOp.MAX:
        op_type = 'c_reduce_max'
    elif op == ReduceOp.MIN:
        op_type = 'c_reduce_min'
    elif op == ReduceOp.PROD:
        op_type = 'c_reduce_prod'

    helper = LayerHelper(op_type, **locals())
830 831 832 833 834 835 836 837
    helper.append_op(type=op_type,
                     inputs={'X': [tensor]},
                     outputs={'Out': [tensor]},
                     attrs={
                         'ring_id': ring_id,
                         'use_calc_stream': use_calc_stream,
                         'root_id': gdst,
                     })
838 839


K
kuizhiqing 已提交
840
def all_gather(tensor_list, tensor, group=None, use_calc_stream=True):
841 842
    """

843 844 845 846 847 848 849 850 851
    Gather tensors from all participators and all get the result. As shown
    below, 4 GPUs each start 4 processes and the data on each GPU is represnted
    by the GPU number. 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
852 853 854 855 856 857

    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 or int64.
        tensor (Tensor): The Tensor to send. Its data type
            should be float16, float32, float64, int32 or int64.
K
kuizhiqing 已提交
858
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
859 860
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
861 862 863 864 865 866 867

    Returns:
        None.

    Examples:
        .. code-block:: python

868
            # required: distributed
869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887
            import numpy as np
            import paddle
            from paddle.distributed import init_parallel_env

            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
            init_parallel_env()
            tensor_list = []
            if paddle.distributed.ParallelEnv().local_rank == 0:
                np_data1 = np.array([[4, 5, 6], [4, 5, 6]])
                np_data2 = np.array([[4, 5, 6], [4, 5, 6]])
                data1 = paddle.to_tensor(np_data1)
                data2 = paddle.to_tensor(np_data2)
                paddle.distributed.all_gather(tensor_list, data1)
            else:
                np_data1 = np.array([[1, 2, 3], [1, 2, 3]])
                np_data2 = np.array([[1, 2, 3], [1, 2, 3]])
                data1 = paddle.to_tensor(np_data1)
                data2 = paddle.to_tensor(np_data2)
                paddle.distributed.all_gather(tensor_list, data2)
888
    """
K
kuizhiqing 已提交
889 890 891
    if group is not None and not group.is_member():
        return

L
lilong12 已提交
892
    if in_dygraph_mode():
893
        group = _get_default_group() if group is None else group
894 895 896 897 898 899
        if len(tensor_list) == 0:
            tensor_shape = list(tensor.shape)
            tensor_shape[0] *= group.nranks
            out = paddle.empty(tensor_shape, tensor.dtype)
        else:
            out = paddle.concat(tensor_list, axis=0)
900 901 902 903 904 905
        task = group.process_group.all_gather(tensor, out)
        task.wait()
        tensor_list.clear()
        tensor_list.extend(paddle.split(out, group.nranks, 0))
        return

K
kuizhiqing 已提交
906 907 908
    ring_id = 0 if group is None else group.id
    nranks = _get_global_group().nranks if group is None else group.nranks

J
Jiabin Yang 已提交
909
    if _non_static_mode():
910 911
        out = _C_ops.c_allgather(tensor, 'use_calc_stream', use_calc_stream,
                                 'ring_id', ring_id, 'nranks', nranks)
912
    else:
913 914 915
        op_type = 'c_allgather'
        helper = LayerHelper(op_type, **locals())
        out = helper.create_variable_for_type_inference(dtype=tensor.dtype)
916 917 918 919 920 921 922 923 924 925 926
        if not isinstance(tensor_list, list):
            raise ValueError("The type of 'tensor_list' for all_gather "
                             "should be list.")
        for elem in tensor_list:
            check_variable_and_dtype(
                elem, 'tensor_list',
                ['float16', 'float32', 'float64', 'int32', 'int64'],
                'all_gather')
        check_variable_and_dtype(
            tensor, 'tensor',
            ['float16', 'float32', 'float64', 'int32', 'int64'], 'all_gather')
927 928 929 930 931 932 933 934
        helper.append_op(type=op_type,
                         inputs={'X': [tensor]},
                         outputs={'Out': [out]},
                         attrs={
                             'ring_id': ring_id,
                             'use_calc_stream': use_calc_stream,
                             'nranks': nranks
                         })
935

K
kuizhiqing 已提交
936
    tensor_list.extend(paddle.split(out, nranks, 0))
937 938


K
kuizhiqing 已提交
939
def scatter(tensor, tensor_list=None, src=0, group=None, use_calc_stream=True):
940 941
    """

942 943 944 945 946 947 948
    Scatter a tensor to all participators. As shown below, 4 GPUs each start 4 processes 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
949 950 951 952

    Args:
        tensor (Tensor): The output Tensor. Its data type
            should be float16, float32, float64, int32 or int64.
953
        tensor_list (list|tuple): A list/tuple of Tensors to scatter. Every element in the list must be a Tensor whose data type
K
kuizhiqing 已提交
954 955
            should be float16, float32, float64, int32 or int64. Default value is None.
        src (int): The source rank id. Default value is 0.
K
kuizhiqing 已提交
956
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
957 958
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
959 960 961 962 963 964 965

    Returns:
        None.

    Examples:
        .. code-block:: python

966
            # required: distributed
967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985
            import numpy as np
            import paddle
            from paddle.distributed import init_parallel_env

            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
            init_parallel_env()
            if paddle.distributed.ParallelEnv().local_rank == 0:
                np_data1 = np.array([7, 8, 9])
                np_data2 = np.array([10, 11, 12])
            else:
                np_data1 = np.array([1, 2, 3])
                np_data2 = np.array([4, 5, 6])
            data1 = paddle.to_tensor(np_data1)
            data2 = paddle.to_tensor(np_data2)
            if paddle.distributed.ParallelEnv().local_rank == 0:
                paddle.distributed.scatter(data1, src=1)
            else:
                paddle.distributed.scatter(data1, tensor_list=[data1, data2], src=1)
            out = data1.numpy()
986
    """
K
kuizhiqing 已提交
987 988 989 990 991 992
    if group is not None and not group.is_member():
        return

    if not isinstance(src, int):
        raise ValueError("src should be int.")

L
lilong12 已提交
993
    if in_dygraph_mode():
994 995 996 997 998 999 1000 1001 1002
        group = _get_default_group() if group is None else group
        gsrc = group.get_group_rank(src)
        rank = group.rank
        nranks = group.nranks
    else:
        ring_id = 0 if group is None else group.id
        gsrc = src if group is None else group.get_group_rank(src)
        rank = _get_global_group().rank if group is None else group.rank
        nranks = _get_global_group().nranks if group is None else group.nranks
K
kuizhiqing 已提交
1003
    assert gsrc >= 0, ("src rank out of group, need global rank")
K
kuizhiqing 已提交
1004 1005

    if rank != gsrc:
1006 1007 1008 1009
        tensor_list = []
        for _ in range(nranks):
            tensor_list.append(tensor)
    temp = paddle.concat(tensor_list, axis=0)
L
lilong12 已提交
1010
    if in_dygraph_mode():
1011 1012 1013 1014 1015 1016 1017
        task = group.process_group.scatter(temp, tensor, gsrc)
        if use_calc_stream:
            task.wait()
            return None
        else:
            return task

L
lilong12 已提交
1018
    if _non_static_mode():
W
wanghuancoder 已提交
1019 1020 1021
        return _C_ops.c_scatter(temp, tensor, 'use_calc_stream',
                                use_calc_stream, 'ring_id', ring_id, 'nranks',
                                nranks, 'root', gsrc)
W
wanghuancoder 已提交
1022
    op_type = 'c_scatter'
1023 1024 1025 1026
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'scatter')
    helper = LayerHelper(op_type, **locals())
1027 1028 1029 1030 1031 1032 1033 1034 1035
    helper.append_op(type=op_type,
                     inputs={'X': [temp]},
                     outputs={'Out': [tensor]},
                     attrs={
                         'ring_id': ring_id,
                         'root': gsrc,
                         'use_calc_stream': use_calc_stream,
                         'nranks': nranks,
                     })
1036 1037


1038
def _c_identity(tensor, group=None):
L
lilong12 已提交
1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049
    """
    Return a copy of the tensor, mainly used with model parallel.

    Args:
        tensor (Tensor): The input Tensor. Its data type
            should be float16, float32, float64, int32 or int64.
        group (int): The id of the process group to work on.

    Returns:
        Tensor.
    """
1050 1051 1052 1053
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

J
Jiabin Yang 已提交
1054
    if _non_static_mode():
W
wanghuancoder 已提交
1055 1056
        return _C_ops.c_identity(tensor, 'use_calc_stream', True, 'ring_id',
                                 ring_id, 'use_model_parallel', True)
L
lilong12 已提交
1057 1058 1059
    op_type = 'c_identity'
    helper = LayerHelper(op_type, **locals())
    out = helper.create_variable_for_type_inference(dtype=tensor.dtype)
1060

L
lilong12 已提交
1061 1062 1063
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        '_c_identity')
1064

1065 1066 1067 1068 1069 1070 1071 1072
    helper.append_op(type=op_type,
                     inputs={'X': tensor},
                     outputs={'Out': out},
                     attrs={
                         'ring_id': ring_id,
                         'use_calc_stream': True,
                         'use_model_parallel': True,
                     })
L
lilong12 已提交
1073 1074 1075
    return out


1076
def _c_concat(tensor, group=None):
1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089
    """
    Return allgather of the tensor, mainly used with model parallel.

    Args:
        tensor (Tensor): The input Tensor. Its data type
            should be float16, float32, float64, int32 or int64.
        group (int): The id of the process group to work on.

    Returns:
        Tensor.
    """
    if group is not None and not group.is_member():
        return
1090 1091
    group = _get_default_group() if group is None else group
    ring_id = group.id
1092

1093
    global_rank = _get_global_env().rank
1094 1095
    rank = group.rank
    nranks = group.nranks
1096

J
Jiabin Yang 已提交
1097
    if _non_static_mode():
W
wanghuancoder 已提交
1098 1099 1100
        return _C_ops.c_concat(tensor, 'ring_id', ring_id, 'use_calc_stream',
                               True, 'rank', rank, 'nranks', nranks,
                               'use_model_parallel', True)
1101 1102 1103 1104 1105 1106 1107 1108 1109

    op_type = 'c_concat'
    helper = LayerHelper(op_type, **locals())
    out = helper.create_variable_for_type_inference(dtype=tensor.dtype)

    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        '_c_concat')

1110 1111 1112 1113 1114 1115 1116 1117 1118 1119
    helper.append_op(type=op_type,
                     inputs={'X': tensor},
                     outputs={'Out': out},
                     attrs={
                         'ring_id': ring_id,
                         'use_calc_stream': True,
                         'use_model_parallel': True,
                         'nranks': nranks,
                         'rank': rank
                     })
1120 1121 1122
    return out


1123
def _c_split(tensor, group=None):
L
lilong12 已提交
1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135
    """
    Split tensor evenly among all members, mainly used with model parallel.

    Args:
        tensor (Tensor): The input Tensor. Its data type
            should be float16, float32, float64, int32 or int64.
        rank (int): The rank of the current process.
        group (int): The id of the process group to work on.

    Returns:
        Tensor.
    """
1136 1137 1138 1139
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

1140 1141 1142 1143
    global_rank = _get_global_env().rank
    rank = global_rank if group is None else group.get_group_rank(global_rank)
    nranks = _get_global_env().world_size if group is None else group.nranks

J
Jiabin Yang 已提交
1144
    if _non_static_mode():
W
wanghuancoder 已提交
1145 1146 1147
        return _C_ops.c_split(tensor, 'use_calc_stream', True, 'ring_id',
                              ring_id, 'rank', rank, 'nranks', nranks,
                              'use_model_parallel', True)
1148

L
lilong12 已提交
1149 1150 1151
    op_type = 'c_split'
    helper = LayerHelper(op_type, **locals())
    out = helper.create_variable_for_type_inference(dtype=tensor.dtype)
1152

L
lilong12 已提交
1153 1154 1155
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        '_c_split')
1156

1157 1158 1159 1160 1161 1162 1163 1164 1165 1166
    helper.append_op(type=op_type,
                     inputs={'X': tensor},
                     outputs={'Out': out},
                     attrs={
                         'ring_id': ring_id,
                         'use_calc_stream': True,
                         'rank': rank,
                         'nranks': nranks,
                         'use_model_parallel': True,
                     })
L
lilong12 已提交
1167 1168 1169
    return out


1170 1171 1172 1173 1174
def _mp_allreduce(tensor,
                  op=ReduceOp.SUM,
                  group=None,
                  use_calc_stream=True,
                  use_model_parallel=True):
1175
    """[it is same as allreduce above, but it supports model parallel. And it support inplace startegy]
1176 1177 1178 1179 1180
    """
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

1181 1182 1183 1184 1185 1186
    if in_dygraph_mode():
        assert op == ReduceOp.SUM, "Unknown parameter: {}.".format(op)

        from paddle.autograd import EagerPyLayer

        class mp_allreduce_eager(EagerPyLayer):
1187

1188 1189 1190 1191
            @staticmethod
            def forward(ctx, tensor, use_calc_stream, ring_id,
                        use_model_parallel):
                ctx.ring_id = ring_id
1192 1193 1194 1195
                return _C_ops.c_allreduce_sum_(tensor, 'use_calc_stream',
                                               use_calc_stream, 'ring_id',
                                               ring_id, "use_model_parallel",
                                               use_model_parallel)
1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206

            @staticmethod
            def backward(ctx, dy):
                return _C_ops.c_identity(dy, 'use_calc_stream', True, 'ring_id',
                                         ctx.ring_id, 'use_model_parallel',
                                         True)

        return mp_allreduce_eager.apply(tensor, use_calc_stream, ring_id,
                                        use_model_parallel)

    elif _in_legacy_dygraph():
1207
        if op == ReduceOp.SUM:
1208 1209 1210 1211
            return _C_ops.c_allreduce_sum_(tensor, 'use_calc_stream',
                                           use_calc_stream, 'ring_id', ring_id,
                                           "use_model_parallel",
                                           use_model_parallel)
1212 1213
        else:
            raise ValueError("Unknown parameter: {}.".format(op))
1214 1215 1216 1217 1218 1219 1220 1221 1222

    op_type = 'c_allreduce_sum'
    helper = LayerHelper(op_type, **locals())
    out = helper.create_variable_for_type_inference(dtype=tensor.dtype)

    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        op_type)

1223 1224 1225 1226 1227 1228 1229 1230
    helper.append_op(type=op_type,
                     inputs={'X': tensor},
                     outputs={'Out': out},
                     attrs={
                         'ring_id': ring_id,
                         'use_calc_stream': use_calc_stream,
                         'use_model_parallel': use_model_parallel,
                     })
1231
    return out
1232 1233


1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247
def _c_lookup_table(table, index, start_index=0, name=None):
    """
    Lookup table according to index.

    Args:
        table (Tensor): The input Tensor. Its data type
            should be float16, float32, float64.
        index (Tensor): The index to lookup table.
        start_index (int): The initial index for table range.
        name (string): The name of the api

    Returns:
        Tensor.
    """
J
Jiabin Yang 已提交
1248
    if _non_static_mode():
W
wanghuancoder 已提交
1249
        return _C_ops.c_embedding(table, index, "start_index", start_index)
1250

1251 1252 1253 1254 1255
    op_type = 'c_embedding'
    helper = LayerHelper(op_type, **locals())
    dtype = helper.input_dtype(input_param_name='table')
    check_variable_and_dtype(index, 'input', ['int32', 'int64'], op_type)
    tmp = helper.create_variable_for_type_inference(dtype)
1256 1257 1258 1259 1260 1261 1262
    helper.append_op(type='c_embedding',
                     inputs={
                         'Ids': index,
                         'W': table
                     },
                     outputs={'Out': tmp},
                     attrs={"start_index": start_index})
1263 1264
    return tmp

1265

B
Baibaifan 已提交
1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280
class _Linear(layers.Layer):
    """
    Linear
    """

    def __init__(self,
                 in_features,
                 out_features,
                 weight_attr=None,
                 bias_attr=None,
                 name=None):
        super(_Linear, self).__init__()
        self._dtype = self._helper.get_default_dtype()
        self._weight_attr = weight_attr
        self._bias_attr = bias_attr
1281 1282 1283 1284 1285 1286 1287 1288
        self.weight = self.create_parameter(shape=[in_features, out_features],
                                            attr=self._weight_attr,
                                            dtype=self._dtype,
                                            is_bias=False)
        self.bias = self.create_parameter(shape=[out_features],
                                          attr=self._bias_attr,
                                          dtype=self._dtype,
                                          is_bias=True)
B
Baibaifan 已提交
1289 1290 1291
        self.name = name

    def forward(self, input):
1292 1293 1294 1295
        out = _linear(x=input,
                      weight=self.weight,
                      bias=self.bias,
                      name=self.name)
B
Baibaifan 已提交
1296 1297 1298 1299 1300 1301 1302 1303
        return out

    def extra_repr(self):
        name_str = ', name={}'.format(self.name) if self.name else ''
        return 'in_features={}, out_features={}, dtype={}{}'.format(
            self.weight.shape[0], self.weight.shape[1], self._dtype, name_str)


1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323
def _c_softmax_with_cross_entropy(logits,
                                  label,
                                  group=None,
                                  return_softmax=False):
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id
    global_rank = _get_global_env().rank
    rank = global_rank if group is None else group.get_group_rank(global_rank)
    nranks = _get_global_env().world_size if group is None else group.nranks

    input_dims = len(list(logits.shape))
    label_dims = len(list(label.shape))
    if input_dims - 1 != label_dims and input_dims != label_dims:
        raise ValueError(
            'Expected nput_dims - 1 = label_dims or input_dims == label_dims\
             (got nput_dims{}, label_dims{})'.format(input_dims, label_dims))
    if input_dims - 1 == label_dims:
        label = paddle.unsqueeze(label, axis=-1)

J
Jiabin Yang 已提交
1324
    if _non_static_mode():
W
wanghuancoder 已提交
1325
        softmax, loss = _C_ops.c_softmax_with_cross_entropy(
1326 1327 1328 1329 1330 1331
            logits, label, 'ring_id', ring_id, 'rank', rank, 'nranks', nranks)
        if not return_softmax:
            return loss
        else:
            return loss, softmax

W
WangXi 已提交
1332 1333 1334 1335 1336 1337 1338 1339
    attrs = {
        'ring_id': ring_id,
        'rank': rank,
        'nranks': nranks,
    }
    helper = LayerHelper('c_softmax_with_cross_entropy', **locals())
    softmax = helper.create_variable_for_type_inference(dtype=logits.dtype)
    loss = helper.create_variable_for_type_inference(dtype=logits.dtype)
1340 1341 1342 1343 1344 1345 1346 1347 1348 1349
    helper.append_op(type='c_softmax_with_cross_entropy',
                     inputs={
                         'Logits': logits,
                         'Label': label
                     },
                     outputs={
                         'Softmax': softmax,
                         'Loss': loss
                     },
                     attrs=attrs)
W
WangXi 已提交
1350 1351 1352 1353 1354 1355

    if return_softmax:
        return loss, softmax

    return loss

1356

B
Baibaifan 已提交
1357 1358 1359 1360
def _linear(x, weight, bias=None, name=None):
    """
    Fuction Linear
    """
J
Jiabin Yang 已提交
1361
    if _non_static_mode():
B
Baibaifan 已提交
1362
        pre_bias = _varbase_creator(dtype=x.dtype)
W
wanghuancoder 已提交
1363 1364
        _C_ops.matmul(x, weight, pre_bias, 'transpose_X', False, 'transpose_Y',
                      False, "alpha", 1)
1365 1366 1367
        return dygraph_utils._append_bias_in_dygraph(pre_bias,
                                                     bias,
                                                     axis=len(x.shape) - 1)
B
Baibaifan 已提交
1368 1369 1370
    else:
        helper = LayerHelper('linear', **locals())
        dtype = x.dtype
B
Baibaifan 已提交
1371 1372
        assert len(
            x.shape) < 4, "X latitude is not supported greater than 3 now."
B
Baibaifan 已提交
1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384

        check_variable_and_dtype(x, 'x', ['float16', 'float32', 'float64'],
                                 'linear')
        check_dtype(dtype, 'dtype', ['float16', 'float32', 'float64'], 'linear')

        inputs = {'X': [x], 'Y': [weight]}
        attrs = {
            'transpose_X': False,
            'transpose_Y': False,
            'alpha': 1,
        }
        tmp = helper.create_variable_for_type_inference(dtype)
1385 1386 1387 1388
        helper.append_op(type='matmul_v2',
                         inputs=inputs,
                         outputs={'Out': tmp},
                         attrs=attrs)
B
Baibaifan 已提交
1389 1390
        if bias is not None:
            res = helper.create_variable_for_type_inference(dtype)
1391 1392 1393 1394 1395 1396 1397
            helper.append_op(type='elementwise_add',
                             inputs={
                                 'X': [tmp],
                                 'Y': [bias]
                             },
                             outputs={'Out': [res]},
                             attrs={'axis': len(x.shape) - 1})
B
Baibaifan 已提交
1398 1399 1400 1401 1402
        else:
            res = tmp
        return res


1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415
def _set_var_distributed(var):
    if var is None:
        return

    var.is_distributed = True

    # NOTE: use current_block and find_var_recursive to support while_loop
    startup_block = paddle.static.default_startup_program().current_block()
    main_block = paddle.static.default_main_program().current_block()
    startup_block._find_var_recursive(var.name).is_distributed = True
    main_block._find_var_recursive(var.name).is_distributed = True


L
lilong12 已提交
1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426
def _parallel_linear(x,
                     num_rows,
                     num_cols,
                     axis,
                     param_attr,
                     bias_attr,
                     gather_out,
                     inner_rank,
                     nranks,
                     split_tensor,
                     name,
1427
                     group=None):
1428 1429
    """
    Parallel Linear
1430 1431 1432

    axis the dimension of the parameter of linear layer. 
    axis = 0: the row dimension
1433
    axis = 1: the col dimension
1434
    
1435
    """
1436 1437 1438 1439
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

L
lilong12 已提交
1440 1441
    if axis == 0:
        if split_tensor:
1442
            x = _c_split(x, group=group)
1443
    else:
L
lilong12 已提交
1444 1445
        x = _c_identity(x, group=group)

1446 1447 1448 1449 1450
    linear = paddle.nn.Linear(num_rows,
                              num_cols,
                              weight_attr=param_attr,
                              bias_attr=bias_attr,
                              name=name)
1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462

    # NOTE: npu linear function use matmul_v2 but linear use matmul
    linear_function = _linear if core.is_compiled_with_npu()\
        else paddle.nn.functional.linear
    linear_out = linear_function(
        x,
        linear.weight,
        # NOTE(wangxi): row split, bias need add after allreduce
        None if axis == 0 else linear.bias,
        linear.name)

    _set_var_distributed(linear.weight)
1463 1464 1465 1466
    # set is_distributed for splited bias
    # if a linear layer is splited by row, each rank would hold a complete bias and they should be the same in each rank.
    # if a linear layer is splited by col, the bias would also be split into each rank as its weight
    if axis == 1 and linear._bias_attr != False:
1467
        _set_var_distributed(linear.bias)
L
lilong12 已提交
1468 1469 1470 1471 1472

    if not gather_out: return linear_out

    out_shape = list(linear_out.shape)
    out_shape[0] *= 1 if axis == 0 else nranks
1473
    main_block = paddle.static.default_main_program().current_block()
L
lilong12 已提交
1474 1475 1476 1477 1478 1479 1480 1481 1482
    out = main_block.create_var(
        shape=out_shape,
        dtype=linear_out.dtype,
        type=linear_out.type,
        lod_level=linear_out.lod_level,
        persistable=False,
        is_data=False,
        need_check_feed=linear_out.desc.need_check_feed())
    if axis == 0:
1483 1484 1485 1486 1487 1488 1489 1490
        main_block.append_op(type='c_allreduce_sum',
                             inputs={'X': linear_out},
                             outputs={'Out': out},
                             attrs={
                                 'ring_id': ring_id,
                                 'use_calc_stream': True,
                                 'use_model_parallel': True
                             })
1491 1492
        if linear.bias is not None:
            out = out + linear.bias
L
lilong12 已提交
1493
    else:
1494 1495 1496 1497 1498 1499 1500 1501 1502 1503
        main_block.append_op(type='c_concat',
                             inputs={'X': linear_out},
                             outputs={'Out': out},
                             attrs={
                                 'rank': inner_rank,
                                 'ring_id': ring_id,
                                 'nranks': nranks,
                                 'use_calc_stream': True,
                                 'use_model_parallel': True
                             })
L
lilong12 已提交
1504
    return out
1505 1506


L
lilong12 已提交
1507 1508 1509 1510 1511 1512 1513
def _parallel_embedding(x,
                        per_part_embeddings,
                        origin_size,
                        param_attr,
                        inner_rank,
                        num_partitions,
                        name,
1514
                        group=None):
1515 1516 1517
    """
    Parallel Embedding
    """
1518 1519 1520 1521
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

1522 1523 1524 1525 1526 1527 1528 1529 1530
    helper = LayerHelper("_parallel_embedding", **locals())

    per_part_size = per_part_embeddings
    rank = inner_rank

    vocab_start_index = rank * per_part_size
    dtype = helper.get_default_dtype()
    size = [per_part_size, origin_size[1]]

1531 1532 1533 1534
    weight = helper.create_parameter(attr=param_attr,
                                     shape=size,
                                     dtype=dtype,
                                     is_bias=False)
1535 1536

    if num_partitions == 1:
1537 1538 1539 1540 1541
        return paddle.nn.functional.embedding(x,
                                              weight=weight,
                                              padding_idx=None,
                                              sparse=False,
                                              name=name)
1542

1543 1544
    startup_block = paddle.static.default_startup_program().global_block()
    main_block = paddle.static.default_main_program().global_block()
1545 1546 1547 1548 1549
    startup_block.vars[weight.name].is_distributed = True
    main_block.vars[weight.name].is_distributed = True

    output_parallel = paddle.distributed.collective._c_lookup_table(
        weight, x, start_index=vocab_start_index, name=name)
1550 1551 1552 1553
    out = paddle.distributed.collective._mp_allreduce(output_parallel,
                                                      group=group,
                                                      use_calc_stream=True,
                                                      use_model_parallel=True)
L
lilong12 已提交
1554
    return out
1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577


def split(x,
          size,
          operation,
          axis=0,
          num_partitions=1,
          gather_out=True,
          weight_attr=None,
          bias_attr=None,
          name=None):
    """

    Split the weight of the specified operation into multiple devices
    and do the computation in parallel.

    Now the following three cases are supported.

    Case 1: Parallel Embedding
        The weight of the embedding operation is a NxM matrix with N rows and M columns.
        With parallel embedding, the weight is split into num_partitions partitions, each
        of which is a matrix with (N/num_partitions + 1) rows and M column where the last
        row as the padding idx.
K
kuizhiqing 已提交
1578

1579 1580 1581 1582 1583 1584 1585 1586 1587
        Suppose we split the NxM weight into two partitons on device_0 and device_1
        respectively. Then, one each device, the final weight has (N/2 + 1) rows with the
        index range from 0 to N/2. On device_0, all values in the input within [0, N/2 -1]
        keep unchanged and all other values are changed to N/2 which is the padding index and
        are mapped to all zeros after embedding. In the same way, on device_1, the value V in the
        input within [N/2, N-1] will be changed to (V - N/2), and all other values are changed
        to N/2 and are mapped to all zeros after embedding. Finally, the results on the two
        devices are sum-reduced.

1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602
        The Embedding put on single card is as shown below:

        .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/split_embedding_single.png
            :width: 800
            :height: 350
            :alt: single_embedding
            :align: center

        Parallel Embedding is shown as below:

        .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/split_embedding_split.png
            :width: 800
            :alt: split_embedding
            :align: center

1603 1604 1605 1606 1607
    Case 2: Row Parallel Linear
        The weight of the linear operation is a NxM matrix with N rows and M columns.
        With row parallel linear, the weight is split into num_partitions partitions, each
        of which is a matrix with N/num_partitions rows and M column.

1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625
        The linear layer put on single card is shown as below, the input variable is represented by X,
        the weight matrix is represented by W and the output vaiable is O. The linear layer on single card is 
        simple matrix multiplication operation, O = X * W.

        .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/split_single.png
            :width: 800
            :alt: single_linear
            :align: center

        Row Parallel Linear is shown as below. As the name suggests, Row Parallel Linear splits the weight matrix W into
        [[W_row1], [W_row2]] along the row. And accordingly the input is splitted along the column into [X_col1, X_col2] and multiply their
        respective weight matrices. Finally apply AllReduce on the output from each card to get the final output.

        .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/split_row.png
            :width: 800
            :alt: split_row
            :align: center

1626 1627 1628 1629 1630
    Case 3: Column Parallel Linear
        The weight of the linear operation is a NxM matrix with N rows and M columns.
        With column parallel linear, the weight is split into num_paratitions partitions, each
        of which is a matrix with N rows and M/num_partitions column.

1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646 1647
        The linear layer put on single card has been illustrated on case 2 and Column Parallel Linear
        is shown as below. The Column Parallel Linear splits the weight matrix W into [W_col1, W_col2] along the column and 
        these splitted matrices respectively multiply the input. Finally apply AllGather on the output from each card to get the final output. 

        .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/split_col.png
            :width: 800
            :alt: split_col
            :align: center
    
    As observed, the column parallel linear and row parallel linear can be combined to skip one ALLGATHER communication
    operator. Furthermore the Attention and MLP can be combined to imporve the performance as shown below.

    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/split_col_row.png
            :width: 800
            :alt: split_col_row
            :align: center

1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666 1667
    Args:
        x (Tensor): Input tensor. It's data type should be float16, float32, float64, int32 or int64.
        size (list|tuple): A list or tuple with two elements indicating the shape of the weight.
        operation (str): The name of the operation. The supported operations are 'linear' and 'embedding'.
        axis (int, Optional): Indicate along which axis to split the weight. Default: 0.
        num_partitions (int, Optional): How many parts the weight is partitioned. Default: 1.
        gather_out (bool, Optional): Whether to gather the output after computation. By default, the output
            on each partitions will be gathered after computation. Default: True.
        weight_attr (ParamAttr, Optional): The parameter attribute for the learnable
            weights(Parameter) of the specified operation. Default: None.
        bias_attr (ParamAttr, Optional): The parameter attribute for the bias
            of the specified operation. Default: None.
        name (str, Optional): The default value is None. Normally there is no need for user to set this
            property. Default: None. For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        Tensor.

    Examples:
        .. code-block:: python
1668

1669
            # required: distributed
1670
            import paddle
1671
            import paddle.distributed.fleet as fleet
1672

1673
            paddle.enable_static()
1674
            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
1675
            fleet.init(is_collective=True)
1676
            data = paddle.randint(0, 8, shape=[10,4])
1677
            emb_out = paddle.distributed.split(
1678 1679 1680 1681
                data,
                (8, 8),
                operation="embedding",
                num_partitions=2)
1682

1683
    """
1684 1685 1686 1687
    assert isinstance(
        size,
        (list, tuple)), ("The type of size for "
                         "paddle.distributed.split must be list or tuple.")
1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699
    assert len(size) == 2, ("Number of elements in size of "
                            "paddle.distributed.split must be two.")
    assert isinstance(operation, str), ("The type of operation for "
                                        "paddle.distributed.split must be str.")
    supported_operations = [
        'linear',
        'embedding',
    ]
    assert operation in supported_operations, (
        "The operation for "
        "paddle.distributed.split must be one of {}.".format(
            supported_operations))
J
Jiabin Yang 已提交
1700
    if _non_static_mode():
L
lilong12 已提交
1701 1702 1703 1704
        raise ValueError(
            "paddle.distributed.split cannot be used in dynamic "
            "graph mode, plese use ParallelEmbedding, ParallelRowLinear, "
            "ParallelColumnLinear instead.")
1705
    else:
1706
        from .fleet import fleet
1707 1708 1709 1710 1711 1712 1713 1714 1715 1716 1717
        assert fleet._role_maker, ("To use paddle.distributed.split, "
                                   "you must call fleet.init() firstly.")
        rank = fleet.worker_index()
        nranks = fleet.worker_num()

    # rank within a model parallel group
    inner_rank = rank % num_partitions

    if operation == "embedding":
        assert axis == 0, ("We only support to split the weight of embedding "
                           "along the first axis now.")
1718 1719 1720
        assert size[0] % num_partitions == 0, \
            "The length of the vocabulary must be divisible by num_partitions " \
            "but received vocabulary={} num_partitions={}".format(size[0], num_partitions)
1721

1722
        per_part_size = size[0] // num_partitions
1723 1724 1725 1726 1727 1728 1729 1730
        emb_out = _parallel_embedding(x,
                                      per_part_size,
                                      size,
                                      weight_attr,
                                      inner_rank,
                                      num_partitions,
                                      name,
                                      group=None)
B
Baibaifan 已提交
1731
        return emb_out
1732
    else:
L
lilong12 已提交
1733
        should_split = False
1734 1735 1736
        if axis == 0:
            assert size[0] % num_partitions == 0, (
                "Number of rows of the weight for linear ({}) must be"
1737 1738
                " divisible by num_partitions ({})".format(
                    size[0], num_partitions))
1739 1740
            per_part_size = size[0] // num_partitions
            linear_size = (per_part_size, size[1])
L
lilong12 已提交
1741
            if x.shape[-1] == size[0]: should_split = True
1742 1743 1744 1745

        elif axis == 1:
            assert size[1] % num_partitions == 0, (
                "Number of column of the weight for linear ({}) must be"
1746 1747
                " divisible by num_partitions ({})".format(
                    size[1], num_partitions))
1748 1749 1750 1751 1752 1753
            per_part_size = size[1] // num_partitions
            linear_size = (size[0], per_part_size)
        else:
            raise ValueError("The value of axis must be 0 or 1, but the value "
                             "given is {}.".format(axis))

1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765
        linear_out = _parallel_linear(x,
                                      linear_size[0],
                                      linear_size[1],
                                      axis,
                                      weight_attr,
                                      bias_attr,
                                      gather_out,
                                      inner_rank,
                                      num_partitions,
                                      should_split,
                                      name=name,
                                      group=None)
1766
        return linear_out
L
lilong12 已提交
1767 1768


L
lilong12 已提交
1769 1770
def alltoall(in_tensor_list, out_tensor_list, group=None, use_calc_stream=True):
    """
1771 1772 1773 1774 1775 1776 1777 1778 1779 1780
    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

L
lilong12 已提交
1781 1782 1783 1784 1785 1786 1787
    Args:
        in_tensor_list (list): A list of input Tensors. Every element in the list must be a Tensor whose data type
            should be float16, float32, float64, int32 or int64.
        out_tensor_list (Tensor): A list of output Tensors. The data type of its elements should be the same as the
            data type of the input Tensors.
        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
        use_calc_stream (bool, optional): Wether to use calculation stream (True) or communication stream. Default: True.
1788
    
L
lilong12 已提交
1789 1790
    Returns:
        None.
1791
    
L
lilong12 已提交
1792 1793
    Examples:
        .. code-block:: python
1794

L
lilong12 已提交
1795 1796 1797 1798
            # required: distributed
            import numpy as np
            import paddle
            from paddle.distributed import init_parallel_env
1799
            
L
lilong12 已提交
1800 1801 1802 1803 1804 1805 1806 1807 1808 1809
            init_parallel_env()
            out_tensor_list = []
            if paddle.distributed.ParallelEnv().rank == 0:
                np_data1 = np.array([[1, 2, 3], [4, 5, 6]])
                np_data2 = np.array([[7, 8, 9], [10, 11, 12]])
            else:
                np_data1 = np.array([[13, 14, 15], [16, 17, 18]])
                np_data2 = np.array([[19, 20, 21], [22, 23, 24]])
            data1 = paddle.to_tensor(np_data1)
            data2 = paddle.to_tensor(np_data2)
李季 已提交
1810
            paddle.distributed.alltoall([data1, data2], out_tensor_list)
L
lilong12 已提交
1811 1812 1813 1814 1815 1816
            # out for rank 0: [[[1, 2, 3], [4, 5, 6]], [[13, 14, 15], [16, 17, 18]]]
            # out for rank 1: [[[7, 8, 9], [10, 11, 12]], [[19, 20, 21], [22, 23, 24]]]
    """
    if group is not None and not group.is_member():
        return

L
lilong12 已提交
1817
    if in_dygraph_mode():
1818 1819 1820 1821
        group = _get_default_group() if group is None else group
    else:
        ring_id = 0 if group is None else group.id

L
lilong12 已提交
1822
    temp = paddle.concat(in_tensor_list, axis=0)
李季 已提交
1823
    nranks = len(in_tensor_list)
L
lilong12 已提交
1824
    if in_dygraph_mode():
1825 1826 1827 1828 1829 1830
        if len(out_tensor_list) == 0:
            tensor_shape = list(in_tensor_list[0].shape)
            tensor_shape[0] *= nranks
            out = paddle.empty(tensor_shape, in_tensor_list[0].dtype)
        else:
            out = paddle.concat(out_tensor_list, axis=0)
1831 1832 1833 1834 1835 1836
        task = group.process_group.alltoall(temp, out)
        task.wait()
        out_tensor_list.clear()
        out_tensor_list.extend(paddle.split(out, nranks, 0))
        return

J
Jiabin Yang 已提交
1837
    if _non_static_mode():
李季 已提交
1838 1839
        out = _C_ops.alltoall(temp, 'use_calc_stream', use_calc_stream,
                              'ring_id', ring_id)
L
lilong12 已提交
1840
    else:
W
wanghuancoder 已提交
1841 1842 1843 1844 1845
        op_type = 'alltoall'
        helper = LayerHelper(op_type, **locals())
        out = helper.create_variable_for_type_inference(
            dtype=in_tensor_list[0].dtype)

L
lilong12 已提交
1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859
        if not isinstance(in_tensor_list, list):
            raise ValueError("The type of 'in_tensor_list' for all_to_all "
                             "should be list.")
        for elem in in_tensor_list:
            check_variable_and_dtype(
                elem, 'in_tensor_list',
                ['float16', 'float32', 'float64', 'int32', 'int64'],
                'all_to_all')
        if not isinstance(out_tensor_list, list):
            raise ValueError("The type of 'out_tensor_list' for all_to_all "
                             "should be list.")
        if len(out_tensor_list) != 0:
            raise ValueError("The 'out_tensor_list' for all_to_all "
                             "must be an empty list.")
1860 1861 1862 1863 1864 1865 1866
        helper.append_op(type=op_type,
                         inputs={'X': [temp]},
                         outputs={'Out': [out]},
                         attrs={
                             'ring_id': ring_id,
                             'use_calc_stream': use_calc_stream,
                         })
L
lilong12 已提交
1867 1868 1869
    out_tensor_list.extend(paddle.split(out, nranks, 0))


L
lilong12 已提交
1870 1871 1872 1873 1874 1875 1876 1877
def send(tensor, dst=0, group=None, use_calc_stream=True):
    """
    Send a tensor to the receiver.

    Args:
        tensor (Tensor): The Tensor to send. Its data type
            should be float16, float32, float64, int32 or int64.
        dst (int): The destination rank id.
L
lilong12 已提交
1878 1879
        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
        use_calc_stream (bool, optional): Whether to use calculate stream or communication stream. Default: True.
1880
    
L
lilong12 已提交
1881 1882 1883 1884 1885
    Returns:
        None.

    Examples:
        .. code-block:: python
1886

L
lilong12 已提交
1887
            # required: distributed
L
lilong12 已提交
1888
            import paddle
L
lilong12 已提交
1889
            from paddle.distributed import init_parallel_env
1890

L
lilong12 已提交
1891 1892 1893 1894 1895 1896 1897 1898
            init_parallel_env()
            if paddle.distributed.ParallelEnv().rank == 0:
                data = paddle.to_tensor([7, 8, 9])
                paddle.distributed.send(data, dst=1)
            else:
                data = paddle.to_tensor([1,2,3])
                paddle.distributed.recv(data, src=0)
            out = data.numpy()
L
lilong12 已提交
1899 1900 1901
    """
    if group is not None and not group.is_member():
        return
1902

L
lilong12 已提交
1903
    if in_dygraph_mode():
1904 1905 1906 1907 1908 1909 1910 1911
        group = _get_default_group() if group is None else group
        task = group.process_group.send(tensor, dst)
        if use_calc_stream:
            task.wait()
            return None
        else:
            return task

L
lilong12 已提交
1912 1913
    ring_id = 0 if group is None else group.id

J
Jiabin Yang 已提交
1914
    if _non_static_mode():
W
wanghuancoder 已提交
1915 1916
        return _C_ops.send_v2(tensor, 'use_calc_stream', use_calc_stream,
                              'ring_id', ring_id, 'peer', dst)
W
wanghuancoder 已提交
1917
    op_type = 'send_v2'
L
lilong12 已提交
1918 1919 1920 1921 1922
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'send')

    helper = LayerHelper(op_type, **locals())
1923 1924 1925 1926 1927 1928 1929
    helper.append_op(type=op_type,
                     inputs={'X': [tensor]},
                     attrs={
                         'ring_id': ring_id,
                         'peer': dst,
                         'use_calc_stream': use_calc_stream,
                     })
L
lilong12 已提交
1930 1931 1932 1933 1934 1935 1936 1937 1938 1939


def recv(tensor, src=0, group=None, use_calc_stream=True):
    """
    Receive a tensor to the sender.

    Args:
        tensor (Tensor): The Tensor to receive. Its data type
            should be float16, float32, float64, int32 or int64.
        src (int): The source rank id.
L
lilong12 已提交
1940 1941
        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
        use_calc_stream (bool, optional): Whether to use calculate stream or communication stream. Default: True.
1942
    
L
lilong12 已提交
1943 1944 1945 1946 1947
    Returns:
        None.

    Examples:
        .. code-block:: python
1948

L
lilong12 已提交
1949
            # required: distributed
L
lilong12 已提交
1950
            import paddle
L
lilong12 已提交
1951
            from paddle.distributed import init_parallel_env
1952

L
lilong12 已提交
1953 1954 1955 1956 1957 1958 1959 1960
            init_parallel_env()
            if paddle.distributed.ParallelEnv().rank == 0:
                data = paddle.to_tensor([7, 8, 9])
                paddle.distributed.send(data, dst=1)
            else:
                data = paddle.to_tensor([1,2,3])
                paddle.distributed.recv(data, src=0)
            out = data.numpy()
L
lilong12 已提交
1961 1962 1963
    """
    if group is not None and not group.is_member():
        return
1964

L
lilong12 已提交
1965
    if in_dygraph_mode():
1966 1967 1968 1969 1970 1971 1972 1973
        group = _get_default_group() if group is None else group
        task = group.process_group.recv(tensor, src)
        if use_calc_stream:
            task.wait()
            return None
        else:
            return task

L
lilong12 已提交
1974 1975
    ring_id = 0 if group is None else group.id

J
Jiabin Yang 已提交
1976
    if _non_static_mode():
W
wanghuancoder 已提交
1977 1978 1979
        return _C_ops.recv_v2(tensor, 'use_calc_stream', use_calc_stream,
                              'ring_id', ring_id, 'peer', src, 'dtype',
                              tensor.dtype, 'out_shape', tensor.shape)
W
wanghuancoder 已提交
1980
    op_type = 'recv_v2'
L
lilong12 已提交
1981 1982 1983 1984
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'recv')
    helper = LayerHelper(op_type, **locals())
1985 1986 1987 1988 1989 1990 1991 1992 1993
    helper.append_op(type=op_type,
                     outputs={'Out': [tensor]},
                     attrs={
                         'ring_id': ring_id,
                         'peer': src,
                         'out_shape': tensor.shape,
                         'dtype': tensor.dtype,
                         'use_calc_stream': use_calc_stream,
                     })