collective.py 55.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
#   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
from ..fluid.layer_helper import LayerHelper
18 19 20 21
from ..fluid.framework import Variable
from ..fluid.framework import OpProtoHolder
from ..fluid.framework import in_dygraph_mode
from ..fluid.framework import convert_np_dtype_to_dtype_
J
Jiangxinz 已提交
22
from ..fluid.framework import _varbase_creator
23 24 25 26
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
27 28
from ..fluid.layers.tensor import fill_constant
from ..fluid.layers import utils
B
Baibaifan 已提交
29
from ..fluid.dygraph import layers
30 31
from ..fluid.dygraph.parallel import prepare_context
import paddle
32
from .fleet import fleet
33 34
import paddle.fluid as fluid
import paddle.fluid.core as core
W
wanghuancoder 已提交
35
from paddle import _C_ops
J
Jiangxinz 已提交
36
import paddle.fluid.dygraph_utils as dygraph_utils
37

38
__all__ = []
39 40 41


class ReduceOp:
L
lilong12 已提交
42 43 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
    """
    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]]
    """
73 74 75 76 77 78
    SUM = 0
    MAX = 1
    MIN = 2
    PROD = 3


K
kuizhiqing 已提交
79 80 81 82
class Group():
    """
    The abstract representation of group.
    """
83

K
kuizhiqing 已提交
84
    def __init__(self, rank, rank_num, id=0, ranks=[]):
85 86
        self.rank = rank
        self.nranks = rank_num
K
kuizhiqing 已提交
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
        self.id = id
        self.ranks = ranks

    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

103 104 105 106 107 108 109
    def __repr__(self):
        debug_str = "rank: {}, nranks: {}, id: {}, ranks: ".format(
            self.rank, self.nranks, self.id)
        debug_str += ", ".join(map(str, self.ranks))
        debug_str += ". "
        return debug_str

K
kuizhiqing 已提交
110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129

_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 = {}


def _get_group_map():
    global _group_map
    if not _group_map:
        genv = _get_global_env()
W
WangXi 已提交
130 131
        _group_map[0] = Group(genv.rank, genv.world_size,
                              list(range(genv.world_size)))
K
kuizhiqing 已提交
132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148
    return _group_map


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


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 已提交
149
        id (int): the group id. Default value is 0.
K
kuizhiqing 已提交
150 151 152 153 154 155 156 157 158 159 160 161 162 163

    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 已提交
164
    return gm[id] if id in gm else None
K
kuizhiqing 已提交
165 166


S
ShenLiang 已提交
167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194
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

    ring_id = 0 if group is None else group.id

    temp = fill_constant([1], dtype="int32", value="1")
    if in_dygraph_mode():
W
wanghuancoder 已提交
195
        return _C_ops.barrier(temp, temp, 'ring_id', ring_id)
W
wanghuancoder 已提交
196 197 198

    op_type = 'barrier'

S
ShenLiang 已提交
199 200 201 202 203 204 205 206 207 208
    if not isinstance(ring_id, int):
        raise ValueError("The type of 'group' for barrier must be int.")
    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [temp]},
        outputs={'Out': [temp]},
        attrs={'ring_id': ring_id})


K
kuizhiqing 已提交
209 210 211
def new_group(ranks=None, backend=None):
    """

K
kuizhiqing 已提交
212
    Creates a new distributed communication group.
K
kuizhiqing 已提交
213 214

    Args:
K
kuizhiqing 已提交
215
        ranks (list): The global ranks of group members.
K
kuizhiqing 已提交
216 217 218
        backend (str): The backend used to create group, only nccl is supported now.

    Returns:
K
kuizhiqing 已提交
219
        Group: The group instance.
K
kuizhiqing 已提交
220 221 222 223 224 225 226

    Examples:
        .. code-block:: python

            import paddle

            paddle.distributed.init_parallel_env()
K
kuizhiqing 已提交
227 228 229
            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 已提交
230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246

    """

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

    global _group_map
    if global_rank not in ranks:
        gp = Group(-1, -1, ring_id, ranks)
        _group_map[ring_id] = gp
    else:
247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273
        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)
            else:
                assert False, ("no cuda device found")
        else:
            return gp

    # TODO(shenliang03): This is a temporary solution to solve the problem of 
    # hang caused by cross-creation of new_group
274 275 276
    tmp = paddle.to_tensor(
        [1], dtype="int32") if in_dygraph_mode() else fill_constant(
            [0], dtype="int32", value="1")
