chore: import upstream snapshot with attribution

python/paddle/distributed/collective.py

@@ -0,0 +1,454 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import datetime
+import hashlib
+from typing import (
+    TYPE_CHECKING,
+    Literal,
+    TypeAlias,
+)
+
+import paddle
+
+# (TODO: GhostScreaming) It will be removed later.
+from paddle.base import core
+from paddle.framework import in_dynamic_mode
+
+from .communication.group import Group, _add_new_group, is_initialized
+from .fleet.layers.mpu.mp_ops import (  # noqa: F401
+    _c_concat,
+    _c_identity,
+    _c_lookup_table,
+    _c_softmax_with_cross_entropy,
+    _c_softmax_with_multi_label_cross_entropy,
+    _c_split,
+    _Linear,
+    _linear,
+    _mp_allreduce,
+    _parallel_embedding,
+    _parallel_linear,
+    _set_var_distributed,
+    split,
+)
+
+if TYPE_CHECKING:
+    _BackendList: TypeAlias = Literal["gloo", "nccl", "xccl", "bkcl", "flagcx"]
+
+    from paddle.base.libpaddle import NCCLConfig
+
+__all__ = []
+
+_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 = {}
+_global_env_gid = 0
+
+# group map by name : the map of all groups from their names
+# Dict[name, Group]
+_group_map_by_name = {}
+
+# backend map by group : the map of all backend from their groups
+# Dict[group, backend]
+_group_map_backend = {}
+
+# Name of the default group for init_parallel_env
+_default_group_name = "_default_pg"
+
+_valid_backend_list = ['nccl', 'gloo', 'heter', 'xccl', 'bkcl', 'flagcx']
+_default_store = None  # the default tcp store
+_default_backend = None
+_default_timeout = datetime.timedelta(seconds=1800)
+_start_ring_id = 0
+
+
+def _set_default_backend(backend):
+    global _default_backend
+    _default_backend = backend
+
+
+def _set_default_store(store):
+    global _default_store
+    _default_store = store
+
+
+def _get_group_map():
+    global _group_map
+    if _global_env_gid not in _group_map:
+        genv = _get_global_env()
+        _group_map[_global_env_gid] = Group(
+            genv.rank, 0, list(range(genv.world_size))
+        )
+    return _group_map
+
+
+def _get_global_group():
+    return _get_group_map()[_global_env_gid]
+
+
+def _get_group_map_by_name():
+    global _group_map_by_name
+    return _group_map_by_name
+
+
+def _get_default_group():
+    global _group_map_by_name
+    assert is_initialized(), (
+        "Call paddle.distributed.init_parallel_env first "
+        "to initialize the distributed environment."
+    )
+    return _get_group_map_by_name()[_default_group_name]
+
+
+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
+
+
+def _set_group_map_backend(group, backend):
+    global _group_map_backend
+    assert group not in _group_map_backend
+    _group_map_backend[group] = backend
+
+
+def _new_ring_id():
+    # NOTE(liyurui): For compatible reason, auto parallel and eager mode relay on previous syntax.
+    if in_dynamic_mode():
+        global _start_ring_id
+        _start_ring_id += 1
+        return _start_ring_id + max(_get_global_env().nrings, 9)
+    else:
+        return len(_get_group_map()) + max(_get_global_env().nrings, 9)
+
+
+def _new_process_group_impl(
+    backend,
+    store,
+    rank,
+    world_size,
+    group_name,
+    pg_options,
+    group_id=0,
+    nccl_comm_init_option=0,
+    nccl_config=None,
+):
+    pg = None
+    genv = _get_global_env()
+    assert backend in _valid_backend_list, f"Unsupported backend: {backend}."
+    if backend == "gloo":
+        pg = core.ProcessGroupGloo.create(store, rank, world_size, group_id)
+    elif backend == "nccl":
+        pg = core.ProcessGroupNCCL.create(
+            store,
+            rank,
+            world_size,
+            group_id,
+            genv.pg_timeout,
+            nccl_comm_init_option,
+            nccl_config,
+        )
+    elif backend == "xccl":
+        pg = core.ProcessGroupCustom.create(
+            store, genv.device_type, rank, world_size, group_id
+        )
+    elif backend == "bkcl":
+        pg = core.ProcessGroupBKCL.create(store, rank, world_size, group_id)
+    elif backend == "flagcx":
+        pg = core.ProcessGroupFlagcx.create(
+            store,
+            rank,
+            world_size,
+            group_id,
+            genv.pg_timeout,
+            nccl_comm_init_option,
+        )
+    return pg
+
+
+# _custom_gid provides a way for users to
+# set the group id, which is usually useful
+# to be compatible with the static graph mode.
+_custom_gid = None
+
+
+def _set_custom_gid(gid):
+    global _custom_gid
+    _custom_gid = gid
+
+
+def new_group(
+    ranks: list[int] | None = None,
+    backend: Literal['nccl'] | None = None,
+    timeout: datetime.timedelta = _default_timeout,
+    nccl_comm_init_option: int = 0,
+    nccl_config: NCCLConfig | None = None,
+) -> Group:
+    """
+
+    Creates a new distributed communication group.
+
+    Args:
+        ranks (list): The global ranks of group members.
+        backend (str): The backend used to create group, only nccl is supported now.
+        timeout (datetime.timedelta, optional): The waiting timeout for store relevant options, default is 30 minutes.
+
+    Returns:
+        Group: The group instance.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+
+            >>> paddle.distributed.init_parallel_env()
+            >>> tindata = paddle.randn(shape=[2, 3])
+            >>> gp = paddle.distributed.new_group([2, 4, 6])
+            >>> paddle.distributed.all_reduce(tindata, group=gp, sync_op=False)
+
+    """
+    global _custom_gid
+    global _group_map
+    if in_dynamic_mode():
+        global _default_group_name
+        gid = _custom_gid if _custom_gid else _new_ring_id()
+        group_name = _default_group_name + str(gid)
+        if backend != 'heter' and (ranks is None or len(ranks) > 1):
+            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."
+            )
+        size = len(ranks)
+        ranks = sorted(ranks)
+        if size > 1 and global_rank in ranks:
+            rank = 0 if backend == 'heter' else ranks.index(global_rank)
+            pg = _new_process_group_impl(
+                backend,
+                _default_store,
+                rank,
+                size,
+                group_name,
+                pg_options=None,
+                group_id=gid,
+                nccl_comm_init_option=nccl_comm_init_option,
+                nccl_config=nccl_config,
+            )
+        else:
+            rank = -1
+            pg = None
+        group = Group(rank, gid, ranks, pg=pg, name=group_name)
+        _group_map_by_name[group_name] = group
+        _group_map[gid] = group
+        _group_map_backend[group] = backend
+        # TODO: The method below is a new method for group management, will replace the previous
+        # three in the future.
+        _add_new_group(group)
+        return group
+
+    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, ring_id, ranks)
+        _group_map[ring_id] = gp
+    else:
+        ranks = sorted(ranks)
+        group_rank = ranks.index(global_rank)
+        group_size = len(ranks)
+        gp = Group(group_rank, 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
+                )
+            elif core.is_compiled_with_xpu():
+                place = core.XPUPlace(genv.device_id)
+                core.BKCLParallelContext(strategy, place).init_with_ring_id(
+                    ring_id
+                )
+            else:
+                raise AssertionError("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
+    tmp = (
+        paddle.to_tensor([1], dtype="int32")
+        if in_dynamic_mode()
+        else paddle.full([0], 1, dtype="int32")
+    )
+    paddle.distributed.all_reduce(tmp, sync_op=True)
+    paddle.distributed.wait(tmp)
+    return gp
+
+
+def is_available() -> bool:
+    """
+    Check whether the distributed package is available.
+
+    Returns:
+        Returns True if the distributed package is available, otherwise False.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> print(paddle.distributed.is_available())
+
+    """
+    return core.is_compiled_with_dist()
+
+
+def _init_parallel_env(backend: _BackendList) -> None:
+    store = core.create_or_get_global_tcp_store()
+    global_env = _get_global_env()
+    rank = global_env.rank
+    world_size = global_env.world_size
+    dev_id = global_env.device_id
+
+    if backend == "gloo":
+        core.CommContextManager.create_gloo_comm_context(
+            store, "0", rank, world_size
+        )
+    elif backend == "nccl":
+        endpoints_str = ""
+        for endpoint in global_env.trainer_endpoints:
+            endpoints_str += endpoint
+        endpoints_str += "ring_id:{}".format("0")
+        endpoints_str_hash = hashlib.md5(
+            endpoints_str.encode(encoding='UTF-8')
+        ).hexdigest()
+        core.CommContextManager.set_device_id(dev_id)
+        core.CommContextManager.create_nccl_comm_context(
+            store, "0", rank, world_size, endpoints_str_hash
+        )
+    elif backend == "xccl":
+        dev_type = global_env.device_type
+        paddle.device.set_device(f"{dev_type}:{dev_id}")
+        core.CommContextManager.create_xccl_comm_context(
+            store, "0", rank, world_size, dev_type
+        )
+    elif backend == "bkcl":
+        endpoints_str = ""
+        for endpoint in global_env.trainer_endpoints:
+            endpoints_str += endpoint
+        endpoints_str += "ring_id:{}".format("0")
+        endpoints_str_hash = hashlib.md5(
+            endpoints_str.encode(encoding='UTF-8')
+        ).hexdigest()
+        core.CommContextManager.set_device_id(dev_id)
+        core.CommContextManager.create_bkcl_comm_context(
+            store, "0", rank, world_size, endpoints_str_hash
+        )
+
+
+_shutdown_group_map_by_name = {}
+
+
+def _get_shutdown_group_map_by_name():
+    global _shutdown_group_map_by_name
+    return _shutdown_group_map_by_name
+
+
+def _update_shutdown_group_map_by_name(pg_name, group):
+    global _shutdown_group_map_by_name
+    _shutdown_group_map_by_name[pg_name] = group
+
+
+def _delete_shutdown_group_map_by_name(pg_name):
+    global _shutdown_group_map_by_name
+    del _shutdown_group_map_by_name[pg_name]
+
+
+def _clear_shutdown_group_map_by_name():
+    global _shutdown_group_map_by_name
+    _shutdown_group_map_by_name.clear()
+
+
+def shutdown_process_group(group: Group | None = None) -> None:
+    shutdown_groups = _get_shutdown_group_map_by_name()
+
+    if group is None:
+        global _default_group_name
+        for pg_name, pg in _get_group_map_by_name().items():
+            if (
+                pg.process_group is not None
+                and pg_name not in shutdown_groups
+                and pg_name != _default_group_name
+            ):
+                pg.process_group.shutdown()
+                _update_shutdown_group_map_by_name(pg_name, pg)
+    else:
+        if (
+            group.process_group is not None
+            and group.name not in shutdown_groups
+        ):
+            group.process_group.shutdown()
+            _update_shutdown_group_map_by_name(group.name, group)
+
+
+def restart_process_group(group: Group | None = None) -> None:
+    shutdown_groups = _get_shutdown_group_map_by_name()
+
+    if group is None:
+        for pg in shutdown_groups.values():
+            pg.process_group.restart()
+        _clear_shutdown_group_map_by_name()
+    else:
+        if group.process_group is not None and group.name in shutdown_groups:
+            group.process_group.restart()
+            _delete_shutdown_group_map_by_name(group.name)