From f35dd1407fa32b1ad07b1a1f83e046bb514417d7 Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Thu, 13 Jun 2024 11:09:49 -0700 Subject: [PATCH] custom allreduce cuda kernel (#20703) ### Description Conditionally route to custom AllReduce kernel when buffer size and gpu numbers meet certain requirements. Otherwise, keep using NCCL's AllReduce. ### Motivation and Context --------- Co-authored-by: Ye Wang Co-authored-by: Your Name --- cmake/onnxruntime_rocm_hipify.cmake | 2 + .../cuda/collective/custom_reduce_impl.cu | 626 ++++++++++++++++++ .../cuda/collective/custom_reduce_impl.h | 81 +++ .../contrib_ops/cuda/collective/ipc_utils.cc | 123 ++++ .../contrib_ops/cuda/collective/ipc_utils.h | 89 +++ .../cuda/collective/nccl_kernels.cc | 66 +- .../cuda/collective/nccl_kernels.h | 18 + .../cuda/collective/sharded_moe.cc | 72 +- .../contrib_ops/cuda/collective/sharded_moe.h | 4 +- .../providers/cuda/shared_inc/cuda_utils.h | 11 + ...ortmodule-distributed-test-ci-pipeline.yml | 4 +- 11 files changed, 1052 insertions(+), 44 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu create mode 100644 onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h create mode 100644 onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc create mode 100644 onnxruntime/contrib_ops/cuda/collective/ipc_utils.h diff --git a/cmake/onnxruntime_rocm_hipify.cmake b/cmake/onnxruntime_rocm_hipify.cmake index bd472f08f6..928a58f7b8 100644 --- a/cmake/onnxruntime_rocm_hipify.cmake +++ b/cmake/onnxruntime_rocm_hipify.cmake @@ -92,7 +92,9 @@ set(contrib_ops_excluded_files "bert/group_query_attention.cc" "bert/group_query_attention_impl.h" "bert/group_query_attention_impl.cu" + "collective/custom_*" "collective/distributed_*" + "collective/ipc_*" "collective/shard*" ) diff --git a/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu new file mode 100644 index 0000000000..666ec3a993 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu @@ -0,0 +1,626 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/shared_inc/cuda_utils.h" +#include "core/providers/shared_library/provider_api.h" +#include "custom_reduce_impl.h" +#include +#include +#include +#include + +namespace onnxruntime { +namespace cuda { +namespace collective { + +#if defined(USE_MPI) || defined(USE_NCCL) + +using namespace onnxruntime; +using namespace onnxruntime::cuda; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +static inline __device__ void st_flag_release(uint32_t const &flag, uint32_t *flag_addr) { +#if __CUDA_ARCH__ >= 700 + asm volatile("st.global.release.sys.b32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +#else + __threadfence_system(); + asm volatile("st.global.volatile.b32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +static inline __device__ uint32_t ld_flag_acquire(uint32_t *flag_addr) { + uint32_t flag; +#if __CUDA_ARCH__ >= 700 + asm volatile("ld.global.acquire.sys.b32 %0, [%1];" : "=r"(flag) : "l"(flag_addr)); +#else + asm volatile("ld.global.volatile.b32 %0, [%1];" : "=r"(flag) : "l"(flag_addr)); +#endif + return flag; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Type Converter that packs data format to 128 bits data type +// +using PackedFloat = union { + int4 packed; + float unpacked[4]; +}; + +using PackedHalf = union { + int4 packed; + half2 unpacked[4]; +}; + +template struct PackedOn16Bytes {}; + +template <> struct PackedOn16Bytes { + using Type = PackedFloat; +}; + +template <> struct PackedOn16Bytes { + using Type = PackedHalf; +}; + +// add two 128b data +template inline __device__ int4 add128b(T &a, T &b) { + T c; + c.unpacked[0] = a.unpacked[0] + b.unpacked[0]; + c.unpacked[1] = a.unpacked[1] + b.unpacked[1]; + c.unpacked[2] = a.unpacked[2] + b.unpacked[2]; + c.unpacked[3] = a.unpacked[3] + b.unpacked[3]; + return c.packed; +} + +__inline__ __device__ void multi_gpu_barrier(uint32_t **signals, uint32_t const flag, size_t const local_rank, + size_t const world_size, int const tidx, int const bidx) { + // After this function, at least one block in each GPU has reached the barrier + if (tidx < world_size) { + // we can think of signals having the shape [world_size, world_size] + // Dimension 0 is the "listening" dimension, dimension 2 is "emitting" dimension + + // Block 0 broadcasts its flag (local_rank on emitting dimension) to all receivers + if (bidx == 0) { + signals[tidx][local_rank] = flag; + } + + // All blocks check that corresponding block 0 on other GPUs have set the flag + // No deadlock because block #0 is always the first block started + uint32_t volatile *my_signals = signals[local_rank]; + while (my_signals[tidx] != flag) { + } + } + + __syncthreads(); +} + +__inline__ __device__ void block_barrier(uint32_t **signals, uint32_t const flag, size_t const local_rank, + size_t const world_size, int const tidx, int const bidx) { + // After this function, the block of id == bidx of each GPU has reached the barrier + if (tidx < world_size) { + // we can think of signals having the shape [world_size, num_blocks, world_size] + // (+ an offset on dim 1 to account for flags used in multi_gpu_barrier) + // Dimension 0 is the "listening" dimension, dimension 2 is "emitting" dimension + + // Block broadcast its flag (local_rank on emitting dimension) to all receivers + uint32_t flag_block_offset = world_size + bidx * world_size; + st_flag_release(flag, signals[tidx] + flag_block_offset + local_rank); + + // Blocks check that corresponding blocks on other GPUs have also set the flag + uint32_t *peer_barrier_d = signals[local_rank] + flag_block_offset + tidx; + while (ld_flag_acquire(peer_barrier_d) != flag) { + } + } + + __syncthreads(); +} + +template +static __global__ void oneShotAllReduceKernel(AllReduceParams params) { + // Suppose that two GPUs participate in the AR exchange, and we start four blocks. + // The message is partitioned into chunks as detailed below: + // message + // |-------------------| + // GPU 0 | B0 | B1 | B2 | B3 | + // GPU 1 | B0 | B1 | B2 | B3 | + // + // Here the step-by-step behavior of one block: + // 1. B0 copies the chunk it is responsible for, from local_input to shareable buffer + // 2. B0 on GPU 0 and B0 on GPU 1 wait for each other (block_barrier) + // 3. B0 on GPU 0 pull and sum the chunk from GPU 1, writes the result to local_output + // + // With COPY_INPUT == false, skip step 1. and use gpu_barrier instead of block barrier during step 2. + // We only to know if the other GPU as arrived at the AR kernel, that would mean that data is ready + // + // With PUSH_MODE, we consider that the shared buffer is of size: + // params.peer_comm_buffer_ptrs: [world_size, world_size, message_size] + // + // Here the step-by-step behavior of one block: + // 1. B0 push the chunk is it responsible for into all other GPUs: + // params.peer_comm_buffer_ptrs[:, local_gpu, B0 slice] + // 2. block sync so the block is shared by other GPUs + // 3. Reduce along second dimension params.peer_comm_buffer_ptrs[local_gpu, :, B0 slice] + + int const bidx = blockIdx.x; + int const tidx = threadIdx.x; + + // The number of elements packed into one for comms + static constexpr int PACKED_ELTS = 16 / sizeof(T); + using PackedStruct = typename PackedOn16Bytes::Type; + + [[maybe_unused]] T const *local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + [[maybe_unused]] T *local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); + [[maybe_unused]] T *local_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + + // Start and end offsets of the thread + size_t const chunk_start = bidx * params.elts_per_block + tidx * PACKED_ELTS; + size_t const chunk_end = std::min((bidx + 1) * params.elts_per_block, params.elts_total); + + T *buffers[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + // buffers[0] is always the local buffers. Helps load balancing reads. + int rank = (params.local_rank + ii) % RANKS_PER_NODE; + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + } + + if constexpr (PUSH_MODE || COPY_INPUT) { + // Copy from local buffer to shareable buffer + for (size_t iter_offset = chunk_start; iter_offset < chunk_end; iter_offset += blockDim.x * PACKED_ELTS) { + if constexpr (PUSH_MODE) { +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + *reinterpret_cast(&buffers[ii][params.local_rank * params.elts_total + iter_offset]) = + *reinterpret_cast(&local_input_buffer[iter_offset]); + } + } else { + *reinterpret_cast(&local_shared_buffer[iter_offset]) = + *reinterpret_cast(&local_input_buffer[iter_offset]); + } + } + // wait for equivalent blocks of other GPUs to have copied data to their shareable buffer + block_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); + } else { + // In the non-copy case, we assume that once the kernel has been started, data is ready to be consumed + multi_gpu_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, + bidx); + } + + // Each block accumulates the values from the different GPUs on the same node. + for (size_t iter_offset = chunk_start; iter_offset < chunk_end; iter_offset += blockDim.x * PACKED_ELTS) { + // Iterate over the different ranks/devices on the node to load the values. + PackedStruct vals[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + if constexpr (PUSH_MODE) { + vals[ii].packed = + *reinterpret_cast(&buffers[params.local_rank][ii * params.elts_total + iter_offset]); + } else { + vals[ii].packed = *reinterpret_cast(&buffers[ii][iter_offset]); + } + } + + // Sum the values from the different ranks. + PackedStruct sums; + sums.packed = {0, 0, 0, 0}; +#pragma unroll + for (int rank = 0; rank < RANKS_PER_NODE; ++rank) { + // Always reduce from rank 0 to ensure stable reduce order. + int ii = (rank + RANKS_PER_NODE - params.local_rank) % RANKS_PER_NODE; + sums.packed = add128b(sums, vals[ii]); + } + + // Store to the destination buffer. + *reinterpret_cast(&local_output_buffer[iter_offset]) = sums.packed; + } +} + +template +static __global__ void twoShotAllReduceKernel(AllReduceParams params) { + // Suppose that two GPUs participate in the AR exchange, and we start two blocks. + // The message is partitioned into chunks as detailed below: + // message + // |-------------------| + // |--GPU 0--|--GPU 1--| (GPU responsibility parts) + // GPU 0 | B0 | B1 | B0 | B1 | + // GPU 1 | B0 | B1 | B0 | B1 | + // + // Here the step-by-step behavior of one block: + // 1. B0 copies all chunks is it responsible for, from local_input to shareable buffer + // 2. B0 on GPU 0 and B0 on GPU 1 wait for each other (block_barrier #0) + // 3. B0 on GPU 0 gather and sum the B0 chunks from GPU 1, that are in the GPU 0 responsibility + // part (the first half of the message, see GPU responsibility row above) + // 3bis. Likewise, B0 on GPU 1 copies and sum the chunks for GPU 0, + // where GPU 1 is responsible: the second half of the message. + // 4. B0 on GPU 0 and B0 on GPU 1 wait for each other (block_barrier #1) + // 5. B0 writes result to local_output. It gathers each chunk from its responsible GPU. + // For example, here it reads the first chunk from GPU 0 and second chunk from GPU 1. + // + // With COPY_INPUT == false, skip step 1. and use gpu_barrier instead of block barrier during step 2. + // We only to know if the other GPU as arrived at the AR kernel, that would mean that data is ready + // to be read. + // + // Note that compared to one-shot, one block (CTA) writes multiple input chunks and write multiple output chunks. + // However, it's only responsible for the summation of a single chunk. + // + // With PUSH_MODE, we consider that the shared buffer is of size: + // params.peer_comm_buffer_ptrs: [world_size, world_size, message_size / world_size] + // + // Here the step-by-step behavior of one block: + // 1. B0 push the chunks is it responsible for into the corresponding GPUs: + // params.peer_comm_buffer_ptrs[target_gpu, local_gpu, current B0 slice] + // 2. block sync so the blocks have been shared by other GPUs + // 3. Reduce along second dimension params.peer_comm_buffer_ptrs[local_gpu, :, B0 slice] + // 4. block barrier (corresponding blocks have finished reduction) + // 5. pull and write on local buffer, by reading params.peer_comm_buffer_ptrs[:, 0, B0 slice] (reduction result is + // written at index 0 of 2nd dim) + + int const bidx = blockIdx.x; + int const tidx = threadIdx.x; + + // The number of elements packed into one for comms + static constexpr int PACKED_ELTS = 16 / sizeof(T); + using PackedType = typename PackedOn16Bytes::Type; + + [[maybe_unused]] T const *local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + [[maybe_unused]] T *local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); + [[maybe_unused]] T *local_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + + size_t const chunk_start = bidx * params.elts_per_block + tidx * PACKED_ELTS; + size_t const chunk_end = min(chunk_start + params.elts_per_block, params.elts_per_rank); + + T *buffers[RANKS_PER_NODE]; + int ranks[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + // A mapping of the ranks to scatter reads as much as possible + int rank = (params.local_rank + ii) % RANKS_PER_NODE; + ranks[ii] = rank; + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + } + + if constexpr (PUSH_MODE || COPY_INPUT) { + // Copy all blocks from local buffer to shareable buffer + for (size_t local_offset = chunk_start; local_offset < chunk_end; local_offset += blockDim.x * PACKED_ELTS) { +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + size_t offset_rank = ii * params.elts_per_rank + local_offset; + if (offset_rank >= params.elts_total) { + continue; + } + + if constexpr (PUSH_MODE) { + *reinterpret_cast(&buffers[ii][params.local_rank * params.elts_per_rank + local_offset]) = + *reinterpret_cast(&local_input_buffer[offset_rank]); + } else { + *reinterpret_cast(&local_shared_buffer[offset_rank]) = + *reinterpret_cast(&local_input_buffer[offset_rank]); + } + } + } + block_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); + } else { + // In the non-copy case, we assume that once the kernel has been started, data is ready to be consumed + multi_gpu_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, + bidx); + } + + // Each block accumulates the values from the different GPUs on the same node. + for (size_t local_offset = chunk_start; local_offset < chunk_end; local_offset += blockDim.x * PACKED_ELTS) { + size_t const responsible_block_offset = local_offset + params.rank_offset; + + // Iterate over the different ranks/devices on the node to load the values. + PackedType vals[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + if constexpr (PUSH_MODE) { + vals[ii].packed = + *reinterpret_cast(&local_shared_buffer[ii * params.elts_per_rank + local_offset]); + } else { + vals[ii].packed = *reinterpret_cast(&buffers[ii][responsible_block_offset]); + } + } + + // Sum the values from the different ranks. + PackedType sums; + sums.packed = {0, 0, 0, 0}; +#pragma unroll + for (int rank = 0; rank < RANKS_PER_NODE; ++rank) { + // Always reduce from rank 0 to ensure stable reduce order. + int ii = (rank + RANKS_PER_NODE - params.local_rank) % RANKS_PER_NODE; + sums.packed = add128b(sums, vals[ii]); + } + + // Store to the local buffer. + if constexpr (PUSH_MODE) { + *reinterpret_cast(&local_shared_buffer[local_offset]) = sums.packed; + } else { + *reinterpret_cast(&local_shared_buffer[responsible_block_offset]) = sums.packed; + } + } + + block_barrier(params.peer_barrier_ptrs_out, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); + + // Gather all needed elts from other intra-node ranks + for (size_t local_offset = chunk_start; local_offset < chunk_end; local_offset += blockDim.x * PACKED_ELTS) { +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + // use round-robin gathering from other ranks + size_t offset_rank = ranks[ii] * params.elts_per_rank + local_offset; + if (offset_rank >= params.elts_total) { + continue; + } + + if constexpr (PUSH_MODE) { + *reinterpret_cast(&local_output_buffer[offset_rank]) = + *reinterpret_cast(&buffers[ii][local_offset]); + } else { + *reinterpret_cast(&local_output_buffer[offset_rank]) = + *reinterpret_cast(&buffers[ii][offset_rank]); + } + } + } +} + +bool ConfigurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t world_size, + onnxruntime::MLDataType type) { + size_t elts_per_thread = 16 / type->Size(); + int const msg_align = (algo == AllReduceStrategyType::TWOSHOT) ? world_size * elts_per_thread : elts_per_thread; + bool supported_algo = (algo == AllReduceStrategyType::ONESHOT || algo == AllReduceStrategyType::TWOSHOT); + return supported_algo && (msg_size % msg_align == 0); +} + +std::tuple kernelLaunchConfig(AllReduceStrategyType algo, AllReduceParams ¶m, size_t elts_per_thread) { + int blocks_per_grid = 1, threads_per_block = DEFAULT_BLOCK_SIZE; + + switch (algo) { + case AllReduceStrategyType::ONESHOT: { + ORT_ENFORCE(param.elts_total % elts_per_thread == 0); + size_t const total_threads = roundUp(param.elts_total / elts_per_thread, WARP_SIZE); + threads_per_block = std::min(DEFAULT_BLOCK_SIZE, total_threads); + blocks_per_grid = std::min(static_cast(MAX_ALL_REDUCE_BLOCKS), + divUp(total_threads, static_cast(threads_per_block))); + param.elts_per_block = roundUp(divUp(param.elts_total, static_cast(blocks_per_grid)), elts_per_thread); + break; + } + case AllReduceStrategyType::TWOSHOT: { + ORT_ENFORCE(param.elts_total % (elts_per_thread * param.ranks_per_node) == 0); + size_t const total_threads = roundUp(param.elts_total / (elts_per_thread * param.ranks_per_node), WARP_SIZE); + + /* + threads_per_block = std::min(DEFAULT_BLOCK_SIZE, total_threads); + blocks_per_grid = std::min(static_cast(MAX_ALL_REDUCE_BLOCKS), divUp(total_threads, threads_per_block)); + */ + + while (total_threads % blocks_per_grid != 0 || total_threads / blocks_per_grid > DEFAULT_BLOCK_SIZE) { + blocks_per_grid += 1; + } + + threads_per_block = total_threads / blocks_per_grid; + + // NOTE: need to adjust here + if (static_cast(blocks_per_grid) > MAX_ALL_REDUCE_BLOCKS) { + size_t iter_factor = 1; + while (blocks_per_grid / iter_factor > MAX_ALL_REDUCE_BLOCKS || blocks_per_grid % iter_factor) { + iter_factor += 1; + } + blocks_per_grid /= iter_factor; + } + param.elts_per_rank = param.elts_total / param.ranks_per_node; + param.rank_offset = param.local_rank * param.elts_per_rank; + param.elts_per_block = + roundUp(divUp(param.elts_per_rank, static_cast(blocks_per_grid)), elts_per_thread); + break; + } + default: + ORT_THROW("Algorithm not supported here."); + } + + return std::make_tuple(blocks_per_grid, threads_per_block); +} + +template +void AllReduceDispatchMemcpy(AllReduceStrategyType algo, AllReduceStrategyConfig config, AllReduceParams ¶m, + cudaStream_t stream) { + ORT_ENFORCE(!(USE_MEMCPY && PUSH_MODE), "Memcpy cannot be used with PUSH_MODE."); + size_t elts_per_thread = 16 / sizeof(T); + auto [blocks_per_grid, threads_per_block] = kernelLaunchConfig(algo, param, elts_per_thread); + + if (USE_MEMCPY) { + cudaMemcpyAsync(param.peer_comm_buffer_ptrs[param.local_rank], param.local_input_buffer_ptr, + param.elts_total * sizeof(T), cudaMemcpyDeviceToDevice, stream); + } + + if (algo == AllReduceStrategyType::ONESHOT) { + oneShotAllReduceKernel + <<>>(param); + } else { + twoShotAllReduceKernel + <<>>(param); + } +} + +template +void AllReduceDispatchPushMode(AllReduceStrategyType algo, AllReduceStrategyConfig config, AllReduceParams ¶m, + cudaStream_t stream) { + if (static_cast>(config) & + static_cast>(AllReduceStrategyConfig::USE_MEMCPY)) { + AllReduceDispatchMemcpy(algo, config, param, stream); + } else { + AllReduceDispatchMemcpy(algo, config, param, stream); + } +} + +template //, bool USE_MEMCPY = false, bool PUSH_MODE = false> +void AllReduceDispatchRanksPerNode(AllReduceStrategyType algo, AllReduceStrategyConfig config, AllReduceParams ¶m, + cudaStream_t stream) { + if (static_cast>(config) & + static_cast>(AllReduceStrategyConfig::PUSH_MODE)) { + AllReduceDispatchPushMode(algo, config, param, stream); + } else { + AllReduceDispatchPushMode(algo, config, param, stream); + } +} + +template +void AllReduceDispatchType(AllReduceParams ¶m, AllReduceStrategyType strategy, AllReduceStrategyConfig config, + cudaStream_t stream) { + switch (param.ranks_per_node) { + case 2: + AllReduceDispatchRanksPerNode(strategy, config, param, stream); + break; + case 4: + AllReduceDispatchRanksPerNode(strategy, config, param, stream); + break; + case 6: + AllReduceDispatchRanksPerNode(strategy, config, param, stream); + break; + case 8: + AllReduceDispatchRanksPerNode(strategy, config, param, stream); + break; + default: + ORT_THROW("Custom all reduce only supported on {2, 4, 6, 8} GPUs per node."); + } +} + +AllReduceParams AllReduceParams::deserialize(const int32_t *buffer, size_t tp_size, size_t tp_rank, uint32_t flag) { + void *const *buffer_ptrs = reinterpret_cast(buffer); + AllReduceParams params; + + for (size_t i = 0; i < tp_size; ++i) { + params.peer_comm_buffer_ptrs[i] = buffer_ptrs[i]; + } + for (size_t i = 0; i < tp_size; ++i) { + params.peer_barrier_ptrs_in[i] = reinterpret_cast(buffer_ptrs[tp_size + i]); + } + for (size_t i = 0; i < tp_size; ++i) { + params.peer_barrier_ptrs_out[i] = reinterpret_cast(buffer_ptrs[2 * tp_size + i]); + } + params.barrier_flag = flag; + params.ranks_per_node = tp_size; + params.rank = tp_rank; + params.local_rank = tp_rank; + + return params; +} + +void CustomAllReduce(AllReduceParams ¶ms, onnxruntime::MLDataType data_type, AllReduceStrategyType strategy, + AllReduceStrategyConfig config, cudaStream_t stream) { + ORT_ENFORCE(ConfigurationSupported(strategy, params.elts_total, params.ranks_per_node, data_type), + "Custom all-reduce configuration unsupported"); + if (data_type == onnxruntime::DataTypeImpl::GetType()) { + AllReduceDispatchType(params, strategy, config, stream); + } else if (data_type == onnxruntime::DataTypeImpl::GetType()) { + AllReduceDispatchType(params, strategy, config, stream); + } else { + ORT_THROW("Unsupported data type for CustomAllReduce"); + } +} + +size_t GetMaxRequiredWorkspaceSize(int world_size) { + if (world_size <= 2) { + return 16 * 1000 * 1000; + } + return 8 * 1000 * 1000; +} + +Status SetPeerAccess(int rank, int world_size, bool enable, int &can_access_peer) { + const int src_node = rank; + + for (int dst_node = 0; dst_node < world_size; dst_node++) { + if (dst_node == src_node) { + continue; + } + + CUDA_RETURN_IF_ERROR(cudaDeviceCanAccessPeer(&can_access_peer, src_node, dst_node)); + + if (!can_access_peer) { + return Status::OK(); + } + + if (enable) { + cudaDeviceEnablePeerAccess(dst_node, 0); + } else { + cudaDeviceDisablePeerAccess(dst_node); + } + + auto const error = cudaGetLastError(); + if (error != cudaErrorPeerAccessAlreadyEnabled && error != cudaErrorPeerAccessNotEnabled) { + CUDA_RETURN_IF_ERROR(error); + } + } + + return Status::OK(); +} + +AllReduceStrategyType SelectImplementation(size_t message_size, int rank, int world_size, + onnxruntime::MLDataType type) { + AllReduceStrategyType strategy = AllReduceStrategyType::NCCL; + if (type != onnxruntime::DataTypeImpl::GetType() && + type != onnxruntime::DataTypeImpl::GetType()) { + return strategy; + } + + if (world_size != 2 && world_size != 4 && world_size != 6 && world_size != 8) { + return strategy; + } + + int can_access_peer = 0; + ORT_ENFORCE(SetPeerAccess(rank, world_size, true, can_access_peer) == Status::OK()); + // If P2P is not enabled, we cannot use the custom allreduce, so default to NCCL. + if (!can_access_peer) { + return strategy; + } + + const size_t maxWorkspaceSize = GetMaxRequiredWorkspaceSize(world_size); + const size_t message_size_bytes = message_size * type->Size(); + + if (message_size_bytes <= maxWorkspaceSize) { + if (world_size <= 2) { + strategy = AllReduceStrategyType::ONESHOT; + } else if (world_size <= 4) { + if (message_size_bytes < 1 * 1000 * 1000) { + strategy = AllReduceStrategyType::ONESHOT; + } else { + strategy = AllReduceStrategyType::TWOSHOT; + } + } else { + if (message_size_bytes < 500 * 1000) { + strategy = AllReduceStrategyType::ONESHOT; + } else { + strategy = AllReduceStrategyType::TWOSHOT; + } + } + } + + if (!ConfigurationSupported(strategy, message_size, world_size, type)) { + strategy = AllReduceStrategyType::NCCL; + } + + return strategy; +} + +#endif + +} // namespace collective +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h new file mode 100644 index 0000000000..3ca3c1dd16 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h @@ -0,0 +1,81 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_kernel.h" +#include + +namespace onnxruntime { +namespace cuda { +namespace collective { + +#if defined(USE_MPI) || defined(USE_NCCL) + +constexpr size_t WARP_SIZE = 32; +constexpr size_t MAX_ALL_REDUCE_BLOCKS = 24; +constexpr size_t MAX_RANKS_PER_NODE = 8; +constexpr size_t DEFAULT_BLOCK_SIZE = 1024; + +// Warning: python definition is in tensorrt_llm/functional.py +// they must be kept in sync +enum class AllReduceStrategyType : int8_t { + NCCL = 0, + ONESHOT = 1, + TWOSHOT = 2, + AUTO = 3, +}; + +enum class AllReduceStrategyConfig : int8_t { + USE_MEMCPY = 1 << 0, + PUSH_MODE = 1 << 1, +}; + +struct AllReduceParams { + size_t elts_total; + size_t elts_per_rank; + size_t elts_per_block; + size_t rank_offset; + size_t ranks_per_node, rank, local_rank; + uint32_t barrier_flag; + uint32_t* peer_barrier_ptrs_in[MAX_RANKS_PER_NODE]; + uint32_t* peer_barrier_ptrs_out[MAX_RANKS_PER_NODE]; + void* peer_comm_buffer_ptrs[MAX_RANKS_PER_NODE]; + void* local_output_buffer_ptr; + const void* local_input_buffer_ptr; + + static AllReduceParams deserialize(const int32_t* buffer, size_t tp_size, size_t tp_rank, uint32_t flag); +}; + +bool ConfigurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t world_size, + onnxruntime::MLDataType type); + +void CustomAllReduce(AllReduceParams& params, onnxruntime::MLDataType data_type, AllReduceStrategyType strategy, + AllReduceStrategyConfig config, cudaStream_t stream); + +size_t GetMaxRequiredWorkspaceSize(int world_size); + +Status SetPeerAccess(int rank, int world_size, bool enable, int& can_access_peer); + +AllReduceStrategyType SelectImplementation(size_t message_size, int rank, int world_size, onnxruntime::MLDataType type); + +#endif + +} // namespace collective +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc new file mode 100644 index 0000000000..b4e602228e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc @@ -0,0 +1,123 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "ipc_utils.h" +#include "mpi_include.h" + +namespace onnxruntime { +namespace cuda { +namespace collective { + +#if defined(USE_MPI) || defined(USE_NCCL) + +using namespace onnxruntime; + +IpcMemory::IpcMemory(int rank, int world_size, std::size_t buffer_size) + : rank_(rank), world_size_(world_size), m_comm_ptrs_(world_size), mbuffer_size_(buffer_size) { + ORT_ENFORCE(AllocateIpcMemory() == Status::OK()); +} + +Status IpcMemory::AllocateIpcMemory() { + void* m_buffer_ptr; + CUDA_RETURN_IF_ERROR(cudaMalloc(&m_buffer_ptr, mbuffer_size_)); + m_buffer_uptr_ = CudaMemPtrT{m_buffer_ptr, CudaDeleter()}; + + CUDA_RETURN_IF_ERROR(cudaMemset(m_buffer_uptr_.get(), 0, mbuffer_size_)); + + cudaIpcMemHandle_t local_handle; + CUDA_RETURN_IF_ERROR(cudaIpcGetMemHandle(&local_handle, m_buffer_uptr_.get())); + + // Assume no pipeline parallelism. + InlinedVector serial_handles(CUDA_IPC_HANDLE_SIZE * world_size_, 0); + +#ifdef USE_MPI + MPI_CHECK(MPI_Allgather(local_handle.reserved, CUDA_IPC_HANDLE_SIZE, MPI_BYTE, serial_handles.data(), + CUDA_IPC_HANDLE_SIZE, MPI_BYTE, MPI_COMM_WORLD)); +#else + // TODO: Implement this for NCCL. + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Please compile ORT with USE_MPI."); +#endif + + InlinedVector handles(world_size_); + for (size_t i = 0; i < handles.size(); ++i) { + memcpy(handles[i].reserved, &serial_handles[i * CUDA_IPC_HANDLE_SIZE], CUDA_IPC_HANDLE_SIZE); + } + + for (size_t node_id = 0; node_id < handles.size(); node_id++) { + if ((int)node_id == rank_) { + m_comm_ptrs_[node_id] = m_buffer_uptr_.get(); + } else { + uint8_t* foreign_buffer; + CUDA_RETURN_IF_ERROR(cudaIpcOpenMemHandle( + reinterpret_cast(&foreign_buffer), handles[node_id], cudaIpcMemLazyEnablePeerAccess)); + m_ipc_uptrs_.emplace_back(foreign_buffer, IpcDeleter()); + m_comm_ptrs_[node_id] = foreign_buffer; + } + } + + return Status::OK(); +} + +IpcMemory::~IpcMemory() { +} + +Status GetCustomAllReduceWorkspace(int rank, int world_size, size_t input_size, + IPCMemoryResourcePack& ipc_mem_res_pack) { + if (input_size <= ipc_mem_res_pack.max_input_size) { + return Status::OK(); + } + + const std::size_t buffer_size = world_size * input_size; + + InlinedVector>& m_ipc_memory_handles = ipc_mem_res_pack.m_ipc_momery_handles; + const size_t handles_size{m_ipc_memory_handles.size()}; + constexpr size_t k_num_handles{3}; + + m_ipc_memory_handles.emplace_back(std::make_unique(rank, world_size, buffer_size)); + m_ipc_memory_handles.emplace_back( + std::make_unique(rank, world_size, IpcMemory::FLAGS_SIZE * world_size)); + m_ipc_memory_handles.emplace_back( + std::make_unique(rank, world_size, IpcMemory::FLAGS_SIZE * world_size)); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + + InlinedVector& m_comm_ptrs = ipc_mem_res_pack.m_comm_ptrs; + m_comm_ptrs.reserve(k_num_handles * world_size); + m_comm_ptrs.resize(k_num_handles * world_size); + + for (size_t mem_idx = handles_size; mem_idx < m_ipc_memory_handles.size(); mem_idx++) { + const auto& mem_comm_ptrs = m_ipc_memory_handles[mem_idx]->GetCommPtrsTensor(); + for (size_t tpIdx = 0; tpIdx < static_cast(world_size); tpIdx++) { + m_comm_ptrs[(mem_idx - handles_size) * world_size + tpIdx] = mem_comm_ptrs[tpIdx]; + } + } + + ipc_mem_res_pack.max_input_size = input_size; + + return Status::OK(); +} + +IPCMemoryResourcePack& IPCMemoryResourcePack::GetGlobalInstance() { + static IPCMemoryResourcePack ipc_mem_res_pack; + return ipc_mem_res_pack; +} + +#endif + +} // namespace collective +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/ipc_utils.h b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.h new file mode 100644 index 0000000000..cda0f3437b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.h @@ -0,0 +1,89 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "custom_reduce_impl.h" + +namespace onnxruntime { +namespace cuda { +namespace collective { + +#if defined(USE_MPI) || defined(USE_NCCL) + +struct CudaDeleter { + void operator()(void* ptr) const noexcept { + if (ptr != nullptr) { + cudaFree(ptr); + } + } +}; + +struct IpcDeleter { + void operator()(void* ptr) const noexcept { + if (ptr != nullptr) { + cudaIpcCloseMemHandle(ptr); + } + } +}; + +class IpcMemory { + public: + size_t static constexpr FLAGS_SIZE = (MAX_ALL_REDUCE_BLOCKS + 1) * sizeof(uint32_t); + + IpcMemory(int rank, int world_size, std::size_t buffer_size); + ~IpcMemory(); + + const InlinedVector& GetCommPtrsTensor() const { + return m_comm_ptrs_; + } + + private: + Status AllocateIpcMemory(); + + int rank_; + int world_size_; + InlinedVector m_comm_ptrs_; + std::size_t mbuffer_size_; + + using CudaMemPtrT = std::unique_ptr; + CudaMemPtrT m_buffer_uptr_; + + using IpcMemPtrT = std::unique_ptr; + InlinedVector m_ipc_uptrs_; +}; + +// A global resource pack for IPC memory used in custom reduce kernel. +// Resource retrieval and deserialization are made atomic to thread safety of accessing it. +struct IPCMemoryResourcePack { + InlinedVector> m_ipc_momery_handles; + InlinedVector m_comm_ptrs; + size_t max_input_size{0}; + uint32_t counter{0}; + + static IPCMemoryResourcePack& GetGlobalInstance(); +}; + +Status +GetCustomAllReduceWorkspace(int rank, int world_size, size_t input_size, IPCMemoryResourcePack& ipc_mem_res_pack); + +#endif + +} // namespace collective +} // namespace cuda +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc index 0f42363bca..feb6613690 100644 --- a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc @@ -251,18 +251,28 @@ AllReduce::AllReduce(const OpKernelInfo& info) : NcclKernel(info) { } Status AllReduce::ComputeInternal(OpKernelContext* context) const { - ncclComm_t comm = nccl_->Comm(); - auto input_tensor = context->Input(0); const void* input_data = input_tensor->DataRaw(); + const auto in_shape = input_tensor->Shape(); int64_t input_count = in_shape.Size(); void* output_data = context->Output(0, in_shape)->MutableDataRaw(); +#ifndef USE_ROCM + return FuncCustomAllReduce(nccl_, + Stream(context), + input_data, + output_data, + input_count, + input_tensor->DataType(), + onnxruntime::cuda::collective::IPCMemoryResourcePack::GetGlobalInstance()); +#else + ncclComm_t comm = nccl_->Comm(); ncclDataType_t dtype = GetNcclDataType(input_tensor->DataType()); NCCL_RETURN_IF_ERROR(ncclAllReduce(input_data, output_data, input_count, dtype, ncclSum, comm, Stream(context))); return Status::OK(); +#endif } AllGather::AllGather(const OpKernelInfo& info) : NcclKernel(info) { @@ -418,6 +428,58 @@ Status FuncAllReduce( return Status::OK(); } +#ifndef USE_ROCM +Status FuncCustomAllReduce( + NcclContext* nccl, + cudaStream_t stream, + const void* input_data, + void* output_data, + int64_t input_count, + onnxruntime::MLDataType data_type, + onnxruntime::cuda::collective::IPCMemoryResourcePack& ipc_mem_res_pack) { + int rank = nccl->Rank(); + int world_size = nccl->Size(); + + onnxruntime::cuda::collective::AllReduceStrategyType runtime_strategy = + onnxruntime::cuda::collective::SelectImplementation(input_count, rank, world_size, data_type); + + if (runtime_strategy == onnxruntime::cuda::collective::AllReduceStrategyType::NCCL) { + ncclDataType_t dtype = GetNcclDataType(data_type); + NCCL_RETURN_IF_ERROR(ncclAllReduce(input_data, output_data, input_count, dtype, ncclSum, nccl->Comm(), stream)); + + return Status::OK(); + } + + onnxruntime::cuda::collective::AllReduceStrategyConfig m_config = + onnxruntime::cuda::collective::AllReduceStrategyConfig::USE_MEMCPY; + + static std::mutex s_mutex; + std::unique_lock lock(s_mutex); + ORT_RETURN_IF_ERROR(onnxruntime::cuda::collective::GetCustomAllReduceWorkspace(rank, + world_size, + input_count * data_type->Size(), + ipc_mem_res_pack)); + + onnxruntime::cuda::collective::AllReduceParams params = onnxruntime::cuda::collective::AllReduceParams::deserialize( + reinterpret_cast(ipc_mem_res_pack.m_comm_ptrs.data()), + world_size, + rank, + ++ipc_mem_res_pack.counter); + lock.unlock(); + + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + + params.local_output_buffer_ptr = output_data; + params.local_input_buffer_ptr = input_data; + params.elts_total = input_count; + + onnxruntime::cuda::collective::CustomAllReduce(params, data_type, runtime_strategy, m_config, stream); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + + return Status::OK(); +} +#endif + static std::vector CalculatePermToSwapAxes( const int64_t axis, const int64_t another_axis, diff --git a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h index 9ea61f2bd9..10b6f7dd56 100644 --- a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h +++ b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h @@ -5,6 +5,13 @@ #include "core/providers/cuda/cuda_kernel.h" +#if defined(ORT_USE_NCCL) || defined(ORT_USE_MPI) +#ifndef USE_ROCM +#include "custom_reduce_impl.h" +#include "ipc_utils.h" +#endif +#endif + #if defined(ORT_USE_NCCL) #include #include @@ -100,6 +107,17 @@ Status FuncAllReduce( const Tensor* input, Tensor* output); +#ifndef USE_ROCM +Status FuncCustomAllReduce( + NcclContext* nccl, + cudaStream_t stream, + const void* input_data, + void* output_data, + int64_t input_count, + onnxruntime::MLDataType data_type, + onnxruntime::cuda::collective::IPCMemoryResourcePack& ipc_mem_res_pack); +#endif + void FuncAllGather( const NcclKernel* nccl_kernel, OpKernelContext* ctx, diff --git a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc index 1dbbe8c4e7..013b7e1779 100644 --- a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc +++ b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include "core/common/safeint.h" @@ -18,6 +19,18 @@ namespace cuda { #if defined(ORT_USE_NCCL) +#define CHECK_CUDA(res) \ + if (res != cudaSuccess) { \ + cuda_result = res; \ + return; \ + } + +#define CHECK_NCCL(res) \ + if (res != ncclSuccess) { \ + nccl_result = res; \ + return; \ + } + #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ ShardedMoE, kMSDomain, 1, T, kCudaExecutionProvider, \ @@ -32,8 +45,9 @@ ShardedMoE::ShardedMoE(const OpKernelInfo& op_kernel_info) : NcclKernel(op_ke ORT_ENFORCE(op_kernel_info.GetAttr("tensor_shards", &tensor_shards_).IsOK()); ORT_ENFORCE(op_kernel_info.GetAttr("local_experts_start_index", &local_experts_start_index_).IsOK()); rank_to_experts_start_index_.resize(nccl_->Size()); - // Initialize rank_to_experts_start_index_[0] to a value to convey that it is not initialized. - rank_to_experts_start_index_[0] = std::numeric_limits::min(); + + auto allocator = op_kernel_info.GetAllocator(OrtMemTypeDefault); + ORT_ENFORCE(SynchronizeExpertsStartIndex(allocator) == Status::OK()); } template @@ -47,10 +61,6 @@ Status ShardedMoE::ComputeInternal(OpKernelContext* context) const { AllocatorPtr allocator; ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); - // Create a {Rank, ExpertsStartIndex} map on Host. - AutoDestoryCudaEvent cuda_event; - cudaEvent_t& copy_event = cuda_event.Get(); - const Tensor* input = context->Input(0); const Tensor* router_probs = context->Input(1); const Tensor* fc1_experts_weights = context->Input(2); @@ -68,10 +78,6 @@ Status ShardedMoE::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(moe_params.num_experts % nccl_->Size() == 0, "num_experts should be divisible by world_size"); - if (moe_params.parallel_type == MoEParallelType::EP || moe_params.parallel_type == MoEParallelType::EPAndTP) { - ORT_RETURN_IF_ERROR(SynchronizeExpertsStartIndex(allocator, context, copy_event)); - } - ort_fastertransformer::CutlassMoeFCRunner moe_runner(sm, fc3_experts_weights_optional != nullptr, normalize_routing_weights_); @@ -133,11 +139,14 @@ Status ShardedMoE::ComputeInternal(OpKernelContext* context) const { if (moe_params.parallel_type == MoEParallelType::TP) { ORT_ENFORCE(moe_params.tensor_shards == nccl_->Size()); - NCCL_RETURN_IF_ERROR(ncclGroupStart()); - NCCL_RETURN_IF_ERROR(ncclAllReduce(reinterpret_cast(fc2_output.get()), - reinterpret_cast(fc2_output_bc.get()), fc2_output_size / sizeof(CudaT), - GetNcclDataType(input->DataType()), ncclSum, nccl_->Comm(), Stream(context))); - NCCL_RETURN_IF_ERROR(ncclGroupEnd()); + + ORT_RETURN_IF_ERROR(FuncCustomAllReduce(nccl_, + Stream(context), + fc2_output.get(), + fc2_output_bc.get(), + static_cast(fc2_output_size / sizeof(CudaT)), + input->DataType(), + collective::IPCMemoryResourcePack::GetGlobalInstance())); } if (moe_params.parallel_type == MoEParallelType::EP) { @@ -145,9 +154,7 @@ Status ShardedMoE::ComputeInternal(OpKernelContext* context) const { size_t stride_bytes = stride_count * sizeof(CudaT); int64_t total_past_rows = 0; int64_t total_covered_rows = 0; - if (copy_event != nullptr) { - CUDA_RETURN_IF_ERROR(cudaEventSynchronize(copy_event)); - } + NCCL_RETURN_IF_ERROR(ncclGroupStart()); for (int rank = 0; rank < nccl_->Size(); ++rank) { int64_t experts_start_index = rank_to_experts_start_index_[rank]; @@ -175,36 +182,25 @@ Status ShardedMoE::ComputeInternal(OpKernelContext* context) const { } template -Status ShardedMoE::SynchronizeExpertsStartIndex(AllocatorPtr& allocator, OpKernelContext* context, - cudaEvent_t& cuda_event) const { - if (rank_to_experts_start_index_[0] != std::numeric_limits::min()) { - return Status::OK(); - } - - auto stream = context->GetComputeStream(); - +Status ShardedMoE::SynchronizeExpertsStartIndex(AllocatorPtr& allocator) const { using IndexType = int64_t; size_t IndexTypeSize = sizeof(IndexType); IAllocatorUniquePtr experts_start_index_d = - IAllocator::MakeUniquePtr(allocator, 1, false, stream); + IAllocator::MakeUniquePtr(allocator, 1, false); IAllocatorUniquePtr rank_to_experts_start_index_d = - IAllocator::MakeUniquePtr(allocator, nccl_->Size(), false, stream); + IAllocator::MakeUniquePtr(allocator, nccl_->Size(), false); - // Only happens in the first run. - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(experts_start_index_d.get(), &local_experts_start_index_, IndexTypeSize, - cudaMemcpyHostToDevice, Stream(context))); + CUDA_RETURN_IF_ERROR(cudaMemcpy(experts_start_index_d.get(), &local_experts_start_index_, IndexTypeSize, + cudaMemcpyHostToDevice)); NCCL_RETURN_IF_ERROR(ncclAllGather(reinterpret_cast(experts_start_index_d.get()), reinterpret_cast(rank_to_experts_start_index_d.get()), 1, GetNcclDataType(DataTypeImpl::GetType()), nccl_->Comm(), - Stream(context))); - // The const_cast<> violates the const modifier to make sure the synchronization happens only once per session. - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(const_cast(rank_to_experts_start_index_.data()), - rank_to_experts_start_index_d.get(), nccl_->Size() * IndexTypeSize, - cudaMemcpyDeviceToHost, Stream(context))); + nullptr)); - CUDA_RETURN_IF_ERROR(cudaEventCreateWithFlags(&cuda_event, cudaEventDisableTiming)); - CUDA_RETURN_IF_ERROR(cudaEventRecord(cuda_event, Stream(context))); + CUDA_RETURN_IF_ERROR(cudaMemcpy(const_cast(rank_to_experts_start_index_.data()), + rank_to_experts_start_index_d.get(), nccl_->Size() * IndexTypeSize, + cudaMemcpyDeviceToHost)); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h index 827283a794..a0f1032834 100644 --- a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h +++ b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h @@ -23,11 +23,11 @@ class ShardedMoE final : public NcclKernel, public MoEBase { Status ComputeInternal(OpKernelContext* ctx) const override; private: - Status SynchronizeExpertsStartIndex(AllocatorPtr& alloc, OpKernelContext* ctx, cudaEvent_t& cuda_event) const; + Status SynchronizeExpertsStartIndex(AllocatorPtr& alloc) const; int64_t local_experts_start_index_; int64_t tensor_shards_; - std::vector rank_to_experts_start_index_; + InlinedVector rank_to_experts_start_index_; }; #endif diff --git a/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h b/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h index 54c024793f..2df0a38d22 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h +++ b/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h @@ -194,5 +194,16 @@ struct Channels { static constexpr size_t W = 3; }; +// Calculates ceil(a / b). User must be careful to ensure that there +// is no overflow or underflow in the calculation. +template +constexpr T divUp(T a, T b) { return (a + b - (T)1) / b; } + +// Rounds a up to the next highest multiple of b. User must be careful +// to ensure that there is no overflow or underflow in the calculation +// of divUp. +template +constexpr T roundUp(T a, T b) { return divUp(a, b) * b; } + } // namespace cuda } // namespace onnxruntime diff --git a/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-ortmodule-distributed-test-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-ortmodule-distributed-test-ci-pipeline.yml index 2c6b6183a9..82aa7b24e7 100644 --- a/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-ortmodule-distributed-test-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-ortmodule-distributed-test-ci-pipeline.yml @@ -122,8 +122,8 @@ stages: --volume $(Build.BinariesDirectory):/build \ --volume /mnist:/mnist \ onnxruntime_ortmodule_distributed_tests_image \ - bash -c "rm -rf /build/RelWithDebInfo/onnxruntime/ && python3 -m pip install mpi4py onnxscript && python3 -m pip install /build/RelWithDebInfo/dist/onnxruntime*.whl && mpirun -n 4 -x NCCL_DEBUG=INFO python /onnxruntime_src/onnxruntime/test/python/onnxruntime_test_collective.py && mpirun -n 2 -x NCCL_DEBUG=INFO python /onnxruntime_src/onnxruntime/test/python/onnxruntime_test_distributed.py" \ - displayName: 'Run onnxruntime_test_collective.py' + bash -c "rm -rf /build/RelWithDebInfo/onnxruntime/ && python3 -m pip install mpi4py onnxscript && python3 -m pip install /build/RelWithDebInfo/dist/onnxruntime*.whl && mpirun -n 4 -x NCCL_DEBUG=INFO python /onnxruntime_src/onnxruntime/test/python/onnxruntime_test_collective.py && mpirun -n 2 -x NCCL_DEBUG=INFO python /onnxruntime_src/onnxruntime/test/python/onnxruntime_test_distributed.py && mpirun -n 2 -x NCCL_DEBUG=INFO python /onnxruntime_src/onnxruntime/test/python/transformers/sharded_moe/test_sharded_moe.py" \ + displayName: 'Run onnxruntime_test_collective.py, onnxruntime_test_distributed.py and test_sharded_moe.py' condition: succeededOrFailed() timeoutInMinutes: 30