277 278
    paddle.distributed.all_reduce(tmp, use_calc_stream=True)
    paddle.distributed.wait(tmp)
K
kuizhiqing 已提交
279 280
    return gp

281

K
kuizhiqing 已提交
282 283 284 285 286 287 288 289
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 已提交
290 291
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
K
kuizhiqing 已提交
292 293 294 295 296 297 298 299 300 301

    Returns:
        None.

    Examples:
        .. code-block:: python

            import paddle

            paddle.distributed.init_parallel_env()
K
kuizhiqing 已提交
302
            tindata = paddle.randn(shape=[2, 3])
K
kuizhiqing 已提交
303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321
            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):

    if in_dygraph_mode():
W
wanghuancoder 已提交
322
        return _C_ops.c_sync_calc_stream(tensor, tensor)
K
kuizhiqing 已提交
323 324 325 326 327 328 329 330

    op_type = 'c_sync_calc_stream'

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

332

K
kuizhiqing 已提交
333
def _sync_comm_stream(tensor, ring_id=0):
334

K
kuizhiqing 已提交
335
    if in_dygraph_mode():
W
wanghuancoder 已提交
336
        return _C_ops.c_sync_comm_stream([tensor], [tensor], 'ring_id', ring_id)
337

K
kuizhiqing 已提交
338
    op_type = 'c_sync_comm_stream'
339

K
kuizhiqing 已提交
340 341 342 343 344 345 346 347 348
    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [tensor]},
        outputs={'Out': [tensor]},
        attrs={'ring_id': ring_id}, )


def broadcast(tensor, src, group=None, use_calc_stream=True):
349 350 351 352 353 354 355 356
    """

    Broadcast a tensor from the source to all others.

    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 已提交
357
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
358 359
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
360 361 362 363 364 365 366

    Returns:
        None.

    Examples:
        .. code-block:: python

367 368 369 370 371 372 373 374 375 376 377 378 379 380
            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]]
381
    """
K
kuizhiqing 已提交
382 383 384 385 386 387 388 389 390

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

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

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

393
    if in_dygraph_mode():
W
wanghuancoder 已提交
394 395 396
        return _C_ops.c_broadcast(tensor, tensor, 'root', gsrc,
                                  'use_calc_stream', use_calc_stream, 'ring_id',
                                  ring_id)
397 398 399 400 401 402 403 404 405 406 407 408

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

    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [tensor]},
        outputs={'Out': [tensor]},
        attrs={
K
kuizhiqing 已提交
409 410 411
            'root': gsrc,
            'use_calc_stream': use_calc_stream,
            'ring_id': ring_id,
412 413 414
        })


K
kuizhiqing 已提交
415
def all_reduce(tensor, op=ReduceOp.SUM, group=None, use_calc_stream=True):
416 417 418 419 420 421 422
    """

    Reduce a tensor over all ranks so that all get the result.

    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 已提交
423
        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.Min|ReduceOp.PROD): Optional. The operation used. Default value is ReduceOp.SUM.
K
kuizhiqing 已提交
424
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
425 426
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
427 428 429 430 431 432 433

    Returns:
        None.

    Examples:
        .. code-block:: python

434 435 436 437 438 439 440 441 442 443 444 445 446 447 448
            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]]
449
    """
K
kuizhiqing 已提交
450 451 452 453
    if group is not None and not group.is_member():
        return

    ring_id = 0 if group is None else group.id
454 455
    if in_dygraph_mode():
        if op == ReduceOp.SUM:
W
wanghuancoder 已提交
456 457
            return _C_ops.c_allreduce_sum_(tensor, 'use_calc_stream',
                                           use_calc_stream, 'ring_id', ring_id)
458
        elif op == ReduceOp.MAX:
W
wanghuancoder 已提交
459 460
            return _C_ops.c_allreduce_max_(tensor, 'use_calc_stream',
                                           use_calc_stream, 'ring_id', ring_id)
461
        elif op == ReduceOp.MIN:
W
wanghuancoder 已提交
462 463
            return _C_ops.c_allreduce_min_(tensor, 'use_calc_stream',
                                           use_calc_stream, 'ring_id', ring_id)
464
        elif op == ReduceOp.PROD:
W
wanghuancoder 已提交
465 466
            return _C_ops.c_allreduce_prod_(tensor, 'use_calc_stream',
                                            use_calc_stream, 'ring_id', ring_id)
