wehub-resource-sync
381 lines
14 KiB
Python
381 lines
14 KiB
Python
import logging
|
|
from types import ModuleType
|
|
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
|
|
|
|
import ray
|
|
from ray.exceptions import RayChannelError
|
|
from ray.experimental.channel.accelerator_context import AcceleratorContext
|
|
from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
|
|
from ray.experimental.util.types import ReduceOp
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
|
|
|
|
# Logger for this module. It should be configured at the entry point
|
|
# into the program using Ray. Ray provides a default configuration at
|
|
# entry/init points.
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _NcclGroup(Communicator):
|
|
"""
|
|
Represents an actor's NCCL communicator. This is the default NCCL communicator
|
|
to be used in Compiled Graph if a custom communicator is not provided.
|
|
|
|
This class is not thread-safe.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
world_size: int,
|
|
comm_id: tuple,
|
|
rank: Optional[int],
|
|
actor_handles: List["ray.actor.ActorHandle"],
|
|
cuda_stream: Optional["torch.cuda.Stream"],
|
|
use_communication_streams: bool = False,
|
|
):
|
|
"""
|
|
Initialize a NCCL communicator that can be used to communicate p2p with
|
|
other GPU actors.
|
|
|
|
This method blocks until the same call has been made on all other
|
|
actors in the group, with the same arguments for world_size and
|
|
comm_id.
|
|
|
|
NOTE: A concurrent NCCL group can coexist with this one but using the
|
|
two groups concurrently on different CUDA streams may cause deadlock.
|
|
See
|
|
https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html
|
|
#using-multiple-nccl-communicators-concurrently.
|
|
|
|
If the user can guarantee that all involved actors execute the same ops
|
|
in the same order, then the other NCCL group should use the given
|
|
`cuda_stream`, and there will not be a concurrency issue. Otherwise,
|
|
the other stream needs to synchronize with the given `cuda_stream`
|
|
before and after it launches NCCL ops, e.g., at the beginning and end
|
|
of a DAG task.
|
|
|
|
Args:
|
|
world_size: The number of participating actors/devices.
|
|
comm_id: A unique communicator ID returned by
|
|
cupy.cuda.nccl.get_unique_id().
|
|
rank: The rank of this actor. If None, then the caller is not a
|
|
participant of the NCCL group.
|
|
actor_handles: A list of actor handles, in rank order.
|
|
cuda_stream: A raw CUDA stream to dispatch NCCL ops to. If rank is
|
|
specified, then this must be specified too.
|
|
use_communication_streams: Whether to use dedicated send and recv
|
|
streams for communication. If True, communication and computation
|
|
can be overlapped to improve performance.
|
|
"""
|
|
self._world_size = world_size
|
|
self._rank: Optional[int] = rank
|
|
self.nccl_util: Optional[ModuleType] = None
|
|
self._actor_handles = actor_handles
|
|
self._use_communication_streams = use_communication_streams
|
|
|
|
if rank is not None:
|
|
assert ray.get_gpu_ids(), "NCCL actor has no GPUs assigned"
|
|
assert cuda_stream is not None, "NCCL actor must specify cuda_stream"
|
|
|
|
expected_rank = self.get_rank(ray.get_runtime_context().current_actor)
|
|
assert (
|
|
rank == expected_rank
|
|
), f"NCCL actor's rank {rank} does not match expected rank {expected_rank}"
|
|
|
|
from ray.util.collective.collective_group import nccl_util
|
|
|
|
self.nccl_util = nccl_util
|
|
self._comm = self.nccl_util.NcclCommunicator(world_size, comm_id, rank)
|
|
else:
|
|
# Driver does not have a rank.
|
|
self._comm = None
|
|
|
|
self._cuda_stream: Optional["torch.cuda.Stream"] = None
|
|
self._send_stream: Optional["torch.cuda.Stream"] = None
|
|
self._recv_stream: Optional["torch.cuda.Stream"] = None
|
|
if cuda_stream is not None:
|
|
assert rank is not None, "NCCL actor has no rank assigned"
|
|
self._cuda_stream = cuda_stream
|
|
|
|
if use_communication_streams:
|
|
import torch
|
|
|
|
# TODO(swang): Allow default device to be overridden.
|
|
device = AcceleratorContext.get().get_accelerator_devices()[0]
|
|
|
|
self._send_stream = torch.cuda.Stream(device=device)
|
|
self._recv_stream = torch.cuda.Stream(device=device)
|
|
else:
|
|
self._send_stream = self._cuda_stream
|
|
self._recv_stream = self._cuda_stream
|
|
|
|
self._closed = False
|
|
|
|
def initialize(self, rank: int) -> None:
|
|
# No additional initialization is needed.
|
|
pass
|
|
|
|
def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
|
|
return self._actor_handles
|
|
|
|
def get_rank(self, actor: ray.actor.ActorHandle) -> int:
|
|
"""Return the given actor's rank in the NCCL communicator.
|
|
|
|
Args:
|
|
actor: The actor handle to look up.
|
|
|
|
Returns:
|
|
The rank of ``actor`` within the NCCL group.
|
|
"""
|
|
actor_ids = [a._ray_actor_id for a in self._actor_handles]
|
|
try:
|
|
rank = actor_ids.index(actor._ray_actor_id)
|
|
except ValueError:
|
|
raise ValueError("Actor is not in the NCCL group.")
|
|
return rank
|
|
|
|
def get_self_rank(self) -> Optional[int]:
|
|
"""
|
|
Return this actor's rank.
|
|
"""
|
|
return self._rank
|
|
|
|
def get_world_size(self) -> int:
|
|
"""
|
|
Return the number of ranks in the NCCL communicator.
|
|
"""
|
|
return self._world_size
|
|
|
|
def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
|
|
"""
|
|
Send a torch.Tensor to a peer.
|
|
|
|
This returns when the send kernel has been queued, but the kernel may
|
|
not have completed. Therefore, the caller should ensure that there are
|
|
no concurrent writes to the sent `buf` until the send has finished.
|
|
That is, either all writes should be submitted on the current stream
|
|
(self._cuda_stream) or, if on a different stream, that stream should
|
|
synchronize with the current stream.
|
|
|
|
Args:
|
|
buf: The torch.Tensor to send. It should already be on this
|
|
actor's default device.
|
|
peer_rank: The rank of the actor to send to.
|
|
"""
|
|
if self._closed:
|
|
raise RayChannelError("NCCL group has been destroyed.")
|
|
|
|
if self._use_communication_streams:
|
|
# We observed that if all recv/compute/send operations run on GPU,
|
|
# since there is no synchronization, the CPU execution loop may be
|
|
# far ahead of the GPU operations and lead to runtime failures.
|
|
# To avoid that, we synchronize on the send stream.
|
|
# TODO(rui): find a better approach
|
|
self._send_stream.synchronize()
|
|
|
|
# TODO(swang): Handle send/recv async NCCL errors such as network
|
|
# failures.
|
|
self._comm.send(
|
|
self.nccl_util.get_tensor_ptr(buf),
|
|
buf.numel(),
|
|
self.nccl_util.get_nccl_tensor_dtype(buf),
|
|
peer_rank,
|
|
self._send_stream.cuda_stream,
|
|
)
|
|
|
|
def recv(
|
|
self,
|
|
shape: Tuple[int],
|
|
dtype: "torch.dtype",
|
|
peer_rank: int,
|
|
allocator: Optional[TorchTensorAllocator] = None,
|
|
) -> "torch.Tensor":
|
|
"""Receive a torch.Tensor from a peer and synchronize the current stream.
|
|
|
|
After this call returns, the receive buffer is safe to read from from
|
|
any stream. An RayChannelError will be raised if an error occurred (e.g.,
|
|
remote actor died), and the buffer is not safe to read.
|
|
|
|
Args:
|
|
shape: The shape of the tensor to receive.
|
|
dtype: The dtype of the tensor to receive.
|
|
peer_rank: The rank of the actor to receive from.
|
|
allocator: A function used to allocate the receive buffer.
|
|
|
|
Returns:
|
|
The tensor received from ``peer_rank``.
|
|
"""
|
|
if self._closed:
|
|
raise RayChannelError("NCCL group has been destroyed.")
|
|
assert allocator is not None, "NCCL group requires a tensor allocator"
|
|
buf = allocator(shape, dtype)
|
|
|
|
if self._use_communication_streams:
|
|
# We observed that if all recv/compute/send operations run on GPU,
|
|
# since there is no synchronization, the CPU execution loop may be
|
|
# far ahead of the GPU operations and lead to runtime failures.
|
|
# To avoid that, we synchronize on the recv stream.
|
|
# TODO(rui): find a better approach
|
|
self._recv_stream.synchronize()
|
|
|
|
self._comm.recv(
|
|
self.nccl_util.get_tensor_ptr(buf),
|
|
buf.numel(),
|
|
self.nccl_util.get_nccl_tensor_dtype(buf),
|
|
peer_rank,
|
|
self._recv_stream.cuda_stream,
|
|
)
|
|
else:
|
|
self._comm.recv(
|
|
self.nccl_util.get_tensor_ptr(buf),
|
|
buf.numel(),
|
|
self.nccl_util.get_nccl_tensor_dtype(buf),
|
|
peer_rank,
|
|
self._recv_stream.cuda_stream,
|
|
)
|
|
|
|
# Buffer values are undefined if NCCL ops are aborted. Therefore, we
|
|
# need to synchronize here and check that the channel is still open to
|
|
# ensure that the receive buffer is valid.
|
|
# TODO(swang): Avoid CUDA synchronization.
|
|
self._cuda_stream.synchronize()
|
|
|
|
if self._closed:
|
|
raise RayChannelError("NCCL group has been destroyed.")
|
|
return buf
|
|
|
|
def _exec_collective(
|
|
self,
|
|
send_buf: "torch.Tensor",
|
|
recv_buf: "torch.Tensor",
|
|
operation: "Callable[..., None]",
|
|
*operation_args,
|
|
):
|
|
if self._closed:
|
|
raise RayChannelError("NCCL group has been destroyed.")
|
|
|
|
assert send_buf.dtype == recv_buf.dtype, (
|
|
"Ray Compiled Graph derived the dtype of recv_buf from send_buf, "
|
|
"so send_buf and recv_buf must have the same dtype. "
|
|
"If you see this error, please file an issue at Ray repository."
|
|
)
|
|
|
|
operation(*operation_args)
|
|
|
|
# Buffer values are undefined if NCCL ops are aborted. Therefore, we
|
|
# need to synchronize here and check that the channel is still open to
|
|
# ensure that the receive buffer is valid.
|
|
# TODO(swang): Avoid CUDA synchronization.
|
|
# TODO(wxdeng): This synchronize will be optional after merging the unify PR.
|
|
self._cuda_stream.synchronize()
|
|
if self._closed:
|
|
raise RayChannelError(
|
|
"NCCL group has been destroyed during allreduce operation. "
|
|
"There may be a dtype mismatch between input tensors from "
|
|
"different ranks."
|
|
)
|
|
|
|
def allgather(
|
|
self,
|
|
send_buf: "torch.Tensor",
|
|
recv_buf: "torch.Tensor",
|
|
):
|
|
operation_args = [
|
|
self.nccl_util.get_tensor_ptr(send_buf),
|
|
self.nccl_util.get_tensor_ptr(recv_buf),
|
|
send_buf.numel(),
|
|
self.nccl_util.get_nccl_tensor_dtype(send_buf),
|
|
self._cuda_stream.cuda_stream,
|
|
]
|
|
self._exec_collective(
|
|
send_buf,
|
|
recv_buf,
|
|
self._comm.allGather,
|
|
*operation_args,
|
|
)
|
|
|
|
def allreduce(
|
|
self,
|
|
send_buf: "torch.Tensor",
|
|
recv_buf: "torch.Tensor",
|
|
op: ReduceOp = ReduceOp.SUM,
|
|
):
|
|
operation_args = [
|
|
self.nccl_util.get_tensor_ptr(send_buf),
|
|
self.nccl_util.get_tensor_ptr(recv_buf),
|
|
send_buf.numel(),
|
|
self.nccl_util.get_nccl_tensor_dtype(send_buf),
|
|
op.value,
|
|
self._cuda_stream.cuda_stream,
|
|
]
|
|
self._exec_collective(
|
|
send_buf,
|
|
recv_buf,
|
|
self._comm.allReduce,
|
|
*operation_args,
|
|
)
|
|
|
|
def reducescatter(
|
|
self,
|
|
send_buf: "torch.Tensor",
|
|
recv_buf: "torch.Tensor",
|
|
op: ReduceOp = ReduceOp.SUM,
|
|
):
|
|
operation_args = [
|
|
self.nccl_util.get_tensor_ptr(send_buf),
|
|
self.nccl_util.get_tensor_ptr(recv_buf),
|
|
recv_buf.numel(),
|
|
self.nccl_util.get_nccl_tensor_dtype(send_buf),
|
|
op.value,
|
|
self._cuda_stream.cuda_stream,
|
|
]
|
|
self._exec_collective(
|
|
send_buf,
|
|
recv_buf,
|
|
self._comm.reduceScatter,
|
|
*operation_args,
|
|
)
|
|
|
|
@property
|
|
def recv_stream(self):
|
|
import torch
|
|
|
|
return torch.cuda.StreamContext(self._recv_stream)
|
|
|
|
@property
|
|
def send_stream(self):
|
|
import torch
|
|
|
|
return torch.cuda.StreamContext(self._send_stream)
|
|
|
|
def destroy(self) -> None:
|
|
"""
|
|
Destroy the NCCL group.
|
|
"""
|
|
if self._closed:
|
|
return
|
|
|
|
self._closed = True
|
|
|
|
if self._comm is not None:
|
|
logger.info(
|
|
"Destructing NCCL group on actor: "
|
|
f"{ray.get_runtime_context().current_actor}"
|
|
)
|
|
# Abort *after* setting the _closed flag. This ensures that NCCL
|
|
# ops that were blocked on a remote peer will see that the _closed
|
|
# flag is True when they exit from the abort.
|
|
self._comm.abort()
|
|
self._comm.destroy()
|
|
|
|
def get_transport_name(self) -> str:
|
|
return "accelerator"
|
|
|
|
@classmethod
|
|
def generate_communicator_id(cls) -> str:
|
|
from cupy.cuda import nccl
|
|
|
|
return nccl.get_unique_id()
|