ikawrakow/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-04-28 18:32:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP: it blocks on ncclAllReduce

Browse Source
This commit is contained in:
Kawrakow
2025-12-19 17:13:22 +00:00
parent 158f2f03f1
commit 0d552e9b38
7 changed files with 293 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
ggml/include/ggml.h
View File

@@ -1068,16 +1068,6 @@ extern "C" {
struct ggml_tensor * a,
enum ggml_op op);
GGML_API struct ggml_tensor * ggml_reduce(
struct ggml_context * ctx,
struct ggml_tensor * a,
enum ggml_op op);
GGML_API struct ggml_tensor * ggml_reduce_inplace(
struct ggml_context * ctx,
struct ggml_tensor * a,
enum ggml_op op);
GGML_API struct ggml_tensor * ggml_add_cast(
struct ggml_context * ctx,
struct ggml_tensor * a,
@@ -3058,6 +3048,18 @@ extern "C" {
int split_dim,
struct ggml_tensor * tensor);
GGML_API struct ggml_tensor * ggml_reduce(
struct ggml_context * ctx,
struct ggml_tensor * a,
ggml_split_tensor_t * b,
enum ggml_op op);
GGML_API struct ggml_tensor * ggml_reduce_inplace(
struct ggml_context * ctx,
struct ggml_tensor * a,
ggml_split_tensor_t * b,
enum ggml_op op);
#ifdef __cplusplus
}
#endif

47
ggml/src/ggml-backend.cpp
View File

@@ -14,6 +14,7 @@
#include <set>
#include <array>
#include <chrono>
#include <thread>
#define IK_PRINT_TIMING 0
@@ -1421,6 +1422,15 @@ static void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct gg
for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i];
int * node_backend_id = &tensor_backend_id(node);
if (node->op == GGML_OP_REDUCE) {
auto extra = (const ggml_split_tensor_t *)node->extra;
GGML_ASSERT(extra);
for (int j = extra->n_device-1; j >= 0; --j) {
if (extra->splits[j]) {
*node_backend_id = j; break;
}
}
}
// do not overwrite user assignments
if (*node_backend_id == -1) {
*node_backend_id = ggml_backend_sched_backend_id_from_cur(sched, node);
@@ -1652,6 +1662,7 @@ static void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct gg
// check if we should start a new split based on the sources of the current node
bool need_new_split = false;
if ((node->op == GGML_OP_ADD && node->op_params[0] == 0xff) ||
node->op == GGML_OP_REDUCE ||
node->op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t) - 1] == 0xff) {
need_new_split = true;
}
@@ -2184,13 +2195,17 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
}
}
//auto compute = [sched, &needs_sync, &own_cpy] (int my_backend_id) {
struct ggml_backend_sched_split * splits = sched->splits;
//bool is_cpu = ggml_backend_is_cpu(sched->backends[my_backend_id]);
std::vector<int32_t> ids;
std::vector<uint32_t> unique_ids;
ggml_tensor * last_ids_tensor = nullptr;
for (int i = 0; i < sched->n_splits; i++) {
//printf("Thread %d: split %d\n", my_backend_id, i);
#if IK_PRINT_TIMING
int64_t tim1 = ggml_time_us();
#endif
@@ -2198,6 +2213,19 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
int split_backend_id = split->backend_id;
ggml_backend_t split_backend = sched->backends[split_backend_id];
printf("Split %d on backend %d\n", i, split_backend_id);
auto node = split->graph.nodes[0];
//if (node->op == GGML_OP_REDUCE && split_backend_id != my_backend_id && !is_cpu) {
// printf("%s: triggering reduce for %s on backend %d\n", __func__, node->name, my_backend_id);
// auto graph = split->graph;
// graph.n_nodes = 1;
// auto ec = ggml_backend_graph_compute_async(sched->backends[my_backend_id], &graph);
// if (ec != GGML_STATUS_SUCCESS) return ec;
//}
//if (split_backend_id != my_backend_id) continue;
// copy the input tensors to the split backend
ggml_backend_sched_copy_inputs(sched, split, needs_sync, ids, unique_ids, last_ids_tensor);
@@ -2257,6 +2285,17 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
j0 = j1;
}
}
if (node->op == GGML_OP_REDUCE) {
for (int ib = 0; ib < sched->n_backends; ++ib) {
if (ib != split_backend_id && !ggml_backend_is_cpu(sched->backends[ib])) {
printf("%s: triggering reduce for %s on backend %d\n", __func__, node->name, ib);
auto graph = split->graph;
graph.n_nodes = 1;
auto ec = ggml_backend_graph_compute_async(sched->backends[ib], &graph);
if (ec != GGML_STATUS_SUCCESS) return ec;
}
}
}
// record the event of this copy
if (split->n_inputs > 0) {
@@ -2265,6 +2304,14 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
}
}
}
//};
//std::vector<std::thread> workers;
//workers.reserve(sched->n_backends);
//for (int i = 0; i < sched->n_backends; ++i) {
// workers.emplace_back(compute, i);
//}
//for (auto & w : workers) w.join();
sched->cur_copy = (sched->cur_copy + 1) % sched->n_copies;

