diff --git a/python/mscclpp/_core/comm.py b/python/mscclpp/_core/comm.py index d42349dd..8e05ca16 100644 --- a/python/mscclpp/_core/comm.py +++ b/python/mscclpp/_core/comm.py @@ -54,14 +54,21 @@ class CommGroup: import torch import torch.distributed as dist + backend = str(dist.get_backend(torch_group)).lower() + device = torch.device("cuda", torch.cuda.current_device()) if "nccl" in backend else torch.device("cpu") if rank == 0: - uniq_id_global = uniq_id pickled_data = pickle.dumps(uniq_id) - data_tensor = torch.frombuffer(bytearray(pickled_data), dtype=torch.uint8).clone() + size_tensor = torch.tensor([len(pickled_data)], dtype=torch.int64, device=device) else: - data_tensor = torch.zeros(256, dtype=torch.uint8) + size_tensor = torch.zeros(1, dtype=torch.int64, device=device) + dist.broadcast(size_tensor, src=0, group=torch_group) + payload_size = int(size_tensor.item()) + if rank == 0: + data_tensor = torch.frombuffer(bytearray(pickled_data), dtype=torch.uint8).clone().to(device) + else: + data_tensor = torch.zeros(payload_size, dtype=torch.uint8, device=device) dist.broadcast(data_tensor, src=0, group=torch_group) - uniq_id_global = pickle.loads(data_tensor.numpy().tobytes()) + uniq_id_global = pickle.loads(data_tensor.cpu().numpy().tobytes()) self.bootstrap.initialize(uniq_id_global) elif not interfaceIpPortTrio == "": assert rank >= 0 and size >= 1 diff --git a/python/mscclpp/ext/ep/communicator.py b/python/mscclpp/ext/ep/communicator.py index 41bd548a..92fda5f9 100644 --- a/python/mscclpp/ext/ep/communicator.py +++ b/python/mscclpp/ext/ep/communicator.py @@ -26,11 +26,13 @@ Current status (see ``src/ext/ep/README.md``): from __future__ import annotations import os +import pickle from dataclasses import dataclass -from typing import List, Optional +from typing import Optional +import numpy as np import torch -import torch.distributed as dist +from mscclpp._core import CommGroup try: import mscclpp_ep_cpp as _cpp # type: ignore[import-not-found] @@ -45,11 +47,59 @@ Config = _cpp.Config EventHandle = _cpp.EventHandle +def _send_bytes(comm: CommGroup, data: bytes, peer: int, tag: int) -> None: + size = np.array([len(data)], dtype=np.int64) + comm.send(size, peer, tag) + if data: + payload = np.frombuffer(data, dtype=np.uint8).copy() + comm.send(payload, peer, tag + 1) + + +def _recv_bytes(comm: CommGroup, peer: int, tag: int) -> bytes: + size = np.empty(1, dtype=np.int64) + comm.recv(size, peer, tag) + payload_size = int(size[0]) + if payload_size == 0: + return b"" + payload = np.empty(payload_size, dtype=np.uint8) + comm.recv(payload, peer, tag + 1) + return payload.tobytes() + + +def _all_gather_object(comm: CommGroup, obj, tag: int): + rank = comm.my_rank + world_size = comm.nranks + if rank == 0: + values = [obj] + for peer in range(1, world_size): + values.append(pickle.loads(_recv_bytes(comm, peer, tag))) + payload = pickle.dumps(values) + for peer in range(1, world_size): + _send_bytes(comm, payload, peer, tag + 2) + return values + + _send_bytes(comm, pickle.dumps(obj), 0, tag) + return pickle.loads(_recv_bytes(comm, 0, tag + 2)) + + +def _broadcast_object(comm: CommGroup, obj, src: int, tag: int): + rank = comm.my_rank + world_size = comm.nranks + if rank == src: + payload = pickle.dumps(obj) + for peer in range(world_size): + if peer != src: + _send_bytes(comm, payload, peer, tag) + return obj + + return pickle.loads(_recv_bytes(comm, src, tag)) + + @dataclass class MoECommunicatorConfig: """Configuration for the high-level MoE dispatch/combine API.""" - comm: Optional[dist.ProcessGroup] = None + comm: Optional[CommGroup] = None device: Optional[torch.device | int] = None num_experts: int = 0 num_local_experts: Optional[int] = None @@ -120,8 +170,8 @@ class ExpertParallelRuntime: Parameters ---------- - group: - The ``torch.distributed`` process group. Used only for out-of-band + comm: + The :class:`mscclpp.CommGroup`. Used only for out-of-band exchange of IPC handles and the MSCCL++ unique id. num_nvl_bytes: Size of the NVLink-accessible scratch buffer (shared via CUDA IPC). @@ -141,15 +191,15 @@ class ExpertParallelRuntime: def __init__( self, - group: dist.ProcessGroup, + comm: CommGroup, num_nvl_bytes: int = 0, num_rdma_bytes: int = 0, low_latency_mode: bool = False, num_qps_per_rank: int = 12, ) -> None: - self.rank: int = group.rank() - self.group_size: int = group.size() - self.group = group + self.rank: int = comm.my_rank + self.group_size: int = comm.nranks + self.comm = comm self.num_nvl_bytes = num_nvl_bytes self.num_rdma_bytes = num_rdma_bytes self.low_latency_mode = low_latency_mode @@ -160,13 +210,11 @@ class ExpertParallelRuntime: ) # Exchange device IDs + IPC handles + (for RDMA) the MSCCL++ unique id. - device_ids: List[Optional[int]] = [None] * self.group_size local_device_id = self._cpp_buffer.get_local_device_id() - dist.all_gather_object(device_ids, local_device_id, group) + device_ids = _all_gather_object(comm, local_device_id, tag=1000) - ipc_handles: List[Optional[bytes]] = [None] * self.group_size local_ipc_handle = self._cpp_buffer.get_local_ipc_handle() - dist.all_gather_object(ipc_handles, local_ipc_handle, group) + ipc_handles = _all_gather_object(comm, local_ipc_handle, tag=1010) root_unique_id: Optional[bytes] = None # MSCCL++ requires a bootstrapped Communicator even for pure-NVLink @@ -177,9 +225,7 @@ class ExpertParallelRuntime: if self.rank == 0: root_unique_id = self._cpp_buffer.create_unique_id() - broadcast_list = [root_unique_id] - dist.broadcast_object_list(broadcast_list, src=0, group=group) - root_unique_id = broadcast_list[0] + root_unique_id = _broadcast_object(comm, root_unique_id, src=0, tag=1020) assert root_unique_id is not None self._cpp_buffer.connect(root_unique_id) @@ -262,15 +308,13 @@ class MoECommunicator: if config.device is not None: torch.cuda.set_device(config.device) - group = config.comm - if group is None: - if not dist.is_initialized(): - raise ValueError("MoECommunicator requires a process group or an initialized torch.distributed group") - group = dist.group.WORLD + comm = config.comm + if comm is None: + raise ValueError("MoECommunicator requires an mscclpp.CommGroup") - self.group = group - self.rank: int = group.rank() - self.world_size: int = group.size() + self.comm = comm + self.rank: int = comm.my_rank + self.world_size: int = comm.nranks self.local_rank: int = torch.cuda.current_device() self.device = torch.device("cuda", self.local_rank) @@ -328,7 +372,7 @@ class MoECommunicator: self._dispatch_count: Optional[torch.Tensor] = None self._runtime = ExpertParallelRuntime( - group, + comm, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, diff --git a/test/python/ext/ep/test_internode_multirank.py b/test/python/ext/ep/test_internode_multirank.py index c38b0945..51f7901e 100644 --- a/test/python/ext/ep/test_internode_multirank.py +++ b/test/python/ext/ep/test_internode_multirank.py @@ -85,8 +85,11 @@ def inplace_unique(x: torch.Tensor, num_slots: int): def main(): rank, num_ranks, local_rank, group = init_dist() + from mscclpp import CommGroup from mscclpp.ext import ep + ep_group = CommGroup(torch_group=group) + NUM_MAX_NVL_PEERS = _detect_local_world_size() assert ( num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS @@ -162,7 +165,7 @@ def main(): print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True) buf = ep.ExpertParallelRuntime( - group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False + ep_group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False ) print( f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()} " diff --git a/test/python/ext/ep/test_intranode_multirank.py b/test/python/ext/ep/test_intranode_multirank.py index 5860aeb0..ab46a1f8 100644 --- a/test/python/ext/ep/test_intranode_multirank.py +++ b/test/python/ext/ep/test_intranode_multirank.py @@ -65,8 +65,11 @@ def inplace_unique(x: torch.Tensor, num_slots: int): def main(): rank, num_ranks, local_rank, group = init_dist() + from mscclpp import CommGroup from mscclpp.ext import ep + ep_group = CommGroup(torch_group=group) + # Small settings for functional check num_tokens = 128 hidden = 1024 @@ -122,7 +125,7 @@ def main(): ) print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True) - buf = ep.ExpertParallelRuntime(group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False) + buf = ep.ExpertParallelRuntime(ep_group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False) print(f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()}", flush=True) assert buf.is_available() diff --git a/test/python/ext/ep/test_low_latency_multirank.py b/test/python/ext/ep/test_low_latency_multirank.py index 26e09038..822e59fb 100644 --- a/test/python/ext/ep/test_low_latency_multirank.py +++ b/test/python/ext/ep/test_low_latency_multirank.py @@ -78,8 +78,11 @@ def init_dist(): def main(): args = parse_args() rank, num_ranks, local_rank, group = init_dist() + from mscclpp import CommGroup from mscclpp.ext import ep + ep_group = CommGroup(torch_group=group) + # Shrink the "bf16 precision" anchor to keep values small. rank_offset = 128 assert num_ranks - rank_offset < 257, "too many ranks for bf16 precision anchor" @@ -107,7 +110,7 @@ def main(): topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = -1 moe_comm = ep.MoECommunicator( - comm=group, + comm=ep_group, num_experts=num_experts, num_local_experts=num_local_experts, hidden_size=hidden,