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