18
ggml/src/ggml-cuda.cu
View File

@@ -48,6 +48,7 @@
#include "ggml-cuda/argmax.cuh"
#include "ggml-cuda/multiadd.cuh"
#include "ggml-cuda/hadamard.cuh"
#include "ggml-cuda/reduce.cuh"
#include <algorithm>
#include <array>
@@ -2948,6 +2949,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
//printf("%4d %s(%s) on device %d. time = %ld\n", i, ggml_op_name(dst->op), dst->name, ctx.device, ggml_time_us());
switch (dst->op) {
case GGML_OP_REDUCE:
ggml_cuda_op_reduce(ctx, dst);
break;
case GGML_OP_ARGMAX:
ggml_cuda_argmax(ctx, dst);
break;
@@ -3747,12 +3751,12 @@ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx
}
#endif
#ifndef NDEBUG
assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device));
for (int j = 0; j < GGML_MAX_SRC; j++) {
if (node->src[j] != nullptr) {
assert(node->src[j]->buffer);
}
}
//assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device));
//for (int j = 0; j < GGML_MAX_SRC; j++) {
// if (node->src[j] != nullptr) {
// assert(node->src[j]->buffer);
// }
//}
#endif // NDEBUG
bool ok = ggml_cuda_compute_forward(*cuda_ctx, node, cgraph, i);
@@ -3881,6 +3885,8 @@ GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t
GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, const ggml_tensor * op) {
ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
switch (op->op) {
case GGML_OP_REDUCE:
return true;
case GGML_OP_UNARY:
switch (ggml_get_unary_op(op)) {
case GGML_UNARY_OP_GELU:

67
ggml/src/ggml-cuda/reduce.cu Normal file
View File

@@ -0,0 +1,67 @@
//
// Copyright (C) 2023-2024 The ggml authors
// Copyright (C) 2024 Iwan Kawrakow
// MIT license
// SPDX-License-Identifier: MIT
//
#include "reduce.cuh"
#ifdef GGML_USE_NCCL
#include <nccl.h>
#endif
void ggml_cuda_op_reduce(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
auto op = (ggml_op)dst->op_params[0];
GGML_ASSERT(op == GGML_OP_ADD);
GGML_ASSERT(dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_contiguous(dst));
auto extra = (ggml_split_tensor_t *)dst->extra;
GGML_ASSERT(extra && extra->n_device > 1);
printf("============================== %s on device %d\n", __func__, ctx.device);
#ifdef GGML_USE_NCCL
auto & info = ggml_cuda_info();
GGML_ASSERT(info.have_nccl);
GGML_ASSERT(info.device_count >= extra->n_device);
int nhave = 0;
auto type = dst->type;
for (int j = 0; j < extra->n_device; ++j) {
if (extra->splits[j]) {
GGML_ASSERT(extra->splits[j]->type == type);
GGML_ASSERT(ggml_are_same_shape(dst, extra->splits[j]));
++nhave;
}
}
int device = ctx.device;
ncclComm_t this_comm;
if (nhave == info.device_count) {
this_comm = info.nccl_coms[device];
} else {
int color = extra->splits[device] ? 1 : 0;
auto status = ncclCommSplit(info.nccl_coms[0], color, ctx.device, &this_comm, nullptr);
GGML_ASSERT(status == ncclSuccess);
}
GGML_ASSERT(this_comm);
ncclResult_t status;
if (type == GGML_TYPE_F32) {
status = ncclAllReduce(extra->splits[device]->data,
extra->splits[device]->data,
ggml_nelements(extra->splits[device]),
ncclFloat, ncclSum, this_comm, ctx.stream());
} else {
status = ncclAllReduce(extra->splits[device]->data,
extra->splits[device]->data,
ggml_nelements(extra->splits[device]),
ncclHalf, ncclSum, this_comm, ctx.stream());
}
if (status != ncclSuccess) {
fprintf(stderr, "%s: ncclAllReduce failed with status %d\n", __func__, (int)status);
GGML_ABORT("Fatal error");
}
return;
#endif
fprintf(stderr, "%s: not implemented without NCCL\n", __func__);
GGML_ABORT("Fatal error");
}

5
ggml/src/ggml-cuda/reduce.cuh Normal file
View File

@@ -0,0 +1,5 @@
#include "common.cuh"
#define CUDA_REDUCE_BLOCK_SIZE 256
void ggml_cuda_op_reduce(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

101
ggml/src/ggml.c
View File

@@ -5314,7 +5314,7 @@ ggml_split_tensor_t * ggml_new_split(
result->n_device = n_device;
result->split_dim = split_dim;
result->tensor = tensor;
result->splits = (struct ggml_tensor**)(result->tensor + 1);
result->splits = (struct ggml_tensor**)(&result->splits + 1);
for (int i = 0; i < n_device; ++i) {
result->splits[i] = NULL;
}
@@ -6085,34 +6085,46 @@ struct ggml_tensor * ggml_dup_inplace(
static struct ggml_tensor * ggml_reduce_impl(
struct ggml_context * ctx,
struct ggml_tensor * a,
ggml_split_tensor_t * extra,
enum ggml_op op,
bool inplace) {
GGML_ASSERT(op == GGML_OP_ADD); // the only op we currently support
GGML_ASSERT(a->extra);
ggml_split_tensor_t * extra = (ggml_split_tensor_t *)a->extra;
GGML_ASSERT(extra->n_device > 1);
GGML_ASSERT(extra->splits);
int idx[GGML_MAX_SRC];
int nhave = 0;
for (int j = 0; j < extra->n_device; ++j) {
if (extra->splits[j]) ++nhave;
if (extra->splits[j]) {
if (nhave == GGML_MAX_SRC) {
GGML_ABORT("Too many tensors to reduce");
}
idx[nhave++] = j;
}
}
GGML_ASSERT(nhave > 1);
for (int j = 1; j < nhave; ++j) {
GGML_ASSERT(ggml_are_same_shape(extra->splits[idx[j]], extra->splits[idx[0]]));
}
struct ggml_tensor * result;
if (inplace) {
result = ggml_view_tensor(ctx, a);
result->src[0] = a;
result->extra = a->extra;
//for (int j = 0; j < nhave; ++j) result->src[j] = extra->splits[idx[j]];
result->extra = extra;
} else {
result = ggml_new_tensor_4d(ctx, a->type, a->ne[0], a->ne[1], a->ne[2], a->ne[3]);
result = ggml_dup_tensor(ctx, a);
ggml_split_tensor_t * new_extra = ggml_new_split(ctx, extra->n_device, extra->split_dim, result);
result->extra = new_extra;
//int jj = 0;
for (int j = 0; j < extra->n_device; ++j) {
if (extra->splits[j]) {
new_extra->splits[j] = ggml_dup_tensor(ctx, extra->splits[j]);
//result->src[jj++] = new_extra->splits[j];
}
}
}
result->src[0] = a;
result->op = GGML_OP_REDUCE;
result->op_params[0] = (int32_t)op;
return result;
@@ -6121,15 +6133,17 @@ static struct ggml_tensor * ggml_reduce_impl(
struct ggml_tensor * ggml_reduce(
struct ggml_context * ctx,
struct ggml_tensor * a,
ggml_split_tensor_t * extra,
enum ggml_op op) {
return ggml_reduce_impl(ctx, a, op, false);
return ggml_reduce_impl(ctx, a, extra, op, false);
}
struct ggml_tensor * ggml_reduce_inplace(
struct ggml_context * ctx,
struct ggml_tensor * a,
ggml_split_tensor_t * extra,
enum ggml_op op) {
return ggml_reduce_impl(ctx, a, op, true);
return ggml_reduce_impl(ctx, a, extra, op, true);
}
@@ -6140,6 +6154,49 @@ static struct ggml_tensor * ggml_add_impl(
struct ggml_tensor * a,
struct ggml_tensor * b,
bool inplace) {
//if (a->extra) {
// ggml_split_tensor_t * a_extra = (ggml_split_tensor_t *)a->extra;
// ggml_split_tensor_t * b_extra = (ggml_split_tensor_t *)b->extra;
// GGML_ASSERT(a_extra->n_device > 1);
// if (b_extra) {
// GGML_ASSERT(b_extra->n_device == a_extra->n_device);
// }
// int nhave = 0;
// for (int j = 0; j < a_extra->n_device; ++j) if (a_extra->splits[j]) ++nhave;
// GGML_ASSERT(nhave > 1);
// ggml_split_tensor_t * new_extra = ggml_new_split(ctx, a_extra->n_device, -1, NULL);
// struct ggml_tensor * last = NULL;
// for (int j = 0; j < a_extra->n_device; ++j) {
// if (!a_extra->splits[j]) continue;
// struct ggml_tensor * aj = a_extra->splits[j];
// struct ggml_tensor * bj = b;
// if (b_extra) {
// GGML_ASSERT(b_extra->splits[j]);
// bj = b_extra->splits[j];
// }
// GGML_ASSERT(ggml_are_same_shape(aj, bj));
// GGML_ASSERT(!aj->extra && !bj->extra);
// struct ggml_tensor * abj = inplace ? ggml_view_tensor(ctx, aj) : ggml_dup_tensor(ctx, aj);
// abj->op = GGML_OP_ADD;
// abj->src[0] = aj;
// abj->src[1] = bj;
// abj->src[2] = a;
// abj->src[3] = b;
// new_extra->splits[j] = abj;
// last = abj;
// }
// GGML_ASSERT(last);
// struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
// result->op = GGML_OP_ADD;
// result->src[0] = a;
// result->src[1] = b;
// new_extra->tensor = result;
// result->extra = new_extra;
// ggml_set_input(result);
// return result;
//}
GGML_ASSERT(ggml_can_repeat(b, a));
bool is_node = false;
@@ -8495,6 +8552,32 @@ struct ggml_tensor * ggml_get_rows(
GGML_ASSERT(b->ne[3] == 1);
GGML_ASSERT(b->type == GGML_TYPE_I32);
if (a->extra) {
ggml_split_tensor_t * extra = (ggml_split_tensor_t *)a->extra;
int nhave = 0;
for (int j = 0; j < extra->n_device; ++j) if (extra->splits[j]) ++nhave;
if (nhave > 1) {
ggml_split_tensor_t * new_extra = ggml_new_split(ctx, extra->n_device, -1, NULL);
struct ggml_tensor * last = NULL;
ggml_split_tensor_t * b_extra = (ggml_split_tensor_t *)b->extra;
for (int j = 0; j < extra->n_device; ++j) {
if (!extra->splits[j]) continue;
struct ggml_tensor * aj = extra->splits[j];
struct ggml_tensor * bj = b_extra && b_extra->splits[j] ? b_extra->splits[j] : b;
GGML_ASSERT(!aj->extra && !bj->extra);
new_extra->splits[j] = ggml_get_rows(ctx, aj, bj);
last = new_extra->splits[j];
}
struct ggml_tensor * result = ggml_view_tensor(ctx, last);
result->src[0] = a;
result->src[1] = b;
result->op = GGML_OP_GET_ROWS;
new_extra->tensor = result;
result->extra = new_extra;
return result;
}
}
bool is_node = false;
if (a->grad || b->grad) {