diff --git a/include/mscclpp/executor.hpp b/include/mscclpp/executor.hpp index 4109844f..b2c6b55b 100644 --- a/include/mscclpp/executor.hpp +++ b/include/mscclpp/executor.hpp @@ -87,6 +87,12 @@ class Executor { void execute(int rank, void* sendbuff, void* recvBuff, size_t sendBuffSize, size_t recvBuffSize, DataType dataType, const ExecutionPlan& plan, cudaStream_t stream, PacketType packetType = PacketType::LL16); + /// Release cached execution contexts while retaining the communicator and default scratch buffer. + /// + /// The caller must ensure that no execution is in flight and that all participating ranks reset collectively before + /// launching another execution. + void reset(); + private: struct Impl; std::unique_ptr impl_; diff --git a/python/csrc/executor_py.cpp b/python/csrc/executor_py.cpp index 350a1e7a..43113b51 100644 --- a/python/csrc/executor_py.cpp +++ b/python/csrc/executor_py.cpp @@ -35,5 +35,7 @@ void register_executor(nb::module_& m) { }, nb::arg("rank"), nb::arg("send_buff"), nb::arg("recv_buff"), nb::arg("send_buff_size"), nb::arg("recv_buff_size"), nb::arg("data_type"), nb::arg("plan"), nb::arg("stream"), - nb::arg("packet_type") = PacketType::LL16); + nb::arg("packet_type") = PacketType::LL16) + .def("reset", &Executor::reset, + "Release cached execution contexts while retaining the communicator and default scratch buffer."); } diff --git a/python/mscclpp/default_algos/__init__.py b/python/mscclpp/default_algos/__init__.py index 32de1ac7..f613d0c7 100644 --- a/python/mscclpp/default_algos/__init__.py +++ b/python/mscclpp/default_algos/__init__.py @@ -3,5 +3,12 @@ from mscclpp.default_algos.allgather_multi_nodes import allgather_multi_nodes from mscclpp.default_algos.allreduce_multi_nodes import allreduce_multi_nodes +from mscclpp.default_algos.reducescatter_multi_nodes import ( + reducescatter_multi_nodes, +) -__all__ = ["allgather_multi_nodes", "allreduce_multi_nodes"] +__all__ = [ + "allgather_multi_nodes", + "allreduce_multi_nodes", + "reducescatter_multi_nodes", +] diff --git a/python/mscclpp/default_algos/reducescatter_multi_nodes.py b/python/mscclpp/default_algos/reducescatter_multi_nodes.py new file mode 100644 index 00000000..9a13a9d6 --- /dev/null +++ b/python/mscclpp/default_algos/reducescatter_multi_nodes.py @@ -0,0 +1,194 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Hierarchical multi-node ReduceScatter for the low-latency packet protocol.""" + +from mscclpp.language.channel import MemoryChannel, PortChannel +from mscclpp.language.collectives import ReduceScatter +from mscclpp.language.program import CollectiveProgram +from mscclpp.language.rank import Buffer, Rank +from mscclpp.language.thread_block_group import ThreadBlockGroup +from mscclpp.language.utils import AlgoSpec + + +def reducescatter_multi_nodes( + spec: AlgoSpec, + thread_block_group_size: int = 1, +) -> CollectiveProgram: + """Build a hierarchical ReduceScatter across nodes and local GPUs.""" + if not isinstance(spec.collective, ReduceScatter): + raise ValueError("reducescatter_multi_nodes requires a ReduceScatter collective") + if spec.protocol != "LL": + raise ValueError("reducescatter_multi_nodes requires protocol='LL'") + if spec.world_size % spec.nranks_per_node != 0: + raise ValueError("world_size must be divisible by nranks_per_node") + if spec.collective.chunk_factor != 1: + raise ValueError("reducescatter_multi_nodes requires chunk_factor=1") + if not spec.in_place or not spec.collective.inplace: + raise ValueError("reducescatter_multi_nodes requires in-place buffers") + if thread_block_group_size <= 0: + raise ValueError("thread_block_group_size must be positive") + + num_nodes = spec.world_size // spec.nranks_per_node + gpus_per_node = spec.nranks_per_node + total_gpus = spec.world_size + + with CollectiveProgram.from_spec(spec) as prog: + local_receive_slots = (gpus_per_node - 1) * num_nodes + local_send_offset = local_receive_slots + remote_receive_offset = local_send_offset + num_nodes + local_owner_offset = remote_receive_offset + num_nodes - 1 + scratch_slots = local_owner_offset + 1 + scratch_buffers = [Buffer(rank, scratch_slots) for rank in range(total_gpus)] + logical_thread_blocks = (gpus_per_node - 1) + num_nodes + thread_block_groups = [ + ThreadBlockGroup( + tb_list=[ + logical_block * thread_block_group_size + group_offset + for group_offset in range(thread_block_group_size) + ] + ) + for logical_block in range(logical_thread_blocks) + ] + + intra_node_channels: dict[tuple[int, int], MemoryChannel] = {} + for node_id in range(num_nodes): + for src_local_rank in range(gpus_per_node): + for dst_local_rank in range(gpus_per_node): + if src_local_rank == dst_local_rank: + continue + src_rank = src_local_rank + node_id * gpus_per_node + dst_rank = dst_local_rank + node_id * gpus_per_node + intra_node_channels[(dst_rank, src_rank)] = MemoryChannel( + dst_rank, + src_rank, + ) + + inter_node_channels: dict[tuple[int, int], PortChannel] = {} + for reducer_local_rank in range(gpus_per_node): + for chunk_offset in range(num_nodes): + owner_rank = reducer_local_rank * num_nodes + chunk_offset + owner_node_id = owner_rank // gpus_per_node + for src_node_id in range(num_nodes): + if src_node_id == owner_node_id: + continue + src_rank = reducer_local_rank + src_node_id * gpus_per_node + inter_node_channels[(owner_rank, src_rank)] = PortChannel( + owner_rank, + src_rank, + ) + + # Each local GPU reduces one contiguous M / gpus_per_node group. Exchange + # those groups with one packet operation per local peer. + for node_id in range(num_nodes): + for src_local_rank in range(gpus_per_node): + src_rank = src_local_rank + node_id * gpus_per_node + src_input = Rank(src_rank).get_input_buffer() + for dst_local_rank in range(gpus_per_node): + if src_local_rank == dst_local_rank: + continue + dst_rank = dst_local_rank + node_id * gpus_per_node + local_peer_slot = src_local_rank if src_local_rank < dst_local_rank else src_local_rank - 1 + dst_peer_slot = dst_local_rank - 1 if src_local_rank < dst_local_rank else dst_local_rank + chunk_index = dst_local_rank * num_nodes + scratch_slot = local_peer_slot * num_nodes + intra_node_channels[(dst_rank, src_rank)].put_packets( + scratch_buffers[dst_rank][scratch_slot : scratch_slot + num_nodes], + src_input[chunk_index : chunk_index + num_nodes], + tb_group=thread_block_groups[dst_peer_slot], + ) + + # Reduce each contiguous group locally. Remote nodes send their partials + # directly to the standard owner while the owner node transfers its one + # local partial over NVLink. + local_reduce_offset = gpus_per_node - 1 + for src_node_id in range(num_nodes): + for reducer_local_rank in range(gpus_per_node): + src_rank = reducer_local_rank + src_node_id * gpus_per_node + rank = Rank(src_rank) + input_buffer = rank.get_input_buffer() + + for chunk_offset in range(num_nodes): + thread_block_group = thread_block_groups[local_reduce_offset + chunk_offset] + chunk_index = reducer_local_rank * num_nodes + chunk_offset + owner_rank = chunk_index + owner_node_id = owner_rank // gpus_per_node + local_packets = [] + for peer_local_rank in range(gpus_per_node): + if peer_local_rank == reducer_local_rank: + continue + local_peer_slot = ( + peer_local_rank if peer_local_rank < reducer_local_rank else peer_local_rank - 1 + ) + scratch_slot = local_peer_slot * num_nodes + chunk_offset + local_packets.append(scratch_buffers[src_rank][scratch_slot : scratch_slot + 1]) + + local_reduced_chunk = input_buffer[chunk_index : chunk_index + 1] + if local_packets: + rank.reduce( + local_reduced_chunk, + local_packets, + tb_group=thread_block_group, + packet=True, + ) + + if src_node_id == owner_node_id: + if src_rank != owner_rank: + intra_node_channels[(owner_rank, src_rank)].put_packets( + scratch_buffers[owner_rank][local_owner_offset : local_owner_offset + 1], + local_reduced_chunk, + tb_group=thread_block_group, + ) + continue + + local_packet_slot = local_send_offset + chunk_offset + rank.copy_packets( + scratch_buffers[src_rank][local_packet_slot : local_packet_slot + 1], + local_reduced_chunk, + tb_group=thread_block_group, + ) + remote_node_slot = src_node_id if src_node_id < owner_node_id else src_node_id - 1 + inter_node_channels[(owner_rank, src_rank)].read_put_packets( + scratch_buffers[owner_rank][ + remote_receive_offset + remote_node_slot : remote_receive_offset + remote_node_slot + 1 + ], + scratch_buffers[src_rank][local_packet_slot : local_packet_slot + 1], + tb_group=thread_block_group, + ) + + if num_nodes == 1: + return prog + + # Every rank receives one standard shard. The owner-node handoff and + # direct IB transfers from the other nodes can progress concurrently. + for owner_rank in range(total_gpus): + owner = Rank(owner_rank) + owner_input = owner.get_input_buffer() + owner_node_id = owner_rank // gpus_per_node + reducer_local_rank = owner_rank // num_nodes + chunk_offset = owner_rank % num_nodes + reducer_rank = reducer_local_rank + owner_node_id * gpus_per_node + thread_block_group = thread_block_groups[local_reduce_offset + chunk_offset] + owner_chunk = owner_input[owner_rank : owner_rank + 1] + + if reducer_rank != owner_rank: + owner.unpack_packets( + owner_chunk, + scratch_buffers[owner_rank][local_owner_offset : local_owner_offset + 1], + tb_group=thread_block_group, + ) + + remote_packets = [ + scratch_buffers[owner_rank][ + remote_receive_offset + remote_node_slot : remote_receive_offset + remote_node_slot + 1 + ] + for remote_node_slot in range(num_nodes - 1) + ] + owner.reduce( + owner_chunk, + remote_packets, + tb_group=thread_block_group, + packet=True, + ) + + return prog diff --git a/python/mscclpp/language/channel.py b/python/mscclpp/language/channel.py index 190e4b25..fc148dde 100644 --- a/python/mscclpp/language/channel.py +++ b/python/mscclpp/language/channel.py @@ -737,7 +737,13 @@ def put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int): get_program().add_operation(self.src_rank, tb, op) - def read_put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int): + def read_put_packets( + self, + dst_chunk: Chunk, + src_chunk: Chunk, + tb: int = None, + tb_group: ThreadBlockGroup = None, + ): """Transfer data in packet format from local to remote scratch buffer. Performs a specialized put operation that transfers data in packet format @@ -747,11 +753,14 @@ def read_put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int): Args: dst_chunk (Chunk): The destination scratch chunk on the destination rank. src_chunk (Chunk): The source scratch chunk on the source rank. - tb (int): The thread block ID that will execute this operation. + tb (int, optional): The thread block ID that will execute this operation. + tb_group (ThreadBlockGroup, optional): The thread block group that will + partition and execute this operation. Raises: RuntimeError: If chunk ranks don't match channel configuration, if chunks are not scratch buffers, or if chunk sizes don't match. + ValueError: If neither a thread block ID nor a thread block group is provided. Example: >>> channel.read_put_packet(dst_chunk, src_chunk, tb=0) @@ -773,20 +782,59 @@ def read_put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int): f"Destination chunk size {dst_chunk.size} does not match source chunk size {src_chunk.size}." ) - remote_chunk = RemoteBuffer(src_chunk.rank, dst_chunk.rank, dst_chunk.buffer, self.channel_type) - tb_chunk_id = get_program().setup_remote_chunk(self.src_rank, tb, remote_chunk, self.channel_type) - tb_channel_ids = get_program().setup_channel(tb, self) + if tb is not None: + tb_list = [tb] + elif tb_group is not None: + tb_list = tb_group.tb_list + else: + raise ValueError( + "Either 'tb' (thread block ID) or 'tb_group' " "(ThreadBlockGroup) must be provided, but both are None." + ) - op = PutOperation( - src_buff=[LocalChunk(src_chunk.buffer, src_chunk.index, src_chunk.size)], - dst_buff=[RemoteChunk(dst_chunk.buffer, dst_chunk.index, dst_chunk.size, tb_chunk_id)], - channel_ids=tb_channel_ids, - channel_type=self.channel_type, - from_packet=True, - to_packet=True, + remote_chunk = RemoteBuffer( + src_chunk.rank, + dst_chunk.rank, + dst_chunk.buffer, + self.channel_type, ) - - get_program().add_operation(self.src_rank, tb, op) + for tb_id in tb_list: + tb_chunk_id = get_program().setup_remote_chunk( + self.src_rank, + tb_id, + remote_chunk, + self.channel_type, + ) + tb_channel_ids = get_program().setup_channel(tb_id, self) + op = PutOperation( + src_buff=[ + LocalChunk( + src_chunk.buffer, + src_chunk.index, + src_chunk.size, + ) + ], + dst_buff=[ + RemoteChunk( + dst_chunk.buffer, + dst_chunk.index, + dst_chunk.size, + tb_chunk_id, + ) + ], + channel_ids=tb_channel_ids, + channel_type=self.channel_type, + tbg_info=( + ThreadBlockGroupInfo( + tb_group.get_internal_id(tb_id), + tb_group.numtb(), + ) + if tb_group is not None + else None + ), + from_packet=True, + to_packet=True, + ) + get_program().add_operation(self.src_rank, tb_id, op) @dataclass diff --git a/python/test/test_mscclpp.py b/python/test/test_mscclpp.py index fc57d083..273e09fa 100644 --- a/python/test/test_mscclpp.py +++ b/python/test/test_mscclpp.py @@ -694,6 +694,25 @@ def test_executor(mpi_group: MpiGroup, filename: str): ) stream.synchronize() assert cp.allclose(sendbuf, expected, atol=1e-3 * mpi_group.comm.size) + + mscclpp_group.barrier() + executor.reset() + mscclpp_group.barrier() + for i in range(nelems_per_rank): + sendbuf[i] = sub_arrays[mpi_group.comm.rank][i] + executor.execute( + mpi_group.comm.rank, + sendbuf.data.ptr, + sendbuf.data.ptr, + sendbuf.nbytes, + sendbuf.nbytes, + DataType.float16, + execution_plan, + stream.ptr, + ) + stream.synchronize() + assert cp.allclose(sendbuf, expected, atol=1e-3 * mpi_group.comm.size) + if npkit_dump_dir is not None: npkit.dump(npkit_dump_dir) npkit.shutdown() diff --git a/src/core/executor/executor.cc b/src/core/executor/executor.cc index c272b17a..d7a81df7 100644 --- a/src/core/executor/executor.cc +++ b/src/core/executor/executor.cc @@ -155,10 +155,10 @@ struct Executor::Impl { } ~Impl() = default; - ExecutionContext setupExecutionContext(int rank, void* sendbuff, void* recvbuff, size_t inputMessageSize, - size_t outputMessageSize, size_t constSrcOffset, size_t constDstOffset, - size_t sendMemRange, size_t recvMemRange, const ExecutionPlan& plan, - std::shared_ptr proxyService) { + const ExecutionContext& setupExecutionContext(int rank, void* sendbuff, void* recvbuff, size_t inputMessageSize, + size_t outputMessageSize, size_t constSrcOffset, size_t constDstOffset, + size_t sendMemRange, size_t recvMemRange, const ExecutionPlan& plan, + std::shared_ptr proxyService) { ExecutionContextKey key = {sendbuff, recvbuff, sendMemRange, recvMemRange, plan.impl_->name}; DeviceExecutionPlanKey devicePlanKey = {inputMessageSize, outputMessageSize, constSrcOffset, constDstOffset}; @@ -166,24 +166,25 @@ struct Executor::Impl { if (plan.impl_->reuseResources) { key = {nullptr, nullptr, 0, 0, plan.impl_->name}; } - if (this->contexts.find(key) != this->contexts.end()) { - auto& devicePlans = this->contexts[key].deviceExecutionPlans; - if (this->contexts[key].currentDevicePlan == devicePlanKey) { - return this->contexts[key]; + auto contextIt = this->contexts.find(key); + if (contextIt != this->contexts.end()) { + auto& context = contextIt->second; + auto& devicePlans = context.deviceExecutionPlans; + if (context.currentDevicePlan == devicePlanKey) { + return context; } else if (devicePlans.find(devicePlanKey) != devicePlans.end()) { - this->contexts[key].currentDevicePlan = devicePlanKey; - return this->contexts[key]; + context.currentDevicePlan = devicePlanKey; + return context; } plan.impl_->operationsReset(); plan.impl_->lightLoadExecutionPlan(inputMessageSize, outputMessageSize, constSrcOffset, constDstOffset); - this->setupDeviceExecutionPlan(this->contexts[key], devicePlanKey, plan); - this->contexts[key].deviceExecutionPlansBuffers[devicePlanKey] = + this->setupDeviceExecutionPlan(context, devicePlanKey, plan); + context.deviceExecutionPlansBuffers[devicePlanKey] = GpuBuffer(devicePlans[devicePlanKey].size() * sizeof(DeviceExecutionPlan)).memory(); - gpuMemcpy(this->contexts[key].deviceExecutionPlansBuffers[devicePlanKey].get(), - (char*)devicePlans[devicePlanKey].data(), + gpuMemcpy(context.deviceExecutionPlansBuffers[devicePlanKey].get(), (char*)devicePlans[devicePlanKey].data(), devicePlans[devicePlanKey].size() * sizeof(DeviceExecutionPlan), cudaMemcpyHostToDevice); - this->contexts[key].currentDevicePlan = devicePlanKey; - return this->contexts[key]; + context.currentDevicePlan = devicePlanKey; + return context; } plan.impl_->reset(); @@ -208,8 +209,8 @@ struct Executor::Impl { (char*)context.deviceExecutionPlans[devicePlanKey].data(), context.deviceExecutionPlans[devicePlanKey].size() * sizeof(DeviceExecutionPlan), cudaMemcpyHostToDevice); context.currentDevicePlan = devicePlanKey; - this->contexts.insert({key, context}); - return context; + auto insertedIt = this->contexts.insert_or_assign(key, std::move(context)).first; + return insertedIt->second; } TransportFlags getTransportFlags(const BufferInfo& info, int rank) { @@ -227,6 +228,24 @@ struct Executor::Impl { return flags; }; + bool usesIbTransport(int rank, const ExecutionPlan& plan) { + auto hasIbPeer = [&](const std::vector& channelInfos) { + for (const auto& info : channelInfos) { + for (int peer : info.connectedPeers) { + if (useIB(rank, peer, this->nranksPerIpcDomain)) { + return true; + } + } + } + return false; + }; + + if (hasIbPeer(plan.impl_->getChannelInfos(ChannelType::PORT))) { + return true; + } + return hasIbPeer(plan.impl_->getUnpairedChannelInfos(this->nranks, ChannelType::PORT)); + } + void setupScratchBuffer(ExecutionContext& context, size_t sendBuffSize, size_t recvBuffSize, const ExecutionPlan& plan) { size_t scratchBufferSize = plan.impl_->calScratchBufferSize(std::min(sendBuffSize, plan.impl_->maxMessageSize), @@ -309,9 +328,10 @@ struct Executor::Impl { size_t recvBufferSize, int rank, const ExecutionPlan& plan) { // Add local src,dst and scratch to registeredMemoryIds context.localMemoryIdBegin = context.proxyService->nextMemoryId(3); + bool registerIb = this->usesIbTransport(rank, plan); for (auto& bufferType : {BufferType::INPUT, BufferType::OUTPUT, BufferType::SCRATCH}) { TransportFlags flags = Transport::CudaIpc; - if (hasIBDevices()) flags |= IBs[rank % this->nranksPerNode]; + if (registerIb) flags |= IBs[rank % this->nranksPerNode]; RegisteredMemory localMemory; auto bufferInfo = getBufferInfo(bufferType, sendbuff, recvbuff, context.scratchBuffer.get(), sendBufferSize, recvBufferSize, context.scratchBufferSize); @@ -469,10 +489,10 @@ struct Executor::Impl { } template - void launchKernelHelper(ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType, + void launchKernelHelper(const ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType, cudaStream_t stream, uint32_t sharedMemSize, const uint32_t& flag) { DeviceExecutionPlanKey key = context.currentDevicePlan; - int nthreadblocks = context.deviceExecutionPlans[key].size(); + int nthreadblocks = context.deviceExecutionPlans.at(key).size(); void* scratchBuffer = context.scratchBuffer.get(); size_t scratchOffset = 0; if (context.doubleScratchBuff && (flag & 0x1) == 0) { @@ -481,23 +501,23 @@ struct Executor::Impl { if (context.reuseResources) { ExecutionKernel::launchKernel( rank, nthreadblocks, context.nthreadsPerBlock, sendbuff, recvbuff, scratchBuffer, scratchOffset, - context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers[key].get(), + context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers.at(key).get(), (DeviceSemaphore*)context.smemaphores.get(), context.localMemoryIdBegin, sharedMemSize, stream, flag); } else { ExecutionKernel::launchKernel( rank, nthreadblocks, context.nthreadsPerBlock, sendbuff, recvbuff, scratchBuffer, scratchOffset, - context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers[key].get(), + context.scratchChunkSize, dataType, (DeviceExecutionPlan*)context.deviceExecutionPlansBuffers.at(key).get(), (DeviceSemaphore*)context.smemaphores.get(), context.localMemoryIdBegin, sharedMemSize, stream, flag); } } - void launchKernel(ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType, + void launchKernel(const ExecutionContext& context, int rank, void* sendbuff, void* recvbuff, DataType dataType, cudaStream_t stream, PacketType packetType) { static uint32_t flag = 0; #if defined(ENABLE_NPKIT) #if defined(MSCCLPP_USE_ROCM) DeviceExecutionPlanKey key = context.currentDevicePlan; - int nthreadblocks = context.deviceExecutionPlans[key].size(); + int nthreadblocks = context.deviceExecutionPlans.at(key).size(); if (nthreadblocks > NPKIT_MAX_NUM_GPU_THREADBLOCKS) { throw Error("Executor plan launching " + std::to_string(nthreadblocks) + " thread blocks, exceeding NPKit support (" + std::to_string(NPKIT_MAX_NUM_GPU_THREADBLOCKS) + @@ -536,12 +556,19 @@ void Executor::execute(int rank, void* sendbuff, void* recvbuff, size_t sendBuff size_t offsetIn = (char*)sendbuff - (char*)sendBasePtr; size_t offsetOut = (char*)recvbuff - (char*)recvBasePtr; - ExecutionContext context = this->impl_->setupExecutionContext( + const ExecutionContext& context = this->impl_->setupExecutionContext( rank, (void*)sendBasePtr, (void*)recvBasePtr, sendBuffSize, recvBuffSize, offsetIn, offsetOut, sendMemRange, recvMemRange, plan, this->impl_->proxyService); this->impl_->launchKernel(context, rank, sendbuff, recvbuff, dataType, stream, packetType); } +void Executor::reset() { + this->impl_->proxyService->stopProxy(); + this->impl_->contexts.clear(); + this->impl_->proxyService = std::make_shared(); + this->impl_->proxyService->startProxy(true); +} + Executor::~Executor() = default; } // namespace mscclpp