467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483
        else:
            raise ValueError("Unknown parameter: {}.".format(op))

    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'all_reduce')
    if not op in [ReduceOp.SUM, ReduceOp.MAX, ReduceOp.MIN, ReduceOp.PROD]:
        raise ValueError("The op for all_reduce must be one of educeOp.PROD, "
                         "ReduceOp.SUM, ReduceOp.MAX, ReduceOp.MIN.")
    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 已提交
484 485
    if not isinstance(ring_id, int):
        raise ValueError("The type of 'ring_id' for all_reduce should be int.")
486 487 488 489 490
    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [tensor]},
        outputs={'Out': [tensor]},
K
kuizhiqing 已提交
491 492
        attrs={'ring_id': ring_id,
               'use_calc_stream': use_calc_stream})
493 494


K
kuizhiqing 已提交
495
def reduce(tensor, dst, op=ReduceOp.SUM, group=None, use_calc_stream=True):
496 497 498 499 500 501 502 503
    """

    Reduce a tensor to the destination from all others.

    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 已提交
504
        op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.Min|ReduceOp.PROD): Optional. The operation used. Default value is ReduceOp.SUM.
K
kuizhiqing 已提交
505
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
506 507
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
508 509 510 511 512 513 514

    Returns:
        None.

    Examples:
        .. code-block:: python

515 516 517 518 519 520 521 522 523 524 525 526 527 528
            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]]
529
    """
K
kuizhiqing 已提交
530 531 532 533 534 535 536 537
    if group is not None and not group.is_member():
        return

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

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

540 541
    if in_dygraph_mode():
        if op == ReduceOp.SUM:
W
wanghuancoder 已提交
542 543 544
            return _C_ops.c_reduce_sum(tensor, tensor, 'use_calc_stream',
                                       use_calc_stream, 'ring_id', ring_id,
                                       'root_id', gdst)
545
        elif op == ReduceOp.MAX:
W
wanghuancoder 已提交
546 547 548
            return _C_ops.c_reduce_max(tensor, tensor, 'use_calc_stream',
                                       use_calc_stream, 'ring_id', ring_id,
                                       'root_id', gdst)
549
        elif op == ReduceOp.MIN:
W
wanghuancoder 已提交
550 551 552
            return _C_ops.c_reduce_min(tensor, tensor, 'use_calc_stream',
                                       use_calc_stream, 'ring_id', ring_id,
                                       'root_id', gdst)
553
        elif op == ReduceOp.PROD:
W
wanghuancoder 已提交
554 555 556
            return _C_ops.c_reduce_prod(tensor, tensor, 'use_calc_stream',
                                        use_calc_stream, 'ring_id', ring_id,
                                        'root_id', gdst)
557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582
        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 not op in [ReduceOp.SUM, ReduceOp.MAX, ReduceOp.MIN, ReduceOp.PROD]:
        raise ValueError("The op for reduce must be one of educeOp.PROD, "
                         "ReduceOp.SUM, ReduceOp.MAX, ReduceOp.MIN.")

    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())
    helper.append_op(
        type=op_type,
        inputs={'X': [tensor]},
        outputs={'Out': [tensor]},
        attrs={
K
kuizhiqing 已提交
583 584 585
            'ring_id': ring_id,
            'use_calc_stream': use_calc_stream,
            'root_id': gdst,
586 587 588
        })


K
kuizhiqing 已提交
589
def all_gather(tensor_list, tensor, group=None, use_calc_stream=True):
590 591 592 593 594 595 596 597 598
    """

    Gather tensors from all participators and all get the result.

    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 已提交
599
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
600 601
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
602 603 604 605 606 607 608

    Returns:
        None.

    Examples:
        .. code-block:: python

609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627
            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)
628
    """
K
kuizhiqing 已提交
629 630 631 632 633 634
    if group is not None and not group.is_member():
        return

    ring_id = 0 if group is None else group.id
    nranks = _get_global_group().nranks if group is None else group.nranks

635
    if in_dygraph_mode():
636 637
        out = _C_ops.c_allgather(tensor, 'use_calc_stream', use_calc_stream,
                                 'ring_id', ring_id, 'nranks', nranks)
638
    else:
639 640 641
        op_type = 'c_allgather'
        helper = LayerHelper(op_type, **locals())
        out = helper.create_variable_for_type_inference(dtype=tensor.dtype)
