chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,21 @@
|
||||
from ray.experimental.collective.collective import (
|
||||
create_collective_group,
|
||||
destroy_all_collective_groups,
|
||||
destroy_collective_group,
|
||||
get_collective_groups,
|
||||
)
|
||||
from ray.experimental.collective.operations import (
|
||||
allgather,
|
||||
allreduce,
|
||||
reducescatter,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"allgather",
|
||||
"allreduce",
|
||||
"reducescatter",
|
||||
"get_collective_groups",
|
||||
"create_collective_group",
|
||||
"destroy_collective_group",
|
||||
"destroy_all_collective_groups",
|
||||
]
|
||||
@@ -0,0 +1,209 @@
|
||||
import threading
|
||||
import uuid
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
import ray
|
||||
import ray.experimental.internal_kv as internal_kv
|
||||
from ray.experimental.collective.communicator import CommunicatorHandle
|
||||
from ray.util.annotations import PublicAPI
|
||||
from ray.util.collective.collective_group.torch_gloo_collective_group import (
|
||||
get_master_address_metadata_key,
|
||||
)
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
_remote_communicator_manager: "Optional[RemoteCommunicatorManager]" = None
|
||||
_remote_communicator_manager_lock = threading.Lock()
|
||||
|
||||
|
||||
class RemoteCommunicatorManager:
|
||||
"""Singleton class to store the mapping between actors and communicators
|
||||
that the actors are a part of.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Handles to communicators that we created. Key is a user-provided
|
||||
# name or UUID.
|
||||
self._remote_communicators: Dict[str, CommunicatorHandle] = {}
|
||||
|
||||
@staticmethod
|
||||
def get() -> "RemoteCommunicatorManager":
|
||||
global _remote_communicator_manager
|
||||
with _remote_communicator_manager_lock:
|
||||
if _remote_communicator_manager is None:
|
||||
_remote_communicator_manager = RemoteCommunicatorManager()
|
||||
return _remote_communicator_manager
|
||||
|
||||
def add_remote_communicator(self, comm_handle: CommunicatorHandle):
|
||||
self._remote_communicators[comm_handle.name] = comm_handle
|
||||
|
||||
def remove_remote_communicator(self, name: str):
|
||||
return self._remote_communicators.pop(name, None)
|
||||
|
||||
def get_collective_groups(
|
||||
self,
|
||||
actors: Optional[List[ray.actor.ActorHandle]] = None,
|
||||
backend: Optional[Backend] = None,
|
||||
):
|
||||
"""
|
||||
Get the collective groups that the given actors are a subset of. Filter by
|
||||
backend if provided.
|
||||
"""
|
||||
actors = actors or []
|
||||
actors = set(actors)
|
||||
|
||||
collectives = []
|
||||
# Find all collective groups that the given actors are a subset
|
||||
# of, with the matching backend if provided.
|
||||
for collective in self._remote_communicators.values():
|
||||
if actors.issubset(set(collective.actors)):
|
||||
if backend is None or collective.backend == backend:
|
||||
collectives.append(collective)
|
||||
return collectives
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def get_collective_groups(
|
||||
actors: List[ray.actor.ActorHandle], backend: Optional[str] = None
|
||||
) -> List[CommunicatorHandle]:
|
||||
"""
|
||||
Get the collective groups that the given actors are a subset of. Filter by
|
||||
backend if provided.
|
||||
|
||||
Args:
|
||||
actors: List of actors. Return handles to all collective groups that
|
||||
these actors are a subset of.
|
||||
backend: An optional backend to filter by. See
|
||||
ray.util.collective.types.Backend for valid backends.
|
||||
|
||||
Returns:
|
||||
A list of communicator handles that the actors are a subset of.
|
||||
"""
|
||||
manager = RemoteCommunicatorManager.get()
|
||||
backend = Backend(backend) if backend is not None else None
|
||||
return manager.get_collective_groups(actors, backend)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def create_collective_group(
|
||||
actors: List[ray.actor.ActorHandle],
|
||||
backend: str,
|
||||
name: Optional[str] = None,
|
||||
) -> CommunicatorHandle:
|
||||
"""Create a collective group on the given list of actors. If this function
|
||||
returns successfully, then the collective group has been initialized on all
|
||||
actors, using the given order of actors as the ranks.
|
||||
|
||||
Currently, an actor can only participate in one collective group per
|
||||
backend at a time. To reuse an actor, destroy its collective group and
|
||||
create a new one.
|
||||
|
||||
Args:
|
||||
actors: The actors to participate in the collective group.
|
||||
backend: The backend to use. See ray.util.collective.types.Backend for
|
||||
valid backends.
|
||||
name: A name to use for the collective group. If None is provided, a
|
||||
random name will be generated.
|
||||
|
||||
Returns:
|
||||
Handle to the communicator.
|
||||
"""
|
||||
manager = RemoteCommunicatorManager.get()
|
||||
|
||||
if name is None:
|
||||
name = str(uuid.uuid4())
|
||||
|
||||
# Validate the backend.
|
||||
backend = Backend(backend)
|
||||
|
||||
world_size = len(actors)
|
||||
|
||||
for actor in actors:
|
||||
if manager.get_collective_groups([actor], backend):
|
||||
raise RuntimeError(
|
||||
f"Actor {actor} already in group for backend {backend}. Actors can currently only participate in at most one group per backend."
|
||||
)
|
||||
|
||||
actor_ids = [actor._ray_actor_id for actor in actors]
|
||||
if len(set(actor_ids)) != len(actor_ids):
|
||||
raise ValueError(f"All actors must be unique, got: {actors}")
|
||||
|
||||
metadata_key = None
|
||||
if backend == Backend.GLOO:
|
||||
metadata_key = get_master_address_metadata_key(name)
|
||||
|
||||
def _do_init_collective_group(self, rank: int):
|
||||
ray.util.collective.init_collective_group(
|
||||
world_size, rank, backend, group_name=name
|
||||
)
|
||||
|
||||
try:
|
||||
init_tasks = [
|
||||
actor.__ray_call__.remote(
|
||||
_do_init_collective_group,
|
||||
rank,
|
||||
)
|
||||
for rank, actor in enumerate(actors)
|
||||
]
|
||||
ray.get(init_tasks)
|
||||
finally:
|
||||
# Clean up the metadata once collective group is initialized
|
||||
# (or failed to initialize).
|
||||
if metadata_key is not None:
|
||||
internal_kv._internal_kv_del(metadata_key)
|
||||
|
||||
# Group was successfully created.
|
||||
comm = CommunicatorHandle(actors, name, backend)
|
||||
manager.add_remote_communicator(comm)
|
||||
return comm
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def destroy_collective_group(group_or_name: Union[CommunicatorHandle, str]):
|
||||
"""
|
||||
Destroy a collective group. If this functions returns successfully, then
|
||||
the actors that were in the collective can be reused to create a new
|
||||
collective group.
|
||||
|
||||
Args:
|
||||
group_or_name: Either a communicator handle or the name of the group to
|
||||
destroy.
|
||||
"""
|
||||
if isinstance(group_or_name, CommunicatorHandle):
|
||||
name = group_or_name.name
|
||||
elif isinstance(group_or_name, str):
|
||||
name = group_or_name
|
||||
else:
|
||||
raise ValueError("Expected CommunicatorHandle or str (group name).")
|
||||
|
||||
manager = RemoteCommunicatorManager.get()
|
||||
group = manager.remove_remote_communicator(name)
|
||||
if group is not None:
|
||||
|
||||
def _do_destroy_collective_group(self):
|
||||
ray.util.collective.destroy_collective_group(name)
|
||||
|
||||
destroy_tasks = [
|
||||
actor.__ray_call__.options(concurrency_group="_ray_system").remote(
|
||||
_do_destroy_collective_group
|
||||
)
|
||||
for actor in group.actors
|
||||
]
|
||||
try:
|
||||
ray.get(destroy_tasks)
|
||||
except ray.exceptions.ActorDiedError:
|
||||
pass
|
||||
else:
|
||||
raise ValueError(f"No group with name {name} found.")
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def destroy_all_collective_groups():
|
||||
"""
|
||||
Destroy all collective groups. This will destroy all collective groups that
|
||||
were previously created by this process. After this function returns, the
|
||||
actors participating in those collective groups can be reused to create a
|
||||
new collective group.
|
||||
"""
|
||||
manager = RemoteCommunicatorManager.get()
|
||||
for collective in manager.get_collective_groups():
|
||||
destroy_collective_group(collective.name)
|
||||
@@ -0,0 +1,63 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
|
||||
import ray
|
||||
from ray.util.collective.types import Backend
|
||||
|
||||
|
||||
@dataclass
|
||||
class Communicator:
|
||||
"""
|
||||
A handle to a communicator that we are a member of.
|
||||
"""
|
||||
|
||||
# The name of the communicator.
|
||||
name: str
|
||||
# Our rank in the collective group.
|
||||
rank: int
|
||||
# A valid backend, as defined by
|
||||
# ray.util.collective.types.Backend.
|
||||
backend: str
|
||||
|
||||
|
||||
class CommunicatorHandle:
|
||||
"""
|
||||
A communicator handle used by the driver to store handles to the
|
||||
actors in the communicator.
|
||||
"""
|
||||
|
||||
def __init__(self, actors: List[ray.actor.ActorHandle], name: str, backend: str):
|
||||
"""
|
||||
Initializes the CommunicatorHandle with the given actor handles.
|
||||
Assumes that the communicator has already been initialized on all actors.
|
||||
|
||||
Args:
|
||||
actors: A list of actor handles to be stored.
|
||||
name: Name of the communicator.
|
||||
backend: Communicator backend. See
|
||||
ray.util.collective.types for valid values.
|
||||
"""
|
||||
self._actors = actors
|
||||
self._name = name
|
||||
self._backend = Backend(backend)
|
||||
|
||||
def get_rank(self, actor: ray.actor.ActorHandle):
|
||||
for i, a in enumerate(self._actors):
|
||||
if a == actor:
|
||||
return i
|
||||
return -1
|
||||
|
||||
@property
|
||||
def actors(self) -> List[ray.actor.ActorHandle]:
|
||||
"""
|
||||
Return all actor handles in this communicator.
|
||||
"""
|
||||
return self._actors[:]
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def backend(self) -> str:
|
||||
return self._backend
|
||||
@@ -0,0 +1,243 @@
|
||||
import uuid
|
||||
from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Type
|
||||
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.experimental.channel.common import ChannelContext
|
||||
from ray.experimental.channel.communicator import (
|
||||
Communicator,
|
||||
ReduceOp,
|
||||
TorchTensorAllocator,
|
||||
)
|
||||
|
||||
|
||||
class AbstractNcclGroup(Communicator):
|
||||
"""
|
||||
A dummy NCCL group for testing.
|
||||
"""
|
||||
|
||||
def __init__(self, actor_handles: List[ray.actor.ActorHandle]):
|
||||
self._actor_handles = actor_handles
|
||||
self._rank = None
|
||||
|
||||
def initialize(self, rank: int) -> None:
|
||||
self._rank = rank
|
||||
|
||||
def get_rank(self, actor: ray.actor.ActorHandle) -> int:
|
||||
return self._actor_handles.index(actor)
|
||||
|
||||
def get_world_size(self) -> int:
|
||||
return len(self._actor_handles)
|
||||
|
||||
def get_self_rank(self) -> Optional[int]:
|
||||
return self._rank
|
||||
|
||||
def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
|
||||
return self._actor_handles
|
||||
|
||||
def send(self, value: "torch.Tensor", peer_rank: int) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def recv(
|
||||
self,
|
||||
shape: Tuple[int],
|
||||
dtype: "torch.dtype",
|
||||
peer_rank: int,
|
||||
allocator: Optional[TorchTensorAllocator] = None,
|
||||
) -> "torch.Tensor":
|
||||
raise NotImplementedError
|
||||
|
||||
def allgather(
|
||||
self,
|
||||
send_buf: "torch.Tensor",
|
||||
recv_buf: "torch.Tensor",
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def allreduce(
|
||||
self,
|
||||
send_buf: "torch.Tensor",
|
||||
recv_buf: "torch.Tensor",
|
||||
op: ReduceOp = ReduceOp.SUM,
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def reducescatter(
|
||||
self,
|
||||
send_buf: "torch.Tensor",
|
||||
recv_buf: "torch.Tensor",
|
||||
op: ReduceOp = ReduceOp.SUM,
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def recv_stream(self):
|
||||
return None
|
||||
|
||||
@property
|
||||
def send_stream(self):
|
||||
return None
|
||||
|
||||
def destroy(self) -> None:
|
||||
pass
|
||||
|
||||
def get_transport_name(self) -> str:
|
||||
return "accelerator"
|
||||
|
||||
@classmethod
|
||||
def generate_communicator_id(cls) -> str:
|
||||
pass
|
||||
|
||||
|
||||
class MockNcclGroupSet:
|
||||
def __init__(self):
|
||||
# Represents a mapping from a NCCL group ID to a set of actors and a custom
|
||||
# NCCL group.
|
||||
self.ids_to_actors_and_custom_comms: Dict[
|
||||
str, Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]]
|
||||
] = {}
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
actors: List["ray.actor.ActorHandle"],
|
||||
custom_nccl_group: Optional[Communicator] = None,
|
||||
use_communication_streams: bool = False,
|
||||
accelerator_module_name: Optional[str] = None,
|
||||
accelerator_communicator_cls: Optional[Type[Communicator]] = None,
|
||||
) -> str:
|
||||
group_id = str(uuid.uuid4())
|
||||
self.ids_to_actors_and_custom_comms[group_id] = (
|
||||
frozenset(actors),
|
||||
custom_nccl_group,
|
||||
)
|
||||
|
||||
if custom_nccl_group is None:
|
||||
ranks = list(range(len(actors)))
|
||||
else:
|
||||
ranks = [custom_nccl_group.get_rank(actor) for actor in actors]
|
||||
init_tasks = [
|
||||
actor.__ray_call__.remote(
|
||||
mock_do_init_nccl_group,
|
||||
group_id,
|
||||
rank,
|
||||
actors,
|
||||
custom_nccl_group,
|
||||
)
|
||||
for rank, actor in zip(ranks, actors)
|
||||
]
|
||||
ray.get(init_tasks, timeout=30)
|
||||
|
||||
ctx = ChannelContext.get_current()
|
||||
if custom_nccl_group is not None:
|
||||
ctx.communicators[group_id] = custom_nccl_group
|
||||
else:
|
||||
ctx.communicators[group_id] = AbstractNcclGroup(actors)
|
||||
|
||||
return group_id
|
||||
|
||||
def mock_destroy_nccl_group(self, group_id: str) -> None:
|
||||
ctx = ChannelContext.get_current()
|
||||
if group_id not in ctx.communicators:
|
||||
return
|
||||
|
||||
actors, _ = self.ids_to_actors_and_custom_comms[group_id]
|
||||
destroy_tasks = [
|
||||
actor.__ray_call__.remote(
|
||||
mock_do_destroy_nccl_group,
|
||||
group_id,
|
||||
)
|
||||
for actor in actors
|
||||
]
|
||||
ray.wait(destroy_tasks, timeout=30)
|
||||
|
||||
if group_id in self.ids_to_actors_and_custom_comms:
|
||||
del self.ids_to_actors_and_custom_comms[group_id]
|
||||
ctx.communicators[group_id].destroy()
|
||||
del ctx.communicators[group_id]
|
||||
|
||||
def check_teardown(self, nccl_group_ids: List[str]) -> None:
|
||||
ctx = ChannelContext.get_current()
|
||||
for nccl_group_id in nccl_group_ids:
|
||||
assert nccl_group_id not in self.ids_to_actors_and_custom_comms
|
||||
assert nccl_group_id not in ctx.communicators
|
||||
|
||||
|
||||
@ray.remote
|
||||
class CPUTorchTensorWorker:
|
||||
def __init__(self):
|
||||
self.device = "cpu"
|
||||
|
||||
def return_tensor(
|
||||
self, size: int, dtype: Optional[torch.dtype] = None
|
||||
) -> torch.Tensor:
|
||||
return torch.ones(size, dtype=dtype, device=self.device)
|
||||
|
||||
def recv(self, tensor: torch.Tensor) -> Tuple[int, int]:
|
||||
assert tensor.device == self.device
|
||||
return tensor.shape, tensor[0]
|
||||
|
||||
def recv_tensors(self, *tensors) -> Tuple[torch.Tensor, ...]:
|
||||
return tuple(tensors)
|
||||
|
||||
|
||||
def mock_do_init_nccl_group(
|
||||
self,
|
||||
group_id: str,
|
||||
rank: int,
|
||||
actors: List[ray.actor.ActorHandle],
|
||||
custom_nccl_group: Optional[Communicator],
|
||||
) -> None:
|
||||
ctx = ChannelContext.get_current()
|
||||
if custom_nccl_group is None:
|
||||
nccl_group = AbstractNcclGroup(actors)
|
||||
nccl_group.initialize(rank)
|
||||
ctx.communicators[group_id] = nccl_group
|
||||
else:
|
||||
custom_nccl_group.initialize(rank)
|
||||
ctx.communicators[group_id] = custom_nccl_group
|
||||
|
||||
|
||||
def mock_do_destroy_nccl_group(self, group_id: str) -> None:
|
||||
ctx = ChannelContext.get_current()
|
||||
if group_id not in ctx.communicators:
|
||||
return
|
||||
ctx.communicators[group_id].destroy()
|
||||
del ctx.communicators[group_id]
|
||||
|
||||
|
||||
def check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag: "ray.dag.DAGNode",
|
||||
actors_and_custom_comms: Set[
|
||||
Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]]
|
||||
],
|
||||
) -> "ray.dag.CompiledDAG":
|
||||
mock_nccl_group_set = MockNcclGroupSet()
|
||||
monkeypatch.setattr(
|
||||
"ray.dag.compiled_dag_node._init_communicator",
|
||||
mock_nccl_group_set,
|
||||
)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
assert (
|
||||
set(mock_nccl_group_set.ids_to_actors_and_custom_comms.values())
|
||||
== actors_and_custom_comms
|
||||
)
|
||||
|
||||
return compiled_dag, mock_nccl_group_set
|
||||
|
||||
|
||||
def check_nccl_group_teardown(
|
||||
monkeypatch,
|
||||
compiled_dag: "ray.dag.CompiledDAG",
|
||||
mock_nccl_group_set: MockNcclGroupSet,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
"ray.dag.compiled_dag_node._destroy_communicator",
|
||||
mock_nccl_group_set.mock_destroy_nccl_group,
|
||||
)
|
||||
|
||||
created_communicator_ids = compiled_dag._actors_to_created_communicator_id.values()
|
||||
compiled_dag.teardown()
|
||||
mock_nccl_group_set.check_teardown(created_communicator_ids)
|
||||
@@ -0,0 +1,203 @@
|
||||
import logging
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import ray
|
||||
from ray.dag.collective_node import CollectiveOutputNode, _CollectiveOperation
|
||||
from ray.dag.constants import (
|
||||
BIND_INDEX_KEY,
|
||||
COLLECTIVE_OPERATION_KEY,
|
||||
IS_CLASS_METHOD_OUTPUT_KEY,
|
||||
PARENT_CLASS_NODE_KEY,
|
||||
)
|
||||
from ray.experimental.channel.torch_tensor_type import Communicator, TorchTensorType
|
||||
from ray.experimental.util.types import (
|
||||
AllGatherOp,
|
||||
AllReduceOp,
|
||||
ReduceOp,
|
||||
ReduceScatterOp,
|
||||
_CollectiveOp,
|
||||
)
|
||||
from ray.util.collective.types import ReduceOp as RayReduceOp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _bind(
|
||||
inputs: Union[List["ray.dag.DAGNode"], List[List["ray.dag.DAGNode"]]],
|
||||
op: _CollectiveOp,
|
||||
transport: Optional[Union[str, Communicator]] = None,
|
||||
):
|
||||
"""
|
||||
Bind inputs (input nodes or lists of input nodes) with a collective operation.
|
||||
The collective operation is applied to each list of input nodes. The output nodes
|
||||
will have the same shape as the input nodes.
|
||||
|
||||
Example of binding a list of input node:
|
||||
with InputNode() as inp:
|
||||
res_comp1 = [actor.comp1.bind(inp) for actor in actors]
|
||||
res_comp2 = [actor.comp2.bind(inp) for actor in actors]
|
||||
res_ar = allreduce.bind([res_comp1, res_comp2])
|
||||
|
||||
Requirements:
|
||||
1. Each input node returns a torch tensor.
|
||||
2. Each input node within a list is from a different actor.
|
||||
3. If lists of input nodes are provided, the order of actors should
|
||||
be the same for each nested list.
|
||||
4. If a custom transport is specified, its actor set matches the actor
|
||||
set of the input nodes.
|
||||
5. If input nodes are provided, then all tensors have the same shape.
|
||||
If lists of input nodes are provided, then all tensors in each
|
||||
list have the same shape.
|
||||
|
||||
Requirements 1-3 are checked in the `CollectiveGroup` constructor.
|
||||
Requirement 4 is not checked yet.
|
||||
|
||||
Args:
|
||||
inputs: A list of DAG nodes or a list of lists of DAG nodes. Each leaf list
|
||||
should contain one object per actor.
|
||||
op: The collective operation.
|
||||
transport: GPU communicator for the collective operation. If not
|
||||
specified, the default ACCELERATOR is used.
|
||||
|
||||
Returns:
|
||||
A list of collective output nodes or a list of lists of collective output nodes,
|
||||
with the same shape as the input nodes. Each output node has the same order and
|
||||
belongs to the same actor as the corresponding input node.
|
||||
"""
|
||||
if isinstance(inputs[0], list) and not isinstance(op, AllReduceOp):
|
||||
raise ValueError(
|
||||
"Currently binding a nested list of dag nodes is only supported for allreduce"
|
||||
)
|
||||
|
||||
# Convert list of DAGNode into nested list for type checking
|
||||
if not isinstance(inputs[0], list):
|
||||
inputs = [inputs]
|
||||
|
||||
if transport is None:
|
||||
transport = TorchTensorType.ACCELERATOR
|
||||
collective_op = _CollectiveOperation(inputs, op, transport)
|
||||
collective_output_nodes: List[CollectiveOutputNode] = []
|
||||
|
||||
if isinstance(op, AllGatherOp):
|
||||
method_name = "allgather"
|
||||
elif isinstance(op, AllReduceOp):
|
||||
method_name = f"allreduce.{op.reduceOp}"
|
||||
elif isinstance(op, ReduceScatterOp):
|
||||
method_name = f"reducescatter.{op.reduceOp}"
|
||||
else:
|
||||
raise ValueError(f"Expected a collective operation, but got {op}")
|
||||
|
||||
for i in range(len(inputs[0])):
|
||||
input_node_list = [l[i] for l in inputs if l]
|
||||
actor_handle: Optional["ray.actor.ActorHandle"] = input_node_list[
|
||||
0
|
||||
]._get_actor_handle()
|
||||
assert actor_handle is not None
|
||||
collective_output_node = CollectiveOutputNode(
|
||||
method_name=method_name,
|
||||
method_args=tuple(input_node_list),
|
||||
method_kwargs=dict(),
|
||||
method_options=dict(),
|
||||
other_args_to_resolve={
|
||||
PARENT_CLASS_NODE_KEY: actor_handle,
|
||||
BIND_INDEX_KEY: actor_handle._ray_dag_bind_index,
|
||||
COLLECTIVE_OPERATION_KEY: collective_op,
|
||||
},
|
||||
)
|
||||
actor_handle._ray_dag_bind_index += 1
|
||||
|
||||
if len(input_node_list) > 1:
|
||||
output_nodes: List[CollectiveOutputNode] = []
|
||||
for i in range(len(input_node_list)):
|
||||
output_node = CollectiveOutputNode(
|
||||
f"return_idx_{i}",
|
||||
(collective_output_node, i),
|
||||
dict(),
|
||||
dict(),
|
||||
{
|
||||
BIND_INDEX_KEY: collective_output_node._get_bind_index(),
|
||||
IS_CLASS_METHOD_OUTPUT_KEY: True,
|
||||
PARENT_CLASS_NODE_KEY: actor_handle,
|
||||
},
|
||||
)
|
||||
output_nodes.append(output_node)
|
||||
collective_output_nodes.append(output_nodes)
|
||||
else:
|
||||
collective_output_nodes.append(collective_output_node)
|
||||
return collective_output_nodes
|
||||
|
||||
|
||||
class AllGatherWrapper:
|
||||
"""Wrapper for NCCL all-gather."""
|
||||
|
||||
def bind(
|
||||
self,
|
||||
input_nodes: List["ray.dag.DAGNode"],
|
||||
transport: Optional[Union[str, Communicator]] = None,
|
||||
) -> List[CollectiveOutputNode]:
|
||||
return _bind(input_nodes, AllGatherOp(), transport)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
tensor_list,
|
||||
tensor,
|
||||
group_name: str = "default",
|
||||
):
|
||||
from ray.util.collective.collective import allgather
|
||||
|
||||
return allgather(tensor_list, tensor, group_name)
|
||||
|
||||
|
||||
class AllReduceWrapper:
|
||||
"""Wrapper for NCCL all-reduce."""
|
||||
|
||||
def bind(
|
||||
self,
|
||||
input_nodes: List["ray.dag.DAGNode"],
|
||||
op: ReduceOp = ReduceOp.SUM,
|
||||
transport: Optional[Union[str, Communicator]] = None,
|
||||
) -> List[CollectiveOutputNode]:
|
||||
if not isinstance(op, ReduceOp):
|
||||
raise ValueError(f"Unexpected operation: {op}")
|
||||
|
||||
return _bind(input_nodes, AllReduceOp(reduceOp=op), transport)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
tensor,
|
||||
group_name: str = "default",
|
||||
op: RayReduceOp = RayReduceOp.SUM,
|
||||
):
|
||||
from ray.util.collective.collective import allreduce
|
||||
|
||||
return allreduce(tensor, group_name, op)
|
||||
|
||||
|
||||
class ReduceScatterWrapper:
|
||||
"""Wrapper for NCCL reduce-scatter."""
|
||||
|
||||
def bind(
|
||||
self,
|
||||
input_nodes: List["ray.dag.DAGNode"],
|
||||
op: ReduceOp = ReduceOp.SUM,
|
||||
transport: Optional[Union[str, Communicator]] = None,
|
||||
) -> List[CollectiveOutputNode]:
|
||||
if not isinstance(op, ReduceOp):
|
||||
raise ValueError(f"Unexpected operation: {op}")
|
||||
|
||||
return _bind(input_nodes, ReduceScatterOp(reduceOp=op), transport)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
tensor,
|
||||
group_name: str = "default",
|
||||
op: RayReduceOp = RayReduceOp.SUM,
|
||||
):
|
||||
from ray.util.collective.collective import reducescatter
|
||||
|
||||
return reducescatter(tensor, group_name, op)
|
||||
|
||||
|
||||
allgather = AllGatherWrapper()
|
||||
allreduce = AllReduceWrapper()
|
||||
reducescatter = ReduceScatterWrapper()
|
||||
Reference in New Issue
Block a user