diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 5ef80da30..38197285c 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -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 // diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 5bdade6de..ff3f07998 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -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 ids; - std::vector 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(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;