642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657
        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')
        helper.append_op(
            type=op_type,
            inputs={'X': [tensor]},
            outputs={'Out': [out]},
            attrs={
K
kuizhiqing 已提交
658 659 660
                'ring_id': ring_id,
                'use_calc_stream': use_calc_stream,
                'nranks': nranks
661 662
            })

K
kuizhiqing 已提交
663
    tensor_list.extend(paddle.split(out, nranks, 0))
664 665


K
kuizhiqing 已提交
666
def scatter(tensor, tensor_list=None, src=0, group=None, use_calc_stream=True):
667 668 669 670 671 672 673
    """

    Scatter a tensor to all participators.

    Args:
        tensor (Tensor): The output Tensor. Its data type
            should be float16, float32, float64, int32 or int64.
674
        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 已提交
675 676
            should be float16, float32, float64, int32 or int64. Default value is None.
        src (int): The source rank id. Default value is 0.
K
kuizhiqing 已提交
677
        group (Group): The group instance return by new_group or None for global default group.
K
kuizhiqing 已提交
678 679
        use_calc_stream (bool): Wether to use calculation stream (True) or communication stream (False).
            Default to True.
680 681 682 683 684 685 686

    Returns:
        None.

    Examples:
        .. code-block:: python

687 688 689 690
            import numpy as np
            import paddle
            from paddle.distributed import init_parallel_env

691 692
            # required: gpu

693 694 695 696 697 698 699 700 701 702 703 704 705 706 707
            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()
708
    """
K
kuizhiqing 已提交
709 710 711 712 713 714 715 716
    if group is not None and not group.is_member():
        return

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

    ring_id = 0 if group is None else group.id
    gsrc = src if group is None else group.get_group_rank(src)
K
kuizhiqing 已提交
717
    assert gsrc >= 0, ("src rank out of group, need global rank")
K
kuizhiqing 已提交
718 719 720 721
    rank = _get_global_group().rank if group is None else group.rank
    nranks = _get_global_group().nranks if group is None else group.nranks

    if rank != gsrc:
722 723 724 725 726
        tensor_list = []
        for _ in range(nranks):
            tensor_list.append(tensor)
    temp = paddle.concat(tensor_list, axis=0)
    if in_dygraph_mode():
W
wanghuancoder 已提交
727 728 729
        return _C_ops.c_scatter(temp, tensor, 'use_calc_stream',
                                use_calc_stream, 'ring_id', ring_id, 'nranks',
                                nranks, 'root', gsrc)
W
wanghuancoder 已提交
730
    op_type = 'c_scatter'
731 732 733 734 735 736 737 738 739
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'scatter')
    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [temp]},
        outputs={'Out': [tensor]},
        attrs={
K
kuizhiqing 已提交
740 741 742
            'ring_id': ring_id,
            'root': gsrc,
            'use_calc_stream': use_calc_stream,
743 744 745 746
            'nranks': nranks,
        })


747
def _c_identity(tensor, group=None):
L
lilong12 已提交
748 749 750 751 752 753 754 755 756 757 758
    """
    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.
    """
759 760 761 762 763
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

    if in_dygraph_mode():
W
wanghuancoder 已提交
764 765
        return _C_ops.c_identity(tensor, 'use_calc_stream', True, 'ring_id',
                                 ring_id, 'use_model_parallel', True)
L
lilong12 已提交
766 767 768
    op_type = 'c_identity'
    helper = LayerHelper(op_type, **locals())
    out = helper.create_variable_for_type_inference(dtype=tensor.dtype)
769

L
lilong12 已提交
770 771 772
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        '_c_identity')
773

L
lilong12 已提交
774 775 776 777 778
    helper.append_op(
        type=op_type,
        inputs={'X': tensor},
        outputs={'Out': out},
        attrs={
779
            'ring_id': ring_id,
L
lilong12 已提交
780 781 782 783 784 785
            'use_calc_stream': True,
            'use_model_parallel': True,
        })
    return out


