diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 2999f2e88..f54c87a2a 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -761,6 +761,10 @@ static bool ggml_is_view_op(enum ggml_op op) { #define GGML_SCHED_MAX_COPIES 4 #endif +#ifndef GGML_SCHED_MAX_PREFETCH_SLOTS +#define GGML_SCHED_MAX_PREFETCH_SLOTS 8 +#endif + struct ggml_backend_sched_split { int backend_id; int i_start; @@ -824,10 +828,11 @@ struct ggml_backend_sched { // they overlap compute, alternating between two staging slots bool prefetch_experts; ggml_backend_t prefetch_backend; - ggml_backend_buffer_t prefetch_slots[2]; - ggml_backend_event_t prefetch_ready[2]; - ggml_backend_event_t prefetch_free[2]; - bool prefetch_used[2]; + int prefetch_n_slots; + ggml_backend_buffer_t prefetch_slots[GGML_SCHED_MAX_PREFETCH_SLOTS]; + ggml_backend_event_t prefetch_ready[GGML_SCHED_MAX_PREFETCH_SLOTS]; + ggml_backend_event_t prefetch_free[GGML_SCHED_MAX_PREFETCH_SLOTS]; + bool prefetch_used[GGML_SCHED_MAX_PREFETCH_SLOTS]; int prefetch_cur; int debug; @@ -1550,6 +1555,40 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { return true; } +static void ggml_backend_sched_prefetch_disable(ggml_backend_sched_t sched, ggml_backend_t split_backend) { + sched->prefetch_experts = false; + if (sched->prefetch_backend) { + ggml_backend_synchronize(split_backend); + ggml_backend_synchronize(sched->prefetch_backend); + } + for (int i = 0; i < sched->prefetch_n_slots; i++) { + ggml_backend_buffer_free(sched->prefetch_slots[i]); + sched->prefetch_slots[i] = NULL; + sched->prefetch_used[i] = false; + } +} + +// slots are sized once for the largest offloaded expert tensor in the current graph so +// that they never need to grow mid-eval +static size_t ggml_backend_sched_prefetch_max_size(ggml_backend_sched_t sched) { + size_t max_size = 0; + for (int split_id = 0; split_id < sched->n_splits; split_id++) { + struct ggml_backend_sched_split * split = &sched->splits[split_id]; + if (split->graph.n_nodes == 0 || split->graph.nodes[0]->op != GGML_OP_MUL_MAT_ID) { + continue; + } + for (int input_id = 0; input_id < split->n_inputs; input_id++) { + const ggml_tensor * input = split->inputs[input_id]; + if (input->buffer && + ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(input->buffer)) { + max_size = std::max(max_size, ggml_nbytes(input)); + } + } + } + return max_size; +} + static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_backend_t split_backend, size_t size) { if (sched->prefetch_backend == NULL) { ggml_backend_dev_t dev = split_backend->device; @@ -1564,7 +1603,7 @@ static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_ba sched->prefetch_experts = false; return false; } - for (int i = 0; i < 2; i++) { + for (int i = 0; i < sched->prefetch_n_slots; i++) { sched->prefetch_ready[i] = ggml_backend_event_new(dev); sched->prefetch_free[i] = ggml_backend_event_new(dev); if (sched->prefetch_ready[i] == NULL || sched->prefetch_free[i] == NULL) { @@ -1574,18 +1613,31 @@ static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_ba } } + size = std::max(size, ggml_backend_sched_prefetch_max_size(sched)); + ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(split_backend); - for (int i = 0; i < 2; i++) { + for (int i = 0; i < sched->prefetch_n_slots; i++) { if (sched->prefetch_slots[i] == NULL || ggml_backend_buffer_get_size(sched->prefetch_slots[i]) < size) { - ggml_backend_synchronize(split_backend); - ggml_backend_synchronize(sched->prefetch_backend); - ggml_backend_buffer_free(sched->prefetch_slots[i]); - sched->prefetch_slots[i] = ggml_backend_buft_alloc_buffer(buft, size); - sched->prefetch_used[i] = false; - if (sched->prefetch_slots[i] == NULL) { - sched->prefetch_experts = false; + // allocate before freeing so a failure leaves the old slot intact + ggml_backend_buffer_t new_buf = ggml_backend_buft_alloc_buffer(buft, size); + if (new_buf == NULL) { + // overlap needs at least 2 slots, otherwise run with what fits + if (i >= 2 && sched->prefetch_slots[0] != NULL && + ggml_backend_buffer_get_size(sched->prefetch_slots[0]) >= size) { + sched->prefetch_n_slots = i; + sched->prefetch_cur = 0; + return true; + } + ggml_backend_sched_prefetch_disable(sched, split_backend); return false; } + if (sched->prefetch_slots[i] != NULL) { + ggml_backend_synchronize(split_backend); + ggml_backend_synchronize(sched->prefetch_backend); + ggml_backend_buffer_free(sched->prefetch_slots[i]); + } + sched->prefetch_slots[i] = new_buf; + sched->prefetch_used[i] = false; } } return true; @@ -1604,6 +1656,9 @@ 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]; int split_prefetch_slot = -1; + ggml_tensor * prefetch_input_cpy = NULL; + ggml_backend_buffer_t prefetch_saved_buffer = NULL; + void * prefetch_saved_data = NULL; // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { @@ -1633,11 +1688,17 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (ids->ne[0]*ids->ne[1] >= 2*n_expert && ggml_backend_sched_prefetch_init(sched, split_backend, ggml_nbytes(input))) { const int slot = sched->prefetch_cur; - sched->prefetch_cur ^= 1; + sched->prefetch_cur = (sched->prefetch_cur + 1) % sched->prefetch_n_slots; // wait for the previous user of this slot to finish computing if (sched->prefetch_used[slot]) { ggml_backend_event_wait(sched->prefetch_backend, sched->prefetch_free[slot]); } + // point the staging copy at the slot only for the duration of + // this split, so a fallback to the regular path on a later + // eval can never see a dangling slot pointer + prefetch_input_cpy = input_cpy; + prefetch_saved_buffer = input_cpy->buffer; + prefetch_saved_data = input_cpy->data; input_cpy->buffer = sched->prefetch_slots[slot]; input_cpy->data = ggml_backend_buffer_get_base(sched->prefetch_slots[slot]); ggml_backend_tensor_set_async(sched->prefetch_backend, input_cpy, input->data, 0, ggml_nbytes(input)); @@ -1759,12 +1820,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); - if (ec != GGML_STATUS_SUCCESS) { - return ec; - } if (split_prefetch_slot != -1) { + // the kernels have captured the slot address at launch, safe to restore ggml_backend_event_record(sched->prefetch_free[split_prefetch_slot], split_backend); sched->prefetch_used[split_prefetch_slot] = true; + prefetch_input_cpy->buffer = prefetch_saved_buffer; + prefetch_input_cpy->data = prefetch_saved_data; + } + if (ec != GGML_STATUS_SUCCESS) { + return ec; } } else { // similar to ggml_backend_compare_graph_backend @@ -1875,8 +1939,14 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->op_offload = op_offload; + // GGML_SCHED_PREFETCH_EXPERTS=1 enables the default slot count, higher values set it + // directly; more slots let uploads run further ahead of compute at the cost of one + // max-sized expert tensor of device memory per slot const char * GGML_SCHED_PREFETCH_EXPERTS = getenv("GGML_SCHED_PREFETCH_EXPERTS"); - sched->prefetch_experts = op_offload && GGML_SCHED_PREFETCH_EXPERTS && atoi(GGML_SCHED_PREFETCH_EXPERTS); + const int prefetch_n_slots = GGML_SCHED_PREFETCH_EXPERTS ? atoi(GGML_SCHED_PREFETCH_EXPERTS) : 0; + sched->prefetch_experts = op_offload && prefetch_n_slots > 0; + // default of 3 covers the gate/up/down expert tensors of one MoE layer + sched->prefetch_n_slots = prefetch_n_slots <= 1 ? 3 : std::min(prefetch_n_slots, GGML_SCHED_MAX_PREFETCH_SLOTS); ggml_backend_sched_reset(sched); @@ -1894,7 +1964,8 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { } if (sched->prefetch_backend) { ggml_backend_synchronize(sched->prefetch_backend); - for (int i = 0; i < 2; i++) { + // the slot count may have been reduced after a failed allocation, free everything + for (int i = 0; i < GGML_SCHED_MAX_PREFETCH_SLOTS; i++) { ggml_backend_event_free(sched->prefetch_ready[i]); ggml_backend_event_free(sched->prefetch_free[i]); ggml_backend_buffer_free(sched->prefetch_slots[i]);