From 9261bf2974d88929743525f48f340ef72e4cf650 Mon Sep 17 00:00:00 2001 From: Keshav Pradeep <32313895+keshprad@users.noreply.github.com> Date: Sun, 19 Oct 2025 19:23:48 -0400 Subject: [PATCH] bring over changes from my sc reproducer clone --- csrc/all_gather.cpp | 51 ++++++++++- csrc/all_gather.h | 7 +- csrc/all_reduce.cpp | 50 ++++++++++ csrc/all_reduce.h | 29 ++++++ csrc/common.cu | 4 +- csrc/common.h | 4 +- csrc/pccl_mpi_extension.cpp | 92 +++++++++++++++++++ csrc/reduce_scatter.cpp | 10 +- csrc/reduce_scatter.h | 4 +- pccl/all_gather.py | 102 ++++++++++++++++++--- pccl/all_reduce.py | 178 ++++++++++++++++++++++++++++++++++++ pccl/reduce_scatter.py | 24 ++++- 12 files changed, 525 insertions(+), 30 deletions(-) create mode 100644 csrc/all_reduce.cpp create mode 100644 csrc/all_reduce.h create mode 100644 pccl/all_reduce.py diff --git a/csrc/all_gather.cpp b/csrc/all_gather.cpp index ce91847..a1753d6 100644 --- a/csrc/all_gather.cpp +++ b/csrc/all_gather.cpp @@ -18,7 +18,7 @@ // - comm: MPI communicator (default MPI_COMM_WORLD). void recursiveDoublingAllGatherGPU(void* output, const void* input, - int total_elems, + int64_t total_elems, void* recv_buf, // Same as output size MPI_Comm comm) { @@ -27,7 +27,7 @@ void recursiveDoublingAllGatherGPU(void* output, MPI_Comm_size(comm, &size); assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes"); - int block_size = total_elems / size; + int64_t block_size = total_elems / size; auto stream = at::cuda::getCurrentCUDAStream(); @@ -71,3 +71,50 @@ void recursiveDoublingAllGatherGPU(void* output, CUDA_CHECK(cudaEventDestroy(stream_sync_event)); } + +void ringAllGatherGPU(void* output, + const void* input, + int64_t total_elems, + MPI_Comm comm) { + int rank, size; + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &size); + + assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes"); + int64_t block_size = total_elems / size; + // printf("[Rank %d] block_size = %d\n", rank, block_size); + + auto stream = at::cuda::getCurrentCUDAStream(); + cudaEvent_t stream_sync_event; + + // Copy local input into its designated block in the output buffer. + CUDA_CHECK(cudaMemcpyAsync(static_cast(output) + rank * block_size, + input, + block_size, + cudaMemcpyDeviceToDevice, + stream)); + + CUDA_CHECK(cudaEventCreateWithFlags(&stream_sync_event, cudaEventDisableTiming)); + + // P-1 rounds each sending N/P data (where P is num processes, N is total data size) + for (int step = 0; step < size - 1; step++) { + // Compute block indices + int send_idx = (rank - step + size) % size; + int recv_idx = (rank - step - 1 + size) % size; + int send_peer = (rank + 1) % size; + int recv_peer = (rank - 1 + size) % size; + + // Record an event on the cuda stream. + CUDA_CHECK(cudaEventRecord(stream_sync_event, stream)); + // Wait for the copy to complete. + CUDA_CHECK(cudaEventSynchronize(stream_sync_event)); + + // Send the block to the right neighbor and receive from the left neighbor. + MPI_Sendrecv(static_cast(output) + send_idx * block_size, block_size, MPI_BYTE, send_peer, 0, + static_cast(output) + recv_idx * block_size, block_size, MPI_BYTE, recv_peer, 0, + comm, MPI_STATUS_IGNORE); + } + + // destroy event + CUDA_CHECK(cudaEventDestroy(stream_sync_event)); +} diff --git a/csrc/all_gather.h b/csrc/all_gather.h index 3bb4b4e..20f4a5f 100644 --- a/csrc/all_gather.h +++ b/csrc/all_gather.h @@ -11,8 +11,13 @@ void recursiveDoublingAllGatherGPU(void* output, const void* input, - int total_elems, + int64_t total_elems, void* recv_buf, MPI_Comm comm = MPI_COMM_WORLD); +void ringAllGatherGPU(void* output, + const void* input, + int64_t total_elems, + MPI_Comm comm); + #endif // ALL_GATHER_H diff --git a/csrc/all_reduce.cpp b/csrc/all_reduce.cpp new file mode 100644 index 0000000..7d96b11 --- /dev/null +++ b/csrc/all_reduce.cpp @@ -0,0 +1,50 @@ +// Copyright 2025 Parallel Software and Systems Group, University of Maryland. +// See the top-level LICENSE file for details. +// +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + + +#include +#include + +#include "all_gather.h" +#include "all_reduce.h" +#include "common.h" +#include "reduce_scatter.h" + + +// Performs an all-reduce on GPU tensors via recursive-halving reduce-scatter followed by recursive-doubling all-gather. +// - output: CUDA device pointer where the final gathered tensor will be stored. +// - input: CUDA device pointer to the local block of size block_size. +// - total_elems: total number of elements in output (P * block_size). +// - buf: main working buffer for the algorithm +// - recv_buf: buffer to receive data before it is processed +// - comm: MPI communicator (default MPI_COMM_WORLD). +void recursiveHalvingDoublingAllReduceGPU(float* output, + const float* input, + int64_t total_elems, + float* buf, // Same as input size + float* recv_buf, // Same as input size + float* intermediate_buf, // Input size / world size + MPI_Comm comm) { + recursiveHalvingReduceScatterGPU(intermediate_buf, input, total_elems, buf, recv_buf, comm); + + // allgather uses void* so multiply total_elems by size of float dtype + recursiveDoublingAllGatherGPU(output, intermediate_buf, total_elems*sizeof(float), recv_buf, comm); +} + +// Performs an all-reduce on GPU tensors via ring reduce-scatter followed by ring all-gather. +void ringAllReduceGPU(float* output, + const float* input, + int64_t total_elems, + float* intermediate_buf, // Input size / world size + float* d_buf, // Input size + float* d_send, // Input size / world size + float* d_tmp, // Input size / world size + MPI_Comm comm) { + ringReduceScatterGPU(intermediate_buf, input, total_elems, d_buf, d_send, d_tmp); + + // allgather uses void* so multiply total_elems by size of float dtype + ringAllGatherGPU(output, intermediate_buf, total_elems*sizeof(float), comm); +} + diff --git a/csrc/all_reduce.h b/csrc/all_reduce.h new file mode 100644 index 0000000..be37f07 --- /dev/null +++ b/csrc/all_reduce.h @@ -0,0 +1,29 @@ +// Copyright 2025 Parallel Software and Systems Group, University of Maryland. +// See the top-level LICENSE file for details. +// +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + + +#ifndef ALL_REDUCE_H +#define ALL_REDUCE_H + +#include + +void recursiveHalvingDoublingAllReduceGPU(float* output, + const float* input, + int64_t total_elems, + float* buf, + float* recv_buf, + float* intermediate_buf, + MPI_Comm comm = MPI_COMM_WORLD); + +void ringAllReduceGPU(float* output, + const float* input, + int64_t total_elems, + float* intermediate_buf, // Input size / world size + float* d_buf, // Input size + float* d_send, // Input size / world size + float* d_tmp, // Input size / world size + MPI_Comm comm = MPI_COMM_WORLD); + +#endif // ALL_REDUCE_H diff --git a/csrc/common.cu b/csrc/common.cu index 0a6ac62..efe4a13 100644 --- a/csrc/common.cu +++ b/csrc/common.cu @@ -8,7 +8,7 @@ // Kernel for vector addition. -__global__ void vectorAddKernel(float* a, const float* b, int n) { +__global__ void vectorAddKernel(float* a, const float* b, int64_t n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { a[idx] += b[idx]; @@ -16,7 +16,7 @@ __global__ void vectorAddKernel(float* a, const float* b, int n) { } // Function to launch the kernel. -void vectorAdd(float* a, const float* b, int n, cudaStream_t stream) { +void vectorAdd(float* a, const float* b, int64_t n, cudaStream_t stream) { int threads = 256; int blocks = (n + threads - 1) / threads; diff --git a/csrc/common.h b/csrc/common.h index 5b2a977..ca5d4f9 100644 --- a/csrc/common.h +++ b/csrc/common.h @@ -21,9 +21,9 @@ } while(0) // Kernel for vector addition. -__global__ void vectorAddKernel(float* a, const float* b, int n); +__global__ void vectorAddKernel(float* a, const float* b, int64_t n); // Function to launch the kernel. -void vectorAdd(float* a, const float* b, int n, cudaStream_t stream); +void vectorAdd(float* a, const float* b, int64_t n, cudaStream_t stream); #endif // COMMON_H diff --git a/csrc/pccl_mpi_extension.cpp b/csrc/pccl_mpi_extension.cpp index a83cdf2..e4cf363 100644 --- a/csrc/pccl_mpi_extension.cpp +++ b/csrc/pccl_mpi_extension.cpp @@ -10,6 +10,7 @@ #include #include "reduce_scatter.h" #include "all_gather.h" +#include "all_reduce.h" namespace py = pybind11; @@ -142,13 +143,104 @@ void all_gather_mpi(const torch::Tensor& output_tensor, tmp_wrkspace_tensor_1.data_ptr(), //tmp_wrkspace_tensor_2.data_ptr(), comm); + } else if (algorithm == "ring") { + // always use torch tensors. do NOT use malloc. + // malloc's have high overheads and will slow your communication down + // torch mallocs memory in advance and manages it internally. + // therefore these calls are low overheads + // auto tmp_wrkspace_tensor_1 = torch::empty_like(output_tensor); + // No workspace tensors needed for ring all-gather algorithm + ringAllGatherGPU(output_ptr, + input_ptr, + total_elems * dtype_size, + comm); } else { TORCH_CHECK(false, "Unknown algorithm specified for all_gather_mpi: ", algorithm); } } +void all_reduce_mpi(const torch::Tensor& output_tensor, + const torch::Tensor& input_tensor, + py::object py_comm, + const std::string& algorithm = "recursive") +{ + TORCH_CHECK(output_tensor.is_contiguous(), "output tensor must be contiguous."); + TORCH_CHECK(input_tensor.is_contiguous(), "input tensor must be contiguous."); + + // Ensure 1D tensors. + TORCH_CHECK(output_tensor.dim() == 1, "output tensor must be 1D"); + TORCH_CHECK(input_tensor.dim() == 1, "input tensor must be 1D"); + + // Ensure input and output dtypes are the same + TORCH_CHECK(input_tensor.dtype() == output_tensor.dtype(), + "Input and output tensors must have the same dtype."); + + // Get MPI rank/size. + int rank, size; + // Get reference to base communicator + MPI_Comm comm = ((PyMPIIntracommObject*)(py_comm.ptr()))->__pyx_base.ob_mpi; + + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &size); + + // Input tensor has one block + int64_t block_size = input_tensor.numel(); + int64_t total_elems = block_size; + // Ensure output tensor is same size as input tensor. + TORCH_CHECK(output_tensor.numel() == block_size, + "Output tensor must have same size as input tensor"); + + // Ensure input tensor divisible by world size. + TORCH_CHECK(block_size % size == 0, + "Input tensor size must be divisible by world_size for recursive halving algorithm"); + + // Get raw device pointers (assumes tensors reside on GPU). + float* output_ptr = output_tensor.data_ptr(); + const float* input_ptr = input_tensor.data_ptr(); + + // Call the corresponding GPU reduce-scatter algorithm. + if (algorithm == "recursive") { + // always use torch tensors. do NOT use malloc. + // malloc's have high overheads and will slow your communication down + // torch mallocs memory in advance and manages it internally. + // therefore these calls are low overheads + auto tmp_wrkspace_tensor_1 = torch::empty_like(input_tensor); + auto tmp_wrkspace_tensor_2 = torch::empty_like(input_tensor); + auto tmp_wrkspace_tensor_4 = torch::empty({block_size / size}, input_tensor.options()); + + recursiveHalvingDoublingAllReduceGPU(output_ptr, + input_ptr, + total_elems, + tmp_wrkspace_tensor_1.data_ptr(), + tmp_wrkspace_tensor_2.data_ptr(), + tmp_wrkspace_tensor_4.data_ptr(), + comm); + } else if (algorithm == "ring") { + // always use torch tensors. do NOT use malloc. + // malloc's have high overheads and will slow your communication down + // torch mallocs memory in advance and manages it internally. + // therefore these calls are low overheads + auto tmp_wrkspace_tensor_1 = torch::empty({block_size / size}, input_tensor.options()); + auto tmp_wrkspace_tensor_2 = torch::empty_like(input_tensor); + auto tmp_wrkspace_tensor_3 = torch::empty({block_size / size}, input_tensor.options()); + auto tmp_wrkspace_tensor_4 = torch::empty({block_size / size}, input_tensor.options()); + + ringAllReduceGPU(output_ptr, + input_ptr, + total_elems, + tmp_wrkspace_tensor_1.data_ptr(), + tmp_wrkspace_tensor_2.data_ptr(), + tmp_wrkspace_tensor_3.data_ptr(), + tmp_wrkspace_tensor_4.data_ptr(), + comm); + } else { + TORCH_CHECK(false, "Unknown algorithm specified for all_reduce_mpi: ", algorithm); + } +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("reduce_scatter_mpi", reduce_scatter_mpi); m.def("all_gather_mpi", all_gather_mpi); + m.def("all_reduce_mpi", all_reduce_mpi); } diff --git a/csrc/reduce_scatter.cpp b/csrc/reduce_scatter.cpp index 89852ce..dd8691e 100644 --- a/csrc/reduce_scatter.cpp +++ b/csrc/reduce_scatter.cpp @@ -19,7 +19,7 @@ // The reduction operation is elementwise addition. void recursiveHalvingReduceScatterGPU(float* output, const float* input, - int total_elems, + int64_t total_elems, float* buf, // same as size of input float* recv_buf, // same as size of input MPI_Comm comm) { @@ -29,7 +29,7 @@ void recursiveHalvingReduceScatterGPU(float* output, MPI_Comm_size(comm, &size); assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes"); - int block_size = total_elems / size; + int64_t block_size = total_elems / size; auto stream = at::cuda::getCurrentCUDAStream(); // copy the input into buf @@ -55,7 +55,7 @@ void recursiveHalvingReduceScatterGPU(float* output, // The current buffer holds 'current_blocks' contiguous blocks. int half = current_blocks / 2; // Number of elements to send/receive in this round. - int count = half * block_size; + int64_t count = half * block_size; if ((rank % group_size) < (group_size / 2)) { // Lower half: keep the lower half and send the upper half. @@ -102,7 +102,7 @@ void recursiveHalvingReduceScatterGPU(float* output, void ringReduceScatterGPU(float* output, const float* input, - int total_elems, + int64_t total_elems, float* d_buf, // same as size of input float* d_send, // same as size of output float* d_tmp, // same as size of output @@ -113,7 +113,7 @@ void ringReduceScatterGPU(float* output, MPI_Comm_size(comm, &size); assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes"); - int block_size = total_elems / size; + int64_t block_size = total_elems / size; auto stream = at::cuda::getCurrentCUDAStream(); cudaEvent_t stream_sync_event; diff --git a/csrc/reduce_scatter.h b/csrc/reduce_scatter.h index d661f3c..a0675b2 100644 --- a/csrc/reduce_scatter.h +++ b/csrc/reduce_scatter.h @@ -11,14 +11,14 @@ void recursiveHalvingReduceScatterGPU(float* output, const float* input, - int total_elems, + int64_t total_elems, float* buf, float* recv_buf, MPI_Comm comm = MPI_COMM_WORLD); void ringReduceScatterGPU(float* output, const float* input, - int total_elems, + int64_t total_elems, float* d_buf, float* d_send, float* d_tmp, diff --git a/pccl/all_gather.py b/pccl/all_gather.py index 2dc6a28..417cfb4 100644 --- a/pccl/all_gather.py +++ b/pccl/all_gather.py @@ -109,12 +109,81 @@ def recursive_doubling_allgather_mpi( # At the end, output_tensor holds blocks from rank 0, 1, ..., size-1 in order. +def ring_allgather_mpi(output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: Optional[MPI.Comm] = None, + async_op: bool = False): + """ + Performs a ring-based reduce-scatter using torch tensors. + + Each process holds a 1D tensor of shape (N/P, P), where P is the number of + processes and N is the block size. Each row of the tensor corresponds to a + block. The goal is to gather all blocks on every process so that at the + end, each process gets the full tensor of size (N, P). + + The algorithm performs P-1 steps. At each step, each process: + - Sends a designated block to its right neighbor. + - Receives a block from its left neighbor. + - The received block is places at index equal to the peer's rank + + Parameters: + output_tensor : torch.Tensor + Pre-allocated tensor of shape (P * block_size,) on a CUDA device. + input_tensor : torch.Tensor + 1D tensor of shape (block_size,) representing local data. + group : Optional[MPI.Comm] + MPI communicator; defaults to MPI.COMM_WORLD. + async_op : bool + Non-blocking operations are not supported in this implementation. + """ + assert not async_op, "non-blocking primitives not supported" + comm = MPI.COMM_WORLD if group is None else group + rank = comm.Get_rank() + size = comm.Get_size() + + # Determine block size and ensure output_tensor is large enough. + block_size = input_tensor.numel() + assert output_tensor.numel() == size * block_size, "Output tensor has incorrect size" + + # Copy local data into the proper slot. + output_tensor[rank * block_size : (rank + 1) * block_size].copy_(input_tensor) + + # Make sure any CUDA work is complete before we start communication. + # torch.cuda.current_stream().synchronize() + + # working buffers + # tmp = torch.empty(block_size, dtype=input_tensor.dtype, device=input_tensor.device) + + for step in range(size-1): + # Adjusted indices + send_idx = (rank - step) % size + recv_idx = (rank - step - 1) % size + # Identify neighbors in the ring + send_peer = (rank + 1) % size + recv_peer = (rank - 1) % size + + # copy block to be sent + # send_data = output_tensor[send_idx * block_size : (send_idx+1) * block_size].clone() + + torch.cuda.current_stream().synchronize() + + comm.Sendrecv( + sendbuf=output_tensor[send_idx * block_size : (send_idx+1) * block_size], + dest=send_peer, + sendtag=0, + recvbuf=output_tensor[recv_idx * block_size : (recv_idx+1) * block_size], + source=recv_peer, + recvtag=0 + ) + + # At the end, output_tensor holds blocks from rank 0, 1, ..., size-1 in order. def _all_gather( output_tensor: torch.Tensor, input_tensor: torch.Tensor, group: Optional[Union[dist.ProcessGroup, MPI.Comm]] = None, async_op: bool = False, + directly_call_mpi: bool = False, use_rd: bool = False, use_pccl_cpp_backend: bool = False, ) -> Optional[Request]: @@ -128,27 +197,32 @@ def _all_gather( # Case 2: mpi4py.MPI.Comm elif isinstance(group, MPI.Comm): # make sure that the cpu is synchronized with the current stream - if use_rd: - if use_pccl_cpp_backend: - request = pccl_mpi_extension.all_gather_mpi( - output_tensor, input_tensor, group, "recursive" - ) - else: - request = recursive_doubling_allgather_mpi( - output_tensor, input_tensor, group, async_op - ) + if use_pccl_cpp_backend: + request = pccl_mpi_extension.all_gather_mpi( + output_tensor, input_tensor, group, "recursive" if use_rd else "ring" + ) else: - torch.cuda.current_stream().synchronize() - if async_op: - request = group.Iallgather(input_tensor, output_tensor) + if not directly_call_mpi: + if use_rd: + request = recursive_doubling_allgather_mpi( + output_tensor, input_tensor, group, async_op + ) + else: + request = ring_allgather_mpi( + output_tensor, input_tensor, group, async_op + ) else: - request = group.Allgather(input_tensor, output_tensor) + torch.cuda.current_stream().synchronize() + if async_op: + request = group.Iallgather(input_tensor, output_tensor) + else: + request = group.Allgather(input_tensor, output_tensor) else: raise TypeError( f"Unsupported group type: {type(group)}. " "Expected torch.distributed.ProcessGroup or mpi4py.MPI.Comm." ) - return request + return request def all_gather_2D( diff --git a/pccl/all_reduce.py b/pccl/all_reduce.py new file mode 100644 index 0000000..f322d78 --- /dev/null +++ b/pccl/all_reduce.py @@ -0,0 +1,178 @@ +# Copyright 2025 Parallel Software and Systems Group, University of Maryland. +# See the top-level LICENSE file for details. +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + + +import torch +import torch.distributed as dist +from mpi4py import MPI +from typing import List, Optional, Union +from .request import Request +from .process_groups import ProcessGroups +import numpy as np +import pccl_mpi_extension +from .all_gather import all_gather_2D, recursive_doubling_allgather_mpi, ring_allgather_mpi +from .reduce_scatter import reduce_scatter_2D, recursive_halving_reduce_scatter_mpi, ring_reduce_scatter_mpi + +def recursive_halving_doubling_allreduce_mpi( + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: Optional[MPI.Comm] = None, + async_op: bool = False, +): + """ + Performs a recursive halving and doubling based all-reduce on CUDA tensors using MPI point-to-point + Sendrecv operations. + + Each process starts with a 1D input_tensor (block_size) and the final output_tensor + is a 1D tensor of size (block_size). The goal is to reduce (using op, e.g. torch.add) + the block over all processes so that all processes end with the fully reduced block. + + Parameters: + output_tensor : torch.Tensor + Pre-allocated tensor of shape (P * block_size,) on a CUDA device. + input_tensor : torch.Tensor + 1D tensor of shape (P * block_size,) representing local data. + group : Optional[MPI.Comm] + MPI communicator; defaults to MPI.COMM_WORLD. + async_op : bool + Non-blocking operations are not supported in this implementation. + """ + assert not async_op, "Non-blocking operations not supported" + + world_size = group.Get_size() + + output_intermediate = torch.empty( + input_tensor.size(0) // world_size, + device=input_tensor.device, + dtype=input_tensor.dtype + ) + recursive_halving_reduce_scatter_mpi(output_intermediate, input_tensor, group, async_op) + recursive_doubling_allgather_mpi(output_tensor, output_intermediate, group, async_op) + +def ring_allreduce_mpi(output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: Optional[MPI.Comm] = None, + async_op: bool = False): + """ + Performs a ring based all-reduce on CUDA tensors using MPI point-to-point Sendrecv operations. + + Each process starts with a 1D input_tensor (block_size) and the final output_tensor + is a 1D tensor of size (block_size). The goal is to reduce (using op, e.g. torch.add) + the block over all processes so that all processes end with the fully reduced block. + + Parameters: + output_tensor : torch.Tensor + Pre-allocated tensor of shape (P * block_size,) on a CUDA device. + input_tensor : torch.Tensor + 1D tensor of shape (P * block_size,) representing local data. + group : Optional[MPI.Comm] + MPI communicator; defaults to MPI.COMM_WORLD. + async_op : bool + Non-blocking operations are not supported in this implementation. + """ + assert not async_op, "Non-blocking operations not supported" + + world_size = group.Get_size() + + output_intermediate = torch.empty( + input_tensor.size(0) // world_size, + device=input_tensor.device, + dtype=input_tensor.dtype + ) + ring_reduce_scatter_mpi(output_intermediate, input_tensor, group, async_op) + ring_allgather_mpi(output_tensor, output_intermediate, group, async_op) + +def _all_reduce( + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: Optional[Union[dist.ProcessGroup, MPI.Comm]] = None, + async_op: bool = False, + directly_call_mpi: bool = False, + use_rh_and_rd: bool = False, + use_pccl_cpp_backend: bool = False, +) -> Optional[Request]: + + # Case 1: torch.distributed.ProcessGroup + if group is None or isinstance(group, dist.ProcessGroup): + # Delegate to torch.distributed.all_reduce + + # all_reduce_into_tensor doesn't exist... + # Copy input tensor to output tensor, then perform in-place all_reduce + output_tensor.copy_(input_tensor) + request = dist.all_reduce( + output_tensor, + group=group, + async_op=async_op + ) + + # Case 2: mpi4py.MPI.Comm + elif isinstance(group, MPI.Comm): + if use_pccl_cpp_backend: + request = pccl_mpi_extension.all_reduce_mpi( + output_tensor, input_tensor, group, "recursive" if use_rh_and_rd else "ring" + ) + else: + if not directly_call_mpi: + if use_rh_and_rd: + request = recursive_halving_doubling_allreduce_mpi( + output_tensor, input_tensor, group, async_op + ) + else: + request = ring_allreduce_mpi( + output_tensor, input_tensor, group, async_op + ) + else: + torch.cuda.current_stream().synchronize() + if async_op: + request = group.Iallreduce(input_tensor, output_tensor) + else: + request = group.Allreduce(input_tensor, output_tensor) + + return request + +def all_reduce_2D(output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: Optional[ProcessGroups] = None, + async_op: bool = False, + use_rh_and_rd: bool = False, + use_pccl_cpp_backend: bool = False): + + assert not async_op, "Non blocking version not implemented" + + assert input_tensor.dim() == 1 and output_tensor.dim() == 1, "all_gather_2D only admits 1D tensors" + + # # TESTING cpp allreduce + # output_intermediate = torch.empty(input_tensor.size(0), device=input_tensor.device, dtype=input_tensor.dtype) + # # Step-1 inter-node all-reduce + # _all_reduce(output_intermediate, input_tensor, group.get_outer_group(), async_op=False, use_rh_and_rd=True, use_pccl_cpp_backend=True, directly_call_mpi=True) + # # Step-2 intra-node all-reduce + # _all_reduce(output_tensor, output_intermediate, group.get_inner_group(), async_op=False, use_rh_and_rd=True, use_pccl_cpp_backend=True, directly_call_mpi=True) + + intra_node_group_size, inter_node_group_size = group.get_world_size() + world_size = intra_node_group_size * inter_node_group_size + output_intermediate = torch.empty( + input_tensor.size(0) // world_size, + device=input_tensor.device, + dtype=input_tensor.dtype + ) + + # Step-1 2-dim reduce-scatter + reduce_scatter_2D(output_tensor=output_intermediate, + input_tensor=input_tensor, + group=group, + async_op=async_op, + use_rh=use_rh_and_rd, + use_pccl_cpp_backend=use_pccl_cpp_backend + ) + + # Step-2 2-dim all-gather + all_gather_2D(output_tensor=output_tensor, + input_tensor=output_intermediate, + group=group, + async_op=async_op, + # directly_call_mpi= TODO: standardize API for all_gather_2D, reduce_scatter_2D, all_gather_2D + use_rd=use_rh_and_rd, + use_pccl_cpp_backend=use_pccl_cpp_backend + ) diff --git a/pccl/reduce_scatter.py b/pccl/reduce_scatter.py index e3dbc53..6293260 100644 --- a/pccl/reduce_scatter.py +++ b/pccl/reduce_scatter.py @@ -39,7 +39,7 @@ def recursive_halving_reduce_scatter_mpi( At each round the process splits its current buffer (which initially has p blocks) into two equal halves. Then: - - If the process’s position within its current group is in the lower half, it keeps + - If the process's position within its current group is in the lower half, it keeps the lower half (which should contain the blocks destined for lower-ranked processes) and sends the upper half. - Otherwise it keeps the upper half and sends the lower half. @@ -47,6 +47,16 @@ def recursive_halving_reduce_scatter_mpi( After log2(P) rounds, only one block remains—and it is naturally ordered (process i gets block i). Assumes that P (the number of processes) is a power of 2. + + Parameters: + output_tensor : torch.Tensor + Pre-allocated tensor of shape (block_size,) on a CUDA device. + input_tensor : torch.Tensor + 1D tensor of shape (P * block_size,) representing local data. + group : Optional[MPI.Comm] + MPI communicator; defaults to MPI.COMM_WORLD. + async_op : bool + Non-blocking operations are not supported in this implementation. """ assert not async_op, "Non-blocking operations not supported" comm = MPI.COMM_WORLD if group is None else group @@ -138,7 +148,17 @@ def ring_reduce_scatter_mpi( - Receives a block from its left neighbor. - Immediately reduces the received block into its local copy. - After the loop, the fully reduced block is at index equal to the process’s rank. + After the loop, the fully reduced block is at index equal to the process's rank. + + Parameters: + output_tensor : torch.Tensor + Pre-allocated tensor of shape (block_size,) on a CUDA device. + input_tensor : torch.Tensor + 1D tensor of shape (P * block_size,) representing local data. + group : Optional[MPI.Comm] + MPI communicator; defaults to MPI.COMM_WORLD. + async_op : bool + Non-blocking operations are not supported in this implementation. """ assert not async_op, "non-blocking primitives not supported" comm = MPI.COMM_WORLD if group is None else group