786
def _c_concat(tensor, group=None):
787 788 789 790 791 792 793 794 795 796 797 798 799 800 801
    """
    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
    ring_id = 0 if group is None else group.id

802 803 804 805
    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

806
    if in_dygraph_mode():
W
wanghuancoder 已提交
807 808 809
        return _C_ops.c_concat(tensor, 'ring_id', ring_id, 'use_calc_stream',
                               True, 'rank', rank, 'nranks', nranks,
                               'use_model_parallel', True)
810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826

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

    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,
827 828
            'nranks': nranks,
            'rank': rank
829 830 831 832
        })
    return out


833
def _c_split(tensor, group=None):
L
lilong12 已提交
834 835 836 837 838 839 840 841 842 843 844 845
    """
    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.
    """
846 847 848 849
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

850 851 852 853
    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

854
    if in_dygraph_mode():
W
wanghuancoder 已提交
855 856 857
        return _C_ops.c_split(tensor, 'use_calc_stream', True, 'ring_id',
                              ring_id, 'rank', rank, 'nranks', nranks,
                              'use_model_parallel', True)
858

L
lilong12 已提交
859 860 861
    op_type = 'c_split'
    helper = LayerHelper(op_type, **locals())
    out = helper.create_variable_for_type_inference(dtype=tensor.dtype)
862

L
lilong12 已提交
863 864 865
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        '_c_split')
866

L
lilong12 已提交
867 868 869 870 871
    helper.append_op(
        type=op_type,
        inputs={'X': tensor},
        outputs={'Out': out},
        attrs={
872
            'ring_id': ring_id,
L
lilong12 已提交
873 874 875 876 877 878 879 880
            'use_calc_stream': True,
            'rank': rank,
            'nranks': nranks,
            'use_model_parallel': True,
        })
    return out


881 882 883 884 885 886 887 888 889 890 891 892 893
def _mp_allreduce(tensor,
                  op=ReduceOp.SUM,
                  group=None,
                  use_calc_stream=True,
                  use_model_parallel=True):
    """[it is same as allreduce above, but it suuports model parallel. And it support inplace startegy]
    """
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

    if in_dygraph_mode():
        if op == ReduceOp.SUM:
W
wanghuancoder 已提交
894
            return _C_ops.c_allreduce_sum_(
895 896 897 898
                tensor, 'use_calc_stream', use_calc_stream, 'ring_id', ring_id,
                "use_model_parallel", use_model_parallel)
        else:
            raise ValueError("Unknown parameter: {}.".format(op))
899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917

    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)

    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,
        })
    return out
918 919


920 921 922 923 924 925 926 927 928 929 930 931 932 933 934
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.
    """
    if in_dygraph_mode():
W
wanghuancoder 已提交
935
        return _C_ops.c_embedding(table, index, "start_index", start_index)
936

937 938 939 940 941 942 943 944 945 946 947 948 949
    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)
    helper.append_op(
        type='c_embedding',
        inputs={'Ids': index,
                'W': table},
        outputs={'Out': tmp},
        attrs={"start_index": start_index})
    return tmp

950

B
Baibaifan 已提交
951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988
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
        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)
        self.name = name

    def forward(self, input):
        out = _linear(
            x=input, weight=self.weight, bias=self.bias, name=self.name)
        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)


989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009
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)

    if in_dygraph_mode():
W
wanghuancoder 已提交
1010
        softmax, loss = _C_ops.c_softmax_with_cross_entropy(
1011 1012 1013 1014 1015 1016
            logits, label, 'ring_id', ring_id, 'rank', rank, 'nranks', nranks)
        if not return_softmax:
            return loss
        else:
            return loss, softmax

W
WangXi 已提交
1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037
    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)
    helper.append_op(
        type='c_softmax_with_cross_entropy',
        inputs={'Logits': logits,
                'Label': label},
        outputs={'Softmax': softmax,
                 'Loss': loss},
        attrs=attrs)

    if return_softmax:
        return loss, softmax

    return loss

1038

B
Baibaifan 已提交
1039 1040 1041 1042 1043 1044
def _linear(x, weight, bias=None, name=None):
    """
    Fuction Linear
    """
    if in_dygraph_mode():
        pre_bias = _varbase_creator(dtype=x.dtype)
W
wanghuancoder 已提交
1045 1046
        _C_ops.matmul(x, weight, pre_bias, 'transpose_X', False, 'transpose_Y',
                      False, "alpha", 1)
B
Baibaifan 已提交
1047 1048 1049 1050 1051
        return dygraph_utils._append_bias_in_dygraph(
            pre_bias, bias, axis=len(x.shape) - 1)
    else:
        helper = LayerHelper('linear', **locals())
        dtype = x.dtype
B
Baibaifan 已提交
1052 1053
        assert len(
            x.shape) < 4, "X latitude is not supported greater than 3 now."
