mirror of
https://github.com/ggml-org/whisper.cpp.git
synced 2026-10-21 18:19:35 +02:00
ggml: refactor selective expert copying to user code (llama/29943)
* ggml: refactor selective expert copying to user code * tests: enroll two models into selective expert copy test * tests: use deepseek2 as test model * improve comment in ggml-backend.h Co-authored-by: Georgi Gerganov <ggerganov@gmail.com> * cont: fix whitespace * cont : better comments Co-authored-by: Georgi Gerganov <ggerganov@gmail.com> --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
co-authored by
Georgi Gerganov
parent
c244a59428
commit
673f4f8367
@@ -317,6 +317,15 @@ extern "C" {
|
||||
//
|
||||
typedef bool (*ggml_backend_sched_eval_callback)(struct ggml_tensor * t, bool ask, void * user_data);
|
||||
|
||||
// Callback while copying input weights of a split
|
||||
// if the user returns false the scheduler simply copies the entire weight
|
||||
// the callback is called only for input weights in host buffers
|
||||
// the callback is called after all non-weight inputs of the split have been copied
|
||||
// `src` is the tensor in the previous split
|
||||
// `dst` is the copy of `src` in the split
|
||||
// `graph` is the compute graph of the split
|
||||
typedef bool (*ggml_backend_sched_copy_callback)(ggml_backend_t backend, const struct ggml_tensor * src, struct ggml_tensor * dst, struct ggml_cgraph * graph, void * user_data);
|
||||
|
||||
// Initialize a backend scheduler, backends with low index are given priority over backends with high index
|
||||
GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload);
|
||||
GGML_API void ggml_backend_sched_free(ggml_backend_sched_t sched);
|
||||
@@ -355,6 +364,9 @@ extern "C" {
|
||||
// Set a callback to be called for each resulting node during graph compute
|
||||
GGML_API void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data);
|
||||
|
||||
// Set a callback to be called when the inputs weights of a split are being copied
|
||||
GGML_API void ggml_backend_sched_set_copy_callback(ggml_backend_sched_t sched, ggml_backend_sched_copy_callback callback, void * user_data);
|
||||
|
||||
//
|
||||
// Meta backend
|
||||
//
|
||||
|
||||
+65
-126
@@ -966,6 +966,9 @@ struct ggml_backend_sched {
|
||||
ggml_backend_sched_eval_callback callback_eval;
|
||||
void * callback_eval_user_data;
|
||||
|
||||
ggml_backend_sched_copy_callback callback_copy;
|
||||
void * callback_copy_user_data;
|
||||
|
||||
char * context_buffer;
|
||||
size_t context_buffer_size;
|
||||
|
||||
@@ -1799,14 +1802,58 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool ggml_backend_sched_is_host_weight(const struct ggml_tensor * t) {
|
||||
return t->buffer != NULL &&
|
||||
ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
|
||||
ggml_backend_buffer_is_host(t->buffer);
|
||||
}
|
||||
|
||||
static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggml_backend_sched_split * split, struct ggml_tensor * input) {
|
||||
const int split_backend_id = split->backend_id;
|
||||
ggml_backend_t split_backend = sched->backends[split_backend_id];
|
||||
ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, input);
|
||||
struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);
|
||||
|
||||
if (input->flags & GGML_TENSOR_FLAG_INPUT) {
|
||||
// inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
|
||||
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
|
||||
ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
|
||||
} else {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
}
|
||||
ggml_backend_tensor_copy(input, input_cpy);
|
||||
return;
|
||||
}
|
||||
|
||||
// wait for the split backend to finish using the input before overwriting it
|
||||
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
|
||||
ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
|
||||
} else {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
}
|
||||
|
||||
if (sched->callback_copy != NULL && ggml_backend_sched_is_host_weight(input) &&
|
||||
sched->callback_copy(split_backend, input, input_cpy, &split->graph, sched->callback_copy_user_data)) {
|
||||
return;
|
||||
}
|
||||
|
||||
// try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events
|
||||
// TODO: add public function to facilitate this, since applications do not have direct access to the backend interface
|
||||
if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) {
|
||||
ggml_backend_synchronize(input_backend);
|
||||
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
|
||||
ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
|
||||
} else {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
}
|
||||
ggml_backend_tensor_copy(input, input_cpy);
|
||||
}
|
||||
}
|
||||
|
||||
static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
|
||||
GGML_ASSERT(sched);
|
||||
struct ggml_backend_sched_split * splits = sched->splits;
|
||||
|
||||
ggml_tensor * prev_ids_tensor = nullptr;
|
||||
std::vector<int32_t> ids;
|
||||
std::vector<ggml_bitset_t> used_ids;
|
||||
|
||||
int prev_backend_id = -1;
|
||||
|
||||
for (int split_id = 0; split_id < sched->n_splits; split_id++) {
|
||||
@@ -1825,129 +1872,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
|
||||
}
|
||||
|
||||
// copy the input tensors to the split backend
|
||||
// the weights in host memory are copied last, so that the copy callback can read the other inputs of the split
|
||||
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
|
||||
ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]);
|
||||
struct ggml_tensor * input = split->inputs[input_id];
|
||||
struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);
|
||||
|
||||
if (input->flags & GGML_TENSOR_FLAG_INPUT) {
|
||||
// inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
|
||||
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
|
||||
ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
|
||||
} else {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
}
|
||||
ggml_backend_tensor_copy(input, input_cpy);
|
||||
} else {
|
||||
// wait for the split backend to finish using the input before overwriting it
|
||||
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
|
||||
ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
|
||||
} else {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
}
|
||||
|
||||
// when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used
|
||||
ggml_tensor * node = split->graph.nodes[0];
|
||||
if (split->graph.n_nodes > 0 &&
|
||||
ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
|
||||
ggml_backend_buffer_is_host(input->buffer) && (
|
||||
(node->src[0] == input_cpy && node->op == GGML_OP_MUL_MAT_ID)
|
||||
//|| (node->src[1] == input_cpy && node->op == GGML_OP_ADD_ID) /* GGML_OP_ADD_ID weights are small and not worth splitting */
|
||||
)) {
|
||||
|
||||
const int64_t n_expert = node->op == GGML_OP_MUL_MAT_ID ? input->ne[2] : input->ne[1];
|
||||
const size_t expert_size = node->op == GGML_OP_MUL_MAT_ID ? input->nb[2] : input->nb[1];
|
||||
|
||||
ggml_backend_synchronize(input_backend);
|
||||
|
||||
// get the ids
|
||||
ggml_tensor * ids_tensor = node->src[2];
|
||||
ggml_backend_t ids_backend = split_backend;
|
||||
|
||||
if (ggml_nelements(ids_tensor) == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// if the ids tensor is also an input of the split, it may not have been copied yet to the split backend
|
||||
// in that case, we use the original ids tensor
|
||||
for (int i = input_id + 1; i < split->n_inputs; i++) {
|
||||
if (ids_tensor == tensor_copy(split->inputs[i], split_backend_id, sched->cur_copy)) {
|
||||
ids_tensor = split->inputs[i];
|
||||
ids_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[i]);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (ids_tensor != prev_ids_tensor) {
|
||||
ids.resize(ggml_nbytes(ids_tensor) / sizeof(int32_t));
|
||||
ggml_backend_tensor_get_async(ids_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor));
|
||||
ggml_backend_synchronize(ids_backend);
|
||||
|
||||
// find the used experts
|
||||
used_ids.clear();
|
||||
used_ids.resize(ggml_bitset_size(n_expert));
|
||||
for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) {
|
||||
for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) {
|
||||
int32_t id = ids[i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t)];
|
||||
GGML_ASSERT(id >= 0 && id < n_expert);
|
||||
ggml_bitset_set(used_ids.data(), id);
|
||||
}
|
||||
}
|
||||
|
||||
prev_ids_tensor = ids_tensor;
|
||||
}
|
||||
|
||||
// group consecutive experts and copy them together
|
||||
auto copy_experts = [&](int32_t first_id, int32_t last_id) {
|
||||
const size_t expert_offset = first_id * expert_size;
|
||||
const size_t expert_size_copy = (last_id - first_id + 1) * expert_size;
|
||||
const size_t padding = std::min<size_t>(expert_size, 512);
|
||||
const size_t padding_end = last_id < n_expert - 1 ? padding : 0;
|
||||
|
||||
ggml_backend_tensor_set_async(split_backend,
|
||||
input_cpy,
|
||||
(const uint8_t *)input->data + expert_offset, expert_offset,
|
||||
// copy a bit extra at the to ensure there are no NaNs in the padding of the last expert
|
||||
// this is necessary for MMQ in the CUDA backend
|
||||
expert_size_copy + padding_end);
|
||||
};
|
||||
|
||||
int id = 0;
|
||||
while (!ggml_bitset_get(used_ids.data(), id)) {
|
||||
id++;
|
||||
}
|
||||
int32_t first_id = id;
|
||||
int32_t last_id = first_id;
|
||||
|
||||
for (++id; id < n_expert; ++id) {
|
||||
if (!ggml_bitset_get(used_ids.data(), id)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (id == last_id + 1) {
|
||||
last_id = id;
|
||||
continue;
|
||||
}
|
||||
|
||||
copy_experts(first_id, last_id);
|
||||
|
||||
first_id = id;
|
||||
last_id = id;
|
||||
}
|
||||
copy_experts(first_id, last_id);
|
||||
} else {
|
||||
// try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events
|
||||
// TODO: add public function to facilitate this, since applications do not have direct access to the backend interface
|
||||
if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) {
|
||||
ggml_backend_synchronize(input_backend);
|
||||
if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
|
||||
ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
|
||||
} else {
|
||||
ggml_backend_synchronize(split_backend);
|
||||
}
|
||||
ggml_backend_tensor_copy(input, input_cpy);
|
||||
}
|
||||
}
|
||||
if (!ggml_backend_sched_is_host_weight(split->inputs[input_id])) {
|
||||
ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
|
||||
}
|
||||
}
|
||||
for (int input_id = 0; input_id < split->n_inputs; input_id++) {
|
||||
if (ggml_backend_sched_is_host_weight(split->inputs[input_id])) {
|
||||
ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2204,6 +2137,12 @@ void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backe
|
||||
sched->callback_eval_user_data = user_data;
|
||||
}
|
||||
|
||||
void ggml_backend_sched_set_copy_callback(ggml_backend_sched_t sched, ggml_backend_sched_copy_callback callback, void * user_data) {
|
||||
GGML_ASSERT(sched);
|
||||
sched->callback_copy = callback;
|
||||
sched->callback_copy_user_data = user_data;
|
||||
}
|
||||
|
||||
int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) {
|
||||
GGML_ASSERT(sched);
|
||||
return sched->n_splits;
|
||||
|
||||
Reference in New Issue
Block a user