chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,59 @@
|
||||
from ray.util.collective.backend_registry import (
|
||||
register_collective_backend,
|
||||
)
|
||||
from ray.util.collective.collective import (
|
||||
allgather,
|
||||
allgather_multigpu,
|
||||
allreduce,
|
||||
allreduce_multigpu,
|
||||
barrier,
|
||||
broadcast,
|
||||
broadcast_multigpu,
|
||||
create_collective_group,
|
||||
destroy_collective_group,
|
||||
get_collective_group_size,
|
||||
get_group_handle,
|
||||
get_rank,
|
||||
gloo_available,
|
||||
init_collective_group,
|
||||
is_backend_available,
|
||||
is_group_initialized,
|
||||
nccl_available,
|
||||
recv,
|
||||
recv_multigpu,
|
||||
reduce,
|
||||
reduce_multigpu,
|
||||
reducescatter,
|
||||
reducescatter_multigpu,
|
||||
send,
|
||||
send_multigpu,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"nccl_available",
|
||||
"gloo_available",
|
||||
"is_backend_available",
|
||||
"is_group_initialized",
|
||||
"init_collective_group",
|
||||
"destroy_collective_group",
|
||||
"create_collective_group",
|
||||
"get_rank",
|
||||
"get_collective_group_size",
|
||||
"allreduce",
|
||||
"allreduce_multigpu",
|
||||
"barrier",
|
||||
"reduce",
|
||||
"reduce_multigpu",
|
||||
"broadcast",
|
||||
"broadcast_multigpu",
|
||||
"allgather",
|
||||
"allgather_multigpu",
|
||||
"reducescatter",
|
||||
"reducescatter_multigpu",
|
||||
"send",
|
||||
"send_multigpu",
|
||||
"recv",
|
||||
"recv_multigpu",
|
||||
"get_group_handle",
|
||||
"register_collective_backend",
|
||||
]
|
||||
@@ -0,0 +1,122 @@
|
||||
from typing import Dict, Type
|
||||
|
||||
from .collective_group.base_collective_group import BaseGroup
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
|
||||
class BackendRegistry:
|
||||
_instance = None
|
||||
_map: Dict[str, Type[BaseGroup]]
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super(BackendRegistry, cls).__new__(cls)
|
||||
cls._instance._map = {}
|
||||
return cls._instance
|
||||
|
||||
def put(self, name: str, group_cls: Type[BaseGroup]) -> None:
|
||||
if not issubclass(group_cls, BaseGroup):
|
||||
raise TypeError(f"{group_cls} is not a subclass of BaseGroup")
|
||||
if name.upper() in self._map:
|
||||
raise ValueError(f"Backend {name.upper()} already registered")
|
||||
self._map[name.upper()] = group_cls
|
||||
|
||||
def get(self, name: str) -> Type[BaseGroup]:
|
||||
name = name.upper()
|
||||
if name not in self._map:
|
||||
raise ValueError(f"Backend {name} not registered")
|
||||
return self._map[name]
|
||||
|
||||
def is_registered(self, name: str) -> bool:
|
||||
"""Check if a backend is registered (regardless of availability)."""
|
||||
return name.upper() in self._map
|
||||
|
||||
def check(self, name: str) -> bool:
|
||||
"""Check if a backend is both registered and available."""
|
||||
try:
|
||||
cls = self.get(name)
|
||||
return cls.check_backend_availability()
|
||||
except (ValueError, AttributeError):
|
||||
return False
|
||||
|
||||
|
||||
_global_registry = BackendRegistry()
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def register_collective_backend(name: str, group_cls: Type[BaseGroup]):
|
||||
"""Register a custom collective backend with Ray.
|
||||
|
||||
This function registers a custom backend class that can be used for
|
||||
collective operations. The backend must be a subclass of
|
||||
:class:`~ray.util.collective.collective_group.base_collective_group.BaseGroup`
|
||||
and implement all required collective operations.
|
||||
|
||||
Important: The backend must be registered on both the driver and all
|
||||
actors before creating collective groups. This is because each process
|
||||
(driver and each actor) needs to know about your backend class to
|
||||
instantiate it.
|
||||
|
||||
Args:
|
||||
name: The name of the backend (e.g., "MY_BACKEND"). This will be
|
||||
automatically added to the Backend enum as Backend.MY_BACKEND.
|
||||
group_cls: The backend class, which must be a subclass of
|
||||
:class:`~ray.util.collective.collective_group.base_collective_group.BaseGroup`.
|
||||
|
||||
Example:
|
||||
>>> import ray
|
||||
>>> from ray.util.collective import create_collective_group, init_collective_group
|
||||
>>> from ray.util.collective.backend_registry import register_collective_backend
|
||||
>>> from ray.util.collective.collective_group.base_collective_group import BaseGroup
|
||||
>>>
|
||||
>>> class MyCustomBackend(BaseGroup):
|
||||
... def __init__(self, world_size, rank, group_name):
|
||||
... super().__init__(world_size, rank, group_name)
|
||||
... @classmethod
|
||||
... def backend(cls):
|
||||
... return "MY_BACKEND"
|
||||
... @classmethod
|
||||
... def check_backend_availability(cls) -> bool:
|
||||
... return True
|
||||
... def allreduce(self, tensor, allreduce_options=None):
|
||||
... pass
|
||||
... def broadcast(self, tensor, broadcast_options=None):
|
||||
... pass
|
||||
... def barrier(self, barrier_options=None):
|
||||
... pass
|
||||
>>>
|
||||
>>> # Register on the driver
|
||||
>>> register_collective_backend("MY_BACKEND", MyCustomBackend)
|
||||
>>>
|
||||
>>> ray.init()
|
||||
>>>
|
||||
>>> @ray.remote
|
||||
... class Worker:
|
||||
... def __init__(self, rank):
|
||||
... self.rank = rank
|
||||
... def setup(self, world_size):
|
||||
... # IMPORTANT: Register on each worker too
|
||||
... register_collective_backend("MY_BACKEND", MyCustomBackend)
|
||||
... init_collective_group(
|
||||
... world_size=world_size,
|
||||
... rank=self.rank,
|
||||
... backend="MY_BACKEND",
|
||||
... group_name="default",
|
||||
... )
|
||||
>>>
|
||||
>>> actors = [Worker.remote(rank=i) for i in range(2)]
|
||||
>>> create_collective_group(
|
||||
... actors=actors,
|
||||
... world_size=2,
|
||||
... ranks=[0, 1],
|
||||
... backend="MY_BACKEND",
|
||||
... group_name="default",
|
||||
... )
|
||||
>>> ray.get([a.setup.remote(2) for a in actors])
|
||||
"""
|
||||
_global_registry.put(name, group_cls)
|
||||
from . import types
|
||||
|
||||
upper_name = name.upper()
|
||||
if not hasattr(types.Backend, upper_name):
|
||||
setattr(types.Backend, upper_name, upper_name)
|
||||
@@ -0,0 +1,816 @@
|
||||
"""APIs exposed under the namespace ray.util.collective."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import threading
|
||||
from typing import Any, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
import ray
|
||||
from . import types
|
||||
from ray._common.network_utils import find_free_port, is_ipv6
|
||||
from ray.util.collective.backend_registry import (
|
||||
_global_registry,
|
||||
register_collective_backend,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from ray.util.collective.collective_group.nccl_collective_group import NCCLGroup
|
||||
|
||||
register_collective_backend("NCCL", NCCLGroup)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from ray.util.collective.collective_group.torch_gloo_collective_group import (
|
||||
TorchGLOOGroup,
|
||||
)
|
||||
|
||||
register_collective_backend("GLOO", TorchGLOOGroup)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def nccl_available():
|
||||
return is_backend_available("NCCL")
|
||||
|
||||
|
||||
def gloo_available():
|
||||
return is_backend_available("GLOO")
|
||||
|
||||
|
||||
def is_backend_available(backend: str) -> bool:
|
||||
"""Check if a collective backend is available.
|
||||
|
||||
Args:
|
||||
backend: The name of the backend to check (e.g., "NCCL", "GLOO").
|
||||
|
||||
Returns:
|
||||
True if the backend is available, False otherwise.
|
||||
"""
|
||||
return _global_registry.check(backend)
|
||||
|
||||
|
||||
def get_address_and_port() -> Tuple[str, int]:
|
||||
"""Returns the IP address and a free port on this node."""
|
||||
addr = ray.util.get_node_ip_address()
|
||||
port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET)
|
||||
return addr, port
|
||||
|
||||
|
||||
class GroupManager(object):
|
||||
"""Use this class to manage the collective groups we created so far.
|
||||
|
||||
Each process will have an instance of `GroupManager`. Each process
|
||||
could belong to multiple collective groups. The membership information
|
||||
and other metadata are stored in the global `_group_mgr` object.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._name_group_map = {}
|
||||
self._registry = _global_registry
|
||||
|
||||
def create_collective_group(
|
||||
self, backend, world_size, rank, group_name, gloo_timeout=None
|
||||
):
|
||||
"""The entry to create new collective groups in the manager.
|
||||
|
||||
Put the registration and the group information into the manager
|
||||
metadata as well.
|
||||
"""
|
||||
backend = backend.upper()
|
||||
backend_cls = self._registry.get(backend)
|
||||
|
||||
if not backend_cls.check_backend_availability():
|
||||
raise RuntimeError(
|
||||
f"Backend {backend} is not available. Please check the installation."
|
||||
)
|
||||
|
||||
if backend == "GLOO":
|
||||
logger.debug(
|
||||
"Creating torch.distributed GLOO group: '{}'...".format(group_name)
|
||||
)
|
||||
g = backend_cls(world_size, rank, group_name, gloo_timeout)
|
||||
else:
|
||||
g = backend_cls(world_size, rank, group_name)
|
||||
|
||||
self._name_group_map[group_name] = g
|
||||
return self._name_group_map[group_name]
|
||||
|
||||
def is_group_exist(self, group_name):
|
||||
return group_name in self._name_group_map
|
||||
|
||||
def get_group_by_name(self, group_name):
|
||||
"""Get the collective group handle by its name."""
|
||||
if not self.is_group_exist(group_name):
|
||||
logger.warning("The group '{}' is not initialized.".format(group_name))
|
||||
return None
|
||||
return self._name_group_map[group_name]
|
||||
|
||||
def destroy_collective_group(self, group_name):
|
||||
"""Group destructor."""
|
||||
if not self.is_group_exist(group_name):
|
||||
logger.warning("The group '{}' does not exist.".format(group_name))
|
||||
return
|
||||
|
||||
# release the collective group resource
|
||||
g = self._name_group_map[group_name]
|
||||
# clean up the dicts
|
||||
del self._name_group_map[group_name]
|
||||
# Release the communicator resources
|
||||
g.destroy_group()
|
||||
|
||||
# Release the detached actors spawned by `create_collective_group()`
|
||||
name = "info_" + group_name
|
||||
try:
|
||||
store = ray.get_actor(name)
|
||||
ray.kill(store)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
|
||||
_group_mgr = GroupManager()
|
||||
# This lock is used to make external calls to the _group_mgr thread-safe.
|
||||
_group_mgr_lock = threading.Lock()
|
||||
|
||||
|
||||
def is_group_initialized(group_name):
|
||||
"""Check if the group is initialized in this process by the group name."""
|
||||
global _group_mgr
|
||||
global _group_mgr_lock
|
||||
with _group_mgr_lock:
|
||||
return _group_mgr.is_group_exist(group_name)
|
||||
|
||||
|
||||
def init_collective_group(
|
||||
world_size: int,
|
||||
rank: int,
|
||||
backend: types.Backend = types.Backend.NCCL,
|
||||
group_name: str = "default",
|
||||
gloo_timeout: int = 30000,
|
||||
):
|
||||
"""Initialize a collective group inside an actor process.
|
||||
|
||||
Args:
|
||||
world_size: the total number of processes in the group.
|
||||
rank: the rank of the current process.
|
||||
backend: the CCL backend to use, NCCL or GLOO.
|
||||
group_name: the name of the collective group.
|
||||
gloo_timeout: timeout in milliseconds for GLOO operations.
|
||||
"""
|
||||
_check_inside_actor()
|
||||
|
||||
global _group_mgr
|
||||
global _group_mgr_lock
|
||||
|
||||
# TODO(Hao): implement a group auto-counter.
|
||||
if not group_name:
|
||||
raise ValueError("group_name '{}' needs to be a string.".format(group_name))
|
||||
|
||||
with _group_mgr_lock:
|
||||
if _group_mgr.is_group_exist(group_name):
|
||||
raise RuntimeError("Trying to initialize a group a second time.")
|
||||
|
||||
assert world_size > 0
|
||||
assert rank >= 0
|
||||
assert rank < world_size
|
||||
_group_mgr.create_collective_group(
|
||||
backend, world_size, rank, group_name, gloo_timeout
|
||||
)
|
||||
|
||||
|
||||
def create_collective_group(
|
||||
actors: List[Any],
|
||||
world_size: int,
|
||||
ranks: List[int],
|
||||
backend: types.Backend = types.Backend.NCCL,
|
||||
group_name: str = "default",
|
||||
gloo_timeout: int = 30000,
|
||||
):
|
||||
"""Declare a list of actors as a collective group.
|
||||
|
||||
Note: This function should be called in a driver process.
|
||||
|
||||
Args:
|
||||
actors: a list of actors to be set in a collective group.
|
||||
world_size: the total number of processes in the group.
|
||||
ranks: the rank of each actor.
|
||||
backend: the CCL backend to use, NCCL or GLOO.
|
||||
group_name: the name of the collective group.
|
||||
gloo_timeout: timeout in milliseconds for GLOO operations.
|
||||
"""
|
||||
|
||||
name = "info_" + group_name
|
||||
try:
|
||||
ray.get_actor(name)
|
||||
raise RuntimeError("Trying to initialize a group twice.")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if len(ranks) != len(actors):
|
||||
raise RuntimeError(
|
||||
"Each actor should correspond to one rank. Got '{}' "
|
||||
"ranks but '{}' actors".format(len(ranks), len(actors))
|
||||
)
|
||||
|
||||
if set(ranks) != set(range(len(ranks))):
|
||||
raise RuntimeError(
|
||||
"Ranks must be a permutation from 0 to '{}'. Got '{}'.".format(
|
||||
len(ranks), "".join([str(r) for r in ranks])
|
||||
)
|
||||
)
|
||||
|
||||
if world_size <= 0:
|
||||
raise RuntimeError(
|
||||
"World size must be greater than zero. Got '{}'.".format(world_size)
|
||||
)
|
||||
if not all(ranks) >= 0:
|
||||
raise RuntimeError("Ranks must be non-negative.")
|
||||
if not all(ranks) < world_size:
|
||||
raise RuntimeError("Ranks cannot be greater than world_size.")
|
||||
|
||||
# Check if backend is registered and available
|
||||
backend_upper = backend.upper()
|
||||
if not _global_registry.is_registered(backend_upper):
|
||||
raise RuntimeError(
|
||||
f"Backend {backend_upper} is not registered. "
|
||||
f"Please register it using register_collective_backend('{backend_upper}', YourBackendClass)."
|
||||
)
|
||||
if not _global_registry.check(backend_upper):
|
||||
raise RuntimeError(
|
||||
f"Backend {backend_upper} is registered but not available. "
|
||||
f"Please check the installation requirements for this backend."
|
||||
)
|
||||
|
||||
# avoid a circular dependency
|
||||
from ray.util.collective.util import Info
|
||||
|
||||
# store the information into a NamedActor that can be accessed later.
|
||||
name = "info_" + group_name
|
||||
actors_id = [a._ray_actor_id for a in actors]
|
||||
# TODO (Dacheng): how do we recycle this name actor?
|
||||
info = Info.options(name=name, lifetime="detached").remote()
|
||||
ray.get([info.set_info.remote(actors_id, world_size, ranks, backend, gloo_timeout)])
|
||||
|
||||
|
||||
# TODO (we need a declarative destroy() API here.)
|
||||
def destroy_collective_group(group_name: str = "default") -> None:
|
||||
"""Destroy a collective group given its group name."""
|
||||
_check_inside_actor()
|
||||
global _group_mgr
|
||||
global _group_mgr_lock
|
||||
with _group_mgr_lock:
|
||||
_group_mgr.destroy_collective_group(group_name)
|
||||
|
||||
|
||||
def get_rank(group_name: str = "default") -> int:
|
||||
"""Return the rank of this process in the given group.
|
||||
|
||||
Args:
|
||||
group_name: the name of the group to query
|
||||
|
||||
Returns:
|
||||
the rank of this process in the named group,
|
||||
-1 if the group does not exist or the process does
|
||||
not belong to the group.
|
||||
"""
|
||||
_check_inside_actor()
|
||||
|
||||
global _group_mgr
|
||||
global _group_mgr_lock
|
||||
with _group_mgr_lock:
|
||||
if not _group_mgr.is_group_exist(group_name):
|
||||
return -1
|
||||
g = _group_mgr.get_group_by_name(group_name)
|
||||
return g.rank
|
||||
|
||||
|
||||
def get_collective_group_size(group_name: str = "default") -> int:
|
||||
"""Return the size of the collective group with the given name.
|
||||
|
||||
Args:
|
||||
group_name: the name of the group to query
|
||||
|
||||
Returns:
|
||||
The world size of the collective group, -1 if the group does
|
||||
not exist or the process does not belong to the group.
|
||||
"""
|
||||
_check_inside_actor()
|
||||
global _group_mgr
|
||||
global _group_mgr_lock
|
||||
with _group_mgr_lock:
|
||||
if not _group_mgr.is_group_exist(group_name):
|
||||
return -1
|
||||
g = _group_mgr.get_group_by_name(group_name)
|
||||
return g.world_size
|
||||
|
||||
|
||||
def allreduce(
|
||||
tensor: Any,
|
||||
group_name: str = "default",
|
||||
op: types.ReduceOp = types.ReduceOp.SUM,
|
||||
):
|
||||
"""Collective allreduce the tensor across the group.
|
||||
|
||||
Args:
|
||||
tensor: the tensor to be all-reduced on this process.
|
||||
group_name: the collective group name to perform allreduce.
|
||||
op: The reduce operation.
|
||||
"""
|
||||
_check_single_tensor_input(tensor)
|
||||
g = get_group_handle(group_name)
|
||||
opts = types.AllReduceOptions
|
||||
opts.reduceOp = op
|
||||
g.allreduce([tensor], opts)
|
||||
|
||||
|
||||
def allreduce_multigpu(
|
||||
tensor_list: list,
|
||||
group_name: str = "default",
|
||||
op: types.ReduceOp = types.ReduceOp.SUM,
|
||||
):
|
||||
"""Collective allreduce a list of tensors across the group.
|
||||
|
||||
Args:
|
||||
tensor_list: list of tensors to be allreduced, each on a GPU.
|
||||
group_name: the collective group name to perform allreduce.
|
||||
op: The reduce operation.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
||||
_check_tensor_list_input(tensor_list)
|
||||
g = get_group_handle(group_name)
|
||||
opts = types.AllReduceOptions
|
||||
opts.reduceOp = op
|
||||
g.allreduce(tensor_list, opts)
|
||||
|
||||
|
||||
def barrier(group_name: str = "default"):
|
||||
"""Barrier all processes in the collective group.
|
||||
|
||||
Args:
|
||||
group_name: the name of the group to barrier.
|
||||
"""
|
||||
g = get_group_handle(group_name)
|
||||
g.barrier()
|
||||
|
||||
|
||||
def reduce(
|
||||
tensor: Any,
|
||||
dst_rank: int = 0,
|
||||
group_name: str = "default",
|
||||
op: types.ReduceOp = types.ReduceOp.SUM,
|
||||
):
|
||||
"""Reduce the tensor across the group to the destination rank.
|
||||
|
||||
Args:
|
||||
tensor: the tensor to be reduced on this process.
|
||||
dst_rank: the rank of the destination process.
|
||||
group_name: the collective group name to perform reduce.
|
||||
op: The reduce operation.
|
||||
"""
|
||||
_check_single_tensor_input(tensor)
|
||||
g = get_group_handle(group_name)
|
||||
|
||||
# check dst rank
|
||||
_check_rank_valid(g, dst_rank)
|
||||
opts = types.ReduceOptions()
|
||||
opts.reduceOp = op
|
||||
opts.root_rank = dst_rank
|
||||
opts.root_tensor = 0
|
||||
g.reduce([tensor], opts)
|
||||
|
||||
|
||||
def reduce_multigpu(
|
||||
tensor_list: list,
|
||||
dst_rank: int = 0,
|
||||
dst_tensor: int = 0,
|
||||
group_name: str = "default",
|
||||
op: types.ReduceOp = types.ReduceOp.SUM,
|
||||
):
|
||||
"""Reduce the tensor across the group to the destination rank
|
||||
and destination tensor.
|
||||
|
||||
Args:
|
||||
tensor_list: the list of tensors to be reduced on this process;
|
||||
each tensor located on a GPU.
|
||||
dst_rank: the rank of the destination process.
|
||||
dst_tensor: the index of GPU at the destination.
|
||||
group_name: the collective group name to perform reduce.
|
||||
op: The reduce operation.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
||||
_check_tensor_list_input(tensor_list)
|
||||
g = get_group_handle(group_name)
|
||||
|
||||
# check dst rank
|
||||
_check_rank_valid(g, dst_rank)
|
||||
_check_root_tensor_valid(len(tensor_list), dst_tensor)
|
||||
opts = types.ReduceOptions()
|
||||
opts.reduceOp = op
|
||||
opts.root_rank = dst_rank
|
||||
opts.root_tensor = dst_tensor
|
||||
g.reduce(tensor_list, opts)
|
||||
|
||||
|
||||
def broadcast(tensor: Any, src_rank: int = 0, group_name: str = "default"):
|
||||
"""Broadcast the tensor from a source process to all others.
|
||||
|
||||
Args:
|
||||
tensor: the tensor to be broadcasted (src) or received (destination).
|
||||
src_rank: the rank of the source process.
|
||||
group_name: the collective group name to perform broadcast.
|
||||
"""
|
||||
_check_single_tensor_input(tensor)
|
||||
g = get_group_handle(group_name)
|
||||
|
||||
# check src rank
|
||||
_check_rank_valid(g, src_rank)
|
||||
opts = types.BroadcastOptions()
|
||||
opts.root_rank = src_rank
|
||||
opts.root_tensor = 0
|
||||
g.broadcast([tensor], opts)
|
||||
|
||||
|
||||
def broadcast_multigpu(
|
||||
tensor_list: list,
|
||||
src_rank: int = 0,
|
||||
src_tensor: int = 0,
|
||||
group_name: str = "default",
|
||||
):
|
||||
"""Broadcast the tensor from a source GPU to all other GPUs.
|
||||
|
||||
Args:
|
||||
tensor_list: the tensors to broadcast (src) or receive (dst).
|
||||
src_rank: the rank of the source process.
|
||||
src_tensor: the index of the source GPU on the source process.
|
||||
group_name: the collective group name to perform broadcast.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
||||
_check_tensor_list_input(tensor_list)
|
||||
g = get_group_handle(group_name)
|
||||
|
||||
# check src rank
|
||||
_check_rank_valid(g, src_rank)
|
||||
_check_root_tensor_valid(len(tensor_list), src_tensor)
|
||||
opts = types.BroadcastOptions()
|
||||
opts.root_rank = src_rank
|
||||
opts.root_tensor = src_tensor
|
||||
g.broadcast(tensor_list, opts)
|
||||
|
||||
|
||||
def allgather(tensor_list: list, tensor: Any, group_name: str = "default"):
|
||||
"""Allgather tensors from each process of the group into a list.
|
||||
|
||||
Args:
|
||||
tensor_list: the results, stored as a list of tensors.
|
||||
tensor: the tensor (to be gathered) in the current process
|
||||
group_name: the name of the collective group.
|
||||
"""
|
||||
_check_single_tensor_input(tensor)
|
||||
_check_tensor_list_input(tensor_list)
|
||||
g = get_group_handle(group_name)
|
||||
if len(tensor_list) != g.world_size:
|
||||
# Typically CLL lib requires len(tensor_list) >= world_size;
|
||||
# Here we make it more strict: len(tensor_list) == world_size.
|
||||
raise RuntimeError(
|
||||
"The length of the tensor list operands to allgather "
|
||||
"must be equal to world_size."
|
||||
)
|
||||
opts = types.AllGatherOptions()
|
||||
g.allgather([tensor_list], [tensor], opts)
|
||||
|
||||
|
||||
def allgather_multigpu(
|
||||
output_tensor_lists: list, input_tensor_list: list, group_name: str = "default"
|
||||
):
|
||||
"""Allgather tensors from each gpus of the group into lists.
|
||||
|
||||
Args:
|
||||
output_tensor_lists: gathered results, with shape
|
||||
must be num_gpus * world_size * shape(tensor).
|
||||
input_tensor_list: a list of tensors, with shape
|
||||
num_gpus * shape(tensor).
|
||||
group_name: the name of the collective group.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
||||
_check_tensor_lists_input(output_tensor_lists)
|
||||
_check_tensor_list_input(input_tensor_list)
|
||||
g = get_group_handle(group_name)
|
||||
opts = types.AllGatherOptions()
|
||||
g.allgather(output_tensor_lists, input_tensor_list, opts)
|
||||
|
||||
|
||||
def reducescatter(
|
||||
tensor: Any,
|
||||
tensor_list: list,
|
||||
group_name: str = "default",
|
||||
op: types.ReduceOp = types.ReduceOp.SUM,
|
||||
):
|
||||
"""Reducescatter a list of tensors across the group.
|
||||
|
||||
Reduce the list of the tensors across each process in the group, then
|
||||
scatter the reduced list of tensors -- one tensor for each process.
|
||||
|
||||
Args:
|
||||
tensor: the resulted tensor on this process.
|
||||
tensor_list: The list of tensors to be reduced and scattered.
|
||||
group_name: the name of the collective group.
|
||||
op: The reduce operation.
|
||||
"""
|
||||
_check_single_tensor_input(tensor)
|
||||
_check_tensor_list_input(tensor_list)
|
||||
g = get_group_handle(group_name)
|
||||
opts = types.ReduceScatterOptions()
|
||||
opts.reduceOp = op
|
||||
if len(tensor_list) != g.world_size:
|
||||
raise RuntimeError(
|
||||
"The length of the tensor list operands to reducescatter "
|
||||
"must not be equal to world_size."
|
||||
)
|
||||
g.reducescatter([tensor], [tensor_list], opts)
|
||||
|
||||
|
||||
def reducescatter_multigpu(
|
||||
output_tensor_list: list,
|
||||
input_tensor_lists: list,
|
||||
group_name: str = "default",
|
||||
op: types.ReduceOp = types.ReduceOp.SUM,
|
||||
):
|
||||
"""Reducescatter a list of tensors across all GPUs.
|
||||
|
||||
Args:
|
||||
output_tensor_list: the resulted list of tensors, with
|
||||
shape: num_gpus * shape(tensor).
|
||||
input_tensor_lists: the original tensors, with shape:
|
||||
num_gpus * world_size * shape(tensor).
|
||||
group_name: the name of the collective group.
|
||||
op: The reduce operation.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
||||
_check_tensor_lists_input(input_tensor_lists)
|
||||
_check_tensor_list_input(output_tensor_list)
|
||||
g = get_group_handle(group_name)
|
||||
opts = types.ReduceScatterOptions()
|
||||
opts.reduceOp = op
|
||||
g.reducescatter(output_tensor_list, input_tensor_lists, opts)
|
||||
|
||||
|
||||
def send(tensor: Any, dst_rank: int, group_name: str = "default"):
|
||||
"""Send a tensor to a remote process synchronously.
|
||||
|
||||
Args:
|
||||
tensor: the tensor to send.
|
||||
dst_rank: the rank of the destination process.
|
||||
group_name: the name of the collective group.
|
||||
"""
|
||||
_check_single_tensor_input(tensor)
|
||||
g = get_group_handle(group_name)
|
||||
_check_rank_valid(g, dst_rank)
|
||||
if dst_rank == g.rank:
|
||||
raise RuntimeError("The destination rank '{}' is self.".format(dst_rank))
|
||||
opts = types.SendOptions()
|
||||
opts.dst_rank = dst_rank
|
||||
g.send([tensor], opts)
|
||||
|
||||
|
||||
def send_multigpu(
|
||||
tensor: Any,
|
||||
dst_rank: int,
|
||||
dst_gpu_index: int,
|
||||
group_name: str = "default",
|
||||
n_elements: int = 0,
|
||||
):
|
||||
"""Send a tensor to a remote GPU synchronously.
|
||||
|
||||
The function assumes each process owns >1 GPUs, and the sender
|
||||
process and receiver process has equal number of GPUs.
|
||||
|
||||
Args:
|
||||
tensor: the tensor to send, located on a GPU.
|
||||
dst_rank: the rank of the destination process.
|
||||
dst_gpu_index: the destination gpu index.
|
||||
group_name: the name of the collective group.
|
||||
n_elements: if specified, send the next n elements
|
||||
from the starting address of tensor.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("send_multigpu call requires NCCL.")
|
||||
_check_single_tensor_input(tensor)
|
||||
g = get_group_handle(group_name)
|
||||
_check_rank_valid(g, dst_rank)
|
||||
if dst_rank == g.rank:
|
||||
raise RuntimeError(
|
||||
"The dst_rank '{}' is self. Considering "
|
||||
"doing GPU to GPU memcpy instead?".format(dst_rank)
|
||||
)
|
||||
if n_elements < 0:
|
||||
raise RuntimeError("The n_elements '{}' should >= 0.".format(n_elements))
|
||||
opts = types.SendOptions()
|
||||
opts.dst_rank = dst_rank
|
||||
opts.dst_gpu_index = dst_gpu_index
|
||||
opts.n_elements = n_elements
|
||||
g.send([tensor], opts)
|
||||
|
||||
|
||||
def recv(tensor: Any, src_rank: int, group_name: str = "default"):
|
||||
"""Receive a tensor from a remote process synchronously.
|
||||
|
||||
Args:
|
||||
tensor: the received tensor.
|
||||
src_rank: the rank of the source process.
|
||||
group_name: the name of the collective group.
|
||||
"""
|
||||
_check_single_tensor_input(tensor)
|
||||
g = get_group_handle(group_name)
|
||||
_check_rank_valid(g, src_rank)
|
||||
if src_rank == g.rank:
|
||||
raise RuntimeError("The destination rank '{}' is self.".format(src_rank))
|
||||
opts = types.RecvOptions()
|
||||
opts.src_rank = src_rank
|
||||
g.recv([tensor], opts)
|
||||
|
||||
|
||||
def recv_multigpu(
|
||||
tensor: Any,
|
||||
src_rank: int,
|
||||
src_gpu_index: int,
|
||||
group_name: str = "default",
|
||||
n_elements: int = 0,
|
||||
):
|
||||
"""Receive a tensor from a remote GPU synchronously.
|
||||
|
||||
The function asssume each process owns >1 GPUs, and the sender
|
||||
process and receiver process has equal nubmer of GPUs.
|
||||
|
||||
Args:
|
||||
tensor: The received tensor, located on a GPU.
|
||||
src_rank: The rank of the source process.
|
||||
src_gpu_index: The index of the source GPU on the src process.
|
||||
group_name: The name of the collective group.
|
||||
n_elements: if specified, receive the next n elements
|
||||
from the starting address of tensor.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("recv_multigpu call requires NCCL.")
|
||||
_check_single_tensor_input(tensor)
|
||||
g = get_group_handle(group_name)
|
||||
_check_rank_valid(g, src_rank)
|
||||
if src_rank == g.rank:
|
||||
raise RuntimeError(
|
||||
"The dst_rank '{}' is self. Considering "
|
||||
"doing GPU to GPU memcpy instead?".format(src_rank)
|
||||
)
|
||||
if n_elements < 0:
|
||||
raise RuntimeError("The n_elements '{}' should be >= 0.".format(n_elements))
|
||||
opts = types.RecvOptions()
|
||||
opts.src_rank = src_rank
|
||||
opts.src_gpu_index = src_gpu_index
|
||||
opts.n_elements = n_elements
|
||||
g.recv([tensor], opts)
|
||||
|
||||
|
||||
def synchronize(gpu_id: int):
|
||||
"""Synchronize the current process to a give device.
|
||||
|
||||
Args:
|
||||
gpu_id: the GPU device id to synchronize.
|
||||
"""
|
||||
if not types.cupy_available():
|
||||
raise RuntimeError("synchronize call requires CUDA and NCCL.")
|
||||
import cupy as cp
|
||||
|
||||
cp.cuda.Device(gpu_id).synchronize()
|
||||
|
||||
|
||||
def get_group_handle(group_name: str = "default"):
|
||||
"""Check if the group is initialized and return the group handle.
|
||||
|
||||
Args:
|
||||
group_name: the name of the collective group.
|
||||
|
||||
Returns:
|
||||
The collective group handle.
|
||||
"""
|
||||
_check_inside_actor()
|
||||
global _group_mgr
|
||||
global _group_mgr_lock
|
||||
with _group_mgr_lock:
|
||||
if not _group_mgr.is_group_exist(group_name):
|
||||
# try loading from remote info store
|
||||
try:
|
||||
# if the information is stored in an Info object,
|
||||
# get and create the group.
|
||||
name = "info_" + group_name
|
||||
mgr = ray.get_actor(name=name)
|
||||
ids, world_size, rank, backend, gloo_timeout = ray.get(
|
||||
mgr.get_info.remote()
|
||||
)
|
||||
worker = ray._private.worker.global_worker
|
||||
id_ = worker.core_worker.get_actor_id()
|
||||
r = rank[ids.index(id_)]
|
||||
_group_mgr.create_collective_group(
|
||||
backend, world_size, r, group_name, gloo_timeout
|
||||
)
|
||||
except ValueError as exc:
|
||||
# check if this group is initialized using options()
|
||||
if (
|
||||
"collective_group_name" in os.environ
|
||||
and os.environ["collective_group_name"] == group_name
|
||||
):
|
||||
rank = int(os.environ["collective_rank"])
|
||||
world_size = int(os.environ["collective_world_size"])
|
||||
backend = os.environ["collective_backend"]
|
||||
gloo_timeout = int(os.getenv("collective_gloo_timeout", 30000))
|
||||
_group_mgr.create_collective_group(
|
||||
backend, world_size, rank, group_name, gloo_timeout
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"The collective group '{}' is not "
|
||||
"initialized in the process.".format(group_name)
|
||||
) from exc
|
||||
g = _group_mgr.get_group_by_name(group_name)
|
||||
return g
|
||||
|
||||
|
||||
def _check_single_tensor_input(tensor):
|
||||
"""Check if the tensor is with a supported type."""
|
||||
if isinstance(tensor, np.ndarray):
|
||||
return
|
||||
if types.cupy_available():
|
||||
if isinstance(tensor, types.cp.ndarray):
|
||||
return
|
||||
if types.torch_available():
|
||||
if isinstance(tensor, types.th.Tensor):
|
||||
return
|
||||
raise RuntimeError(
|
||||
"Unrecognized tensor type '{}'. Supported types are: "
|
||||
"np.ndarray, torch.Tensor, cupy.ndarray.".format(type(tensor))
|
||||
)
|
||||
|
||||
|
||||
def _check_inside_actor():
|
||||
"""Check if currently it is inside a Ray actor/task."""
|
||||
worker = ray._private.worker.global_worker
|
||||
if worker.mode == ray.WORKER_MODE:
|
||||
return
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"The collective APIs shall be only used inside a Ray actor or task."
|
||||
)
|
||||
|
||||
|
||||
def _check_rank_valid(g, rank: int):
|
||||
"""Check the rank: 0 <= rank < world_size."""
|
||||
if rank < 0:
|
||||
raise ValueError("rank '{}' is negative.".format(rank))
|
||||
if rank >= g.world_size:
|
||||
raise ValueError(
|
||||
"rank '{}' must be less than world size '{}'".format(rank, g.world_size)
|
||||
)
|
||||
|
||||
|
||||
def _check_tensor_list_input(tensor_list):
|
||||
"""Check if the input is a list of supported tensor types."""
|
||||
if not isinstance(tensor_list, list):
|
||||
raise RuntimeError(
|
||||
"The input must be a list of tensors. "
|
||||
"Got '{}'.".format(type(tensor_list))
|
||||
)
|
||||
if not tensor_list:
|
||||
raise RuntimeError("Got an empty list of tensors.")
|
||||
for t in tensor_list:
|
||||
_check_single_tensor_input(t)
|
||||
|
||||
|
||||
def _check_tensor_lists_input(tensor_lists):
|
||||
"""Check if the input is a list of lists of supported tensor types."""
|
||||
if not isinstance(tensor_lists, list):
|
||||
raise RuntimeError(
|
||||
"The input must be a list of lists of tensors. "
|
||||
"Got '{}'.".format(type(tensor_lists))
|
||||
)
|
||||
if not tensor_lists:
|
||||
raise RuntimeError(f"Did not receive tensors. Got: {tensor_lists}")
|
||||
for t in tensor_lists:
|
||||
_check_tensor_list_input(t)
|
||||
|
||||
|
||||
def _check_root_tensor_valid(length, root_tensor):
|
||||
"""Check the root_tensor device is 0 <= root_tensor < length"""
|
||||
if root_tensor < 0:
|
||||
raise ValueError("root_tensor '{}' is negative.".format(root_tensor))
|
||||
if root_tensor >= length:
|
||||
raise ValueError(
|
||||
"root_tensor '{}' is greater than the number of GPUs: "
|
||||
"'{}'".format(root_tensor, length)
|
||||
)
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Abstract class for collective groups."""
|
||||
from abc import ABCMeta, abstractmethod
|
||||
|
||||
from ray.util.collective.types import (
|
||||
AllGatherOptions,
|
||||
AllReduceOptions,
|
||||
BarrierOptions,
|
||||
BroadcastOptions,
|
||||
RecvOptions,
|
||||
ReduceOptions,
|
||||
ReduceScatterOptions,
|
||||
SendOptions,
|
||||
)
|
||||
|
||||
|
||||
class BaseGroup(metaclass=ABCMeta):
|
||||
def __init__(self, world_size: int, rank: int, group_name: str):
|
||||
"""Init the process group with basic information.
|
||||
|
||||
Args:
|
||||
world_size: The total number of processes in the group.
|
||||
rank: The rank of the current process.
|
||||
group_name: The group name.
|
||||
"""
|
||||
self._world_size = world_size
|
||||
self._rank = rank
|
||||
self._group_name = group_name
|
||||
|
||||
@property
|
||||
def rank(self):
|
||||
"""Return the rank of the current process."""
|
||||
return self._rank
|
||||
|
||||
@property
|
||||
def world_size(self):
|
||||
"""Return the number of processes in this group."""
|
||||
return self._world_size
|
||||
|
||||
@property
|
||||
def group_name(self):
|
||||
"""Return the group name of this group."""
|
||||
return self._group_name
|
||||
|
||||
def destroy_group(self):
|
||||
"""GC the communicators."""
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def backend(cls):
|
||||
"""The backend of this collective group."""
|
||||
raise NotImplementedError()
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def check_backend_availability(cls) -> bool:
|
||||
"""Check if the backend is available."""
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def allreduce(self, tensor, allreduce_options=AllReduceOptions()):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def barrier(self, barrier_options=BarrierOptions()):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def reduce(self, tensor, reduce_options=ReduceOptions()):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def allgather(self, tensor_list, tensor, allgather_options=AllGatherOptions()):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def broadcast(self, tensor, broadcast_options=BroadcastOptions()):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def reducescatter(
|
||||
self, tensor, tensor_list, reducescatter_options=ReduceScatterOptions()
|
||||
):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def send(self, tensor, send_options: SendOptions):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def recv(self, tensor, recv_options: RecvOptions):
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,95 @@
|
||||
import logging
|
||||
import threading
|
||||
|
||||
import cupy
|
||||
|
||||
from ray.util.collective.collective_group import nccl_util
|
||||
from ray.util.collective.const import ENV
|
||||
|
||||
NCCL_STREAM_POOL_SIZE = 32
|
||||
MAX_GPU_PER_ACTOR = 16
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StreamPool:
|
||||
"""The class that represents a stream pool associated with a GPU.
|
||||
|
||||
When multistream is enabled, we will allocate a pool of streams for each
|
||||
GPU, and get available stream from this pool when a collective kernel is
|
||||
initialized. This enables overlapping computation/communication kernels
|
||||
using multiple CUDA streams, given that the streams a appropriately
|
||||
synchronized. The class is thread-safe.
|
||||
|
||||
|
||||
Args:
|
||||
device_idx: the absolute index of the device for this pool.
|
||||
"""
|
||||
|
||||
def __init__(self, device_idx: int):
|
||||
self.device_idx = device_idx
|
||||
|
||||
self._initialized = False
|
||||
self._initialized_lock = threading.Lock()
|
||||
|
||||
self._pool = [None] * NCCL_STREAM_POOL_SIZE
|
||||
self._counter = 0
|
||||
self._pool_lock = threading.Lock()
|
||||
|
||||
def get_stream(self):
|
||||
"""Get an available stream from the pool.
|
||||
|
||||
The function locks the stream pool and releases the lock before
|
||||
returning.
|
||||
|
||||
Returns:
|
||||
stream (cupy.cuda.Stream): the returned stream from pool.
|
||||
"""
|
||||
|
||||
# check the flag
|
||||
self._initialized_lock.acquire()
|
||||
if not self._initialized:
|
||||
self._init_once()
|
||||
self._initialized_lock.release()
|
||||
|
||||
# Get the stream from the pool.
|
||||
self._pool_lock.acquire()
|
||||
stream = self._pool[self._counter]
|
||||
self._counter = (self._counter + 1) % NCCL_STREAM_POOL_SIZE
|
||||
self._pool_lock.release()
|
||||
return stream
|
||||
|
||||
def _init_once(self):
|
||||
"""Initialize the stream pool only for once."""
|
||||
with nccl_util.Device(self.device_idx):
|
||||
for i in range(NCCL_STREAM_POOL_SIZE):
|
||||
# this is the only place where self._pool will be written.
|
||||
if ENV.NCCL_USE_MULTISTREAM.val:
|
||||
logger.debug("NCCL multistream enabled.")
|
||||
self._pool[i] = cupy.cuda.Stream(null=False, non_blocking=False)
|
||||
else:
|
||||
logger.debug("NCCL multistream disabled.")
|
||||
self._pool[i] = cupy.cuda.Stream.null
|
||||
self._init_flag = True
|
||||
|
||||
|
||||
# This is a map from GPU index to its stream pool.
|
||||
# It is supposed to be READ-ONLY out of this file
|
||||
_device_stream_pool_map = dict()
|
||||
|
||||
|
||||
def _init_stream_pool():
|
||||
global _device_stream_pool_map
|
||||
for i in range(MAX_GPU_PER_ACTOR):
|
||||
_device_stream_pool_map[i] = StreamPool(i)
|
||||
|
||||
|
||||
def get_stream_pool(device_idx):
|
||||
"""Get the CUDA stream pool of a GPU device."""
|
||||
# In case there will be multiple threads writing to the pool.
|
||||
lock = threading.Lock()
|
||||
lock.acquire()
|
||||
if not _device_stream_pool_map:
|
||||
_init_stream_pool()
|
||||
lock.release()
|
||||
return _device_stream_pool_map[device_idx]
|
||||
@@ -0,0 +1,850 @@
|
||||
import datetime
|
||||
import logging
|
||||
import time
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.util.collective.collective_group import nccl_util
|
||||
from ray.util.collective.collective_group.base_collective_group import BaseGroup
|
||||
from ray.util.collective.collective_group.cuda_stream import get_stream_pool
|
||||
from ray.util.collective.const import ENV, get_store_name
|
||||
from ray.util.collective.types import (
|
||||
AllGatherOptions,
|
||||
AllReduceOptions,
|
||||
Backend,
|
||||
BarrierOptions,
|
||||
BroadcastOptions,
|
||||
RecvOptions,
|
||||
ReduceOptions,
|
||||
ReduceScatterOptions,
|
||||
SendOptions,
|
||||
torch_available,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
import cupy
|
||||
import torch
|
||||
|
||||
_NCCL_AVAILABLE = True
|
||||
_LOG_NCCL_WARNING = False
|
||||
except ImportError:
|
||||
_NCCL_AVAILABLE = False
|
||||
_LOG_NCCL_WARNING = True
|
||||
|
||||
|
||||
class Rendezvous:
|
||||
"""A rendezvous class for different actor/task processes to meet.
|
||||
|
||||
To initialize an NCCL collective communication group, different
|
||||
actors/tasks spawned in Ray in a collective group needs to meet
|
||||
each other to synchronize the NCCLUniqueID. This class guarantees
|
||||
they meet via the NCCLUniqueIDStore, initialized on the rank=0
|
||||
process.
|
||||
|
||||
Args:
|
||||
store_key: the unique store key, usually as a concatanation
|
||||
of group_name and communicator key. See `get_nccl_communicator`
|
||||
for more details.
|
||||
"""
|
||||
|
||||
def __init__(self, store_key: str):
|
||||
if not store_key:
|
||||
raise ValueError(
|
||||
"Invalid store_key. The store_key is a concatenation of "
|
||||
"'group_name' and the 'communicator_key'. See the "
|
||||
"docstring of `get_nccl_communicator` for details."
|
||||
)
|
||||
self._store_key = store_key
|
||||
self._store_name = None
|
||||
self._store = None
|
||||
|
||||
def meet(self, timeout_s: int = 180):
|
||||
"""Meet at the named actor store.
|
||||
|
||||
Args:
|
||||
timeout_s: timeout in seconds.
|
||||
"""
|
||||
if timeout_s <= 0:
|
||||
raise ValueError(
|
||||
"The 'timeout' argument must be positive. "
|
||||
"Got '{}'.".format(timeout_s)
|
||||
)
|
||||
self._store_name = get_store_name(self._store_key)
|
||||
timeout_delta = datetime.timedelta(seconds=timeout_s)
|
||||
elapsed = datetime.timedelta(seconds=0)
|
||||
start_time = datetime.datetime.now()
|
||||
while elapsed < timeout_delta:
|
||||
try:
|
||||
logger.debug(
|
||||
"Trying to meet at the store '{}'".format(self._store_name)
|
||||
)
|
||||
self._store = ray.get_actor(self._store_name)
|
||||
except ValueError:
|
||||
logger.debug(
|
||||
"Failed to meet at the store '{}'."
|
||||
"Trying again...".format(self._store_name)
|
||||
)
|
||||
time.sleep(1)
|
||||
elapsed = datetime.datetime.now() - start_time
|
||||
continue
|
||||
logger.debug("Successful rendezvous!")
|
||||
break
|
||||
if not self._store:
|
||||
raise RuntimeError(
|
||||
"Unable to meet other processes "
|
||||
"at the rendezvous store. If you are using "
|
||||
"P2P communication, please check if tensors "
|
||||
"are put in the correct GPU. "
|
||||
)
|
||||
|
||||
@property
|
||||
def store(self):
|
||||
return self._store
|
||||
|
||||
def get_nccl_id(self, timeout_s: int = 180):
|
||||
"""Get the NCCLUniqueID from the store through Ray.
|
||||
|
||||
Args:
|
||||
timeout_s: timeout in seconds.
|
||||
|
||||
Returns:
|
||||
uid: the NCCLUniqueID if successful.
|
||||
"""
|
||||
if not self._store:
|
||||
raise ValueError("Rendezvous store is not setup.")
|
||||
try:
|
||||
uid = ray.get(self._store.wait_and_get_id.remote(), timeout=timeout_s)
|
||||
except ray.exceptions.GetTimeoutError:
|
||||
raise RuntimeError(
|
||||
f"Unable to get the NCCLUniqueID from the store within {timeout_s} seconds."
|
||||
) from None
|
||||
return uid
|
||||
|
||||
|
||||
class NCCLGroup(BaseGroup):
|
||||
def __init__(self, world_size: int, rank: int, group_name: str):
|
||||
"""Init an NCCL collective group."""
|
||||
super(NCCLGroup, self).__init__(world_size, rank, group_name)
|
||||
|
||||
# communicator and stream cache.
|
||||
# TODO (Hao): we need a lock here...
|
||||
self._dev_comm_map = {}
|
||||
self._dev_streams_map = {}
|
||||
|
||||
# record the used GPU IDs.
|
||||
self._used_gpu_indices = set()
|
||||
|
||||
# TODO(Fu): might need an event map
|
||||
self._dev_event_map = {}
|
||||
|
||||
if nccl_util.get_nccl_build_version() < 2000:
|
||||
raise RuntimeError("NCCL in Ray requires NCCL >= 2.0.")
|
||||
if nccl_util.get_nccl_runtime_version() < 2704:
|
||||
logger.warning("NCCL send/recv calls requires NCCL>=2.7.4")
|
||||
|
||||
def destroy_group(self):
|
||||
"""Destroy the group and release NCCL communicators."""
|
||||
if len(self._dev_comm_map.keys()) > 0:
|
||||
# TODO(Hao): check this barrier call
|
||||
# self.barrier()
|
||||
|
||||
# Destroy the communicators and streams.
|
||||
for comm_key, comms in self._dev_comm_map.items():
|
||||
for c in comms:
|
||||
c.destroy()
|
||||
self._dev_comm_map[comm_key] = None
|
||||
|
||||
if self.rank == 0:
|
||||
for comm_key in self._dev_comm_map:
|
||||
assert not self._dev_comm_map[comm_key]
|
||||
group_key = self._generate_group_key(comm_key)
|
||||
self._destroy_store(group_key)
|
||||
self._barrier_tensor = None
|
||||
self._dev_comm_map = None
|
||||
self._dev_streams_map = None
|
||||
super(NCCLGroup, self).destroy_group()
|
||||
|
||||
@classmethod
|
||||
def backend(cls):
|
||||
return Backend.NCCL
|
||||
|
||||
@classmethod
|
||||
def check_backend_availability(cls) -> bool:
|
||||
global _LOG_NCCL_WARNING, _NCCL_AVAILABLE
|
||||
if _LOG_NCCL_WARNING:
|
||||
logger.warning(
|
||||
"NCCL is not available. Please install Cupy "
|
||||
"following the guide at: "
|
||||
"https://docs.cupy.dev/en/stable/install.html."
|
||||
)
|
||||
_LOG_NCCL_WARNING = False
|
||||
return _NCCL_AVAILABLE
|
||||
|
||||
def allreduce(
|
||||
self,
|
||||
tensors: list,
|
||||
allreduce_options: AllReduceOptions = AllReduceOptions(),
|
||||
):
|
||||
"""AllReduce tensors across the collective group following options.
|
||||
|
||||
Args:
|
||||
tensors: the list of tensors to be reduced. Each tensor must
|
||||
reside on one GPU of the current process.
|
||||
allreduce_options: allreduce options.
|
||||
"""
|
||||
|
||||
def collective_fn(input_tensor, output_tensor, comm, stream):
|
||||
comm.allReduce(
|
||||
nccl_util.get_tensor_ptr(input_tensor),
|
||||
nccl_util.get_tensor_ptr(output_tensor),
|
||||
nccl_util.get_tensor_n_elements(input_tensor),
|
||||
nccl_util.get_nccl_tensor_dtype(input_tensor),
|
||||
nccl_util.get_nccl_reduce_op(allreduce_options.reduceOp),
|
||||
stream.ptr,
|
||||
)
|
||||
|
||||
self._collective(tensors, tensors, collective_fn)
|
||||
|
||||
def barrier(self, barrier_options: BarrierOptions = BarrierOptions()):
|
||||
"""Blocks until all processes reach this barrier.
|
||||
|
||||
Args:
|
||||
barrier_options: barrier options.
|
||||
"""
|
||||
# Get the device list.
|
||||
if self._used_gpu_indices:
|
||||
devices = list(self._used_gpu_indices)
|
||||
else:
|
||||
devices = list(range(nccl_util.get_num_gpus()))
|
||||
barrier_tensors = [None] * len(devices)
|
||||
for i, d in enumerate(devices):
|
||||
with nccl_util.Device(d):
|
||||
barrier_tensors[i] = cupy.array([1])
|
||||
self.allreduce(barrier_tensors)
|
||||
|
||||
def reduce(self, tensors: list, reduce_options: ReduceOptions = ReduceOptions()):
|
||||
"""Reduce tensors to a destination gpu following options.
|
||||
|
||||
Args:
|
||||
tensors: the list of tensors to be reduced, each tensor
|
||||
must reside on one gpu of the current process.
|
||||
reduce_options: reduce options.
|
||||
"""
|
||||
root_rank = len(tensors) * reduce_options.root_rank + reduce_options.root_tensor
|
||||
|
||||
def collective_fn(input_tensor, output_tensor, comm, stream):
|
||||
comm.reduce(
|
||||
nccl_util.get_tensor_ptr(input_tensor),
|
||||
nccl_util.get_tensor_ptr(output_tensor),
|
||||
nccl_util.get_tensor_n_elements(input_tensor),
|
||||
nccl_util.get_nccl_tensor_dtype(input_tensor),
|
||||
nccl_util.get_nccl_reduce_op(reduce_options.reduceOp),
|
||||
root_rank,
|
||||
stream.ptr,
|
||||
)
|
||||
|
||||
self._collective(tensors, tensors, collective_fn)
|
||||
|
||||
def broadcast(
|
||||
self,
|
||||
tensors: list,
|
||||
broadcast_options: BroadcastOptions = BroadcastOptions(),
|
||||
):
|
||||
"""Broadcast tensors to all other gpus following options.
|
||||
|
||||
Args:
|
||||
tensors: tensors to be broadcast or received.
|
||||
broadcast_options: broadcast options.
|
||||
"""
|
||||
root_rank = (
|
||||
len(tensors) * broadcast_options.root_rank + broadcast_options.root_tensor
|
||||
)
|
||||
|
||||
def collective_fn(input_tensor, output_tensor, comm, stream):
|
||||
comm.broadcast(
|
||||
nccl_util.get_tensor_ptr(input_tensor),
|
||||
nccl_util.get_tensor_ptr(output_tensor),
|
||||
nccl_util.get_tensor_n_elements(input_tensor),
|
||||
nccl_util.get_nccl_tensor_dtype(input_tensor),
|
||||
root_rank,
|
||||
stream.ptr,
|
||||
)
|
||||
|
||||
self._collective(tensors, tensors, collective_fn)
|
||||
|
||||
def allgather(
|
||||
self,
|
||||
tensor_lists: list,
|
||||
tensors: list,
|
||||
allgather_options: AllGatherOptions = AllGatherOptions(),
|
||||
):
|
||||
"""Allgather tensors across gpus into a list of tensors.
|
||||
|
||||
Args:
|
||||
tensor_lists: allgathered tensors.
|
||||
tensors: the list of tensors to allgather across the group.
|
||||
Each tensor must lolcate on a GPU of the process.
|
||||
allgather_options: allgather options.
|
||||
"""
|
||||
|
||||
def collective_fn(input_tensor, output_tensor, comm, stream):
|
||||
comm.allGather(
|
||||
nccl_util.get_tensor_ptr(input_tensor),
|
||||
nccl_util.get_tensor_ptr(output_tensor),
|
||||
nccl_util.get_tensor_n_elements(input_tensor),
|
||||
nccl_util.get_nccl_tensor_dtype(input_tensor),
|
||||
stream.ptr,
|
||||
)
|
||||
|
||||
_check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
|
||||
output_flattened = [
|
||||
_flatten_for_scatter_gather(tensor_list, copy=False)
|
||||
for tensor_list in tensor_lists
|
||||
]
|
||||
|
||||
def postprocess_fn(stream):
|
||||
# TODO(Hao): designate a copy stream.
|
||||
for i, tensor_list in enumerate(tensor_lists):
|
||||
for j, tensor in enumerate(tensor_list):
|
||||
nccl_util.copy_tensor(tensor, output_flattened[i][j])
|
||||
|
||||
self._collective(
|
||||
tensors, output_flattened, collective_fn, postprocess_fn=postprocess_fn
|
||||
)
|
||||
|
||||
def reducescatter(
|
||||
self,
|
||||
tensors: list,
|
||||
tensor_lists: list,
|
||||
reducescatter_options: ReduceScatterOptions = ReduceScatterOptions(),
|
||||
):
|
||||
"""Reduce then scatter a list of tensors across the group.
|
||||
|
||||
Args:
|
||||
tensors: the output tensors (could be unspecified), each
|
||||
located on a GPU of the current process.
|
||||
tensor_lists: the list of tensors to be reduced then
|
||||
scattered.
|
||||
reducescatter_options: reduce-scatter options.
|
||||
"""
|
||||
|
||||
def collective_fn(input_tensor, output_tensor, comm, stream):
|
||||
comm.reduceScatter(
|
||||
nccl_util.get_tensor_ptr(input_tensor),
|
||||
nccl_util.get_tensor_ptr(output_tensor),
|
||||
nccl_util.get_tensor_n_elements(output_tensor),
|
||||
nccl_util.get_nccl_tensor_dtype(output_tensor),
|
||||
nccl_util.get_nccl_reduce_op(reducescatter_options.reduceOp),
|
||||
stream.ptr,
|
||||
)
|
||||
|
||||
_check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
|
||||
input_flattened = [
|
||||
_flatten_for_scatter_gather(tensor_list, copy=False)
|
||||
for tensor_list in tensor_lists
|
||||
]
|
||||
|
||||
def preprocess_fn(stream):
|
||||
for i, tensor_list in enumerate(tensor_lists):
|
||||
for j, tensor in enumerate(tensor_list):
|
||||
nccl_util.copy_tensor(input_flattened[i][j], tensor)
|
||||
|
||||
self._collective(
|
||||
input_flattened, tensors, collective_fn, preprocess_fn=preprocess_fn
|
||||
)
|
||||
|
||||
def send(self, tensors: list, send_options: SendOptions = SendOptions()):
|
||||
"""Send a tensor to a destination gpu in the group.
|
||||
|
||||
Args:
|
||||
tensors: the tensor to send.
|
||||
send_options: send options.
|
||||
"""
|
||||
|
||||
def p2p_fn(tensor, comm, stream, peer):
|
||||
comm.send(
|
||||
nccl_util.get_tensor_ptr(tensor),
|
||||
send_options.n_elements
|
||||
if send_options.n_elements > 0
|
||||
else nccl_util.get_tensor_n_elements(tensor),
|
||||
nccl_util.get_nccl_tensor_dtype(tensor),
|
||||
peer,
|
||||
stream.ptr,
|
||||
)
|
||||
|
||||
self._point2point(
|
||||
tensors, p2p_fn, send_options.dst_rank, send_options.dst_gpu_index
|
||||
)
|
||||
|
||||
def recv(self, tensors: list, recv_options: RecvOptions = RecvOptions()):
|
||||
"""Receive a tensor from a source gpu in the group.
|
||||
|
||||
Args:
|
||||
tensors: the received tensor.
|
||||
recv_options: Receive options.
|
||||
"""
|
||||
|
||||
def p2p_fn(tensor, comm, stream, peer):
|
||||
comm.recv(
|
||||
nccl_util.get_tensor_ptr(tensor),
|
||||
recv_options.n_elements
|
||||
if recv_options.n_elements > 0
|
||||
else nccl_util.get_tensor_n_elements(tensor),
|
||||
nccl_util.get_nccl_tensor_dtype(tensor),
|
||||
peer,
|
||||
stream.ptr,
|
||||
)
|
||||
|
||||
self._point2point(
|
||||
tensors, p2p_fn, recv_options.src_rank, recv_options.src_gpu_index
|
||||
)
|
||||
|
||||
def _get_nccl_collective_communicator(self, comm_key: str, device_list: list):
|
||||
"""Create or retrieve an NCCL communicator from cache.
|
||||
|
||||
If the communicator is found in cache, return the communicator. If not,
|
||||
a communicator and a stream will be created and put in cache.
|
||||
TODO(Hao): this function is not thread-safe now.
|
||||
|
||||
Args:
|
||||
comm_key: the key to query the communicator cache.
|
||||
device_list: a list of GPU devices of the current process
|
||||
that participates into the collective.
|
||||
|
||||
Returns:
|
||||
communicator: the NCCL communicator corresponded to the devices.
|
||||
"""
|
||||
if not comm_key:
|
||||
raise RuntimeError("Got empty communicator key.")
|
||||
for d in device_list:
|
||||
self._used_gpu_indices.add(d)
|
||||
|
||||
# TODO(Hao): lock the _dev_comm_map here.
|
||||
if comm_key in self._dev_comm_map:
|
||||
return self._dev_comm_map[comm_key]
|
||||
|
||||
group_key = self._generate_group_key(comm_key)
|
||||
if self.rank == 0:
|
||||
nccl_uid = self._generate_nccl_uid(group_key)
|
||||
else:
|
||||
rendezvous = Rendezvous(group_key)
|
||||
rendezvous.meet()
|
||||
nccl_uid = rendezvous.get_nccl_id()
|
||||
|
||||
# Now create the communicators
|
||||
actual_world_size = len(device_list) * self.world_size
|
||||
comms = [None] * len(device_list)
|
||||
streams = [None] * len(device_list)
|
||||
events = [None] * len(device_list)
|
||||
nccl_util.groupStart()
|
||||
for i, device in enumerate(device_list):
|
||||
actual_rank = self.rank * len(device_list) + i
|
||||
with nccl_util.Device(device):
|
||||
comms[i] = nccl_util.create_nccl_communicator(
|
||||
actual_world_size, nccl_uid, actual_rank
|
||||
)
|
||||
# request a stream from the pool
|
||||
# note the device_idx is absolute index.
|
||||
streams[i] = get_stream_pool(device).get_stream()
|
||||
# TODO(Fu): double check the parameters
|
||||
events[i] = cupy.cuda.Event()
|
||||
nccl_util.groupEnd()
|
||||
# TODO(Fu): lock
|
||||
self._dev_comm_map[comm_key] = comms
|
||||
self._dev_streams_map[comm_key] = streams
|
||||
self._dev_event_map[comm_key] = events
|
||||
return comms
|
||||
|
||||
@staticmethod
|
||||
def _sync_streams(device_list, events, streams):
|
||||
"""Let NCCL streams wait for current streams for every device."""
|
||||
# TODO(Fu): recordStream besides calling this function?
|
||||
if ENV.NCCL_USE_MULTISTREAM.val:
|
||||
for i, device in enumerate(device_list):
|
||||
with nccl_util.Device(device):
|
||||
events[i].record(cupy.cuda.get_current_stream())
|
||||
streams[i].wait_event(events[i])
|
||||
|
||||
def _get_nccl_p2p_communicator(
|
||||
self,
|
||||
comm_key: str,
|
||||
my_gpu_idx: int,
|
||||
peer_rank: int,
|
||||
peer_gpu_idx: int,
|
||||
):
|
||||
"""Create or retrieve an NCCL communicator for p2p tasks.
|
||||
|
||||
Note(Hao): this function is not thread-safe now.
|
||||
|
||||
Args:
|
||||
comm_key: communicator key.
|
||||
my_gpu_idx: the gpu index on the current process.
|
||||
peer_rank: the rank of the destination process.
|
||||
peer_gpu_idx: the gpu index on the peer process.
|
||||
Returns:
|
||||
communicator
|
||||
"""
|
||||
if not comm_key:
|
||||
raise RuntimeError("Got empty communicator key.")
|
||||
|
||||
# TODO(Hao): lock the _dev_comm_map here.
|
||||
if comm_key in self._dev_comm_map:
|
||||
return self._dev_comm_map[comm_key]
|
||||
|
||||
# Note (Hao): This is a bit complex so I decide to take a note here.
|
||||
# Here we need to consider three cases:
|
||||
# Case 1: src_rank != dst_rank, hence the send and recv happen on
|
||||
# different process (actors/tasks); each process makes independent
|
||||
# collective calls and manages corresponding communicators.
|
||||
# Case 2: src_rank == dst_rank, src_gpu_idx == dst_gpu_idx; for
|
||||
# this case, we simply throw a RuntimeError;
|
||||
# Case 3: src_rank == dst_rank, src_gpu_idx != dst_gpu_idx, which
|
||||
# means the send and recv will be called on the same process. We
|
||||
# DO NOT support this case for now. We need to properly scope:
|
||||
# (1) communicators creation, and
|
||||
# (2) send/recv calls
|
||||
# using groupStart(（ and groupEnd() calls to avoid deadlocks.
|
||||
if self.rank < peer_rank:
|
||||
my_p2p_rank = 0
|
||||
elif self.rank > peer_rank:
|
||||
my_p2p_rank = 1
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Send and recv happens on the same process! "
|
||||
"ray.util.collective does not support this case as of now. "
|
||||
"Alternatively, consider doing GPU to GPU memcpy?"
|
||||
)
|
||||
|
||||
group_key = self._generate_group_key(comm_key)
|
||||
if my_p2p_rank == 0:
|
||||
nccl_uid = self._generate_nccl_uid(group_key)
|
||||
else:
|
||||
rendezvous = Rendezvous(group_key)
|
||||
rendezvous.meet()
|
||||
nccl_uid = rendezvous.get_nccl_id()
|
||||
|
||||
# create the p2p communicators
|
||||
with nccl_util.Device(my_gpu_idx):
|
||||
comm = nccl_util.create_nccl_communicator(2, nccl_uid, my_p2p_rank)
|
||||
stream = get_stream_pool(my_gpu_idx).get_stream()
|
||||
event = cupy.cuda.Event()
|
||||
|
||||
# TODO(Fu): lock and might need to add event
|
||||
self._dev_comm_map[comm_key] = [comm]
|
||||
self._dev_streams_map[comm_key] = [stream]
|
||||
self._dev_event_map[comm_key] = [event]
|
||||
return [comm]
|
||||
|
||||
def _generate_group_key(self, comm_key):
|
||||
"""Generate a unique key used to initialize the KV store.
|
||||
|
||||
The group key is a concatenation of the communicator key and
|
||||
the group name, following: [comm_key]@[group_name].
|
||||
"""
|
||||
return comm_key + "@" + self.group_name
|
||||
|
||||
@staticmethod
|
||||
def _destroy_store(group_key: str):
|
||||
"""Destroy the KV store (Ray named actor).
|
||||
|
||||
Args:
|
||||
group_key: the unique key to retrieve the KV store.
|
||||
"""
|
||||
store_name = get_store_name(group_key)
|
||||
store = ray.get_actor(store_name)
|
||||
# ray.get([store.__ray_terminate__.remote()])
|
||||
ray.kill(store)
|
||||
|
||||
def _generate_nccl_uid(self, key: str):
|
||||
"""Generate an NCCL unique ID for initializing communicators.
|
||||
|
||||
The method will also create a KV store using Ray named actor and store
|
||||
the NCCLUniqueID in the store. The store needs to be garbage collected
|
||||
when destroying the collective group.
|
||||
|
||||
Args:
|
||||
key: the key of the .
|
||||
|
||||
Returns:
|
||||
NCCLUniqueID (str): NCCL unique ID.
|
||||
"""
|
||||
group_uid = nccl_util.get_nccl_unique_id()
|
||||
store_name = get_store_name(key)
|
||||
# Avoid a potential circular dependency in ray/actor.py
|
||||
from ray.util.collective.util import NCCLUniqueIDStore
|
||||
|
||||
store = NCCLUniqueIDStore.options(name=store_name, lifetime="detached").remote(
|
||||
store_name
|
||||
)
|
||||
ray.get([store.set_id.remote(group_uid)])
|
||||
return group_uid
|
||||
|
||||
def _collective(
|
||||
self,
|
||||
input_tensors: list,
|
||||
output_tensors: list,
|
||||
collective_fn: Callable,
|
||||
preprocess_fn: Optional[Callable] = None,
|
||||
postprocess_fn: Optional[Callable] = None,
|
||||
):
|
||||
"""A method to encapsulate all collective calls.
|
||||
|
||||
Args:
|
||||
input_tensors: the list of the input tensors.
|
||||
output_tensors: the list of the output tensors.
|
||||
collective_fn: the collective function call.
|
||||
preprocess_fn: preprocess procedures before collective calls.
|
||||
postprocess_fn: postprocess procedures after collective calls.
|
||||
"""
|
||||
_check_gpu_tensors(input_tensors)
|
||||
_check_gpu_tensors(output_tensors)
|
||||
|
||||
devices = nccl_util.get_tensor_device_list(input_tensors)
|
||||
key = _get_comm_key_from_devices(devices)
|
||||
comms = self._get_nccl_collective_communicator(key, devices)
|
||||
streams = self._dev_streams_map[key]
|
||||
events = self._dev_event_map[key]
|
||||
|
||||
# TODO(Hao): sync streams and events
|
||||
self._sync_streams(devices, events, streams)
|
||||
|
||||
# Make the collective call
|
||||
if preprocess_fn:
|
||||
preprocess_fn(streams)
|
||||
|
||||
nccl_util.groupStart()
|
||||
# TODO(Fu): how to recordStreams as there are no library functions
|
||||
# We also need to make sure input tensors are not freed before their
|
||||
# usages on ncclStreams finish. This can be achieved by calling
|
||||
# c10::cuda::CUDACachingAllocator::recordStream, which remembers the
|
||||
# usage stream (ncclStream), creates an event on the usage stream
|
||||
# when GC attempts to free the input tensor, and delays GC until that
|
||||
# event is done.
|
||||
for i, tensor in enumerate(input_tensors):
|
||||
collective_fn(tensor, output_tensors[i], comms[i], streams[i])
|
||||
nccl_util.groupEnd()
|
||||
if postprocess_fn:
|
||||
postprocess_fn(streams)
|
||||
|
||||
def _point2point(
|
||||
self,
|
||||
tensors: list,
|
||||
p2p_fn: Callable,
|
||||
peer_rank: int,
|
||||
peer_gpu_idx: int,
|
||||
):
|
||||
"""A method to encapsulate all peer-to-peer calls (i.e., send/recv).
|
||||
|
||||
Args:
|
||||
tensors: the tensor to send or receive.
|
||||
p2p_fn: the p2p function call.
|
||||
peer_rank: the rank of the peer process.
|
||||
peer_gpu_idx: the index of the gpu on the peer process.
|
||||
"""
|
||||
# check send/recv availability.
|
||||
if nccl_util.get_nccl_runtime_version() < 2704:
|
||||
raise RuntimeError(
|
||||
"P2p send/recv requires NCCL >= 2.7.4. "
|
||||
"Got '{}'.".format(nccl_util.get_nccl_runtime_version())
|
||||
)
|
||||
_check_gpu_tensors(tensors)
|
||||
|
||||
# we currently only support single device to single device send/recv.
|
||||
assert len(tensors) == 1
|
||||
my_gpu_idx = nccl_util.get_tensor_device(tensors[0])
|
||||
comm_key = _get_comm_key_send_recv(
|
||||
self.rank, my_gpu_idx, peer_rank, peer_gpu_idx
|
||||
)
|
||||
comms = self._get_nccl_p2p_communicator(
|
||||
comm_key, my_gpu_idx, peer_rank, peer_gpu_idx
|
||||
)
|
||||
streams = self._dev_streams_map[comm_key]
|
||||
events = self._dev_event_map[comm_key]
|
||||
|
||||
# TODO(Hao): sync streams and events
|
||||
self._sync_streams([my_gpu_idx], events, streams)
|
||||
|
||||
# We have made sure that self.rank != peer_rank during API check.
|
||||
peer_p2p_rank = 0 if self.rank > peer_rank else 1
|
||||
for i, tensor in enumerate(tensors):
|
||||
p2p_fn(tensor, comms[i], streams[i], peer_p2p_rank)
|
||||
# Record the stream to avoid tensor being freed before the send/recv is completed.
|
||||
torch_stream = torch.cuda.ExternalStream(streams[i].ptr)
|
||||
tensor.record_stream(torch_stream)
|
||||
|
||||
|
||||
def _flatten_for_scatter_gather(tensor_list: list, copy: bool = False):
|
||||
"""Flatten the tensor for gather/scatter operations.
|
||||
|
||||
Args:
|
||||
tensor_list: the list of tensors to be scattered/gathered.
|
||||
copy: whether the copy the tensors in tensor_list into the buffer.
|
||||
|
||||
Returns:
|
||||
The flattened tensor buffer.
|
||||
"""
|
||||
if not tensor_list:
|
||||
raise RuntimeError("Received an empty list.")
|
||||
t = tensor_list[0]
|
||||
buffer_shape = [len(tensor_list)] + nccl_util.get_tensor_shape(t)
|
||||
|
||||
# TODO(wuxibin): cupy doesn't support bfloat16 for now,
|
||||
# once it is supported, we can eliminate this if statement.
|
||||
#
|
||||
# Allocate using the same backend as the tensors in `tensor_list`.
|
||||
# Use torch only when the tensors are torch.Tensor; otherwise fall back to CuPy.
|
||||
use_torch = False
|
||||
if torch_available():
|
||||
try:
|
||||
import torch
|
||||
|
||||
use_torch = isinstance(t, torch.Tensor)
|
||||
except ImportError:
|
||||
use_torch = False
|
||||
|
||||
if use_torch:
|
||||
buffer = torch.empty(tuple(buffer_shape), dtype=t.dtype, device=t.device)
|
||||
else:
|
||||
# note we need a cupy dtype here.
|
||||
dtype = nccl_util.get_cupy_tensor_dtype(t)
|
||||
device = nccl_util.get_tensor_device(t)
|
||||
with nccl_util.Device(device):
|
||||
buffer = cupy.empty(buffer_shape, dtype=dtype)
|
||||
|
||||
if copy:
|
||||
for i, tensor in enumerate(tensor_list):
|
||||
nccl_util.copy_tensor(buffer[i], tensor)
|
||||
return buffer
|
||||
|
||||
|
||||
def _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists):
|
||||
"""Check the compatibility between tensor input and tensor list input."""
|
||||
if not tensors or not isinstance(tensors, list):
|
||||
raise RuntimeError("The first argument 'tensors' expects a list of tensors.")
|
||||
if not tensor_lists or not isinstance(tensor_lists, list):
|
||||
raise RuntimeError(
|
||||
"The second argument 'tensor_lists' expects a list of tensor list."
|
||||
)
|
||||
dtype = nccl_util.get_nccl_tensor_dtype(tensors[0])
|
||||
shape = nccl_util.get_tensor_shape(tensors[0])
|
||||
for i, tensor_list in enumerate(tensor_lists):
|
||||
# check all tensor in `tensors` match.
|
||||
dt = nccl_util.get_nccl_tensor_dtype(tensors[i])
|
||||
if dt != dtype:
|
||||
raise RuntimeError(
|
||||
"All tensor operands to scatter/gather must "
|
||||
"have the same dtype. Got '{}' and '{}'.".format(dt, dtype)
|
||||
)
|
||||
# Note: typically CCL libraries only requires they have the same
|
||||
# number of elements; Here we make it more strict -- we require
|
||||
# exact shape match.
|
||||
s = nccl_util.get_tensor_shape(tensors[i])
|
||||
if s != shape:
|
||||
raise RuntimeError(
|
||||
"All tensor operands to scatter/gather must "
|
||||
"have the same shape. Got '{}' and '{}'.".format(s, shape)
|
||||
)
|
||||
# check all tensors in `tensor_lists` match.
|
||||
for t in tensor_lists[i]:
|
||||
# check dtype
|
||||
dt = nccl_util.get_nccl_tensor_dtype(t)
|
||||
if dt != dtype:
|
||||
raise RuntimeError(
|
||||
"All tensor operands to scatter/gather must "
|
||||
"have the same dtype. Got '{}' and '{}'.".format(dt, dtype)
|
||||
)
|
||||
s = nccl_util.get_tensor_shape(t)
|
||||
if s != shape:
|
||||
raise RuntimeError(
|
||||
"All tensor operands to scatter/gather must "
|
||||
"have the same shape. Got '{}' and '{}'.".format(s, shape)
|
||||
)
|
||||
|
||||
|
||||
def _check_gpu_tensors(tensors):
|
||||
"""Check all tensors are distributed on different GPUs."""
|
||||
if not tensors or not isinstance(tensors, list):
|
||||
raise RuntimeError("'tensors' must be a nonempty list.")
|
||||
if len(tensors) > nccl_util.get_num_gpus():
|
||||
raise RuntimeError(
|
||||
"Tensor list cannot be larger than the number"
|
||||
"of available GPUs. Got {} > {}.".format(
|
||||
len(tensors), nccl_util.get_num_gpus()
|
||||
)
|
||||
)
|
||||
t0 = tensors[0]
|
||||
dt = nccl_util.get_nccl_tensor_dtype(t0)
|
||||
s = nccl_util.get_tensor_shape(t0)
|
||||
d = nccl_util.get_tensor_device(t0)
|
||||
for i, t in enumerate(tensors):
|
||||
if i == 0:
|
||||
continue
|
||||
# We need to check the following:
|
||||
# (1) tensor is cuda (already checked during API)
|
||||
# (2) tensor dtype
|
||||
# (3) tensor shape match
|
||||
# (4) each tensor is on a different GPU
|
||||
dtype = nccl_util.get_nccl_tensor_dtype(t)
|
||||
if dt != dtype:
|
||||
raise RuntimeError(
|
||||
"Tensors must have identical dtype. Got: '{}'.".format(dtype)
|
||||
)
|
||||
shape = nccl_util.get_tensor_shape(t)
|
||||
if s != shape:
|
||||
raise RuntimeError(
|
||||
"Tensor must have identical shape. Got: '{}'.".format(shape)
|
||||
)
|
||||
device = nccl_util.get_tensor_device(t)
|
||||
if device == d:
|
||||
raise RuntimeError("Tensor must be on distinct GPUs.")
|
||||
|
||||
|
||||
def _get_comm_key_from_devices(devices: List[int]):
|
||||
"""Return a key from a list of devices for collective calls.
|
||||
|
||||
For example, if the tensors are on gpus 0, 1, 2, 3,
|
||||
then the key would be "0,1,2,3".
|
||||
|
||||
Args:
|
||||
devices: a list of GPU device indices
|
||||
|
||||
Returns:
|
||||
str: a string represents the key to query the communicator cache.
|
||||
|
||||
"""
|
||||
return ",".join([str(d) for d in devices])
|
||||
|
||||
|
||||
def _get_comm_key_send_recv(
|
||||
my_rank: int, my_gpu_idx: int, peer_rank: int, peer_gpu_idx: int
|
||||
):
|
||||
"""Return a key given source and destination ranks for p2p tasks.
|
||||
|
||||
The p2p key is in the following form:
|
||||
[min_rank]_[gpu_index]:[max_rank]_[gpu_index].
|
||||
|
||||
Args:
|
||||
my_rank: the rank of the source process.
|
||||
my_gpu_idx: the source gpu index on the process.
|
||||
peer_rank: the rank of the destination process.
|
||||
peer_gpu_idx: the destination gpu index on the process.
|
||||
|
||||
Returns:
|
||||
comm_key: a string key to query the communication cache.
|
||||
"""
|
||||
if my_rank < peer_rank:
|
||||
lower_key = str(my_rank) + "_" + str(my_gpu_idx)
|
||||
higher_key = str(peer_rank) + "_" + str(peer_gpu_idx)
|
||||
elif my_rank > peer_rank:
|
||||
lower_key = str(peer_rank) + "_" + str(peer_gpu_idx)
|
||||
higher_key = str(my_rank) + "_" + str(my_gpu_idx)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Send and recv happens on the same process. ray.util.collective "
|
||||
"does not support this case as of now. Alternatively, consider "
|
||||
"doing GPU to GPU memcpy?"
|
||||
)
|
||||
comm_key = lower_key + ":" + higher_key
|
||||
return comm_key
|
||||
@@ -0,0 +1,298 @@
|
||||
"""Code to wrap some NCCL API calls."""
|
||||
from typing import Any, List
|
||||
|
||||
import numpy
|
||||
|
||||
try:
|
||||
import cupy
|
||||
from cupy.cuda import (
|
||||
Device, # noqa: F401
|
||||
nccl,
|
||||
)
|
||||
from cupy.cuda.nccl import (
|
||||
NcclCommunicator,
|
||||
get_build_version,
|
||||
get_version,
|
||||
groupEnd, # noqa: F401
|
||||
groupStart, # noqa: F401
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError("NCCL in Ray requires Cupy being available!")
|
||||
|
||||
from ray.util.collective.types import ReduceOp, torch_available
|
||||
|
||||
NCCL_REDUCE_OP_MAP = {
|
||||
ReduceOp.SUM: nccl.NCCL_SUM,
|
||||
ReduceOp.PRODUCT: nccl.NCCL_PROD,
|
||||
ReduceOp.MIN: nccl.NCCL_MIN,
|
||||
ReduceOp.MAX: nccl.NCCL_MAX,
|
||||
}
|
||||
|
||||
# cupy types are the same with numpy types
|
||||
NUMPY_NCCL_DTYPE_MAP = {
|
||||
# INT types
|
||||
numpy.int_: nccl.NCCL_INT64,
|
||||
numpy.uint8: nccl.NCCL_UINT8,
|
||||
numpy.uint32: nccl.NCCL_UINT32,
|
||||
numpy.uint64: nccl.NCCL_UINT64,
|
||||
numpy.int8: nccl.NCCL_INT8,
|
||||
numpy.int32: nccl.NCCL_INT32,
|
||||
numpy.int64: nccl.NCCL_INT64,
|
||||
# FLOAT types
|
||||
numpy.half: nccl.NCCL_HALF,
|
||||
numpy.float16: nccl.NCCL_FLOAT16,
|
||||
numpy.float32: nccl.NCCL_FLOAT32,
|
||||
numpy.float64: nccl.NCCL_FLOAT64,
|
||||
numpy.double: nccl.NCCL_DOUBLE,
|
||||
}
|
||||
|
||||
if torch_available():
|
||||
import torch
|
||||
import torch.utils.dlpack
|
||||
|
||||
TORCH_NCCL_DTYPE_MAP = {
|
||||
torch.bool: nccl.NCCL_INT8,
|
||||
# INT types
|
||||
torch.int: nccl.NCCL_INT,
|
||||
torch.uint8: nccl.NCCL_UINT8,
|
||||
torch.int8: nccl.NCCL_INT8,
|
||||
torch.int32: nccl.NCCL_INT32,
|
||||
torch.int64: nccl.NCCL_INT64,
|
||||
torch.long: nccl.NCCL_INT64,
|
||||
# FLOAT types
|
||||
torch.half: nccl.NCCL_HALF,
|
||||
torch.float: nccl.NCCL_FLOAT,
|
||||
torch.float16: nccl.NCCL_FLOAT16,
|
||||
torch.float32: nccl.NCCL_FLOAT32,
|
||||
torch.float64: nccl.NCCL_FLOAT64,
|
||||
torch.double: nccl.NCCL_DOUBLE,
|
||||
}
|
||||
|
||||
# Older versions of cupy don't support bfloat16.
|
||||
if hasattr(nccl, "NCCL_BFLOAT16"):
|
||||
TORCH_NCCL_DTYPE_MAP[torch.bfloat16] = nccl.NCCL_BFLOAT16
|
||||
|
||||
TORCH_NUMPY_DTYPE_MAP = {
|
||||
# INT types
|
||||
torch.int: numpy.int32,
|
||||
torch.uint8: numpy.uint8,
|
||||
torch.int8: numpy.int8,
|
||||
torch.int32: numpy.int32,
|
||||
torch.int64: numpy.int64,
|
||||
torch.long: numpy.int64,
|
||||
# FLOAT types
|
||||
torch.half: numpy.half,
|
||||
torch.float: numpy.float32,
|
||||
torch.float16: numpy.float16,
|
||||
torch.float32: numpy.float32,
|
||||
torch.float64: numpy.float64,
|
||||
}
|
||||
|
||||
|
||||
def get_num_gpus():
|
||||
"""Returns the number of compute-capable GPUs."""
|
||||
return cupy.cuda.runtime.getDeviceCount()
|
||||
|
||||
|
||||
def get_nccl_build_version():
|
||||
return get_build_version()
|
||||
|
||||
|
||||
def get_nccl_runtime_version():
|
||||
return get_version()
|
||||
|
||||
|
||||
def get_nccl_unique_id():
|
||||
return nccl.get_unique_id()
|
||||
|
||||
|
||||
def create_nccl_communicator(world_size: int, nccl_unique_id: bytes, rank: int):
|
||||
"""Create an NCCL communicator using NCCL APIs.
|
||||
|
||||
Args:
|
||||
world_size: the number of processes of this communicator group.
|
||||
nccl_unique_id: the NCCLUniqueID for this group.
|
||||
rank: the rank of this process.
|
||||
|
||||
Returns:
|
||||
comm: an NCCL communicator.
|
||||
"""
|
||||
comm = NcclCommunicator(world_size, nccl_unique_id, rank)
|
||||
return comm
|
||||
|
||||
|
||||
def get_nccl_reduce_op(reduce_op: ReduceOp):
|
||||
"""Map the reduce op to NCCL reduce op type.
|
||||
|
||||
Args:
|
||||
reduce_op: ReduceOp Enum (SUM/PRODUCT/MIN/MAX).
|
||||
|
||||
Returns:
|
||||
the mapped NCCL reduce op.
|
||||
"""
|
||||
if reduce_op not in NCCL_REDUCE_OP_MAP:
|
||||
raise RuntimeError("NCCL does not support reduce op: '{}'.".format(reduce_op))
|
||||
return NCCL_REDUCE_OP_MAP[reduce_op]
|
||||
|
||||
|
||||
def get_nccl_tensor_dtype(tensor):
|
||||
"""Return the corresponded NCCL dtype given a tensor."""
|
||||
if isinstance(tensor, cupy.ndarray):
|
||||
return NUMPY_NCCL_DTYPE_MAP[tensor.dtype.type]
|
||||
if torch_available():
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
return TORCH_NCCL_DTYPE_MAP[tensor.dtype]
|
||||
raise ValueError(
|
||||
"Unsupported tensor type. Got: {}. Supported "
|
||||
"GPU tensor types are: torch.Tensor, "
|
||||
"cupy.ndarray.".format(type(tensor))
|
||||
)
|
||||
|
||||
|
||||
def get_cupy_tensor_dtype(tensor):
|
||||
"""Return the corresponded Cupy dtype given a tensor."""
|
||||
if isinstance(tensor, cupy.ndarray):
|
||||
return tensor.dtype.type
|
||||
if torch_available():
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
return TORCH_NUMPY_DTYPE_MAP[tensor.dtype]
|
||||
raise ValueError(
|
||||
"Unsupported tensor type. Got: {}. Supported "
|
||||
"GPU tensor types are: torch.Tensor, "
|
||||
"cupy.ndarray.".format(type(tensor))
|
||||
)
|
||||
|
||||
|
||||
def get_tensor_ptr(tensor):
|
||||
"""Return the pointer to the underlying memory storage of a tensor."""
|
||||
if isinstance(tensor, cupy.ndarray):
|
||||
return tensor.data.ptr
|
||||
if isinstance(tensor, numpy.ndarray):
|
||||
return tensor.data
|
||||
if torch_available():
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
if not tensor.is_cuda:
|
||||
raise RuntimeError(
|
||||
"Torch tensor must be on GPU when using NCCL collectives."
|
||||
)
|
||||
return tensor.data_ptr()
|
||||
raise ValueError(
|
||||
"Unsupported tensor type. Got: {}. Supported "
|
||||
"GPU tensor types are: torch.Tensor, "
|
||||
"cupy.ndarray.".format(type(tensor))
|
||||
)
|
||||
|
||||
|
||||
def get_tensor_n_elements(tensor):
|
||||
"""Return the number of elements in a tensor."""
|
||||
if isinstance(tensor, cupy.ndarray) or isinstance(tensor, numpy.ndarray):
|
||||
return tensor.size
|
||||
if torch_available():
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
return torch.numel(tensor)
|
||||
raise ValueError(
|
||||
"Unsupported tensor type. Got: {}. Supported "
|
||||
"GPU tensor types are: torch.Tensor, "
|
||||
"cupy.ndarray.".format(type(tensor))
|
||||
)
|
||||
|
||||
|
||||
def get_tensor_shape(tensor):
|
||||
"""Return the shape of the tensor as a list."""
|
||||
if isinstance(tensor, cupy.ndarray):
|
||||
return list(tensor.shape)
|
||||
if torch_available():
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
return list(tensor.size())
|
||||
raise ValueError(
|
||||
"Unsupported tensor type. Got: {}. Supported "
|
||||
"GPU tensor types are: torch.Tensor, "
|
||||
"cupy.ndarray.".format(type(tensor))
|
||||
)
|
||||
|
||||
|
||||
def get_tensor_strides(tensor):
|
||||
"""Return the strides of the tensor as a list."""
|
||||
if isinstance(tensor, cupy.ndarray):
|
||||
return [int(stride / tensor.dtype.itemsize) for stride in tensor.strides]
|
||||
if torch_available():
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
return list(tensor.stride())
|
||||
raise ValueError(
|
||||
"Unsupported tensor type. Got: {}. Supported "
|
||||
"GPU tensor types are: torch.Tensor, "
|
||||
"cupy.ndarray.".format(type(tensor))
|
||||
)
|
||||
|
||||
|
||||
def get_tensor_device(tensor):
|
||||
"""Return the GPU index of a tensor."""
|
||||
if isinstance(tensor, cupy.ndarray):
|
||||
try:
|
||||
device = tensor.device.id
|
||||
except AttributeError as exec:
|
||||
raise RuntimeError("The tensor is not on a valid GPU.") from exec
|
||||
elif torch_available() and isinstance(tensor, torch.Tensor):
|
||||
device = tensor.device.index
|
||||
if not isinstance(device, int):
|
||||
raise RuntimeError("The tensor is not on a valid GPU.")
|
||||
else:
|
||||
raise ValueError("Unsupported tensor type. Got: {}.".format(type(tensor)))
|
||||
return device
|
||||
|
||||
|
||||
def copy_tensor(dst_tensor: Any, src_tensor: Any):
|
||||
"""Copy the content from src_tensor to dst_tensor.
|
||||
|
||||
Args:
|
||||
dst_tensor: the tensor to copy from.
|
||||
src_tensor: the tensor to copy to.
|
||||
"""
|
||||
copied = True
|
||||
if isinstance(dst_tensor, cupy.ndarray) and isinstance(src_tensor, cupy.ndarray):
|
||||
cupy.copyto(dst_tensor, src_tensor)
|
||||
elif torch_available():
|
||||
if isinstance(dst_tensor, torch.Tensor) and isinstance(
|
||||
src_tensor, torch.Tensor
|
||||
):
|
||||
dst_tensor.copy_(src_tensor)
|
||||
elif isinstance(dst_tensor, torch.Tensor) and isinstance(
|
||||
src_tensor, cupy.ndarray
|
||||
):
|
||||
t = torch.utils.dlpack.from_dlpack(src_tensor.toDlpack())
|
||||
dst_tensor.copy_(t)
|
||||
elif isinstance(dst_tensor, cupy.ndarray) and isinstance(
|
||||
src_tensor, torch.Tensor
|
||||
):
|
||||
t = cupy.fromDlpack(torch.utils.dlpack.to_dlpack(src_tensor))
|
||||
cupy.copyto(dst_tensor, t)
|
||||
else:
|
||||
copied = False
|
||||
else:
|
||||
copied = False
|
||||
if not copied:
|
||||
raise ValueError(
|
||||
"Unsupported tensor type. Got: {} and {}. Supported "
|
||||
"GPU tensor types are: torch.Tensor, cupy.ndarray.".format(
|
||||
type(dst_tensor), type(src_tensor)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_tensor_device_list(tensors: List[Any]):
|
||||
"""Returns the gpu devices of the list of input tensors.
|
||||
|
||||
Args:
|
||||
tensors: a list of tensors, each locates on a GPU.
|
||||
|
||||
Returns:
|
||||
list: the list of GPU devices.
|
||||
|
||||
"""
|
||||
if not isinstance(tensors, list):
|
||||
raise RuntimeError(
|
||||
"Expect a list of tensors each locates on a GPU device. "
|
||||
"Got: '{}'.".format(type(tensors))
|
||||
)
|
||||
devices = [get_tensor_device(t) for t in tensors]
|
||||
return devices
|
||||
@@ -0,0 +1,290 @@
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
from typing import TYPE_CHECKING, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import ray
|
||||
import ray.experimental.internal_kv as internal_kv
|
||||
from ray._common.network_utils import find_free_port, is_ipv6
|
||||
from ray.util.collective.collective_group.base_collective_group import BaseGroup
|
||||
from ray.util.collective.types import (
|
||||
AllGatherOptions,
|
||||
AllReduceOptions,
|
||||
Backend,
|
||||
BarrierOptions,
|
||||
BroadcastOptions,
|
||||
RecvOptions,
|
||||
ReduceOp,
|
||||
ReduceOptions,
|
||||
ReduceScatterOptions,
|
||||
SendOptions,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
import torch.distributed as dist
|
||||
|
||||
_TORCH_DISTRIBUTED_AVAILABLE = True
|
||||
TORCH_REDUCE_OP_MAP = {
|
||||
ReduceOp.SUM: dist.ReduceOp.SUM,
|
||||
ReduceOp.PRODUCT: dist.ReduceOp.PRODUCT,
|
||||
ReduceOp.MIN: dist.ReduceOp.MIN,
|
||||
ReduceOp.MAX: dist.ReduceOp.MAX,
|
||||
}
|
||||
except ImportError:
|
||||
_TORCH_DISTRIBUTED_AVAILABLE = False
|
||||
TORCH_REDUCE_OP_MAP = None
|
||||
|
||||
|
||||
def get_master_address_metadata_key(group_name: str):
|
||||
return f"collective_group_master_address_{group_name}"
|
||||
|
||||
|
||||
def get_address_and_port() -> Tuple[str, int]:
|
||||
"""Returns the IP address and a free port on this node."""
|
||||
addr = ray.util.get_node_ip_address()
|
||||
port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET)
|
||||
return addr, port
|
||||
|
||||
|
||||
class TorchGLOOGroup(BaseGroup):
|
||||
def __init__(
|
||||
self,
|
||||
world_size: int,
|
||||
rank: int,
|
||||
group_name: str,
|
||||
gloo_timeout: Optional[int] = None,
|
||||
):
|
||||
# Initialize the default process group only once per process.
|
||||
if not dist.is_initialized():
|
||||
# Rendezvous: ensure a MASTER_ADDR:MASTER_PORT is published in internal_kv.
|
||||
self._rendezvous(group_name, rank, gloo_timeout)
|
||||
|
||||
metadata_key = get_master_address_metadata_key(group_name)
|
||||
try:
|
||||
metadata = internal_kv._internal_kv_get(metadata_key)
|
||||
except ValueError:
|
||||
raise RuntimeError(
|
||||
f"TorchGLOOGroup expected metadata in internal_kv with name `{metadata_key}`. "
|
||||
"TorchGLOOGroup should not be instantiated directly. "
|
||||
"Use ray.experimental.collective.create_collective_group to create a group."
|
||||
)
|
||||
if metadata is None:
|
||||
raise RuntimeError(
|
||||
f"Missing rendezvous metadata for group `{group_name}` under key `{metadata_key}`."
|
||||
)
|
||||
metadata = metadata.decode()
|
||||
master_addr, master_port = metadata.split(":")
|
||||
os.environ["MASTER_ADDR"] = master_addr
|
||||
os.environ["MASTER_PORT"] = master_port
|
||||
|
||||
dist.init_process_group(
|
||||
backend="gloo", init_method="env://", world_size=world_size, rank=rank
|
||||
)
|
||||
|
||||
# Clean up rendezvous metadata after all ranks have initialized.
|
||||
# dist.init_process_group is synchronous, so all ranks have read the metadata by now.
|
||||
if rank == 0:
|
||||
try:
|
||||
internal_kv._internal_kv_del(metadata_key)
|
||||
except Exception as e:
|
||||
# Ignore errors during cleanup (e.g., key already deleted)
|
||||
logger.warning(
|
||||
f"Failed to delete rendezvous key '{metadata_key}' during init: {e}"
|
||||
)
|
||||
|
||||
super().__init__(world_size, rank, group_name)
|
||||
|
||||
# Create a subgroup for this logical group. For the default group, use WORLD.
|
||||
self._is_default_group = group_name == "default"
|
||||
if self._is_default_group:
|
||||
self._pg = dist.group.WORLD
|
||||
else:
|
||||
# All ranks participate in this subgroup with global ranks [0..world_size-1].
|
||||
ranks = list(range(world_size))
|
||||
self._pg = dist.new_group(ranks=ranks, backend="gloo")
|
||||
|
||||
# Compatibility shim for legacy tests expecting a pygloo context with getTimeout().
|
||||
# Store the rendezvous timeout in milliseconds, defaulting to 30000 if unspecified.
|
||||
class _GlooCompatContext:
|
||||
def __init__(self, timeout_ms: int):
|
||||
self._timeout_ms = timeout_ms
|
||||
|
||||
def getTimeout(self) -> int:
|
||||
return self._timeout_ms
|
||||
|
||||
self._gloo_context = _GlooCompatContext(
|
||||
gloo_timeout if gloo_timeout is not None else 30000
|
||||
)
|
||||
|
||||
def _rendezvous(
|
||||
self, group_name: str, rank: int, gloo_timeout: Optional[int]
|
||||
) -> None:
|
||||
"""Rendezvous: ensure a MASTER_ADDR:MASTER_PORT is published in internal_kv.
|
||||
|
||||
Rank 0 publishes the address and port, other ranks wait for it.
|
||||
"""
|
||||
metadata_key = get_master_address_metadata_key(group_name)
|
||||
if rank == 0:
|
||||
addr, port = get_address_and_port()
|
||||
internal_kv._internal_kv_put(metadata_key, f"{addr}:{port}")
|
||||
else:
|
||||
# Wait until rank 0 publishes the metadata or timeout.
|
||||
deadline_s = time.time() + (gloo_timeout / 1000.0 if gloo_timeout else 30.0)
|
||||
while True:
|
||||
meta = internal_kv._internal_kv_get(metadata_key)
|
||||
if meta is not None:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError(
|
||||
f"Timed out waiting for GLOO rendezvous metadata for group '{group_name}'."
|
||||
)
|
||||
time.sleep(0.05)
|
||||
|
||||
def destroy_group(self):
|
||||
"""GC the communicators."""
|
||||
# Destroy only the subgroup for non-default groups. Allow default to be torn down explicitly.
|
||||
if self._is_default_group:
|
||||
# Destroy default process group to allow re-init in tests that recreate the same group.
|
||||
dist.destroy_process_group()
|
||||
else:
|
||||
# Destroy just this subgroup.
|
||||
if self._pg is not None:
|
||||
dist.destroy_process_group(self._pg)
|
||||
|
||||
@classmethod
|
||||
def backend(cls):
|
||||
"""The backend of this collective group."""
|
||||
return Backend.GLOO
|
||||
|
||||
@classmethod
|
||||
def check_backend_availability(cls) -> bool:
|
||||
return _TORCH_DISTRIBUTED_AVAILABLE
|
||||
|
||||
def _check_tensor_input(self, tensor: List["torch.Tensor"]) -> "torch.Tensor":
|
||||
"""ray.util.collective wraps tensor arguments in a list.
|
||||
Accept a single torch.Tensor or numpy.ndarray and unwrap/convert it.
|
||||
"""
|
||||
assert isinstance(tensor, list) and len(tensor) == 1
|
||||
t = tensor[0]
|
||||
if isinstance(t, torch.Tensor):
|
||||
return t
|
||||
if isinstance(t, np.ndarray):
|
||||
return torch.from_numpy(t)
|
||||
raise ValueError(
|
||||
f"torch_gloo group only accepts torch.Tensor or numpy.ndarray, received {type(t)}"
|
||||
)
|
||||
|
||||
def _check_tensor_list_input(
|
||||
self, tensor_list: List[List["torch.Tensor"]]
|
||||
) -> List["torch.Tensor"]:
|
||||
"""ray.util.collective wraps tensor arguments in a list.
|
||||
Accept a single list containing torch.Tensors or numpy.ndarrays and
|
||||
unwrap/convert items as needed.
|
||||
"""
|
||||
assert isinstance(tensor_list, list) and len(tensor_list) == 1
|
||||
tensor_list = tensor_list[0]
|
||||
converted_tensor_list = []
|
||||
for tensor in tensor_list:
|
||||
if isinstance(tensor, np.ndarray):
|
||||
tensor = torch.from_numpy(tensor)
|
||||
converted_tensor_list.append(tensor)
|
||||
elif isinstance(tensor, torch.Tensor):
|
||||
converted_tensor_list.append(tensor)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"torch_gloo group only accepts torch.Tensor or numpy.ndarray types, received tensor list with value {tensor}"
|
||||
)
|
||||
return converted_tensor_list
|
||||
|
||||
def allreduce(
|
||||
self,
|
||||
tensor: List["torch.Tensor"],
|
||||
allreduce_options: Optional[AllReduceOptions] = None,
|
||||
) -> None:
|
||||
if allreduce_options is None:
|
||||
allreduce_options = AllReduceOptions()
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
torch_reduce_op = TORCH_REDUCE_OP_MAP[allreduce_options.reduceOp]
|
||||
dist.all_reduce(tensor, op=torch_reduce_op, group=self._pg)
|
||||
|
||||
def barrier(self, barrier_options=BarrierOptions()) -> None:
|
||||
dist.barrier(group=self._pg)
|
||||
|
||||
def reduce(
|
||||
self,
|
||||
tensor: List["torch.Tensor"],
|
||||
reduce_options: Optional[ReduceOptions] = None,
|
||||
) -> None:
|
||||
if reduce_options is None:
|
||||
reduce_options = ReduceOptions()
|
||||
t = self._check_tensor_input(tensor)
|
||||
torch_reduce_op = TORCH_REDUCE_OP_MAP[reduce_options.reduceOp]
|
||||
# Avoid mutating non-root ranks' user tensors to match util.collective semantics.
|
||||
if self._rank == reduce_options.root_rank:
|
||||
dist.reduce(
|
||||
t, dst=reduce_options.root_rank, op=torch_reduce_op, group=self._pg
|
||||
)
|
||||
else:
|
||||
tmp = t.detach().clone()
|
||||
dist.reduce(
|
||||
tmp, dst=reduce_options.root_rank, op=torch_reduce_op, group=self._pg
|
||||
)
|
||||
|
||||
def allgather(
|
||||
self,
|
||||
tensor_list: List[List["torch.Tensor"]],
|
||||
tensor: List["torch.Tensor"],
|
||||
allgather_options: Optional[AllGatherOptions] = None,
|
||||
) -> None:
|
||||
if allgather_options is None:
|
||||
allgather_options = AllGatherOptions()
|
||||
tensor_list = self._check_tensor_list_input(tensor_list)
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
dist.all_gather(tensor_list, tensor, group=self._pg)
|
||||
|
||||
def broadcast(
|
||||
self, tensor: List["torch.Tensor"], broadcast_options=BroadcastOptions()
|
||||
) -> None:
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
dist.broadcast(tensor, src=broadcast_options.root_rank, group=self._pg)
|
||||
|
||||
def reducescatter(
|
||||
self,
|
||||
output_tensor: List["torch.Tensor"],
|
||||
tensor_list: List[List["torch.Tensor"]],
|
||||
reducescatter_options: Optional[ReduceScatterOptions] = None,
|
||||
) -> None:
|
||||
if reducescatter_options is None:
|
||||
reducescatter_options = ReduceScatterOptions()
|
||||
tensor_list = self._check_tensor_list_input(tensor_list)
|
||||
output_tensor = self._check_tensor_input(output_tensor)
|
||||
if output_tensor.shape != tensor_list[self._rank].shape:
|
||||
raise ValueError(
|
||||
"Output tensor has wrong shape {output_tensor.shape}, expected {tensor_list[self._rank].shape}"
|
||||
)
|
||||
torch_reduce_op = TORCH_REDUCE_OP_MAP[reducescatter_options.reduceOp]
|
||||
|
||||
# torch.distributed gloo doesn't support reducescatter. Implement a
|
||||
# simple version using allreduce.
|
||||
for tensor in tensor_list:
|
||||
dist.all_reduce(tensor, op=torch_reduce_op, group=self._pg)
|
||||
|
||||
if output_tensor.data_ptr() != tensor_list[self._rank].data_ptr():
|
||||
output_tensor.copy_(tensor_list[self._rank])
|
||||
|
||||
def send(self, tensor: List["torch.Tensor"], send_options: SendOptions) -> None:
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
dist.send(tensor, dst=send_options.dst_rank)
|
||||
|
||||
def recv(self, tensor: List["torch.Tensor"], recv_options: RecvOptions) -> None:
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
dist.recv(tensor, src=recv_options.src_rank)
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
Constants.
|
||||
|
||||
Contains constants used to setup collective groups.
|
||||
"""
|
||||
import hashlib
|
||||
import os
|
||||
from enum import Enum, auto
|
||||
|
||||
|
||||
def get_store_name(group_name: str):
|
||||
"""Generate the unique name for the NCCLUniqueID store (named actor).
|
||||
|
||||
Args:
|
||||
group_name: unique user name for the store.
|
||||
|
||||
Returns:
|
||||
SHA256-hexlified name for the store.
|
||||
"""
|
||||
if not group_name:
|
||||
raise ValueError("group_name is None.")
|
||||
hexlified_name = hashlib.sha256(group_name.encode()).hexdigest()
|
||||
return hexlified_name
|
||||
|
||||
|
||||
class ENV(Enum):
|
||||
"""ray.util.collective environment variables."""
|
||||
|
||||
NCCL_USE_MULTISTREAM = auto(), lambda v: (v or "True") == "True"
|
||||
|
||||
@property
|
||||
def val(self):
|
||||
"""Return the output of the lambda against the system's env value."""
|
||||
_, default_fn = self.value
|
||||
return default_fn(os.getenv(self.name))
|
||||
@@ -0,0 +1,529 @@
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
|
||||
import ray
|
||||
import ray.experimental.internal_kv as internal_kv
|
||||
from ray.util.collective import (
|
||||
allreduce,
|
||||
broadcast,
|
||||
create_collective_group,
|
||||
init_collective_group,
|
||||
)
|
||||
from ray.util.collective.backend_registry import (
|
||||
_global_registry,
|
||||
register_collective_backend,
|
||||
)
|
||||
from ray.util.collective.collective_group.base_collective_group import BaseGroup
|
||||
from ray.util.collective.types import (
|
||||
AllGatherOptions,
|
||||
AllReduceOptions,
|
||||
BarrierOptions,
|
||||
BroadcastOptions,
|
||||
RecvOptions,
|
||||
ReduceOp,
|
||||
ReduceOptions,
|
||||
ReduceScatterOptions,
|
||||
SendOptions,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
|
||||
def _unregister_collective_backend(name: str) -> None:
|
||||
"""Helper function to unregister a backend for testing purposes."""
|
||||
upper_name = name.upper()
|
||||
if upper_name in _global_registry._map:
|
||||
del _global_registry._map[upper_name]
|
||||
|
||||
|
||||
def get_data_key(group_name: str, rank: int, op_name: str):
|
||||
return f"collective_mock_{group_name}_{op_name}_rank_{rank}"
|
||||
|
||||
|
||||
def get_barrier_key(group_name: str, barrier_id: int):
|
||||
return f"collective_mock_{group_name}_barrier_{barrier_id}"
|
||||
|
||||
|
||||
class MockInternalKVGroup(BaseGroup):
|
||||
def __init__(self, world_size: int, rank: int, group_name: str):
|
||||
super().__init__(world_size, rank, group_name)
|
||||
self._barrier_counter = 0
|
||||
|
||||
@classmethod
|
||||
def backend(cls):
|
||||
return "MOCK"
|
||||
|
||||
@classmethod
|
||||
def check_backend_availability(cls) -> bool:
|
||||
return True
|
||||
|
||||
def _check_tensor_input(self, tensor):
|
||||
assert isinstance(tensor, list) and len(tensor) == 1
|
||||
t = tensor[0]
|
||||
if isinstance(t, np.ndarray):
|
||||
return t
|
||||
try:
|
||||
import torch
|
||||
|
||||
if isinstance(t, torch.Tensor):
|
||||
return t
|
||||
except ImportError:
|
||||
pass
|
||||
raise ValueError(
|
||||
f"MockInternalKVGroup only only accepts numpy.ndarray or torch.Tensor, received {type(t)}"
|
||||
)
|
||||
|
||||
def _serialize_tensor(self, tensor):
|
||||
if isinstance(tensor, np.ndarray):
|
||||
return tensor.tobytes(), tensor.shape, tensor.dtype
|
||||
try:
|
||||
import torch
|
||||
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
return (
|
||||
tensor.cpu().numpy().tobytes(),
|
||||
tensor.shape,
|
||||
tensor.cpu().numpy().dtype,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
raise ValueError(f"Unsupported tensor type: {type(tensor)}")
|
||||
|
||||
def _deserialize_tensor(self, data: bytes, shape, dtype, target_tensor):
|
||||
if isinstance(target_tensor, np.ndarray):
|
||||
np_array = np.frombuffer(data, dtype=dtype).reshape(shape)
|
||||
target_tensor[:] = np_array
|
||||
else:
|
||||
try:
|
||||
import torch
|
||||
|
||||
if isinstance(target_tensor, torch.Tensor):
|
||||
np_array = np.frombuffer(data, dtype=dtype).reshape(shape)
|
||||
target_tensor.copy_(torch.from_numpy(np_array))
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
def broadcast(self, tensor, broadcast_options=BroadcastOptions()):
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
root_rank = broadcast_options.root_rank
|
||||
|
||||
data_key = get_data_key(self._group_name, root_rank, "broadcast")
|
||||
|
||||
if self._rank == root_rank:
|
||||
data, shape, dtype = self._serialize_tensor(tensor)
|
||||
internal_kv._internal_kv_put(data_key, data)
|
||||
internal_kv._internal_kv_put(f"{data_key}_shape", str(shape))
|
||||
internal_kv._internal_kv_put(f"{data_key}_dtype", dtype.name)
|
||||
else:
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
data = internal_kv._internal_kv_get(data_key)
|
||||
if data is not None:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError(
|
||||
f"Timed out waiting for broadcast data from rank {root_rank}"
|
||||
)
|
||||
time.sleep(0.01)
|
||||
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
shape_data = internal_kv._internal_kv_get(f"{data_key}_shape")
|
||||
dtype_data = internal_kv._internal_kv_get(f"{data_key}_dtype")
|
||||
if shape_data is not None and dtype_data is not None:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError(
|
||||
f"Timed out waiting for broadcast metadata from rank {root_rank}"
|
||||
)
|
||||
time.sleep(0.01)
|
||||
|
||||
shape_str = shape_data.decode()
|
||||
shape = eval(shape_str)
|
||||
dtype_name = dtype_data.decode()
|
||||
dtype = np.dtype(dtype_name)
|
||||
|
||||
self._deserialize_tensor(data, shape, dtype, tensor)
|
||||
|
||||
def allreduce(self, tensor, allreduce_options=AllReduceOptions()):
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
reduce_op = allreduce_options.reduceOp
|
||||
|
||||
data_key = get_data_key(self._group_name, self._rank, "allreduce")
|
||||
done_key = (
|
||||
f"collective_mock_{self._group_name}_allreduce_done_rank_{self._rank}"
|
||||
)
|
||||
|
||||
data, shape, dtype = self._serialize_tensor(tensor)
|
||||
internal_kv._internal_kv_put(data_key, data)
|
||||
internal_kv._internal_kv_put(f"{data_key}_shape", str(shape))
|
||||
internal_kv._internal_kv_put(f"{data_key}_dtype", dtype.name)
|
||||
internal_kv._internal_kv_put(done_key, b"1")
|
||||
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
all_done = True
|
||||
for r in range(self._world_size):
|
||||
key = f"collective_mock_{self._group_name}_allreduce_done_rank_{r}"
|
||||
if internal_kv._internal_kv_get(key) is None:
|
||||
all_done = False
|
||||
break
|
||||
if all_done:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError(
|
||||
"Timed out waiting for allreduce data from all ranks"
|
||||
)
|
||||
time.sleep(0.01)
|
||||
|
||||
result = None
|
||||
|
||||
for r in range(self._world_size):
|
||||
rank_data_key = get_data_key(self._group_name, r, "allreduce")
|
||||
rank_data = internal_kv._internal_kv_get(rank_data_key)
|
||||
rank_shape_data = internal_kv._internal_kv_get(f"{rank_data_key}_shape")
|
||||
rank_dtype_data = internal_kv._internal_kv_get(f"{rank_data_key}_dtype")
|
||||
|
||||
rank_shape_str = rank_shape_data.decode()
|
||||
rank_shape = eval(rank_shape_str)
|
||||
rank_dtype_name = rank_dtype_data.decode()
|
||||
rank_dtype = np.dtype(rank_dtype_name)
|
||||
|
||||
if isinstance(tensor, np.ndarray):
|
||||
rank_tensor = np.frombuffer(rank_data, dtype=rank_dtype).reshape(
|
||||
rank_shape
|
||||
)
|
||||
else:
|
||||
import torch
|
||||
|
||||
rank_np = np.frombuffer(rank_data, dtype=rank_dtype).reshape(rank_shape)
|
||||
rank_tensor = torch.from_numpy(rank_np)
|
||||
|
||||
if result is None:
|
||||
result = (
|
||||
rank_tensor.copy()
|
||||
if isinstance(rank_tensor, np.ndarray)
|
||||
else rank_tensor.clone()
|
||||
)
|
||||
else:
|
||||
if reduce_op == ReduceOp.SUM:
|
||||
result += rank_tensor
|
||||
elif reduce_op == ReduceOp.PRODUCT:
|
||||
result *= rank_tensor
|
||||
elif reduce_op == ReduceOp.MAX:
|
||||
if isinstance(result, np.ndarray):
|
||||
result = np.maximum(result, rank_tensor)
|
||||
else:
|
||||
import torch
|
||||
|
||||
result = torch.maximum(result, rank_tensor)
|
||||
elif reduce_op == ReduceOp.MIN:
|
||||
if isinstance(result, np.ndarray):
|
||||
result = np.minimum(result, rank_tensor)
|
||||
else:
|
||||
import torch
|
||||
|
||||
result = torch.minimum(result, rank_tensor)
|
||||
|
||||
if isinstance(tensor, np.ndarray):
|
||||
tensor[:] = result
|
||||
else:
|
||||
import torch
|
||||
|
||||
if isinstance(result, np.ndarray):
|
||||
tensor.copy_(torch.from_numpy(result))
|
||||
else:
|
||||
tensor.copy_(result)
|
||||
|
||||
def barrier(self, barrier_options=BarrierOptions()):
|
||||
barrier_id = self._barrier_counter
|
||||
barrier_key = get_barrier_key(self._group_name, barrier_id)
|
||||
rank_key = f"{barrier_key}_rank_{self._rank}"
|
||||
|
||||
internal_kv._internal_kv_put(rank_key, b"1")
|
||||
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
all_arrived = True
|
||||
for r in range(self._world_size):
|
||||
key = f"{barrier_key}_rank_{r}"
|
||||
if internal_kv._internal_kv_get(key) is None:
|
||||
all_arrived = False
|
||||
break
|
||||
if all_arrived:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError("Timed out waiting for barrier")
|
||||
time.sleep(0.01)
|
||||
|
||||
self._barrier_counter += 1
|
||||
|
||||
def reduce(self, tensor, reduce_options=ReduceOptions()):
|
||||
raise NotImplementedError("reduce is not implemented in MockInternalKVGroup")
|
||||
|
||||
def allgather(self, tensor_list, tensor, allgather_options=AllGatherOptions()):
|
||||
raise NotImplementedError("allgather is not implemented in MockInternalKVGroup")
|
||||
|
||||
def reducescatter(
|
||||
self, tensor, tensor_list, reducescatter_options=ReduceScatterOptions()
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"reducescatter is not implemented in MockInternalKVGroup"
|
||||
)
|
||||
|
||||
def send(self, tensor, send_options: SendOptions):
|
||||
raise NotImplementedError("send is not implemented in MockInternalKVGroup")
|
||||
|
||||
def recv(self, tensor, recv_options: RecvOptions):
|
||||
raise NotImplementedError("recv is not implemented in MockInternalKVGroup")
|
||||
|
||||
|
||||
def test_mock_backend_create_group():
|
||||
"""Test using create_collective_group (driver-managed approach).
|
||||
|
||||
In this approach:
|
||||
- Driver calls create_collective_group() to declare the group
|
||||
- Workers only need to register the backend
|
||||
- Workers do NOT call init_collective_group()
|
||||
- The group is automatically initialized when workers call collective ops
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
def setup(self):
|
||||
from ray.util.collective.backend_registry import register_collective_backend
|
||||
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
def broadcast_test(self):
|
||||
if self.rank == 0:
|
||||
tensor = np.array([42.0, 43.0, 44.0], dtype=np.float32)
|
||||
else:
|
||||
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
|
||||
broadcast(tensor, src_rank=0)
|
||||
return tensor.tolist()
|
||||
|
||||
def allreduce_test(self):
|
||||
tensor = np.array([float(self.rank + 1)], dtype=np.float32)
|
||||
allreduce(tensor, op=ReduceOp.SUM)
|
||||
return tensor.item()
|
||||
|
||||
actors = [Worker.remote(rank=i) for i in range(3)]
|
||||
|
||||
create_collective_group(
|
||||
actors=actors,
|
||||
world_size=3,
|
||||
ranks=[0, 1, 2],
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
|
||||
ray.get([a.setup.remote() for a in actors])
|
||||
|
||||
results = ray.get([a.broadcast_test.remote() for a in actors])
|
||||
expected = [[42.0, 43.0, 44.0]] * 3
|
||||
if results == expected:
|
||||
print("Broadcast test passed!")
|
||||
else:
|
||||
print(f"Broadcast test failed! Expected {expected}, got {results}")
|
||||
|
||||
results = ray.get([a.allreduce_test.remote() for a in actors])
|
||||
|
||||
if results == [6.0, 6.0, 6.0]:
|
||||
print("AllReduce test passed!")
|
||||
else:
|
||||
print(f"AllReduce test failed! Expected [6.0, 6.0, 6.0], got {results}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
print("test_mock_backend_create_group completed!")
|
||||
|
||||
|
||||
def test_mock_backend_init_group():
|
||||
"""Test using init_collective_group (worker-managed approach).
|
||||
|
||||
In this approach:
|
||||
- Workers call init_collective_group() inside their setup method
|
||||
- Driver does NOT call create_collective_group()
|
||||
- Each worker explicitly initializes its own group membership
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
def setup(self, world_size):
|
||||
from ray.util.collective.backend_registry import register_collective_backend
|
||||
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
init_collective_group(
|
||||
world_size=world_size,
|
||||
rank=self.rank,
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
|
||||
def broadcast_test(self):
|
||||
if self.rank == 0:
|
||||
tensor = np.array([42.0, 43.0, 44.0], dtype=np.float32)
|
||||
else:
|
||||
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
|
||||
broadcast(tensor, src_rank=0)
|
||||
return tensor.tolist()
|
||||
|
||||
def allreduce_test(self):
|
||||
tensor = np.array([float(self.rank + 1)], dtype=np.float32)
|
||||
allreduce(tensor, op=ReduceOp.SUM)
|
||||
return tensor.item()
|
||||
|
||||
actors = [Worker.remote(rank=i) for i in range(3)]
|
||||
|
||||
# Do NOT call create_collective_group here
|
||||
ray.get([a.setup.remote(3) for a in actors])
|
||||
|
||||
results = ray.get([a.broadcast_test.remote() for a in actors])
|
||||
expected = [[42.0, 43.0, 44.0]] * 3
|
||||
if results == expected:
|
||||
print("Broadcast test passed!")
|
||||
else:
|
||||
print(f"Broadcast test failed! Expected {expected}, got {results}")
|
||||
|
||||
results = ray.get([a.allreduce_test.remote() for a in actors])
|
||||
|
||||
if results == [6.0, 6.0, 6.0]:
|
||||
print("AllReduce test passed!")
|
||||
else:
|
||||
print(f"AllReduce test failed! Expected [6.0, 6.0, 6.0], got {results}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
print("test_mock_backend_init_group completed!")
|
||||
|
||||
|
||||
def test_mock_backend_worker_not_registered():
|
||||
"""Test error handling when backend is not registered in worker.
|
||||
|
||||
This test uses create_collective_group (driver-managed approach).
|
||||
The driver registers the backend, but workers do not.
|
||||
When workers try to call collective ops, they should fail.
|
||||
|
||||
Note: We use world_size=1 to avoid "Unhandled error" messages
|
||||
from multiple workers failing simultaneously.
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
def broadcast_test(self):
|
||||
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
broadcast(tensor, src_rank=0)
|
||||
return tensor.tolist()
|
||||
|
||||
# Use single actor to avoid multiple "Unhandled error" messages
|
||||
actors = [Worker.remote(rank=0)]
|
||||
|
||||
create_collective_group(
|
||||
actors=actors,
|
||||
world_size=1,
|
||||
ranks=[0],
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
|
||||
test_passed = False
|
||||
try:
|
||||
ray.get([a.broadcast_test.remote() for a in actors])
|
||||
print("ERROR: Should have raised an exception for missing registration!")
|
||||
except Exception as e:
|
||||
if "not registered" in str(e) or "not initialized" in str(e):
|
||||
print(
|
||||
"Test passed! Correctly raised error for missing worker registration."
|
||||
)
|
||||
test_passed = True
|
||||
else:
|
||||
print(f"ERROR: Unexpected error: {e}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
|
||||
if not test_passed:
|
||||
print("Test failed!")
|
||||
|
||||
|
||||
def test_mock_backend_driver_not_registered():
|
||||
"""Test error handling when backend is not registered on driver.
|
||||
|
||||
This test uses create_collective_group, but the driver doesn't
|
||||
register the backend first, so it should fail immediately.
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
actors = [Worker.remote(rank=i) for i in range(2)]
|
||||
|
||||
try:
|
||||
create_collective_group(
|
||||
actors=actors,
|
||||
world_size=2,
|
||||
ranks=[0, 1],
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
print("ERROR: Should have raised an exception for missing registration!")
|
||||
except Exception as e:
|
||||
if "not registered" in str(e):
|
||||
print(
|
||||
"Test passed! Correctly raised error for missing driver registration."
|
||||
)
|
||||
else:
|
||||
print(f"ERROR: Unexpected error: {e}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("Test 1: create_collective_group approach (driver-managed)")
|
||||
print("=" * 60)
|
||||
test_mock_backend_create_group()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 2: init_collective_group approach (worker-managed)")
|
||||
print("=" * 60)
|
||||
test_mock_backend_init_group()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 3: Error handling - worker not registered")
|
||||
print("=" * 60)
|
||||
test_mock_backend_worker_not_registered()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 4: Error handling - driver not registered")
|
||||
print("=" * 60)
|
||||
test_mock_backend_driver_not_registered()
|
||||
@@ -0,0 +1,37 @@
|
||||
import cupy as cp
|
||||
|
||||
import ray
|
||||
import ray.util.collective as collective
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
self.send = cp.ones((4,), dtype=cp.float32)
|
||||
self.recv = cp.zeros((4,), dtype=cp.float32)
|
||||
|
||||
def setup(self, world_size, rank):
|
||||
collective.init_collective_group(world_size, rank, "nccl", "default")
|
||||
return True
|
||||
|
||||
def compute(self):
|
||||
collective.allreduce(self.send, "default")
|
||||
return self.send
|
||||
|
||||
def destroy(self):
|
||||
collective.destroy_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
send = cp.ones((4,), dtype=cp.float32)
|
||||
ray.init(num_gpus=2)
|
||||
num_workers = 2
|
||||
workers = []
|
||||
init_rets = []
|
||||
for i in range(num_workers):
|
||||
w = Worker.remote()
|
||||
workers.append(w)
|
||||
init_rets.append(w.setup.remote(num_workers, i))
|
||||
_ = ray.get(init_rets)
|
||||
results = ray.get([w.compute.remote() for w in workers])
|
||||
ray.shutdown()
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
import cupy as cp
|
||||
|
||||
import ray
|
||||
import ray.util.collective as collective
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
self.send = cp.ones((4,), dtype=cp.float32)
|
||||
|
||||
def compute(self):
|
||||
collective.allreduce(self.send, "177")
|
||||
return self.send
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ray.init(num_gpus=2)
|
||||
|
||||
num_workers = 2
|
||||
workers = []
|
||||
for i in range(num_workers):
|
||||
w = Worker.remote()
|
||||
workers.append(w)
|
||||
_options = {
|
||||
"group_name": "177",
|
||||
"world_size": 2,
|
||||
"ranks": [0, 1],
|
||||
"backend": "nccl",
|
||||
}
|
||||
collective.create_collective_group(workers, **_options)
|
||||
results = ray.get([w.compute.remote() for w in workers])
|
||||
ray.shutdown()
|
||||
@@ -0,0 +1,43 @@
|
||||
import cupy as cp
|
||||
from cupy.cuda import Device
|
||||
|
||||
import ray
|
||||
import ray.util.collective as collective
|
||||
|
||||
|
||||
@ray.remote(num_gpus=2)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
with Device(0):
|
||||
self.send1 = cp.ones((4,), dtype=cp.float32)
|
||||
with Device(1):
|
||||
self.send2 = cp.ones((4,), dtype=cp.float32) * 2
|
||||
|
||||
self.recv = cp.zeros((4,), dtype=cp.float32)
|
||||
|
||||
def setup(self, world_size, rank):
|
||||
collective.init_collective_group(world_size, rank, "nccl", "177")
|
||||
return True
|
||||
|
||||
def compute(self):
|
||||
collective.allreduce_multigpu([self.send1, self.send2], "177")
|
||||
return [self.send1, self.send2], self.send1.device, self.send2.device
|
||||
|
||||
def destroy(self):
|
||||
collective.destroy_collective_group("177")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ray.init(address="auto")
|
||||
num_workers = 2
|
||||
workers = []
|
||||
init_rets = []
|
||||
for i in range(num_workers):
|
||||
w = Worker.remote()
|
||||
workers.append(w)
|
||||
init_rets.append(w.setup.remote(num_workers, i))
|
||||
a = ray.get(init_rets)
|
||||
results = ray.get([w.compute.remote() for w in workers])
|
||||
print(results)
|
||||
ray.get([w.destroy.remote() for w in workers])
|
||||
ray.shutdown()
|
||||
@@ -0,0 +1,53 @@
|
||||
import cupy as cp
|
||||
from cupy.cuda import Device
|
||||
|
||||
import ray
|
||||
import ray.util.collective as collective
|
||||
|
||||
|
||||
@ray.remote(num_gpus=2)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
with Device(0):
|
||||
self.send1 = cp.ones((4,), dtype=cp.float32)
|
||||
with Device(1):
|
||||
self.send2 = cp.ones((4,), dtype=cp.float32) * 2
|
||||
|
||||
with Device(0):
|
||||
self.recv1 = cp.zeros((4,), dtype=cp.float32)
|
||||
with Device(1):
|
||||
self.recv2 = cp.zeros((4,), dtype=cp.float32)
|
||||
self.rank = -1
|
||||
|
||||
def setup(self, world_size, rank):
|
||||
self.rank = rank
|
||||
collective.init_collective_group(world_size, rank, "nccl", "8")
|
||||
return True
|
||||
|
||||
def compute(self):
|
||||
if self.rank == 0:
|
||||
with Device(0):
|
||||
collective.send_multigpu(self.send1 * 2, 1, 1, "8")
|
||||
else:
|
||||
# with Device(1):
|
||||
collective.recv_multigpu(self.recv2, 0, 0, "8")
|
||||
return self.recv2
|
||||
|
||||
def destroy(self):
|
||||
collective.destroy_collective_group("8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ray.init(address="auto")
|
||||
num_workers = 2
|
||||
workers = []
|
||||
init_rets = []
|
||||
for i in range(num_workers):
|
||||
w = Worker.remote()
|
||||
workers.append(w)
|
||||
init_rets.append(w.setup.remote(num_workers, i))
|
||||
a = ray.get(init_rets)
|
||||
results = ray.get([w.compute.remote() for w in workers])
|
||||
print(results)
|
||||
ray.get([w.destroy.remote() for w in workers])
|
||||
ray.shutdown()
|
||||
@@ -0,0 +1 @@
|
||||
cupy-cuda100
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Some fixtures for collective tests."""
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
try:
|
||||
from ray.util.collective.collective_group.nccl_collective_group import (
|
||||
_get_comm_key_from_devices,
|
||||
_get_comm_key_send_recv,
|
||||
)
|
||||
except Exception: # Cupy/NCCL may be unavailable on CPU-only setups
|
||||
_get_comm_key_from_devices = None
|
||||
_get_comm_key_send_recv = None
|
||||
from ray.util.collective.const import get_store_name
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel("INFO")
|
||||
|
||||
|
||||
# TODO (Hao): remove this clean_up function as it sometimes crashes Ray.
|
||||
def clean_up():
|
||||
# If NCCL helpers are unavailable (e.g., no cupy), skip cleanup.
|
||||
if _get_comm_key_from_devices is None or _get_comm_key_send_recv is None:
|
||||
return
|
||||
group_names = ["default", "test", "123?34!", "default2", "random"]
|
||||
group_names.extend([str(i) for i in range(10)])
|
||||
max_world_size = 4
|
||||
all_keys = []
|
||||
for name in group_names:
|
||||
devices = [[0], [0, 1], [1, 0]]
|
||||
for d in devices:
|
||||
collective_communicator_key = _get_comm_key_from_devices(d)
|
||||
all_keys.append(collective_communicator_key + "@" + name)
|
||||
for i in range(max_world_size):
|
||||
for j in range(max_world_size):
|
||||
if i < j:
|
||||
p2p_communicator_key = _get_comm_key_send_recv(i, 0, j, 0)
|
||||
all_keys.append(p2p_communicator_key + "@" + name)
|
||||
for group_key in all_keys:
|
||||
store_name = get_store_name(group_key)
|
||||
try:
|
||||
actor = ray.get_actor(store_name)
|
||||
except ValueError:
|
||||
actor = None
|
||||
if actor:
|
||||
logger.debug(
|
||||
"Killing actor with group_key: '{}' and store: '{}'.".format(
|
||||
group_key, store_name
|
||||
)
|
||||
)
|
||||
ray.kill(actor)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_single_node_2_gpus():
|
||||
# Please start this fixture in a cluster with 2 GPUs.
|
||||
address_info = ray.init(num_gpus=2)
|
||||
yield address_info
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
# Hao: this fixture is a bit tricky.
|
||||
# I use a bash script to start a ray cluster on
|
||||
# my own on-premise cluster before run this fixture.
|
||||
@pytest.fixture
|
||||
def ray_start_distributed_2_nodes_4_gpus():
|
||||
# The cluster has a setup of 2 nodes, each node with 2
|
||||
# GPUs. Each actor will be allocated 1 GPU.
|
||||
ray.init("auto")
|
||||
yield
|
||||
clean_up()
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_distributed_multigpu_2_nodes_4_gpus():
|
||||
# The cluster has a setup of 2 nodes, each node with 2
|
||||
# GPUs. Each actor will be allocated 2 GPUs.
|
||||
ray.init("auto")
|
||||
yield
|
||||
clean_up()
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_single_node():
|
||||
address_info = ray.init(num_cpus=8)
|
||||
yield address_info
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_distributed_2_nodes():
|
||||
# The cluster has a setup of 2 nodes.
|
||||
# no GPUs!
|
||||
ray.init("auto")
|
||||
yield
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shutdown_only():
|
||||
yield None
|
||||
ray.shutdown()
|
||||
@@ -0,0 +1,137 @@
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import ray
|
||||
import ray.util.collective as col
|
||||
from ray.util.collective.types import Backend, ReduceOp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ray.remote(num_cpus=1)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
self.buffer = None
|
||||
self.list_buffer = None
|
||||
|
||||
def init_tensors(self):
|
||||
self.buffer = np.ones((10,), dtype=np.float32)
|
||||
self.list_buffer = [np.ones((10,), dtype=np.float32) for _ in range(2)]
|
||||
return True
|
||||
|
||||
def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
|
||||
col.init_collective_group(world_size, rank, backend, group_name)
|
||||
return True
|
||||
|
||||
def set_buffer(self, data):
|
||||
self.buffer = data
|
||||
return self.buffer
|
||||
|
||||
def get_buffer(self):
|
||||
return self.buffer
|
||||
|
||||
def set_list_buffer(self, list_of_arrays, copy=False):
|
||||
if copy:
|
||||
copy_list = []
|
||||
for tensor in list_of_arrays:
|
||||
if isinstance(tensor, np.ndarray):
|
||||
copy_list.append(tensor.copy())
|
||||
elif isinstance(tensor, torch.Tensor):
|
||||
copy_list.append(tensor.clone().detach())
|
||||
self.list_buffer = copy_list
|
||||
else:
|
||||
self.list_buffer = list_of_arrays
|
||||
return self.list_buffer
|
||||
|
||||
def do_allreduce(self, group_name="default", op=ReduceOp.SUM):
|
||||
col.allreduce(self.buffer, group_name, op)
|
||||
return self.buffer
|
||||
|
||||
def do_reduce(self, group_name="default", dst_rank=0, op=ReduceOp.SUM):
|
||||
col.reduce(self.buffer, dst_rank, group_name, op)
|
||||
return self.buffer
|
||||
|
||||
def do_broadcast(self, group_name="default", src_rank=0):
|
||||
col.broadcast(self.buffer, src_rank, group_name)
|
||||
return self.buffer
|
||||
|
||||
def do_allgather(self, group_name="default"):
|
||||
col.allgather(self.list_buffer, self.buffer, group_name)
|
||||
return self.list_buffer
|
||||
|
||||
def do_reducescatter(self, group_name="default", op=ReduceOp.SUM):
|
||||
col.reducescatter(self.buffer, self.list_buffer, group_name, op)
|
||||
return self.buffer
|
||||
|
||||
def do_send(self, group_name="default", dst_rank=0):
|
||||
col.send(self.buffer, dst_rank, group_name)
|
||||
return self.buffer
|
||||
|
||||
def do_recv(self, group_name="default", src_rank=0):
|
||||
col.recv(self.buffer, src_rank, group_name)
|
||||
return self.buffer
|
||||
|
||||
def destroy_group(self, group_name="default"):
|
||||
col.destroy_collective_group(group_name)
|
||||
return True
|
||||
|
||||
def report_rank(self, group_name="default"):
|
||||
rank = col.get_rank(group_name)
|
||||
return rank
|
||||
|
||||
def report_world_size(self, group_name="default"):
|
||||
ws = col.get_collective_group_size(group_name)
|
||||
return ws
|
||||
|
||||
def report_nccl_availability(self):
|
||||
avail = col.nccl_available()
|
||||
return avail
|
||||
|
||||
def report_gloo_availability(self):
|
||||
avail = col.gloo_available()
|
||||
return avail
|
||||
|
||||
def report_is_group_initialized(self, group_name="default"):
|
||||
is_init = col.is_group_initialized(group_name)
|
||||
return is_init
|
||||
|
||||
|
||||
def create_collective_workers(num_workers=2, group_name="default", backend="nccl"):
|
||||
actors = [None] * num_workers
|
||||
for i in range(num_workers):
|
||||
actor = Worker.remote()
|
||||
ray.get([actor.init_tensors.remote()])
|
||||
actors[i] = actor
|
||||
world_size = num_workers
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
return actors, init_results
|
||||
|
||||
|
||||
def init_tensors_for_gather_scatter(
|
||||
actors, array_size=10, dtype=np.float32, tensor_backend="numpy"
|
||||
):
|
||||
world_size = len(actors)
|
||||
for i, a in enumerate(actors):
|
||||
if tensor_backend == "numpy":
|
||||
t = np.ones(array_size, dtype=dtype) * (i + 1)
|
||||
elif tensor_backend == "torch":
|
||||
t = torch.ones(array_size, dtype=torch.float32) * (i + 1)
|
||||
else:
|
||||
raise RuntimeError("Unsupported tensor backend.")
|
||||
ray.get([a.set_buffer.remote(t)])
|
||||
if tensor_backend == "numpy":
|
||||
list_buffer = [np.ones(array_size, dtype=dtype) for _ in range(world_size)]
|
||||
elif tensor_backend == "torch":
|
||||
list_buffer = [
|
||||
torch.ones(array_size, dtype=torch.float32) for _ in range(world_size)
|
||||
]
|
||||
else:
|
||||
raise RuntimeError("Unsupported tensor backend.")
|
||||
ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
|
||||
@@ -0,0 +1,143 @@
|
||||
"""Test the allgather API on a distributed Ray cluster."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_allgather_different_array_size(
|
||||
ray_start_distributed_2_nodes, array_size, tensor_backend, backend
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if tensor_backend == "numpy":
|
||||
assert (
|
||||
results[i][j] == np.ones(array_size, dtype=np.float32) * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(array_size, dtype=torch.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
|
||||
def test_allgather_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == np.ones(10, dtype=dtype) * (j + 1)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("length", [0, 1, 3, 4, 7, 8])
|
||||
def test_unmatched_tensor_list_length(ray_start_distributed_2_nodes, length, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
list_buffer = [np.ones(10, dtype=np.float32) for _ in range(length)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
|
||||
if length != world_size:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
|
||||
def test_unmatched_tensor_shape(ray_start_distributed_2_nodes, shape, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(actors, array_size=10)
|
||||
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
|
||||
ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
|
||||
if shape != 10:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allgather_torch_numpy(ray_start_distributed_2_nodes, backend):
|
||||
world_size = 8
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# tensor is pytorch, list is numpy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)).all()
|
||||
|
||||
# tensor is numpy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (
|
||||
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are numpy
|
||||
for i, a in enumerate(actors):
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32))
|
||||
else:
|
||||
list_buffer.append(np.ones(shape, dtype=np.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
assert (
|
||||
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Test the collective allreduice API on a distributed Ray cluster."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend, ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
|
||||
def test_allreduce_different_name(
|
||||
ray_start_distributed_2_nodes, group_name, world_size, backend
|
||||
):
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
def test_allreduce_different_array_size(
|
||||
ray_start_distributed_2_nodes, array_size, backend
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((array_size,), dtype=np.float32) * world_size).all()
|
||||
assert (results[1] == np.ones((array_size,), dtype=np.float32) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_destroy(
|
||||
ray_start_distributed_2_nodes, backend, group_name="default"
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
|
||||
# destroy the group and try do work, should fail
|
||||
ray.get([a.destroy_group.remote() for a in actors])
|
||||
with pytest.raises(RuntimeError):
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
|
||||
# reinit the same group and all reduce
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (
|
||||
results[0] == np.ones((10,), dtype=np.float32) * world_size * world_size
|
||||
).all()
|
||||
assert (
|
||||
results[1] == np.ones((10,), dtype=np.float32) * world_size * world_size
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_multiple_group(ray_start_distributed_2_nodes, backend, num_groups=5):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
for group_name in range(1, num_groups):
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, str(group_name))
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(num_groups):
|
||||
group_name = "default" if i == 0 else str(i)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (
|
||||
results[0] == np.ones((10,), dtype=np.float32) * (world_size ** (i + 1))
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_different_op(ray_start_distributed_2_nodes, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
|
||||
product = 1
|
||||
for i in range(world_size):
|
||||
product = product * (i + 2)
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * product).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * product).all()
|
||||
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * 2).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * 2).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * 9).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * 9).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
|
||||
def test_allreduce_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait([a.set_buffer.remote(np.ones(10, dtype=dtype)) for a in actors])
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=dtype) * world_size).all()
|
||||
assert (results[1] == np.ones((10,), dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_torch_numpy(ray_start_distributed_2_nodes, backend):
|
||||
# import torch
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((10,)) * world_size).all()
|
||||
|
||||
ray.wait(
|
||||
[
|
||||
actors[0].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
ray.wait([actors[1].set_buffer.remote(np.ones(10, dtype=np.float32))])
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
"""Test the collective group APIs."""
|
||||
from random import shuffle
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import Worker, create_collective_workers
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("world_size", [2, 3, 4])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_init_two_actors(
|
||||
ray_start_distributed_2_nodes, world_size, group_name, backend
|
||||
):
|
||||
actors, results = create_collective_workers(world_size, group_name, backend=backend)
|
||||
for i in range(world_size):
|
||||
assert results[i]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("world_size", [2, 3, 4])
|
||||
def test_init_multiple_groups(ray_start_distributed_2_nodes, world_size, backend):
|
||||
num_groups = 5
|
||||
actors = [Worker.remote() for _ in range(world_size)]
|
||||
for i in range(num_groups):
|
||||
group_name = str(i)
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(
|
||||
world_size, i, group_name=group_name, backend=backend
|
||||
)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for j in range(world_size):
|
||||
assert init_results[j]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
|
||||
def test_get_rank(ray_start_distributed_2_nodes, world_size, backend):
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote())
|
||||
assert actor0_rank == 0
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote())
|
||||
assert actor1_rank == 1
|
||||
|
||||
# create a second group with a different name, and different
|
||||
# orders of ranks.
|
||||
new_group_name = "default2"
|
||||
ranks = list(range(world_size))
|
||||
shuffle(ranks)
|
||||
_ = ray.get(
|
||||
[
|
||||
actor.init_group.remote(
|
||||
world_size, ranks[i], group_name=new_group_name, backend=backend
|
||||
)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
|
||||
assert actor0_rank == ranks[0]
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
|
||||
assert actor1_rank == ranks[1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
|
||||
def test_get_world_size(ray_start_distributed_2_nodes, world_size, backend):
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
actor0_world_size = ray.get(actors[0].report_world_size.remote())
|
||||
actor1_world_size = ray.get(actors[1].report_world_size.remote())
|
||||
assert actor0_world_size == actor1_world_size == world_size
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_is_group_initialized(ray_start_distributed_2_nodes, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
# check group is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
|
||||
assert not actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
|
||||
assert not actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
|
||||
assert not actor1_is_init
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_destroy_group(ray_start_distributed_2_nodes, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
# Now destroy the group at actor0
|
||||
ray.wait([actors[0].destroy_group.remote()])
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert not actor0_is_init
|
||||
|
||||
# should go well as the group `random` does not exist at all
|
||||
ray.wait([actors[0].destroy_group.remote("random")])
|
||||
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("random")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("default")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert not actor1_is_init
|
||||
for i in range(2, world_size):
|
||||
ray.wait([actors[i].destroy_group.remote("default")])
|
||||
|
||||
# Now reconstruct the group using the same name
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend=backend)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert init_results[i]
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Test the broadcast API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
|
||||
def test_broadcast_different_name(
|
||||
ray_start_distributed_2_nodes, group_name, src_rank, backend
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones((10,), dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (src_rank + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
|
||||
def test_broadcast_different_array_size(
|
||||
ray_start_distributed_2_nodes, array_size, src_rank, backend
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == np.ones((array_size,), dtype=np.float32) * (src_rank + 2)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
|
||||
def test_broadcast_torch_numpy(ray_start_distributed_2_nodes, src_rank, backend):
|
||||
import torch
|
||||
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
* world_size
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
if src_rank == 0:
|
||||
assert (results[0] == np.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,))).all()
|
||||
else:
|
||||
assert (results[0] == np.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,)) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_broadcast_invalid_rank(ray_start_distributed_2_nodes, backend, src_rank=9):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
with pytest.raises(ValueError):
|
||||
_ = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Test the reduce API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend, ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
|
||||
def test_reduce_different_name(
|
||||
ray_start_distributed_2_nodes, group_name, backend, dst_rank
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
|
||||
def test_reduce_different_array_size(
|
||||
ray_start_distributed_2_nodes, backend, array_size, dst_rank
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (
|
||||
results[i] == np.ones((array_size,), dtype=np.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((array_size,), dtype=np.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
|
||||
def test_reduce_different_op(ray_start_distributed_2_nodes, backend, dst_rank):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
|
||||
)
|
||||
|
||||
product = 1
|
||||
for i in range(world_size):
|
||||
product = product * (i + 2)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * product).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * 2).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (
|
||||
results[i] == np.ones((10,), dtype=np.float32) * (world_size + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
|
||||
def test_reduce_torch_numpy(ray_start_distributed_2_nodes, backend, dst_rank):
|
||||
import torch
|
||||
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.get(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
if dst_rank == 0:
|
||||
assert (results[0] == np.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,))).all()
|
||||
else:
|
||||
assert (results[0] == np.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,))).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_reduce_invalid_rank(ray_start_distributed_2_nodes, backend, dst_rank=9):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
with pytest.raises(ValueError):
|
||||
_ = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
"""Test the collective reducescatter API on a distributed Ray cluster."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_reducescatter_different_array_size(
|
||||
ray_start_distributed_2_nodes, array_size, tensor_backend, backend
|
||||
):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if tensor_backend == "numpy":
|
||||
assert (
|
||||
results[i] == np.ones(array_size, dtype=np.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i] == torch.ones(array_size, dtype=torch.float32) * world_size
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
|
||||
def test_reducescatter_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i] == np.ones(10, dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_reducescatter_torch_numpy(ray_start_distributed_2_nodes, backend):
|
||||
world_size = 8
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# tensor is pytorch, list is numpy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (results[i] == torch.ones(shape, dtype=torch.float32) * world_size).all()
|
||||
|
||||
# tensor is numpy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are numpy
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
else:
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32))
|
||||
else:
|
||||
list_buffer.append(np.ones(shape, dtype=np.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
|
||||
|
||||
# mixed case
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
else:
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(np.ones(shape, dtype=np.float32))
|
||||
else:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Test the send/recv API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1, 3, 6])
|
||||
@pytest.mark.parametrize("src_rank", [0, 2, 4, 7])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
|
||||
)
|
||||
def test_sendrecv(
|
||||
ray_start_distributed_2_nodes, group_name, array_size, src_rank, dst_rank, backend
|
||||
):
|
||||
if src_rank == dst_rank:
|
||||
return
|
||||
world_size = 8
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
ray.get(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 1))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
refs = []
|
||||
for i in range(world_size):
|
||||
refs.append(actors[i].get_buffer.remote())
|
||||
refs[src_rank] = actors[src_rank].do_send.remote(group_name, dst_rank)
|
||||
refs[dst_rank] = actors[dst_rank].do_recv.remote(group_name, src_rank)
|
||||
results = ray.get(refs)
|
||||
assert (
|
||||
results[src_rank] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
|
||||
).all()
|
||||
assert (
|
||||
results[dst_rank] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
|
||||
).all()
|
||||
ray.get([a.destroy_group.remote(group_name) for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Test the allgather API on a distributed Ray cluster."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_allgather_different_array_size(
|
||||
ray_start_distributed_2_nodes_4_gpus, array_size, tensor_backend
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if tensor_backend == "cupy":
|
||||
assert (
|
||||
results[i][j] == cp.ones(array_size, dtype=cp.float32) * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(array_size, dtype=torch.float32).cuda() * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
|
||||
def test_allgather_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == cp.ones(10, dtype=dtype) * (j + 1)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("length", [0, 1, 3, 4, 7, 8])
|
||||
def test_unmatched_tensor_list_length(ray_start_distributed_2_nodes_4_gpus, length):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
list_buffer = [cp.ones(10, dtype=cp.float32) for _ in range(length)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer) for a in actors])
|
||||
if length != world_size:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
|
||||
def test_unmatched_tensor_shape(ray_start_distributed_2_nodes_4_gpus, shape):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(actors, array_size=10)
|
||||
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
|
||||
ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
|
||||
if shape != 10:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
def test_allgather_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
|
||||
world_size = 4
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# tensor is pytorch, list is cupy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)).all()
|
||||
|
||||
# tensor is cupy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (
|
||||
results[i][j] == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
|
||||
).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are cupy
|
||||
for i, a in enumerate(actors):
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
|
||||
else:
|
||||
list_buffer.append(cp.ones(shape, dtype=cp.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,170 @@
|
||||
"""Test the collective allreduice API on a distributed Ray cluster."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
from ray.util.collective.types import ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("world_size", [2, 3, 4])
|
||||
def test_allreduce_different_name(
|
||||
ray_start_distributed_2_nodes_4_gpus, group_name, world_size
|
||||
):
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
def test_allreduce_different_array_size(
|
||||
ray_start_distributed_2_nodes_4_gpus, array_size
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
|
||||
assert (results[1] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
|
||||
|
||||
|
||||
def test_allreduce_destroy(
|
||||
ray_start_distributed_2_nodes_4_gpus, backend="nccl", group_name="default"
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
|
||||
# destroy the group and try do work, should fail
|
||||
ray.get([a.destroy_group.remote() for a in actors])
|
||||
with pytest.raises(RuntimeError):
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
|
||||
# reinit the same group and all reduce
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (
|
||||
results[0] == cp.ones((10,), dtype=cp.float32) * world_size * world_size
|
||||
).all()
|
||||
assert (
|
||||
results[1] == cp.ones((10,), dtype=cp.float32) * world_size * world_size
|
||||
).all()
|
||||
|
||||
|
||||
def test_allreduce_multiple_group(
|
||||
ray_start_distributed_2_nodes_4_gpus, backend="nccl", num_groups=5
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
for group_name in range(1, num_groups):
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, str(group_name))
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(num_groups):
|
||||
group_name = "default" if i == 0 else str(i)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (
|
||||
results[0] == cp.ones((10,), dtype=cp.float32) * (world_size ** (i + 1))
|
||||
).all()
|
||||
|
||||
|
||||
def test_allreduce_different_op(ray_start_distributed_2_nodes_4_gpus):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 120).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 120).all()
|
||||
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 5).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 5).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
|
||||
def test_allreduce_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait([a.set_buffer.remote(cp.ones(10, dtype=dtype)) for a in actors])
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=dtype) * world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
def test_allreduce_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
|
||||
# import torch
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
).cuda()
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,)) * world_size).all()
|
||||
|
||||
ray.wait(
|
||||
[
|
||||
actors[0].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
cp.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
with pytest.raises(RuntimeError):
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
+122
@@ -0,0 +1,122 @@
|
||||
"""Test the collective group APIs."""
|
||||
from random import shuffle
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import Worker, create_collective_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 3, 4])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_init_two_actors(ray_start_distributed_2_nodes_4_gpus, world_size, group_name):
|
||||
actors, results = create_collective_workers(world_size, group_name)
|
||||
for i in range(world_size):
|
||||
assert results[i]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 3, 4])
|
||||
def test_init_multiple_groups(ray_start_distributed_2_nodes_4_gpus, world_size):
|
||||
num_groups = 1
|
||||
actors = [Worker.remote() for _ in range(world_size)]
|
||||
for i in range(num_groups):
|
||||
group_name = str(i)
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, k, group_name=group_name)
|
||||
for k, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for j in range(world_size):
|
||||
assert init_results[j]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 3, 4])
|
||||
def test_get_rank(ray_start_distributed_2_nodes_4_gpus, world_size):
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote())
|
||||
assert actor0_rank == 0
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote())
|
||||
assert actor1_rank == 1
|
||||
|
||||
# create a second group with a different name, and different
|
||||
# orders of ranks.
|
||||
new_group_name = "default2"
|
||||
ranks = list(range(world_size))
|
||||
shuffle(ranks)
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, ranks[i], group_name=new_group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
|
||||
assert actor0_rank == ranks[0]
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
|
||||
assert actor1_rank == ranks[1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 3, 4])
|
||||
def test_get_collective_group_size(ray_start_distributed_2_nodes_4_gpus, world_size):
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
actor0_world_size = ray.get(actors[0].report_world_size.remote())
|
||||
actor1_world_size = ray.get(actors[1].report_world_size.remote())
|
||||
assert actor0_world_size == actor1_world_size == world_size
|
||||
|
||||
|
||||
def test_is_group_initialized(ray_start_distributed_2_nodes_4_gpus):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
# check group is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
|
||||
assert not actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
|
||||
assert not actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
|
||||
assert not actor1_is_init
|
||||
|
||||
|
||||
def test_destroy_group(ray_start_distributed_2_nodes_4_gpus):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
# Now destroy the group at actor0
|
||||
ray.wait([actors[0].destroy_group.remote()])
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert not actor0_is_init
|
||||
|
||||
# should go well as the group `random` does not exist at all
|
||||
ray.wait([actors[0].destroy_group.remote("random")])
|
||||
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("random")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("default")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert not actor1_is_init
|
||||
for i in [2, 3]:
|
||||
ray.wait([actors[i].destroy_group.remote("default")])
|
||||
|
||||
# Now reconstruct the group using the same name
|
||||
init_results = ray.get(
|
||||
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert init_results[i]
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Test the broadcast API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
|
||||
def test_broadcast_different_name(
|
||||
ray_start_distributed_2_nodes_4_gpus, group_name, src_rank
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones((10,), dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (src_rank + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
|
||||
def test_broadcast_different_array_size(
|
||||
ray_start_distributed_2_nodes_4_gpus, array_size, src_rank
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == cp.ones((array_size,), dtype=cp.float32) * (src_rank + 2)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
def test_broadcast_torch_cupy(ray_start_distributed_2_nodes_4_gpus, src_rank):
|
||||
import torch
|
||||
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
).cuda()
|
||||
* world_size
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
if src_rank == 0:
|
||||
assert (results[0] == cp.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda()).all()
|
||||
else:
|
||||
assert (results[0] == cp.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
|
||||
|
||||
|
||||
def test_broadcast_invalid_rank(ray_start_distributed_2_nodes_4_gpus, src_rank=3):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Test the reduce API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
from ray.util.collective.types import ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
|
||||
def test_reduce_different_name(
|
||||
ray_start_distributed_2_nodes_4_gpus, group_name, dst_rank
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
|
||||
def test_reduce_different_array_size(
|
||||
ray_start_distributed_2_nodes_4_gpus, array_size, dst_rank
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (
|
||||
results[i] == cp.ones((array_size,), dtype=cp.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((array_size,), dtype=cp.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
|
||||
def test_reduce_different_op(ray_start_distributed_2_nodes_4_gpus, dst_rank):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 120).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
|
||||
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 5).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_torch_cupy(ray_start_distributed_2_nodes_4_gpus, dst_rank):
|
||||
import torch
|
||||
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
).cuda()
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
if dst_rank == 0:
|
||||
assert (results[0] == cp.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda()).all()
|
||||
else:
|
||||
assert (results[0] == cp.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
|
||||
|
||||
|
||||
def test_reduce_invalid_rank(ray_start_distributed_2_nodes_4_gpus, dst_rank=7):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
+130
@@ -0,0 +1,130 @@
|
||||
"""Test the collective reducescatter API on a distributed Ray cluster."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_reducescatter_different_array_size(
|
||||
ray_start_distributed_2_nodes_4_gpus, array_size, tensor_backend
|
||||
):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if tensor_backend == "cupy":
|
||||
assert (
|
||||
results[i] == cp.ones(array_size, dtype=cp.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i]
|
||||
== torch.ones(array_size, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
|
||||
def test_reducescatter_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i] == cp.ones(10, dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
def test_reducescatter_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
|
||||
world_size = 4
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# tensor is pytorch, list is cupy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
|
||||
# tensor is cupy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are cupy
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
else:
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
|
||||
else:
|
||||
list_buffer.append(cp.ones(shape, dtype=cp.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
|
||||
|
||||
# mixed case
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
else:
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(cp.ones(shape, dtype=cp.float32))
|
||||
else:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Test the send/recv API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
|
||||
)
|
||||
def test_sendrecv(
|
||||
ray_start_distributed_2_nodes_4_gpus, group_name, array_size, src_rank, dst_rank
|
||||
):
|
||||
if src_rank == dst_rank:
|
||||
return
|
||||
world_size = 4
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
ray.get(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 1))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
refs = []
|
||||
for i in range(world_size):
|
||||
refs.append(actors[i].get_buffer.remote())
|
||||
refs[src_rank] = actors[src_rank].do_send.remote(group_name, dst_rank)
|
||||
refs[dst_rank] = actors[dst_rank].do_recv.remote(group_name, src_rank)
|
||||
results = ray.get(refs)
|
||||
assert (
|
||||
results[src_rank] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
|
||||
).all()
|
||||
assert (
|
||||
results[dst_rank] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
|
||||
).all()
|
||||
ray.get([a.destroy_group.remote(group_name) for a in actors])
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
"""Test the allgather API on a distributed Ray cluster."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import (
|
||||
create_collective_multigpu_workers,
|
||||
init_tensors_for_gather_scatter_multigpu,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_allgather_different_array_size(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, tensor_backend
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
init_tensors_for_gather_scatter_multigpu(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
for k in range(actual_world_size):
|
||||
if tensor_backend == "cupy":
|
||||
assert (
|
||||
results[i][j][k] == cp.ones(array_size, dtype=cp.float32)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j][k]
|
||||
== torch.ones(array_size, dtype=torch.float32).cuda(j)
|
||||
).all()
|
||||
|
||||
|
||||
def test_allgather_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
|
||||
# tensor is pytorch, list is cupy
|
||||
for i, a in enumerate(actors):
|
||||
ray.get(
|
||||
[a.set_buffer.remote(shape, tensor_type0="torch", tensor_type1="torch")]
|
||||
)
|
||||
ray.get(
|
||||
[a.set_list_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")]
|
||||
)
|
||||
results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
for k in range(actual_world_size):
|
||||
assert (results[i][j][k] == cp.ones(shape, dtype=cp.float32)).all()
|
||||
|
||||
# tensor is cupy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
ray.get([a.set_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")])
|
||||
ray.get(
|
||||
[
|
||||
a.set_list_buffer.remote(
|
||||
shape, tensor_type0="torch", tensor_type1="torch"
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
for k in range(actual_world_size):
|
||||
assert (
|
||||
results[i][j][k] == torch.ones(shape, dtype=torch.float32).cuda(j)
|
||||
).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
"""Test the collective allreduice API on a distributed Ray cluster."""
|
||||
import logging
|
||||
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_multigpu_workers
|
||||
from ray.util.collective.types import ReduceOp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel("DEBUG")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_allreduce_multigpu_different_name(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, group_name
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(
|
||||
num_workers=world_size, group_name=group_name
|
||||
)
|
||||
results = ray.get([a.do_allreduce_multigpu.remote(group_name) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
def test_allreduce_multigpu_different_array_size(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
ray.get([a.set_buffer.remote(array_size) for a in actors])
|
||||
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
|
||||
assert (
|
||||
results[0] == cp.ones((array_size,), dtype=cp.float32) * actual_world_size
|
||||
).all()
|
||||
assert (
|
||||
results[1] == cp.ones((array_size,), dtype=cp.float32) * actual_world_size
|
||||
).all()
|
||||
|
||||
|
||||
def test_allreduce_multigpu_destroy(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, backend="nccl", group_name="default"
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
|
||||
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
|
||||
|
||||
# destroy the group and try do work, should fail
|
||||
ray.get([a.destroy_group.remote() for a in actors])
|
||||
with pytest.raises(RuntimeError):
|
||||
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
|
||||
|
||||
# reinit the same group and all reduce
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
|
||||
assert (
|
||||
results[0]
|
||||
== cp.ones((10,), dtype=cp.float32) * actual_world_size * actual_world_size
|
||||
).all()
|
||||
assert (
|
||||
results[1]
|
||||
== cp.ones((10,), dtype=cp.float32) * actual_world_size * actual_world_size
|
||||
).all()
|
||||
|
||||
|
||||
def test_allreduce_multigpu_multiple_group(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, backend="nccl", num_groups=5
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
for group_name in range(1, num_groups):
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, str(group_name))
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(num_groups):
|
||||
group_name = "default" if i == 0 else str(i)
|
||||
results = ray.get([a.do_allreduce_multigpu.remote(group_name) for a in actors])
|
||||
assert (
|
||||
results[0]
|
||||
== cp.ones((10,), dtype=cp.float32) * (actual_world_size ** (i + 1))
|
||||
).all()
|
||||
|
||||
|
||||
def test_allreduce_multigpu_different_op(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
# check product
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
|
||||
results = ray.get(
|
||||
[a.do_allreduce_multigpu.remote(op=ReduceOp.PRODUCT) for a in actors]
|
||||
)
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 120).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 120).all()
|
||||
|
||||
# check min
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
|
||||
results = ray.get([a.do_allreduce_multigpu.remote(op=ReduceOp.MIN) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
|
||||
# check max
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
|
||||
results = ray.get([a.do_allreduce_multigpu.remote(op=ReduceOp.MAX) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 5).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 5).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
|
||||
def test_allreduce_multigpu_different_dtype(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, dtype
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
ray.get([a.set_buffer.remote([10], dtype=dtype) for a in actors])
|
||||
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=dtype) * actual_world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=dtype) * actual_world_size).all()
|
||||
|
||||
|
||||
def test_allreduce_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
# import torch
|
||||
world_size = 2
|
||||
actual_world_size = 4
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
ray.get(actors[0].set_buffer.remote([10]))
|
||||
ray.get(
|
||||
actors[1].set_buffer.remote([10], tensor_type0="torch", tensor_type1="torch")
|
||||
)
|
||||
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,)) * actual_world_size).all()
|
||||
|
||||
ray.get(
|
||||
actors[0].set_buffer.remote([10], tensor_type0="cupy", tensor_type1="torch")
|
||||
)
|
||||
ray.get(
|
||||
actors[1].set_buffer.remote([10], tensor_type0="torch", tensor_type1="cupy")
|
||||
)
|
||||
results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,)) * actual_world_size).all()
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
"""Test the collective group APIs."""
|
||||
from random import shuffle
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_multigpu_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_init_two_actors(ray_start_distributed_multigpu_2_nodes_4_gpus, group_name):
|
||||
world_size = 2
|
||||
actors, results = create_collective_multigpu_workers(world_size, group_name)
|
||||
for i in range(world_size):
|
||||
assert results[i]
|
||||
|
||||
|
||||
def test_report_num_gpus(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
world_size = 2
|
||||
actors, results = create_collective_multigpu_workers(world_size)
|
||||
num_gpus = ray.get([actor.report_num_gpus.remote() for actor in actors])
|
||||
assert num_gpus == [2, 2]
|
||||
|
||||
|
||||
def test_get_rank(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote())
|
||||
assert actor0_rank == 0
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote())
|
||||
assert actor1_rank == 1
|
||||
|
||||
# create a second group with a different name, and different
|
||||
# orders of ranks.
|
||||
new_group_name = "default2"
|
||||
ranks = list(range(world_size))
|
||||
shuffle(ranks)
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, ranks[i], group_name=new_group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
|
||||
assert actor0_rank == ranks[0]
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
|
||||
assert actor1_rank == ranks[1]
|
||||
|
||||
|
||||
def test_is_group_initialized(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
# check group is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
|
||||
assert not actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
|
||||
assert not actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
|
||||
assert not actor1_is_init
|
||||
|
||||
|
||||
def test_destroy_group(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
# Now destroy the group at actor0
|
||||
ray.wait([actors[0].destroy_group.remote()])
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert not actor0_is_init
|
||||
|
||||
# should go well as the group `random` does not exist at all
|
||||
ray.wait([actors[0].destroy_group.remote("random")])
|
||||
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("random")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("default")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert not actor1_is_init
|
||||
|
||||
# Now reconstruct the group using the same name
|
||||
init_results = ray.get(
|
||||
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert init_results[i]
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
+112
@@ -0,0 +1,112 @@
|
||||
"""Test the broadcast API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_multigpu_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
@pytest.mark.parametrize("src_gpu_index", [0, 1])
|
||||
def test_broadcast_different_name(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, group_name, src_rank, src_gpu_index
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actors, _ = create_collective_multigpu_workers(
|
||||
num_workers=world_size, group_name=group_name
|
||||
)
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
|
||||
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_broadcast_multigpu.remote(
|
||||
group_name=group_name, src_rank=src_rank, src_gpu_index=src_gpu_index
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
val = (src_rank + 1) * 2 + src_gpu_index
|
||||
assert (results[i][j] == cp.ones([10], dtype=cp.float32) * val).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
@pytest.mark.parametrize("src_gpu_index", [0, 1])
|
||||
def test_broadcast_different_array_size(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, src_rank, src_gpu_index
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
ray.get(actors[0].set_buffer.remote([array_size], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([array_size], value0=4, value1=5))
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_broadcast_multigpu.remote(
|
||||
src_rank=src_rank, src_gpu_index=src_gpu_index
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
val = (src_rank + 1) * 2 + src_gpu_index
|
||||
assert (
|
||||
results[i][j] == cp.ones((array_size,), dtype=cp.float32) * val
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
@pytest.mark.parametrize("src_gpu_index", [0, 1])
|
||||
def test_broadcast_torch_cupy(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, src_rank, src_gpu_index
|
||||
):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(
|
||||
actors[1].set_buffer.remote(
|
||||
[10], value0=4, value1=5, tensor_type0="torch", tensor_type1="torch"
|
||||
)
|
||||
)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_broadcast_multigpu.remote(
|
||||
src_rank=src_rank, src_gpu_index=src_gpu_index
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
val = (src_rank + 1) * 2 + src_gpu_index
|
||||
if i == 0:
|
||||
assert (results[i][j] == cp.ones([10], dtype=cp.float32) * val).all()
|
||||
else:
|
||||
assert (results[i][j] == torch.ones([10]).cuda(j) * val).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("src_rank", [3, 4])
|
||||
@pytest.mark.parametrize("src_gpu_index", [2, 3])
|
||||
def test_broadcast_invalid_rank(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, src_rank, src_gpu_index
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
with pytest.raises(ValueError):
|
||||
_ = ray.get(
|
||||
[
|
||||
a.do_broadcast_multigpu.remote(
|
||||
src_rank=src_rank, src_gpu_index=src_gpu_index
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
+194
@@ -0,0 +1,194 @@
|
||||
"""Test the reduce API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_multigpu_workers
|
||||
from ray.util.collective.types import ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
|
||||
def test_reduce_different_name(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, group_name, dst_rank, dst_gpu_index
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(
|
||||
num_workers=world_size, group_name=group_name
|
||||
)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce_multigpu.remote(
|
||||
group_name, dst_rank=dst_rank, dst_gpu_index=dst_gpu_index
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
if i == dst_rank and j == dst_gpu_index:
|
||||
assert (
|
||||
results[i][j]
|
||||
== cp.ones((10,), dtype=cp.float32) * actual_world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i][j] == cp.ones((10,), dtype=cp.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
|
||||
def test_reduce_different_array_size(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, dst_rank, dst_gpu_index
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(num_workers=world_size)
|
||||
|
||||
ray.get(actors[0].set_buffer.remote(array_size))
|
||||
ray.get(actors[1].set_buffer.remote(array_size))
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce_multigpu.remote(dst_rank=dst_rank, dst_gpu_index=dst_gpu_index)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
if i == dst_rank and j == dst_gpu_index:
|
||||
assert (
|
||||
results[i][j]
|
||||
== cp.ones((array_size,), dtype=cp.float32) * actual_world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i][j] == cp.ones((array_size,), dtype=cp.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
|
||||
def test_reduce_different_op(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
|
||||
# check product
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce_multigpu.remote(
|
||||
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.PRODUCT
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
if i == dst_rank and j == dst_gpu_index:
|
||||
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 120).all()
|
||||
else:
|
||||
val = (i + 1) * 2 + j
|
||||
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
|
||||
|
||||
# check min
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce_multigpu.remote(
|
||||
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.MIN
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
if i == dst_rank and j == dst_gpu_index:
|
||||
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
else:
|
||||
val = (i + 1) * 2 + j
|
||||
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
|
||||
|
||||
# check max
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce_multigpu.remote(
|
||||
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.MAX
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
if i == dst_rank and j == dst_gpu_index:
|
||||
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 5).all()
|
||||
else:
|
||||
val = (i + 1) * 2 + j
|
||||
assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
|
||||
def test_reduce_torch_cupy(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
|
||||
):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
|
||||
ray.get(
|
||||
actors[1].set_buffer.remote(
|
||||
[10], value0=4, value1=5, tensor_type0="torch", tensor_type1="torch"
|
||||
)
|
||||
)
|
||||
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce_multigpu.remote(dst_rank=dst_rank, dst_gpu_index=dst_gpu_index)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
val = (i + 1) * 2 + j
|
||||
if dst_rank == i and dst_gpu_index == j:
|
||||
if i == 0:
|
||||
assert (results[i][j] == cp.ones([10], dtype=cp.float32) * 14).all()
|
||||
else:
|
||||
assert (results[i][j] == torch.ones([10]).cuda(j) * 14).all()
|
||||
else:
|
||||
if i == 0:
|
||||
assert (
|
||||
results[i][j] == cp.ones([10], dtype=cp.float32) * val
|
||||
).all()
|
||||
else:
|
||||
assert (results[i][j] == torch.ones([10]).cuda(j) * val).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [3, 4])
|
||||
@pytest.mark.parametrize("dst_gpu_index", [2, 3])
|
||||
def test_reduce_invalid_rank(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
with pytest.raises(ValueError):
|
||||
_ = ray.get(
|
||||
[
|
||||
a.do_reduce_multigpu.remote(
|
||||
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index
|
||||
)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
"""Test the collective reducescatter API on a distributed Ray cluster."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import (
|
||||
create_collective_multigpu_workers,
|
||||
init_tensors_for_gather_scatter_multigpu,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_reducescatter_different_array_size(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, tensor_backend
|
||||
):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
|
||||
init_tensors_for_gather_scatter_multigpu(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
if tensor_backend == "cupy":
|
||||
assert (
|
||||
results[i][j]
|
||||
== cp.ones(array_size, dtype=cp.float32) * actual_world_size
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(array_size, dtype=torch.float32).cuda(j)
|
||||
* actual_world_size
|
||||
).all()
|
||||
|
||||
|
||||
def test_reducescatter_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
|
||||
world_size = 2
|
||||
num_gpu_per_worker = 2
|
||||
actual_world_size = world_size * num_gpu_per_worker
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_multigpu_workers(world_size)
|
||||
|
||||
# tensor is pytorch, list is cupy
|
||||
for i, a in enumerate(actors):
|
||||
ray.get(
|
||||
[a.set_buffer.remote(shape, tensor_type0="torch", tensor_type1="torch")]
|
||||
)
|
||||
ray.get(
|
||||
[a.set_list_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")]
|
||||
)
|
||||
results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(shape, dtype=torch.float32).cuda(j) * actual_world_size
|
||||
).all()
|
||||
|
||||
# tensor is cupy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
ray.get([a.set_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")])
|
||||
ray.get(
|
||||
[
|
||||
a.set_list_buffer.remote(
|
||||
shape, tensor_type0="torch", tensor_type1="torch"
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(num_gpu_per_worker):
|
||||
assert (
|
||||
results[i][j] == cp.ones(shape, dtype=cp.float32) * actual_world_size
|
||||
).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
"""Test the send/recv API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_multigpu_workers
|
||||
|
||||
|
||||
# @pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
@pytest.mark.parametrize("dst_gpu_index", [0, 1])
|
||||
@pytest.mark.parametrize("src_gpu_index", [0, 1])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
|
||||
)
|
||||
def test_sendrecv(
|
||||
ray_start_distributed_multigpu_2_nodes_4_gpus,
|
||||
array_size,
|
||||
src_rank,
|
||||
dst_rank,
|
||||
src_gpu_index,
|
||||
dst_gpu_index,
|
||||
):
|
||||
if src_rank == dst_rank:
|
||||
return
|
||||
world_size = 2
|
||||
actors, _ = create_collective_multigpu_workers(num_workers=world_size)
|
||||
|
||||
ray.get(actors[0].set_buffer.remote(array_size, value0=2, value1=3))
|
||||
ray.get(actors[1].set_buffer.remote(array_size, value0=4, value1=5))
|
||||
|
||||
refs = []
|
||||
for i in range(world_size):
|
||||
refs.append(actors[i].get_buffer.remote())
|
||||
refs[src_rank][src_gpu_index] = actors[src_rank].do_send_multigpu.remote(
|
||||
dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, src_gpu_index=src_gpu_index
|
||||
)
|
||||
refs[dst_rank][dst_gpu_index] = actors[dst_rank].do_recv_multigpu.remote(
|
||||
src_rank=src_rank, src_gpu_index=src_gpu_index, dst_gpu_index=dst_gpu_index
|
||||
)
|
||||
results = []
|
||||
results_flattend = ray.get(refs[0] + refs[1])
|
||||
results.append([results_flattend[0], results_flattend[1]])
|
||||
results.append([results_flattend[2], results_flattend[3]])
|
||||
assert (
|
||||
results[src_rank][src_gpu_index]
|
||||
== cp.ones(array_size, dtype=cp.float32) * ((src_rank + 1) * 2 + src_gpu_index)
|
||||
).all()
|
||||
assert (
|
||||
results[dst_rank][dst_gpu_index]
|
||||
== cp.ones(array_size, dtype=cp.float32) * ((src_rank + 1) * 2 + src_gpu_index)
|
||||
).all()
|
||||
ray.get([a.destroy_group.remote() for a in actors])
|
||||
@@ -0,0 +1,143 @@
|
||||
"""Test the collective allgather API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_allgather_different_array_size(
|
||||
ray_start_single_node, array_size, tensor_backend, backend
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if tensor_backend == "numpy":
|
||||
assert (
|
||||
results[i][j] == np.ones(array_size, dtype=np.float32) * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(array_size, dtype=torch.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
|
||||
def test_allgather_different_dtype(ray_start_single_node, dtype, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == np.ones(10, dtype=dtype) * (j + 1)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("length", [0, 1, 2, 3])
|
||||
def test_unmatched_tensor_list_length(ray_start_single_node, length, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
list_buffer = [np.ones(10, dtype=np.float32) for _ in range(length)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
|
||||
if length != world_size:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
|
||||
def test_unmatched_tensor_shape(ray_start_single_node, shape, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(actors, array_size=10)
|
||||
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
|
||||
ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
|
||||
if shape != 10:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allgather_torch_numpy(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# tensor is pytorch, list is numpy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)).all()
|
||||
|
||||
# tensor is numpy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (
|
||||
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are numpy
|
||||
for i, a in enumerate(actors):
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32))
|
||||
else:
|
||||
list_buffer.append(np.ones(shape, dtype=np.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
assert (
|
||||
results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,166 @@
|
||||
"""Test the collective allreduice API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend, ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_allreduce_different_name(ray_start_single_node, group_name, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
def test_allreduce_different_array_size(ray_start_single_node, array_size, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((array_size,), dtype=np.float32) * world_size).all()
|
||||
assert (results[1] == np.ones((array_size,), dtype=np.float32) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_destroy(ray_start_single_node, backend, group_name="default"):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
|
||||
# destroy the group and try do work, should fail
|
||||
ray.get([a.destroy_group.remote() for a in actors])
|
||||
with pytest.raises(RuntimeError):
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
|
||||
# reinit the same group and all reduce
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * world_size * 2).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * world_size * 2).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_multiple_group(ray_start_single_node, backend, num_groups=5):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
for group_name in range(1, num_groups):
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, str(group_name))
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(num_groups):
|
||||
group_name = "default" if i == 0 else str(i)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (
|
||||
results[0] == np.ones((10,), dtype=np.float32) * (world_size ** (i + 1))
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_different_op(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * 6).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * 6).all()
|
||||
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * 2).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * 2).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=np.float32) * 3).all()
|
||||
assert (results[1] == np.ones((10,), dtype=np.float32) * 3).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
|
||||
def test_allreduce_different_dtype(ray_start_single_node, dtype, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait([a.set_buffer.remote(np.ones(10, dtype=dtype)) for a in actors])
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((10,), dtype=dtype) * world_size).all()
|
||||
assert (results[1] == np.ones((10,), dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_allreduce_torch_numpy(ray_start_single_node, backend):
|
||||
# import torch
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == np.ones((10,)) * world_size).all()
|
||||
|
||||
ray.wait(
|
||||
[
|
||||
actors[0].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
ray.wait([actors[1].set_buffer.remote(np.ones(10, dtype=np.float32))])
|
||||
ray.get([a.do_allreduce.remote() for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Test the collective group APIs."""
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import Worker, create_collective_workers
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_init_two_actors(ray_start_single_node, group_name, backend):
|
||||
world_size = 2
|
||||
actors, results = create_collective_workers(world_size, group_name, backend=backend)
|
||||
for i in range(world_size):
|
||||
assert results[i]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_init_multiple_groups(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
num_groups = 10
|
||||
actors = [Worker.remote() for i in range(world_size)]
|
||||
for i in range(num_groups):
|
||||
group_name = str(i)
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(
|
||||
world_size, k, group_name=group_name, backend=backend
|
||||
)
|
||||
for k, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for j in range(world_size):
|
||||
assert init_results[j]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_get_rank(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote())
|
||||
assert actor0_rank == 0
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote())
|
||||
assert actor1_rank == 1
|
||||
|
||||
# create a second group with a different name,
|
||||
# and different order of ranks.
|
||||
new_group_name = "default2"
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(
|
||||
world_size,
|
||||
world_size - 1 - i,
|
||||
group_name=new_group_name,
|
||||
backend=backend,
|
||||
)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
|
||||
assert actor0_rank == 1
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
|
||||
assert actor1_rank == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_get_world_size(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
actor0_world_size = ray.get(actors[0].report_world_size.remote())
|
||||
actor1_world_size = ray.get(actors[1].report_world_size.remote())
|
||||
assert actor0_world_size == actor1_world_size == world_size
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_is_group_initialized(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
# check group is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
|
||||
assert not actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
|
||||
assert not actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
|
||||
assert not actor1_is_init
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_destroy_group(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
# Now destroy the group at actor0
|
||||
ray.wait([actors[0].destroy_group.remote()])
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert not actor0_is_init
|
||||
|
||||
# should go well as the group `random` does not exist at all
|
||||
ray.wait([actors[0].destroy_group.remote("random")])
|
||||
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("random")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("default")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert not actor1_is_init
|
||||
|
||||
# Now reconstruct the group using the same name
|
||||
init_results = ray.get(
|
||||
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert init_results[i]
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Test the broadcast API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
def test_broadcast_different_name(ray_start_single_node, group_name, src_rank, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
ray.get(
|
||||
[
|
||||
a.set_buffer.remote(np.ones((10,), dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (src_rank + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
def test_broadcast_different_array_size(
|
||||
ray_start_single_node, array_size, src_rank, backend
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.get(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == np.ones((array_size,), dtype=np.float32) * (src_rank + 2)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
def test_broadcast_torch_numpy(ray_start_single_node, src_rank, backend):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
* world_size
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
if src_rank == 0:
|
||||
assert (results[0] == np.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,))).all()
|
||||
else:
|
||||
assert (results[0] == np.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,)) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_broadcast_invalid_rank(ray_start_single_node, backend, src_rank=3):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,67 @@
|
||||
import time
|
||||
|
||||
import ray
|
||||
import ray.util.collective as col
|
||||
from ray.util.collective.collective_group.torch_gloo_collective_group import (
|
||||
TorchGLOOGroup as GLOOGroup,
|
||||
)
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def init_gloo_group(
|
||||
self, world_size: int, rank: int, group_name: str, gloo_timeout: int = 30000
|
||||
):
|
||||
col.init_collective_group(
|
||||
world_size, rank, Backend.GLOO, group_name, gloo_timeout
|
||||
)
|
||||
return True
|
||||
|
||||
def get_gloo_timeout(self, group_name: str) -> int:
|
||||
g = col.get_group_handle(group_name)
|
||||
# Check if the group is initialized correctly
|
||||
assert isinstance(g, GLOOGroup)
|
||||
return g._gloo_context.getTimeout()
|
||||
|
||||
|
||||
def test_two_groups_in_one_cluster(ray_start_single_node):
|
||||
name1 = "name_1"
|
||||
name2 = "name_2"
|
||||
time1 = 40000
|
||||
time2 = 60000
|
||||
w1 = Worker.remote()
|
||||
ret1 = w1.init_gloo_group.remote(1, 0, name1, time1)
|
||||
w2 = Worker.remote()
|
||||
ret2 = w2.init_gloo_group.remote(1, 0, name2, time2)
|
||||
assert ray.get(ret1)
|
||||
assert ray.get(ret2)
|
||||
assert ray.get(w1.get_gloo_timeout.remote(name1)) == time1
|
||||
assert ray.get(w2.get_gloo_timeout.remote(name2)) == time2
|
||||
|
||||
|
||||
def test_failure_when_initializing(shutdown_only):
|
||||
# job1
|
||||
ray.init()
|
||||
w1 = Worker.remote()
|
||||
ret1 = w1.init_gloo_group.remote(2, 0, "name_1")
|
||||
ray.wait([ret1], timeout=1)
|
||||
time.sleep(5)
|
||||
ray.shutdown()
|
||||
|
||||
# job2
|
||||
ray.init()
|
||||
w2 = Worker.remote()
|
||||
ret2 = w2.init_gloo_group.remote(1, 0, "name_1")
|
||||
assert ray.get(ret2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,167 @@
|
||||
"""Test the reduce API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend, ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_different_name(ray_start_single_node, group_name, dst_rank, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * world_size).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_different_array_size(
|
||||
ray_start_single_node, array_size, dst_rank, backend
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (
|
||||
results[i] == np.ones((array_size,), dtype=np.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((array_size,), dtype=np.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_multiple_group(ray_start_single_node, dst_rank, backend, num_groups=5):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
for group_name in range(1, num_groups):
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, str(group_name))
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(num_groups):
|
||||
group_name = "default" if i == 0 else str(i)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce.remote(dst_rank=dst_rank, group_name=group_name)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for j in range(world_size):
|
||||
if j == dst_rank:
|
||||
assert (results[j] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
|
||||
else:
|
||||
assert (results[j] == np.ones((10,), dtype=np.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_different_op(ray_start_single_node, dst_rank, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * 6).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
|
||||
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * 2).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * 3).all()
|
||||
else:
|
||||
assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_torch_numpy(ray_start_single_node, dst_rank, backend):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
if dst_rank == 0:
|
||||
assert (results[0] == np.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,))).all()
|
||||
else:
|
||||
assert (results[0] == np.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,)) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_reduce_invalid_rank(ray_start_single_node, backend, dst_rank=3):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Test the collective reducescatter API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_reducescatter_different_array_size(
|
||||
ray_start_single_node, array_size, tensor_backend, backend
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if tensor_backend == "numpy":
|
||||
assert (
|
||||
results[i] == np.ones(array_size, dtype=np.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i] == torch.ones(array_size, dtype=torch.float32) * world_size
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
|
||||
def test_reducescatter_different_dtype(ray_start_single_node, dtype, backend):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i] == np.ones(10, dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_reducescatter_torch_numpy(ray_start_single_node, backend):
|
||||
world_size = 2
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
|
||||
# tensor is pytorch, list is numpy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (results[i] == torch.ones(shape, dtype=torch.float32) * world_size).all()
|
||||
|
||||
# tensor is numpy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are numpy
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
else:
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32))
|
||||
else:
|
||||
list_buffer.append(np.ones(shape, dtype=np.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
|
||||
|
||||
# mixed case
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32) * (i + 1)
|
||||
else:
|
||||
t = np.ones(shape, dtype=np.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(np.ones(shape, dtype=np.float32))
|
||||
else:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Test the send/recv API."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.cpu_util import create_collective_workers
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
|
||||
)
|
||||
def test_reduce_different_name(
|
||||
ray_start_single_node, group_name, array_size, dst_rank, backend
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(
|
||||
num_workers=world_size, group_name=group_name, backend=backend
|
||||
)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 1))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
src_rank = 1 - dst_rank
|
||||
refs = []
|
||||
for i, actor in enumerate(actors):
|
||||
if i != dst_rank:
|
||||
ref = actor.do_send.remote(group_name, dst_rank)
|
||||
else:
|
||||
ref = actor.do_recv.remote(group_name, src_rank)
|
||||
refs.append(ref)
|
||||
results = ray.get(refs)
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_sendrecv_torch_numpy(ray_start_single_node, dst_rank, backend):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
* 2
|
||||
)
|
||||
]
|
||||
)
|
||||
src_rank = 1 - dst_rank
|
||||
|
||||
refs = []
|
||||
for i, actor in enumerate(actors):
|
||||
if i != dst_rank:
|
||||
ref = actor.do_send.remote(dst_rank=dst_rank)
|
||||
else:
|
||||
ref = actor.do_recv.remote(src_rank=src_rank)
|
||||
refs.append(ref)
|
||||
results = ray.get(refs)
|
||||
if dst_rank == 0:
|
||||
assert (results[0] == np.ones((10,)) * 2).all()
|
||||
assert (results[1] == torch.ones((10,)) * 2).all()
|
||||
else:
|
||||
assert (results[0] == np.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,))).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", [Backend.GLOO])
|
||||
def test_sendrecv_invalid_rank(ray_start_single_node, backend, dst_rank=3):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size, backend=backend)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_send.remote(dst_rank=dst_rank) for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Test the collective allgather API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_allgather_different_array_size(
|
||||
ray_start_single_node_2_gpus, array_size, tensor_backend
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if tensor_backend == "cupy":
|
||||
assert (
|
||||
results[i][j] == cp.ones(array_size, dtype=cp.float32) * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(array_size, dtype=torch.float32).cuda() * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
|
||||
def test_allgather_different_dtype(ray_start_single_node_2_gpus, dtype):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == cp.ones(10, dtype=dtype) * (j + 1)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("length", [0, 1, 2, 3])
|
||||
def test_unmatched_tensor_list_length(ray_start_single_node_2_gpus, length):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
list_buffer = [cp.ones(10, dtype=cp.float32) for _ in range(length)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer) for a in actors])
|
||||
if length != world_size:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
|
||||
def test_unmatched_tensor_shape(ray_start_single_node_2_gpus, shape):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(actors, array_size=10)
|
||||
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
|
||||
ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
|
||||
if shape != 10:
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
else:
|
||||
ray.get([a.do_allgather.remote() for a in actors])
|
||||
|
||||
|
||||
def test_allgather_torch_cupy(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# tensor is pytorch, list is cupy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)).all()
|
||||
|
||||
# tensor is cupy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (
|
||||
results[i][j] == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
|
||||
).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are cupy
|
||||
for i, a in enumerate(actors):
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
|
||||
else:
|
||||
list_buffer.append(cp.ones(shape, dtype=cp.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_allgather.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
assert (
|
||||
results[i][j]
|
||||
== torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)
|
||||
).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,170 @@
|
||||
"""Test the collective allreduice API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
from ray.util.collective.types import ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_allreduce_different_name(ray_start_single_node_2_gpus, group_name):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
def test_allreduce_different_array_size(ray_start_single_node_2_gpus, array_size):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
|
||||
assert (results[1] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
|
||||
|
||||
|
||||
def test_allreduce_destroy(
|
||||
ray_start_single_node_2_gpus, backend="nccl", group_name="default"
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
|
||||
# destroy the group and try do work, should fail
|
||||
ray.get([a.destroy_group.remote() for a in actors])
|
||||
with pytest.raises(RuntimeError):
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
|
||||
# reinit the same group and all reduce
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size * 2).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size * 2).all()
|
||||
|
||||
|
||||
def test_allreduce_multiple_group(
|
||||
ray_start_single_node_2_gpus, backend="nccl", num_groups=5
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
for group_name in range(1, num_groups):
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, str(group_name))
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(num_groups):
|
||||
group_name = "default" if i == 0 else str(i)
|
||||
results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
|
||||
assert (
|
||||
results[0] == cp.ones((10,), dtype=cp.float32) * (world_size ** (i + 1))
|
||||
).all()
|
||||
|
||||
|
||||
def test_allreduce_different_op(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 6).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 6).all()
|
||||
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=cp.float32) * 3).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=cp.float32) * 3).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
|
||||
def test_allreduce_different_dtype(ray_start_single_node_2_gpus, dtype):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait([a.set_buffer.remote(cp.ones(10, dtype=dtype)) for a in actors])
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,), dtype=dtype) * world_size).all()
|
||||
assert (results[1] == cp.ones((10,), dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
def test_allreduce_torch_cupy(ray_start_single_node_2_gpus):
|
||||
# import torch
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
).cuda()
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
assert (results[0] == cp.ones((10,)) * world_size).all()
|
||||
|
||||
ray.wait(
|
||||
[
|
||||
actors[0].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
cp.ones(
|
||||
10,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
with pytest.raises(RuntimeError):
|
||||
results = ray.get([a.do_allreduce.remote() for a in actors])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,118 @@
|
||||
"""Test the collective group APIs."""
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import Worker, create_collective_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
def test_init_two_actors(ray_start_single_node_2_gpus, group_name):
|
||||
world_size = 2
|
||||
actors, results = create_collective_workers(world_size, group_name)
|
||||
for i in range(world_size):
|
||||
assert results[i]
|
||||
|
||||
|
||||
def test_init_multiple_groups(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
num_groups = 10
|
||||
actors = [Worker.remote() for i in range(world_size)]
|
||||
for i in range(num_groups):
|
||||
group_name = str(i)
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, k, group_name=group_name)
|
||||
for k, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for j in range(world_size):
|
||||
assert init_results[j]
|
||||
|
||||
|
||||
def test_get_rank(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote())
|
||||
assert actor0_rank == 0
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote())
|
||||
assert actor1_rank == 1
|
||||
|
||||
# create a second group with a different name,
|
||||
# and different order of ranks.
|
||||
new_group_name = "default2"
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(
|
||||
world_size, world_size - 1 - i, group_name=new_group_name
|
||||
)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
|
||||
assert actor0_rank == 1
|
||||
actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
|
||||
assert actor1_rank == 0
|
||||
|
||||
|
||||
def test_get_collective_group_size(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
actor0_world_size = ray.get(actors[0].report_world_size.remote())
|
||||
actor1_world_size = ray.get(actors[1].report_world_size.remote())
|
||||
assert actor0_world_size == actor1_world_size == world_size
|
||||
|
||||
|
||||
def test_is_group_initialized(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
# check group is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
|
||||
assert not actor0_is_init
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
|
||||
assert not actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
|
||||
assert not actor1_is_init
|
||||
|
||||
|
||||
def test_destroy_group(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
# Now destroy the group at actor0
|
||||
ray.wait([actors[0].destroy_group.remote()])
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert not actor0_is_init
|
||||
|
||||
# should go well as the group `random` does not exist at all
|
||||
ray.wait([actors[0].destroy_group.remote("random")])
|
||||
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("random")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
ray.wait([actors[1].destroy_group.remote("default")])
|
||||
actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
|
||||
assert not actor1_is_init
|
||||
|
||||
# Now reconstruct the group using the same name
|
||||
init_results = ray.get(
|
||||
[actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert init_results[i]
|
||||
actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor0_is_init
|
||||
actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
|
||||
assert actor1_is_init
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Test the broadcast API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
def test_broadcast_different_name(ray_start_single_node_2_gpus, group_name, src_rank):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones((10,), dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for i in range(world_size):
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (src_rank + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
def test_broadcast_different_array_size(
|
||||
ray_start_single_node_2_gpus, array_size, src_rank
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == cp.ones((array_size,), dtype=cp.float32) * (src_rank + 2)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("src_rank", [0, 1])
|
||||
def test_broadcast_torch_cupy(ray_start_single_node_2_gpus, src_rank):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
).cuda()
|
||||
* world_size
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
if src_rank == 0:
|
||||
assert (results[0] == cp.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda()).all()
|
||||
else:
|
||||
assert (results[0] == cp.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
|
||||
|
||||
|
||||
def test_broadcast_invalid_rank(ray_start_single_node_2_gpus, src_rank=3):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
|
||||
@@ -0,0 +1,151 @@
|
||||
"""Test the reduce API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
from ray.util.collective.types import ReduceOp
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_different_name(ray_start_single_node_2_gpus, group_name, dst_rank):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * world_size).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_different_array_size(
|
||||
ray_start_single_node_2_gpus, array_size, dst_rank
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (
|
||||
results[i] == cp.ones((array_size,), dtype=cp.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((array_size,), dtype=cp.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_multiple_group(ray_start_single_node_2_gpus, dst_rank, num_groups=5):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
for group_name in range(1, num_groups):
|
||||
ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, "nccl", str(group_name))
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
for i in range(num_groups):
|
||||
group_name = "default" if i == 0 else str(i)
|
||||
results = ray.get(
|
||||
[
|
||||
a.do_reduce.remote(dst_rank=dst_rank, group_name=group_name)
|
||||
for a in actors
|
||||
]
|
||||
)
|
||||
for j in range(world_size):
|
||||
if j == dst_rank:
|
||||
assert (results[j] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
|
||||
else:
|
||||
assert (results[j] == cp.ones((10,), dtype=cp.float32)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_different_op(ray_start_single_node_2_gpus, dst_rank):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# check product
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 6).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
|
||||
|
||||
# check min
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 2).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
|
||||
|
||||
# check max
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
results = ray.get(
|
||||
[a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
|
||||
)
|
||||
for i in range(world_size):
|
||||
if i == dst_rank:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * 3).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_reduce_torch_cupy(ray_start_single_node_2_gpus, dst_rank):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
).cuda()
|
||||
)
|
||||
]
|
||||
)
|
||||
results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
if dst_rank == 0:
|
||||
assert (results[0] == cp.ones((10,)) * world_size).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda()).all()
|
||||
else:
|
||||
assert (results[0] == cp.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
|
||||
|
||||
|
||||
def test_reduce_invalid_rank(ray_start_single_node_2_gpus, dst_rank=3):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Test the collective reducescatter API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import (
|
||||
create_collective_workers,
|
||||
init_tensors_for_gather_scatter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
|
||||
)
|
||||
def test_reducescatter_different_array_size(
|
||||
ray_start_single_node_2_gpus, array_size, tensor_backend
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(
|
||||
actors, array_size=array_size, tensor_backend=tensor_backend
|
||||
)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if tensor_backend == "cupy":
|
||||
assert (
|
||||
results[i] == cp.ones(array_size, dtype=cp.float32) * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (
|
||||
results[i]
|
||||
== torch.ones(array_size, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
|
||||
def test_reducescatter_different_dtype(ray_start_single_node_2_gpus, dtype):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
init_tensors_for_gather_scatter(actors, dtype=dtype)
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
for j in range(world_size):
|
||||
assert (results[i] == cp.ones(10, dtype=dtype) * world_size).all()
|
||||
|
||||
|
||||
def test_reducescatter_torch_cupy(ray_start_single_node_2_gpus):
|
||||
world_size = 2
|
||||
shape = [10, 10]
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
|
||||
# tensor is pytorch, list is cupy
|
||||
for i, a in enumerate(actors):
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
|
||||
# tensor is cupy, list is pytorch
|
||||
for i, a in enumerate(actors):
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = [
|
||||
torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
|
||||
]
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
|
||||
|
||||
# some tensors in the list are pytorch, some are cupy
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
else:
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
|
||||
else:
|
||||
list_buffer.append(cp.ones(shape, dtype=cp.float32))
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
|
||||
|
||||
# mixed case
|
||||
for i, a in enumerate(actors):
|
||||
if i % 2 == 0:
|
||||
t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
|
||||
else:
|
||||
t = cp.ones(shape, dtype=cp.float32) * (i + 1)
|
||||
ray.wait([a.set_buffer.remote(t)])
|
||||
list_buffer = []
|
||||
for j in range(world_size):
|
||||
if j % 2 == 0:
|
||||
list_buffer.append(cp.ones(shape, dtype=cp.float32))
|
||||
else:
|
||||
list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
|
||||
ray.wait([a.set_list_buffer.remote(list_buffer)])
|
||||
results = ray.get([a.do_reducescatter.remote() for a in actors])
|
||||
for i in range(world_size):
|
||||
if i % 2 == 0:
|
||||
assert (
|
||||
results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
|
||||
).all()
|
||||
else:
|
||||
assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Test the send/recv API."""
|
||||
import cupy as cp
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.util.collective.tests.util import create_collective_workers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
@pytest.mark.parametrize(
|
||||
"array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
|
||||
)
|
||||
def test_reduce_different_name(
|
||||
ray_start_single_node_2_gpus, group_name, array_size, dst_rank
|
||||
):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
|
||||
ray.wait(
|
||||
[
|
||||
a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 1))
|
||||
for i, a in enumerate(actors)
|
||||
]
|
||||
)
|
||||
src_rank = 1 - dst_rank
|
||||
refs = []
|
||||
for i, actor in enumerate(actors):
|
||||
if i != dst_rank:
|
||||
ref = actor.do_send.remote(group_name, dst_rank)
|
||||
else:
|
||||
ref = actor.do_recv.remote(group_name, src_rank)
|
||||
refs.append(ref)
|
||||
results = ray.get(refs)
|
||||
for i in range(world_size):
|
||||
assert (
|
||||
results[i] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
|
||||
).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dst_rank", [0, 1])
|
||||
def test_sendrecv_torch_cupy(ray_start_single_node_2_gpus, dst_rank):
|
||||
import torch
|
||||
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
ray.wait(
|
||||
[
|
||||
actors[1].set_buffer.remote(
|
||||
torch.ones(
|
||||
10,
|
||||
).cuda()
|
||||
* 2
|
||||
)
|
||||
]
|
||||
)
|
||||
src_rank = 1 - dst_rank
|
||||
|
||||
refs = []
|
||||
for i, actor in enumerate(actors):
|
||||
if i != dst_rank:
|
||||
ref = actor.do_send.remote(dst_rank=dst_rank)
|
||||
else:
|
||||
ref = actor.do_recv.remote(src_rank=src_rank)
|
||||
refs.append(ref)
|
||||
results = ray.get(refs)
|
||||
if dst_rank == 0:
|
||||
assert (results[0] == cp.ones((10,)) * 2).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda() * 2).all()
|
||||
else:
|
||||
assert (results[0] == cp.ones((10,))).all()
|
||||
assert (results[1] == torch.ones((10,)).cuda()).all()
|
||||
|
||||
|
||||
def test_sendrecv_invalid_rank(ray_start_single_node_2_gpus, dst_rank=3):
|
||||
world_size = 2
|
||||
actors, _ = create_collective_workers(world_size)
|
||||
with pytest.raises(ValueError):
|
||||
ray.get([a.do_send.remote(dst_rank=dst_rank) for a in actors])
|
||||
@@ -0,0 +1,373 @@
|
||||
import logging
|
||||
|
||||
import cupy as cp
|
||||
import torch
|
||||
|
||||
import ray
|
||||
import ray.util.collective as col
|
||||
from ray.util.collective.collective_group.nccl_util import get_num_gpus
|
||||
from ray.util.collective.types import Backend, ReduceOp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
self.buffer = None
|
||||
self.list_buffer = None
|
||||
|
||||
def init_tensors(self):
|
||||
self.buffer = cp.ones((10,), dtype=cp.float32)
|
||||
self.list_buffer = [cp.ones((10,), dtype=cp.float32) for _ in range(2)]
|
||||
cp.cuda.Stream.null.synchronize()
|
||||
return True
|
||||
|
||||
def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
|
||||
col.init_collective_group(world_size, rank, backend, group_name)
|
||||
return True
|
||||
|
||||
def set_buffer(self, data):
|
||||
self.buffer = data
|
||||
return self.buffer
|
||||
|
||||
def get_buffer(self):
|
||||
return self.buffer
|
||||
|
||||
def set_list_buffer(self, list_of_arrays):
|
||||
self.list_buffer = list_of_arrays
|
||||
return self.list_buffer
|
||||
|
||||
def do_allreduce(self, group_name="default", op=ReduceOp.SUM):
|
||||
col.allreduce(self.buffer, group_name, op)
|
||||
return self.buffer
|
||||
|
||||
def do_reduce(self, group_name="default", dst_rank=0, op=ReduceOp.SUM):
|
||||
col.reduce(self.buffer, dst_rank, group_name, op)
|
||||
return self.buffer
|
||||
|
||||
def do_broadcast(self, group_name="default", src_rank=0):
|
||||
col.broadcast(self.buffer, src_rank, group_name)
|
||||
return self.buffer
|
||||
|
||||
def do_allgather(self, group_name="default"):
|
||||
col.allgather(self.list_buffer, self.buffer, group_name)
|
||||
return self.list_buffer
|
||||
|
||||
def do_reducescatter(self, group_name="default", op=ReduceOp.SUM):
|
||||
col.reducescatter(self.buffer, self.list_buffer, group_name, op)
|
||||
return self.buffer
|
||||
|
||||
def do_send(self, group_name="default", dst_rank=0):
|
||||
col.send(self.buffer, dst_rank, group_name)
|
||||
return self.buffer
|
||||
|
||||
def do_recv(self, group_name="default", src_rank=0):
|
||||
col.recv(self.buffer, src_rank, group_name)
|
||||
return self.buffer
|
||||
|
||||
def destroy_group(self, group_name="default"):
|
||||
col.destroy_collective_group(group_name)
|
||||
return True
|
||||
|
||||
def report_rank(self, group_name="default"):
|
||||
rank = col.get_rank(group_name)
|
||||
return rank
|
||||
|
||||
def report_world_size(self, group_name="default"):
|
||||
ws = col.get_collective_group_size(group_name)
|
||||
return ws
|
||||
|
||||
def report_nccl_availability(self):
|
||||
avail = col.nccl_available()
|
||||
return avail
|
||||
|
||||
def report_gloo_availability(self):
|
||||
avail = col.gloo_available()
|
||||
return avail
|
||||
|
||||
def report_is_group_initialized(self, group_name="default"):
|
||||
is_init = col.is_group_initialized(group_name)
|
||||
return is_init
|
||||
|
||||
|
||||
def create_collective_workers(num_workers=2, group_name="default", backend="nccl"):
|
||||
actors = [None] * num_workers
|
||||
for i in range(num_workers):
|
||||
actor = Worker.remote()
|
||||
ray.get([actor.init_tensors.remote()])
|
||||
actors[i] = actor
|
||||
world_size = num_workers
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
return actors, init_results
|
||||
|
||||
|
||||
def init_tensors_for_gather_scatter(
|
||||
actors, array_size=10, dtype=cp.float32, tensor_backend="cupy"
|
||||
):
|
||||
world_size = len(actors)
|
||||
for i, a in enumerate(actors):
|
||||
if tensor_backend == "cupy":
|
||||
t = cp.ones(array_size, dtype=dtype) * (i + 1)
|
||||
elif tensor_backend == "torch":
|
||||
t = torch.ones(array_size, dtype=torch.float32).cuda() * (i + 1)
|
||||
else:
|
||||
raise RuntimeError("Unsupported tensor backend.")
|
||||
ray.get([a.set_buffer.remote(t)])
|
||||
if tensor_backend == "cupy":
|
||||
list_buffer = [cp.ones(array_size, dtype=dtype) for _ in range(world_size)]
|
||||
elif tensor_backend == "torch":
|
||||
list_buffer = [
|
||||
torch.ones(array_size, dtype=torch.float32).cuda()
|
||||
for _ in range(world_size)
|
||||
]
|
||||
else:
|
||||
raise RuntimeError("Unsupported tensor backend.")
|
||||
ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
|
||||
|
||||
|
||||
@ray.remote(num_gpus=2)
|
||||
class MultiGPUWorker:
|
||||
def __init__(self):
|
||||
self.buffer0 = None
|
||||
self.buffer1 = None
|
||||
self.list_buffer0 = None
|
||||
self.list_buffer1 = None
|
||||
|
||||
def __del__(self):
|
||||
self.buffer0 = None
|
||||
self.buffer1 = None
|
||||
self.list_buffer0 = None
|
||||
self.list_buffer1 = None
|
||||
|
||||
def init_tensors(self):
|
||||
with cp.cuda.Device(0):
|
||||
self.buffer0 = cp.ones((10,), dtype=cp.float32)
|
||||
self.list_buffer0 = [cp.ones((10,), dtype=cp.float32) for _ in range(4)]
|
||||
with cp.cuda.Device(1):
|
||||
self.buffer1 = cp.ones((10,), dtype=cp.float32)
|
||||
self.list_buffer1 = [cp.ones((10,), dtype=cp.float32) for _ in range(4)]
|
||||
cp.cuda.Stream.null.synchronize()
|
||||
return True
|
||||
|
||||
def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
|
||||
col.init_collective_group(world_size, rank, backend, group_name)
|
||||
return True
|
||||
|
||||
def set_buffer(
|
||||
self,
|
||||
size,
|
||||
value0=1.0,
|
||||
value1=1.0,
|
||||
dtype=cp.float32,
|
||||
tensor_type0="cupy",
|
||||
tensor_type1="cupy",
|
||||
):
|
||||
if tensor_type0 == "cupy":
|
||||
with cp.cuda.Device(0):
|
||||
self.buffer0 = cp.ones(size, dtype=dtype) * value0
|
||||
elif tensor_type0 == "torch":
|
||||
self.buffer0 = torch.ones(size, dtype=torch.float32).cuda(0) * value0
|
||||
else:
|
||||
raise RuntimeError()
|
||||
|
||||
if tensor_type1 == "cupy":
|
||||
with cp.cuda.Device(1):
|
||||
self.buffer1 = cp.ones(size, dtype=dtype) * value1
|
||||
elif tensor_type1 == "torch":
|
||||
self.buffer1 = torch.ones(size, dtype=torch.float32).cuda(1) * value1
|
||||
else:
|
||||
raise RuntimeError()
|
||||
cp.cuda.Device(0).synchronize()
|
||||
cp.cuda.Device(1).synchronize()
|
||||
# cp.cuda.Stream.null.synchronize()
|
||||
return True
|
||||
|
||||
def set_list_buffer(
|
||||
self,
|
||||
size,
|
||||
value0=1.0,
|
||||
value1=1.0,
|
||||
dtype=cp.float32,
|
||||
tensor_type0="cupy",
|
||||
tensor_type1="cupy",
|
||||
):
|
||||
if tensor_type0 == "cupy":
|
||||
with cp.cuda.Device(0):
|
||||
self.list_buffer0 = [
|
||||
cp.ones(size, dtype=dtype) * value0 for _ in range(4)
|
||||
]
|
||||
elif tensor_type0 == "torch":
|
||||
self.list_buffer0 = [
|
||||
torch.ones(size, dtype=torch.float32).cuda(0) * value0 for _ in range(4)
|
||||
]
|
||||
else:
|
||||
raise RuntimeError()
|
||||
|
||||
if tensor_type1 == "cupy":
|
||||
with cp.cuda.Device(1):
|
||||
self.list_buffer1 = [
|
||||
cp.ones(size, dtype=dtype) * value1 for _ in range(4)
|
||||
]
|
||||
elif tensor_type1 == "torch":
|
||||
self.list_buffer1 = [
|
||||
torch.ones(size, dtype=torch.float32).cuda(1) * value1 for _ in range(4)
|
||||
]
|
||||
else:
|
||||
raise RuntimeError()
|
||||
cp.cuda.Device(0).synchronize()
|
||||
cp.cuda.Device(1).synchronize()
|
||||
return True
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def get_buffer(self):
|
||||
return self.buffer0, self.buffer1
|
||||
|
||||
def do_allreduce_multigpu(self, group_name="default", op=ReduceOp.SUM):
|
||||
col.allreduce_multigpu([self.buffer0, self.buffer1], group_name, op)
|
||||
cp.cuda.Device(0).synchronize()
|
||||
cp.cuda.Device(1).synchronize()
|
||||
return self.buffer0
|
||||
|
||||
def do_reduce_multigpu(
|
||||
self, group_name="default", dst_rank=0, dst_gpu_index=0, op=ReduceOp.SUM
|
||||
):
|
||||
col.reduce_multigpu(
|
||||
[self.buffer0, self.buffer1], dst_rank, dst_gpu_index, group_name, op
|
||||
)
|
||||
cp.cuda.Device(0).synchronize()
|
||||
cp.cuda.Device(1).synchronize()
|
||||
return self.buffer0, self.buffer1
|
||||
|
||||
def do_broadcast_multigpu(self, group_name="default", src_rank=0, src_gpu_index=0):
|
||||
col.broadcast_multigpu(
|
||||
[self.buffer0, self.buffer1], src_rank, src_gpu_index, group_name
|
||||
)
|
||||
return self.buffer0, self.buffer1
|
||||
|
||||
def do_allgather_multigpu(self, group_name="default"):
|
||||
col.allgather_multigpu(
|
||||
[self.list_buffer0, self.list_buffer1],
|
||||
[self.buffer0, self.buffer1],
|
||||
group_name,
|
||||
)
|
||||
cp.cuda.Device(0).synchronize()
|
||||
cp.cuda.Device(1).synchronize()
|
||||
return self.list_buffer0, self.list_buffer1
|
||||
|
||||
def do_reducescatter_multigpu(self, group_name="default", op=ReduceOp.SUM):
|
||||
col.reducescatter_multigpu(
|
||||
[self.buffer0, self.buffer1],
|
||||
[self.list_buffer0, self.list_buffer1],
|
||||
group_name,
|
||||
op,
|
||||
)
|
||||
cp.cuda.Device(0).synchronize()
|
||||
cp.cuda.Device(1).synchronize()
|
||||
return self.buffer0, self.buffer1
|
||||
|
||||
def do_send_multigpu(
|
||||
self, group_name="default", dst_rank=0, dst_gpu_index=0, src_gpu_index=0
|
||||
):
|
||||
if src_gpu_index == 0:
|
||||
col.send_multigpu(self.buffer0, dst_rank, dst_gpu_index, group_name)
|
||||
cp.cuda.Device(0).synchronize()
|
||||
return self.buffer0
|
||||
elif src_gpu_index == 1:
|
||||
col.send_multigpu(self.buffer1, dst_rank, dst_gpu_index, group_name)
|
||||
cp.cuda.Device(1).synchronize()
|
||||
return self.buffer1
|
||||
else:
|
||||
raise RuntimeError()
|
||||
|
||||
def do_recv_multigpu(
|
||||
self, group_name="default", src_rank=0, src_gpu_index=0, dst_gpu_index=0
|
||||
):
|
||||
if dst_gpu_index == 0:
|
||||
col.recv_multigpu(self.buffer0, src_rank, src_gpu_index, group_name)
|
||||
cp.cuda.Device(0).synchronize()
|
||||
return self.buffer0
|
||||
elif dst_gpu_index == 1:
|
||||
col.recv_multigpu(self.buffer1, src_rank, src_gpu_index, group_name)
|
||||
cp.cuda.Device(1).synchronize()
|
||||
return self.buffer1
|
||||
else:
|
||||
raise RuntimeError()
|
||||
|
||||
def destroy_group(self, group_name="default"):
|
||||
col.destroy_collective_group(group_name)
|
||||
return True
|
||||
|
||||
def report_rank(self, group_name="default"):
|
||||
rank = col.get_rank(group_name)
|
||||
return rank
|
||||
|
||||
def report_world_size(self, group_name="default"):
|
||||
ws = col.get_collective_group_size(group_name)
|
||||
return ws
|
||||
|
||||
def report_nccl_availability(self):
|
||||
avail = col.nccl_available()
|
||||
return avail
|
||||
|
||||
def report_gloo_availability(self):
|
||||
avail = col.gloo_available()
|
||||
return avail
|
||||
|
||||
def report_is_group_initialized(self, group_name="default"):
|
||||
is_init = col.is_group_initialized(group_name)
|
||||
return is_init
|
||||
|
||||
def report_num_gpus(self):
|
||||
n_gpus = get_num_gpus()
|
||||
return n_gpus
|
||||
|
||||
|
||||
def create_collective_multigpu_workers(
|
||||
num_workers=2, group_name="default", backend="nccl"
|
||||
):
|
||||
actors = [None] * num_workers
|
||||
for i in range(num_workers):
|
||||
actor = MultiGPUWorker.remote()
|
||||
ray.get([actor.set_buffer.remote([10])], timeout=10)
|
||||
ray.get([actor.set_list_buffer.remote([10])], timeout=10)
|
||||
actors[i] = actor
|
||||
world_size = num_workers
|
||||
init_results = ray.get(
|
||||
[
|
||||
actor.init_group.remote(world_size, i, backend, group_name)
|
||||
for i, actor in enumerate(actors)
|
||||
]
|
||||
)
|
||||
return actors, init_results
|
||||
|
||||
|
||||
def init_tensors_for_gather_scatter_multigpu(
|
||||
actors, array_size=10, tensor_backend="cupy"
|
||||
):
|
||||
for i, a in enumerate(actors):
|
||||
if tensor_backend == "cupy":
|
||||
ray.get([a.set_buffer.remote(array_size)])
|
||||
ray.get([a.set_list_buffer.remote(array_size)])
|
||||
elif tensor_backend == "torch":
|
||||
ray.get(
|
||||
[
|
||||
a.set_buffer.remote(
|
||||
array_size, tensor_type0="torch", tensor_type1="torch"
|
||||
)
|
||||
]
|
||||
)
|
||||
ray.get(
|
||||
[
|
||||
a.set_list_buffer.remote(
|
||||
array_size, tensor_type0="torch", tensor_type1="torch"
|
||||
)
|
||||
]
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Unsupported tensor backend.")
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Types conversion between different backends."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
_NUMPY_AVAILABLE = True
|
||||
_TORCH_AVAILABLE = True
|
||||
_CUPY_AVAILABLE = True
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
try:
|
||||
import torch as th # noqa: F401
|
||||
except ImportError:
|
||||
_TORCH_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import cupy as cp # noqa: F401
|
||||
except ImportError:
|
||||
_CUPY_AVAILABLE = False
|
||||
|
||||
|
||||
def cupy_available():
|
||||
return _CUPY_AVAILABLE
|
||||
|
||||
|
||||
def torch_available():
|
||||
return _TORCH_AVAILABLE
|
||||
|
||||
|
||||
class Backend(object):
|
||||
"""A class to represent different backends."""
|
||||
|
||||
NCCL = "NCCL"
|
||||
GLOO = "GLOO"
|
||||
UNRECOGNIZED = "unrecognized"
|
||||
|
||||
def __new__(cls, name: str):
|
||||
upper_name = name.upper()
|
||||
backend = getattr(Backend, upper_name, Backend.UNRECOGNIZED)
|
||||
if backend == Backend.UNRECOGNIZED:
|
||||
if upper_name == "TORCH_GLOO":
|
||||
return Backend.GLOO
|
||||
raise ValueError(
|
||||
"Unrecognized backend: '{}'. Only NCCL and GLOO are supported".format(
|
||||
name
|
||||
)
|
||||
)
|
||||
return backend
|
||||
|
||||
|
||||
class ReduceOp(Enum):
|
||||
SUM = 0
|
||||
PRODUCT = 1
|
||||
MIN = 2
|
||||
MAX = 3
|
||||
|
||||
|
||||
unset_timeout_ms = timedelta(milliseconds=-1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AllReduceOptions:
|
||||
reduceOp = ReduceOp.SUM
|
||||
timeout_ms = unset_timeout_ms
|
||||
|
||||
|
||||
@dataclass
|
||||
class BarrierOptions:
|
||||
timeout_ms = unset_timeout_ms
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReduceOptions:
|
||||
reduceOp = ReduceOp.SUM
|
||||
root_rank = 0
|
||||
root_tensor = 0 # index for multi-gpu reduce operations
|
||||
timeout_ms = unset_timeout_ms
|
||||
|
||||
|
||||
@dataclass
|
||||
class AllGatherOptions:
|
||||
timeout_ms = unset_timeout_ms
|
||||
|
||||
|
||||
#
|
||||
# @dataclass
|
||||
# class GatherOptions:
|
||||
# root_rank = 0
|
||||
# timeout = unset_timeout
|
||||
|
||||
|
||||
@dataclass
|
||||
class BroadcastOptions:
|
||||
root_rank = 0
|
||||
root_tensor = 0
|
||||
timeout_ms = unset_timeout_ms
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReduceScatterOptions:
|
||||
reduceOp = ReduceOp.SUM
|
||||
timeout_ms = unset_timeout_ms
|
||||
|
||||
|
||||
@dataclass
|
||||
class SendOptions:
|
||||
dst_rank = 0
|
||||
dst_gpu_index = 0
|
||||
n_elements = 0
|
||||
timeout_ms = unset_timeout_ms
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecvOptions:
|
||||
src_rank = 0
|
||||
src_gpu_index = 0
|
||||
n_elements = 0
|
||||
unset_timeout_ms = unset_timeout_ms
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Some utility class for Collectives."""
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
import ray
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ray.remote
|
||||
class NCCLUniqueIDStore:
|
||||
"""NCCLUniqueID Store as a named actor class.
|
||||
|
||||
Args:
|
||||
name: the unique name for this named actor.
|
||||
|
||||
Attributes:
|
||||
name: the unique name for this named actor.
|
||||
nccl_id: the NCCLUniqueID held in this store.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
self.nccl_id = None
|
||||
self.event = asyncio.Event()
|
||||
|
||||
async def set_id(self, uid: bytes):
|
||||
"""
|
||||
Initialize the NCCL unique ID for this store.
|
||||
|
||||
Args:
|
||||
uid: the unique ID generated via the NCCL generate_communicator_id API.
|
||||
|
||||
Returns:
|
||||
The NCCL unique ID set.
|
||||
"""
|
||||
self.nccl_id = uid
|
||||
self.event.set()
|
||||
return uid
|
||||
|
||||
async def wait_and_get_id(self):
|
||||
"""Wait for the NCCL unique ID to be set and return it."""
|
||||
await self.event.wait()
|
||||
return self.nccl_id
|
||||
|
||||
def get_id(self):
|
||||
"""Get the NCCL unique ID held in this store."""
|
||||
if not self.nccl_id:
|
||||
logger.warning(
|
||||
"The NCCL ID has not been set yet for store {}.".format(self.name)
|
||||
)
|
||||
return self.nccl_id
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Info:
|
||||
"""Store the group information created via `create_collective_group`.
|
||||
|
||||
Note: Should be used as a NamedActor.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.ids = None
|
||||
self.world_size = -1
|
||||
self.rank = -1
|
||||
self.backend = None
|
||||
self.gloo_timeout = 30000
|
||||
|
||||
def set_info(self, ids, world_size, rank, backend, gloo_timeout):
|
||||
"""Store collective information."""
|
||||
self.ids = ids
|
||||
self.world_size = world_size
|
||||
self.rank = rank
|
||||
self.backend = backend
|
||||
self.gloo_timeout = gloo_timeout
|
||||
|
||||
def get_info(self):
|
||||
"""Get previously stored collective information."""
|
||||
return (
|
||||
self.ids,
|
||||
self.world_size,
|
||||
self.rank,
|
||||
self.backend,
|
||||
self.gloo_timeout,
|
||||
)
|
||||
Reference in New Issue
Block a user