B
Baibaifan 已提交
1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080

        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)
        helper.append_op(
            type='matmul_v2', inputs=inputs, outputs={'Out': tmp}, attrs=attrs)
        if bias is not None:
            res = helper.create_variable_for_type_inference(dtype)
            helper.append_op(
                type='elementwise_add',
                inputs={'X': [tmp],
                        'Y': [bias]},
                outputs={'Out': [res]},
                attrs={'axis': len(x.shape) - 1})
        else:
            res = tmp
        return res


L
lilong12 已提交
1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091
def _parallel_linear(x,
                     num_rows,
                     num_cols,
                     axis,
                     param_attr,
                     bias_attr,
                     gather_out,
                     inner_rank,
                     nranks,
                     split_tensor,
                     name,
1092
                     group=None):
1093 1094
    """
    Parallel Linear
1095 1096 1097 1098 1099

    axis the dimension of the parameter of linear layer. 
    axis = 0: the row dimension
    axid = 1: the col dimension
    
1100
    """
1101 1102 1103 1104
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

L
lilong12 已提交
1105 1106
    if axis == 0:
        if split_tensor:
1107
            x = _c_split(x, group=group)
1108
    else:
L
lilong12 已提交
1109 1110
        x = _c_identity(x, group=group)

B
Baibaifan 已提交
1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124
    if core.is_compiled_with_npu():
        linear = _Linear(
            num_rows,
            num_cols,
            weight_attr=param_attr,
            bias_attr=bias_attr,
            name=name)
    else:
        linear = paddle.nn.Linear(
            num_rows,
            num_cols,
            weight_attr=param_attr,
            bias_attr=bias_attr,
            name=name)
1125 1126

    linear_out = linear(x)
李季 已提交
1127 1128 1129 1130 1131
    startup_block = paddle.static.default_startup_program().current_block()
    main_block = paddle.static.default_main_program().current_block()
    startup_block._find_var_recursive(linear.weight.name).is_distributed = True
    main_block._find_var_recursive(linear.weight.name).is_distributed = True

1132 1133 1134 1135
    # 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:
李季 已提交
1136 1137 1138
        startup_block._find_var_recursive(
            linear.bias.name).is_distributed = True
        main_block._find_var_recursive(linear.bias.name).is_distributed = True
L
lilong12 已提交
1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158

    if not gather_out: return linear_out

    op_type = 'c_allreduce_sum' if axis == 0 else 'c_concat'
    out_shape = list(linear_out.shape)
    out_shape[0] *= 1 if axis == 0 else nranks
    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:
        main_block.append_op(
            type='c_allreduce_sum',
            inputs={'X': linear_out},
            outputs={'Out': out},
            attrs={
1159
                'ring_id': ring_id,
L
lilong12 已提交
1160 1161 1162 1163 1164 1165 1166 1167 1168
                'use_calc_stream': True,
                'use_model_parallel': True
            })
    else:
        main_block.append_op(
            type='c_concat',
            inputs={'X': linear_out},
            outputs={'Out': out},
            attrs={
1169
                'ring_id': ring_id,
L
lilong12 已提交
1170 1171 1172 1173 1174
                'nranks': nranks,
                'use_calc_stream': True,
                'use_model_parallel': True
            })
    return out
1175 1176


L
lilong12 已提交
1177 1178 1179 1180 1181 1182 1183
def _parallel_embedding(x,
                        per_part_embeddings,
                        origin_size,
                        param_attr,
                        inner_rank,
                        num_partitions,
                        name,
1184
                        group=None):
1185 1186 1187
    """
    Parallel Embedding
    """
1188 1189 1190 1191
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207
    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]]

    weight = helper.create_parameter(
        attr=param_attr, shape=size, dtype=dtype, is_bias=False)

    if num_partitions == 1:
        return paddle.nn.functional.embedding(
            x, weight=weight, padding_idx=None, sparse=False, name=name)

1208 1209
    startup_block = paddle.static.default_startup_program().global_block()
    main_block = paddle.static.default_main_program().global_block()
1210 1211 1212 1213 1214 1215 1216 1217 1218 1219
    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)
    out = paddle.distributed.collective._mp_allreduce(
        output_parallel,
        group=group,
        use_calc_stream=True,
        use_model_parallel=True)
L
lilong12 已提交
1220
    return out
1221 1222


