diff --git a/ggml/src/ggml-cuda/allreduce.cu b/ggml/src/ggml-cuda/allreduce.cu index 2aba8d03b6..c1df8d8d2c 100644 --- a/ggml/src/ggml-cuda/allreduce.cu +++ b/ggml/src/ggml-cuda/allreduce.cu @@ -505,60 +505,70 @@ bool ggml_cuda_ar_allreduce( return false; } - const int64_t ne = ggml_nelements(tensors[0]); - const size_t bytes = (size_t)ne * type_size; + const int64_t ne = ggml_nelements(tensors[0]); if (ne == 0) { return true; } - if (bytes > p->buf_bytes) { + if (p->buf_bytes < type_size) { return false; } - // Cycle through the event pool. On the second pass through the ring, - // synchronise on the slot's ker event before touching arrival ints — - // the event and arrival pools wrap in lock-step so this guarantees the - // kernels which last used this slot have finished. - const int slot = static_cast(p->call_count % GGML_CUDA_AR_POOL_SIZE); - const bool pool_lapped = p->call_count >= GGML_CUDA_AR_POOL_SIZE; - p->call_count++; + const size_t max_chunk_elems = p->buf_bytes / type_size; + GGML_ASSERT(max_chunk_elems > 0); - if (pool_lapped) { - for (int i = 0; i < n; ++i) { - ggml_cuda_set_device(p->devices[i]); - CUDA_CHECK(cudaEventSynchronize(p->ev_pool[i][slot].ker)); - } - } - - // Reset the arrival ints for this slot before any kernel can read them. - for (int i = 0; i < n; ++i) { - *ggml_cuda_ar_arrival_ptr(p, slot, i) = 0; - } - - - // Insert the kernel into each GPU's existing compute stream via events: + // Insert chunked kernels into each GPU's existing compute stream via events: // record(app, compute_stream) — capture "upstream done" // wait(internal_stream, app) — internal stream defers until then - // launch kernel on internal_stream - // record(ker, internal_stream) — capture "kernel done" - // wait(compute_stream, ker) — compute stream resumes after kernel - for (int i = 0; i < n; ++i) { - const int peer = 1 - i; // valid for n == 2 only - ggml_cuda_set_device(p->devices[i]); - auto * cuda_ctx = static_cast(backends[i]->context); - ggml_cuda_ar_event_slot & ev = p->ev_pool[i][slot]; - const bool compute = (tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) != 0; + // launch one or more chunk kernels on internal_stream + // record(ker, internal_stream) — capture "final chunk done" + // wait(compute_stream, ker) — compute stream resumes after reduce + for (int64_t chunk_start = 0; chunk_start < ne; chunk_start += (int64_t) max_chunk_elems) { + const size_t remaining_elems = (size_t) (ne - chunk_start); + const size_t chunk_elems = remaining_elems < max_chunk_elems ? remaining_elems : max_chunk_elems; + const size_t chunk_bytes = chunk_elems * type_size; - CUDA_CHECK(cudaEventRecord(ev.app, cuda_ctx->stream())); - CUDA_CHECK(cudaStreamWaitEvent(p->streams[i], ev.app)); + // Cycle through the event pool. On the second pass through the ring, + // synchronise on the slot's ker event before touching arrival ints — + // the event and arrival pools wrap in lock-step so this guarantees the + // kernels which last used this slot have finished. + const int slot = static_cast(p->call_count % GGML_CUDA_AR_POOL_SIZE); + const bool pool_lapped = p->call_count >= GGML_CUDA_AR_POOL_SIZE; + p->call_count++; - // Match the NCCL and meta-backend semantics: inactive shards - // contribute zeros to the reduction. - if (!compute) { - CUDA_CHECK(cudaMemsetAsync(tensors[i]->data, 0, bytes, p->streams[i])); + if (pool_lapped) { + for (int i = 0; i < n; ++i) { + ggml_cuda_set_device(p->devices[i]); + CUDA_CHECK(cudaEventSynchronize(p->ev_pool[i][slot].ker)); + } } + // Reset the arrival ints for this slot before any kernel can read them. + for (int i = 0; i < n; ++i) { + *ggml_cuda_ar_arrival_ptr(p, slot, i) = 0; + } + + for (int i = 0; i < n; ++i) { + const int peer = 1 - i; // valid for n == 2 only + ggml_cuda_set_device(p->devices[i]); + auto * cuda_ctx = static_cast(backends[i]->context); + ggml_cuda_ar_event_slot & ev = p->ev_pool[i][slot]; + const bool compute = (tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) != 0; + + if (chunk_start == 0) { + CUDA_CHECK(cudaEventRecord(ev.app, cuda_ctx->stream())); + CUDA_CHECK(cudaStreamWaitEvent(p->streams[i], ev.app)); + } + + char * data = static_cast(tensors[i]->data) + chunk_start * (int64_t) type_size; + + // Match the NCCL and meta-backend semantics: inactive shards + // contribute zeros to the reduction. + if (!compute) { + CUDA_CHECK(cudaMemsetAsync(data, 0, chunk_bytes, p->streams[i])); + } + #if GGML_CUDA_AR_WATCHDOG #define GGML_CUDA_AR_WDOG_EXTRA_ARGS , p->debug_ring[i], p->wdog_max_spin, i, slot #else @@ -566,30 +576,32 @@ bool ggml_cuda_ar_allreduce( #endif #define LAUNCH_AR_KERNEL(T) \ - ggml_cuda_ar_kernel<<streams[i]>>>( \ - static_cast(tensors[i]->data), \ - static_cast(tensors[i]->data), \ - reinterpret_cast(p->host_buf[i]), \ - reinterpret_cast(p->host_buf[peer]), \ - static_cast(ne), \ - ggml_cuda_ar_arrival_ptr(p, slot, i), \ - ggml_cuda_ar_arrival_ptr(p, slot, peer) \ - GGML_CUDA_AR_WDOG_EXTRA_ARGS) + ggml_cuda_ar_kernel<<streams[i]>>>( \ + reinterpret_cast(data), \ + reinterpret_cast(data), \ + reinterpret_cast(p->host_buf[i]), \ + reinterpret_cast(p->host_buf[peer]), \ + static_cast(chunk_elems), \ + ggml_cuda_ar_arrival_ptr(p, slot, i), \ + ggml_cuda_ar_arrival_ptr(p, slot, peer) \ + GGML_CUDA_AR_WDOG_EXTRA_ARGS) - switch (type) { - case GGML_TYPE_F32: LAUNCH_AR_KERNEL(float); break; - case GGML_TYPE_F16: LAUNCH_AR_KERNEL(half); break; - case GGML_TYPE_BF16: LAUNCH_AR_KERNEL(__nv_bfloat16); break; - default: GGML_ASSERT(false); - } + switch (type) { + case GGML_TYPE_F32: LAUNCH_AR_KERNEL(float); break; + case GGML_TYPE_F16: LAUNCH_AR_KERNEL(half); break; + case GGML_TYPE_BF16: LAUNCH_AR_KERNEL(__nv_bfloat16); break; + default: GGML_ASSERT(false); + } #undef LAUNCH_AR_KERNEL #undef GGML_CUDA_AR_WDOG_EXTRA_ARGS - CUDA_CHECK(cudaGetLastError()); - - CUDA_CHECK(cudaEventRecord(ev.ker, p->streams[i])); - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), ev.ker)); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaEventRecord(ev.ker, p->streams[i])); + if (chunk_start + (int64_t) chunk_elems == ne) { + CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), ev.ker)); + } + } } return true; diff --git a/ggml/src/ggml-cuda/allreduce.cuh b/ggml/src/ggml-cuda/allreduce.cuh index 77466011a0..4c2be43962 100644 --- a/ggml/src/ggml-cuda/allreduce.cuh +++ b/ggml/src/ggml-cuda/allreduce.cuh @@ -5,9 +5,8 @@ #include -// Maximum tensor size (bytes per GPU) handled by the internal kernel path. -// Tensors larger than this are not yet supported and ggml_cuda_ar_allreduce() -// returns false, allowing the caller to fall back to another provider. +// Maximum chunk size (bytes per GPU) handled by one internal kernel launch. +// Larger tensors are reduced by issuing multiple chunked launches. static constexpr size_t GGML_CUDA_AR_MAX_BYTES = 256 * 1024; // 256 KB // Opaque pipeline context — owns all pinned buffers, streams, and events. diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2dca6b3db4..9a4dd062aa 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1377,13 +1377,6 @@ static ggml_cuda_comm_allreduce_result ggml_backend_cuda_comm_try_allreduce_inte return GGML_CUDA_COMM_ALLREDUCE_SUCCESS; } - const size_t bytes = (size_t) ne * ggml_type_size(type); - if (bytes > GGML_CUDA_AR_MAX_BYTES) { - GGML_LOG_WARN("%s: internal unsupported: ne=%" PRId64 " type=%d bytes=%zu max=%zu\n", - __func__, ne, (int) type, bytes, GGML_CUDA_AR_MAX_BYTES); - return GGML_CUDA_COMM_ALLREDUCE_UNSUPPORTED; - } - for (size_t i = 0; i < n_backends; ++i) { if (tensors[i] == nullptr) { GGML_LOG_ERROR("%s: internal failed: tensor[%zu] is null\n", __func__, i);