chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user