B
Baibaifan 已提交
1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281
def _parallel_embedding_npu(x,
                            per_part_embeddings,
                            origin_size,
                            param_attr,
                            inner_rank,
                            num_partitions,
                            name,
                            group=None):
    """
    NPU Parallel Embedding
    """
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

    origin_num_embeddings = origin_size[0]
    embedding = paddle.nn.Embedding(
        per_part_embeddings,
        origin_size[1],
        padding_idx=per_part_embeddings - 1,
        sparse=False,
        weight_attr=param_attr,
        name=name)

    origin_input_shape = x.shape
    if len(origin_input_shape) == 2:
        x = paddle.unsqueeze(x, axis=-1)
    else:
        assert origin_input_shape[-1] == 1, (
            "The last dimension size of x must be 1.")
    x_shard = paddle.shard_index(x, origin_num_embeddings, num_partitions,
                                 inner_rank, per_part_embeddings - 1)
    if len(origin_input_shape) == 2:
        x_shard = paddle.squeeze(x_shard, axis=-1)
    emb_out = embedding(x_shard)
    startup_block = paddle.static.default_startup_program().global_block()
    main_block = paddle.static.default_main_program().global_block()
    startup_block.vars[embedding.weight.name].is_distributed = True
    main_block.vars[embedding.weight.name].is_distributed = True
    out = main_block.create_var(
        shape=emb_out.shape,
        dtype=emb_out.dtype,
        type=emb_out.type,
        lod_level=emb_out.lod_level,
        persistable=False,
        is_data=False,
        need_check_feed=emb_out.desc.need_check_feed())
    main_block.append_op(
        type='c_allreduce_sum',
        inputs={'X': emb_out},
        outputs={'Out': out},
        attrs={
            'ring_id': ring_id,
            'use_calc_stream': True,
            'use_model_parallel': True
        })
    return out


1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302
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 已提交
1303

1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342
        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.

    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.

    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.

    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
1343

1344
            # required: distributed
1345
            import paddle
1346
            import paddle.distributed.fleet as fleet
1347

1348
            paddle.enable_static()
1349
            paddle.set_device('gpu:%d'%paddle.distributed.ParallelEnv().dev_id)
1350
            fleet.init(is_collective=True)
1351
            data = paddle.randint(0, 8, shape=[10,4])
1352
            emb_out = paddle.distributed.split(
1353 1354 1355 1356
                data,
                (8, 8),
                operation="embedding",
                num_partitions=2)
1357

1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374
    """
    assert isinstance(size, (list, tuple)), (
        "The type of size for "
        "paddle.distributed.split must be list or tuple.")
    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))
    if in_dygraph_mode():
L
lilong12 已提交
1375 1376 1377 1378
        raise ValueError(
            "paddle.distributed.split cannot be used in dynamic "
            "graph mode, plese use ParallelEmbedding, ParallelRowLinear, "
            "ParallelColumnLinear instead.")
1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390
    else:
        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.")
1391 1392 1393
        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)
1394

1395
        per_part_size = size[0] // num_partitions
B
Baibaifan 已提交
1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417
        if core.is_compiled_with_npu():
            emb_out = _parallel_embedding_npu(
                x,
                per_part_size,
                size,
                weight_attr,
                inner_rank,
                num_partitions,
                name,
                group=None)
            return emb_out
        else:
            emb_out = _parallel_embedding(
                x,
                per_part_size,
                size,
                weight_attr,
                inner_rank,
                num_partitions,
                name,
                group=None)
            return emb_out
1418
    else:
L
lilong12 已提交
1419
        should_split = False
1420 1421 1422 1423 1424 1425 1426
        if axis == 0:
            assert size[0] % num_partitions == 0, (
                "Number of rows of the weight for linear ({}) must be"
                " divisible by num_partitions ({})".format(size[0],
                                                           num_partitions))
            per_part_size = size[0] // num_partitions
            linear_size = (per_part_size, size[1])
L
lilong12 已提交
1427
            if x.shape[-1] == size[0]: should_split = True
1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448

        elif axis == 1:
            assert size[1] % num_partitions == 0, (
                "Number of column of the weight for linear ({}) must be"
                " divisible by num_partitions ({})".format(size[1],
                                                           num_partitions))
            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))

        linear_out = _parallel_linear(
            x,
            linear_size[0],
            linear_size[1],
            axis,
            weight_attr,
            bias_attr,
            gather_out,
            inner_rank,
L
lilong12 已提交
1449 1450 1451
            num_partitions,
            should_split,
            name=name,
1452
            group=None)
1453
        return linear_out
L
lilong12 已提交
1454 1455


L
lilong12 已提交
1456 1457 1458
def alltoall(in_tensor_list, out_tensor_list, group=None, use_calc_stream=True):
    """
    Scatter tensors in in_tensor_list to all participators and gather the result tensors in out_tensor_list.
1459
    
L
lilong12 已提交
1460 1461 1462 1463 1464 1465 1466
    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.
1467
    
L
lilong12 已提交
1468 1469
    Returns:
        None.
1470
    
L
lilong12 已提交
1471 1472
    Examples:
        .. code-block:: python
1473

L
lilong12 已提交
1474 1475 1476 1477
            # required: distributed
            import numpy as np
            import paddle
            from paddle.distributed import init_parallel_env
1478
            
L
lilong12 已提交
1479 1480 1481 1482 1483 1484 1485 1486 1487 1488
            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)
李季 已提交
1489
            paddle.distributed.alltoall([data1, data2], out_tensor_list)
L
lilong12 已提交
1490 1491 1492 1493 1494 1495 1496 1497
            # 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

    ring_id = 0 if group is None else group.id
    temp = paddle.concat(in_tensor_list, axis=0)
李季 已提交
1498
    nranks = len(in_tensor_list)
L
lilong12 已提交
1499
    if in_dygraph_mode():
李季 已提交
1500 1501
        out = _C_ops.alltoall(temp, 'use_calc_stream', use_calc_stream,
                              'ring_id', ring_id)
L
lilong12 已提交
1502
    else:
W
wanghuancoder 已提交
1503 1504 1505 1506 1507
        op_type = 'alltoall'
        helper = LayerHelper(op_type, **locals())
        out = helper.create_variable_for_type_inference(
            dtype=in_tensor_list[0].dtype)

L
lilong12 已提交
1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526
        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.")
        helper.append_op(
            type=op_type,
            inputs={'X': [temp]},
            outputs={'Out': [out]},
            attrs={
L
lilong12 已提交
1527
                'ring_id': ring_id,
L
lilong12 已提交
1528 1529 1530 1531 1532
                'use_calc_stream': use_calc_stream,
            })
    out_tensor_list.extend(paddle.split(out, nranks, 0))


L
lilong12 已提交
1533 1534 1535 1536 1537 1538 1539 1540
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 已提交
1541 1542
        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.
1543
    
L
lilong12 已提交
1544 1545 1546 1547 1548
    Returns:
        None.

    Examples:
        .. code-block:: python
1549

L
lilong12 已提交
1550
            # required: distributed
L
lilong12 已提交
1551
            import paddle
L
lilong12 已提交
1552
            from paddle.distributed import init_parallel_env
1553

L
lilong12 已提交
1554 1555 1556 1557 1558 1559 1560 1561
            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 已提交
1562 1563 1564 1565 1566 1567
    """
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

    if in_dygraph_mode():
W
wanghuancoder 已提交
1568 1569
        return _C_ops.send_v2(tensor, 'use_calc_stream', use_calc_stream,
                              'ring_id', ring_id, 'peer', dst)
W
wanghuancoder 已提交
1570
    op_type = 'send_v2'
L
lilong12 已提交
1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'send')

    helper = LayerHelper(op_type, **locals())
    helper.append_op(
        type=op_type,
        inputs={'X': [tensor]},
        attrs={
            'ring_id': ring_id,
            'peer': dst,
            'use_calc_stream': use_calc_stream,
        })


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 已提交
1594 1595
        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.
1596
    
L
lilong12 已提交
1597 1598 1599 1600 1601
    Returns:
        None.

    Examples:
        .. code-block:: python
1602

L
lilong12 已提交
1603
            # required: distributed
L
lilong12 已提交
1604
            import paddle
L
lilong12 已提交
1605
            from paddle.distributed import init_parallel_env
1606

L
lilong12 已提交
1607 1608 1609 1610 1611 1612 1613 1614
            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 已提交
1615 1616 1617 1618 1619 1620
    """
    if group is not None and not group.is_member():
        return
    ring_id = 0 if group is None else group.id

    if in_dygraph_mode():
W
wanghuancoder 已提交
1621 1622 1623
        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 已提交
1624
    op_type = 'recv_v2'
L
lilong12 已提交
1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638
    check_variable_and_dtype(
        tensor, 'tensor', ['float16', 'float32', 'float64', 'int32', 'int64'],
        'recv')
    helper = LayerHelper(op_type, **locals())
    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,
        })