llama: retry pipeline reservation with fewer copies

ggml/src/ggml-backend-impl.h

@@ -105,6 +105,15 @@ extern "C" {
     GGML_API size_t         ggml_backend_meta_n_backends    (ggml_backend_t meta_backend);
     GGML_API ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index);
 
+    // Create a scheduler with an explicit number of pipeline input copies.
+    GGML_API ggml_backend_sched_t ggml_backend_sched_new_with_n_copies(
+        ggml_backend_t * backends,
+        ggml_backend_buffer_type_t * bufts,
+        int n_backends,
+        size_t graph_size,
+        int n_copies,
+        bool op_offload);
+
     //
     // Backend (stream)
     //

ggml/src/ggml-backend.cpp

@@ -2001,15 +2001,17 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
     return GGML_STATUS_SUCCESS;
 }
 
-ggml_backend_sched_t ggml_backend_sched_new(
+ggml_backend_sched_t ggml_backend_sched_new_with_n_copies(
         ggml_backend_t * backends,
         ggml_backend_buffer_type_t * bufts,
         int n_backends,
         size_t graph_size,
-        bool parallel,
+        int n_copies,
         bool op_offload) {
     GGML_ASSERT(n_backends > 0);
     GGML_ASSERT(n_backends <= GGML_SCHED_MAX_BACKENDS);
+    GGML_ASSERT(n_copies >= 1);
+    GGML_ASSERT(n_copies <= GGML_SCHED_MAX_COPIES);
     GGML_ASSERT(ggml_backend_dev_type(ggml_backend_get_device(backends[n_backends - 1])) == GGML_BACKEND_DEVICE_TYPE_CPU);
 
     struct ggml_backend_sched * sched = (ggml_backend_sched *) calloc(1, sizeof(struct ggml_backend_sched));
@@ -2025,7 +2027,7 @@ ggml_backend_sched_t ggml_backend_sched_new(
     sched->debug_realloc = GGML_SCHED_DEBUG_REALLOC ? atoi(GGML_SCHED_DEBUG_REALLOC) : sched->debug_realloc;
 
     sched->n_backends = n_backends;
-    sched->n_copies = parallel ? GGML_SCHED_MAX_COPIES : 1;
+    sched->n_copies = n_copies;
 
     // initialize hash table
     // FIXME: needs to be size*2 to account for leafs (do it in graph_split instead)
@@ -2073,6 +2075,22 @@ ggml_backend_sched_t ggml_backend_sched_new(
     return sched;
 }
 
+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) {
+    return ggml_backend_sched_new_with_n_copies(
+            backends,
+            bufts,
+            n_backends,
+            graph_size,
+            parallel ? GGML_SCHED_MAX_COPIES : 1,
+            op_offload);
+}
+
 void ggml_backend_sched_free(ggml_backend_sched_t sched) {
     if (sched == NULL) {
         return;

src/CMakeLists.txt

@@ -64,7 +64,7 @@ set_source_files_properties(${LLAMA_CORE_SOURCES} ../include/llama.h unicode.h
 
 configure_file(llama-version.h.in ${CMAKE_CURRENT_BINARY_DIR}/llama-version.h @ONLY)
 
-target_include_directories(llama PRIVATE . ${CMAKE_CURRENT_BINARY_DIR})
+target_include_directories(llama PRIVATE . ${CMAKE_CURRENT_BINARY_DIR} ../ggml/src)
 target_include_directories(llama PUBLIC ../include)
 target_compile_features   (llama PRIVATE cxx_std_17) # don't bump
 

src/llama-context.cpp

@@ -1,5 +1,6 @@
 #include "llama-context.h"
 
+#include "ggml-backend-impl.h"
 #include "ggml.h"
 #include "llama-arch.h"
 #include "llama-graph.h"
@@ -675,10 +676,23 @@ void llama_context::sched_reserve() {
                 model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr);
         if (!gf) {
             if (cparams.pipeline_parallel) {
-                LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);
-                cparams.pipeline_parallel = false;
-                sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload));
-                gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
+                const int n_copies = ggml_backend_sched_get_n_copies(sched.get());
+                if (n_copies > 2) {
+                    LLAMA_LOG_WARN("%s: compute buffer allocation failed with %d pipeline copies, retrying with 2 copies\n",
+                            __func__, n_copies);
+                    sched.reset(ggml_backend_sched_new_with_n_copies(
+                            backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, 2, cparams.op_offload));
+                    gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
+                }
+
+                if (!gf && ggml_backend_sched_get_n_copies(sched.get()) > 1) {
+                    const int n_copies_failed = ggml_backend_sched_get_n_copies(sched.get());
+                    LLAMA_LOG_WARN("%s: compute buffer allocation failed with %d pipeline copies, retrying without pipeline parallelism\n",
+                            __func__, n_copies_failed);
+                    cparams.pipeline_parallel = false;
+                    sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload));
+                    gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
+                }
             }
             if (!gf) {
                 throw std::runtime_error("failed to allocate compute pp buffers");

tests/test-alloc.cpp

@@ -15,6 +15,7 @@ uint8_t * const alloc_base = (uint8_t *) 16;
 
 struct dummy_backend_context {
     size_t max_buffer_size = 64;
+    size_t max_total_size  = SIZE_MAX;
     size_t alignment       = 8;
 
     ggml_backend_buffer_i              buffer_interface;
@@ -24,6 +25,8 @@ struct dummy_backend_context {
 
     bool custom_alloc_buffer_n_called   = false;
     bool custom_get_alloc_size_n_called = false;
+    size_t events_created = 0;
+    size_t events_freed   = 0;
 
     size_t allocated_total() const {
         size_t n = 0;
@@ -42,6 +45,10 @@ static const char * dummy_backend_buffer_type_get_name(ggml_backend_buffer_type_
 
 static ggml_backend_buffer_t dummy_backend_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
     dummy_backend_context * ctx    = (dummy_backend_context *) buft->context;
+    const size_t allocated = ctx->allocated_total();
+    if (allocated > ctx->max_total_size || size > ctx->max_total_size - allocated) {
+        return nullptr;
+    }
     ggml_backend_buffer_t & buffer = ctx->buffers.emplace_back();
     buffer                         = ggml_backend_buffer_init(buft, ctx->buffer_interface, ctx, size);
     return buffer;
@@ -126,6 +133,20 @@ static bool dummy_backend_device_supports_buft(ggml_backend_dev_t device, ggml_b
     return device->context == buft->context;
 }
 
+static ggml_backend_event_t dummy_backend_device_event_new(ggml_backend_dev_t dev) {
+    dummy_backend_context * ctx = (dummy_backend_context *) dev->context;
+    ctx->events_created++;
+    return new ggml_backend_event{ dev, ctx };
+}
+
+static void dummy_backend_device_event_free(ggml_backend_dev_t /*dev*/, ggml_backend_event_t event) {
+    dummy_backend_context * ctx = (dummy_backend_context *) event->context;
+    ctx->events_freed++;
+    delete event;
+}
+
+static void dummy_backend_device_event_synchronize(ggml_backend_dev_t /*dev*/, ggml_backend_event_t /*event*/) {}
+
 // ggml_backend interface
 
 static const char * dummy_backend_get_name(ggml_backend_t /*backend*/) {
@@ -139,11 +160,12 @@ struct dummy_backend {
     ggml_backend_buffer_type               buffer_type;
 };
 
-static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment = 8) {
+static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment = 8, size_t max_total_size = SIZE_MAX) {
     dummy_backend b{};
     b.context                  = std::make_unique<dummy_backend_context>();
     b.context->alignment       = alignment;
     b.context->max_buffer_size = max_buffer_size;
+    b.context->max_total_size  = max_total_size;
 
     b.context->buffer_interface.free_buffer   = dummy_backend_buffer_free_buffer;
     b.context->buffer_interface.get_base      = dummy_backend_buffer_get_base;
@@ -157,6 +179,9 @@ static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment
     b.context->device.iface.get_type      = dummy_backend_device_get_type;
     b.context->device.iface.supports_op   = dummy_backend_device_supports_op;
     b.context->device.iface.supports_buft = dummy_backend_device_supports_buft;
+    b.context->device.iface.event_new     = dummy_backend_device_event_new;
+    b.context->device.iface.event_free    = dummy_backend_device_event_free;
+    b.context->device.iface.event_synchronize = dummy_backend_device_event_synchronize;
 
     b.context->backend.context        = b.context.get();
     b.context->backend.device         = &b.context->device;
@@ -677,6 +702,86 @@ static void test_graph_optimize_alloc_dep() {
     GGML_ASSERT(!graph_reuses_allocation(true));
 }
 
+struct sched_alloc_result {
+    bool   allocated;
+    size_t buffer_size;
+    size_t events_created;
+    size_t events_freed;
+};
+
+static sched_alloc_result sched_alloc_with_n_copies(int n_copies, size_t max_total_size = SIZE_MAX) {
+    dummy_backend backend      = dummy_backend_init(SIZE_MAX, 8, max_total_size);
+    auto [ctx, graph, ctx_ptr] = make_context();
+
+    ggml_tensor * input  = make_input_with_size(ctx, 64);
+    ggml_tensor * output = ggml_scale(ctx, input, 2.0f);
+    ggml_set_output(output);
+    ggml_build_forward_expand(graph, output);
+
+    ggml_backend_t             backend_ptr = &backend.context->backend;
+    ggml_backend_buffer_type_t buft        = &backend.buffer_type;
+    ggml_backend_sched_ptr sched(ggml_backend_sched_new_with_n_copies(
+            &backend_ptr, &buft, 1, 8, n_copies, true));
+    GGML_ASSERT(ggml_backend_sched_get_n_copies(sched.get()) == n_copies);
+
+    const bool allocated = ggml_backend_sched_alloc_graph(sched.get(), graph);
+    const size_t buffer_size = backend.context->allocated_total();
+    const size_t events_created = backend.context->events_created;
+
+    sched.reset();
+
+    GGML_ASSERT(backend.context->buffers.empty());
+    GGML_ASSERT(backend.context->events_created == backend.context->events_freed);
+
+    return { allocated, buffer_size, events_created, backend.context->events_freed };
+}
+
+static int sched_default_n_copies(bool parallel) {
+    dummy_backend backend = dummy_backend_init(SIZE_MAX);
+    ggml_backend_t backend_ptr = &backend.context->backend;
+    ggml_backend_buffer_type_t buft = &backend.buffer_type;
+    ggml_backend_sched_ptr sched(ggml_backend_sched_new(&backend_ptr, &buft, 1, 8, parallel, true));
+    const int n_copies = ggml_backend_sched_get_n_copies(sched.get());
+    sched.reset();
+    GGML_ASSERT(backend.context->events_created == backend.context->events_freed);
+    return n_copies;
+}
+
+static void test_sched_copy_count_alloc_and_cleanup() {
+    const int n_copies_parallel = sched_default_n_copies(true);
+    GGML_ASSERT(sched_default_n_copies(false) == 1);
+
+    const sched_alloc_result one = sched_alloc_with_n_copies(1);
+    GGML_ASSERT(one.allocated);
+    GGML_ASSERT(one.events_created == 0);
+    GGML_ASSERT(one.events_freed == 0);
+
+    if (n_copies_parallel >= 2) {
+        const sched_alloc_result two = sched_alloc_with_n_copies(2);
+        GGML_ASSERT(two.allocated);
+        GGML_ASSERT(two.buffer_size > one.buffer_size);
+        GGML_ASSERT(two.events_created == 2);
+        GGML_ASSERT(two.events_freed == 2);
+    }
+
+    if (n_copies_parallel > 2) {
+        const sched_alloc_result parallel = sched_alloc_with_n_copies(n_copies_parallel);
+        GGML_ASSERT(parallel.allocated);
+        GGML_ASSERT(parallel.buffer_size > one.buffer_size);
+
+        const size_t limit = one.buffer_size + (parallel.buffer_size - one.buffer_size + 1) / 2;
+        const sched_alloc_result parallel_limited = sched_alloc_with_n_copies(n_copies_parallel, limit);
+        GGML_ASSERT(!parallel_limited.allocated);
+        GGML_ASSERT(parallel_limited.events_created == (size_t) n_copies_parallel);
+        GGML_ASSERT(parallel_limited.events_freed == (size_t) n_copies_parallel);
+
+        const sched_alloc_result two_limited = sched_alloc_with_n_copies(2, limit);
+        GGML_ASSERT(two_limited.allocated);
+        GGML_ASSERT(two_limited.events_created == 2);
+        GGML_ASSERT(two_limited.events_freed == 2);
+    }
+}
+
 // Check that the size reported by ggml_backend_alloc_ctx_tensors_from_buft_size
 // matches the actual size of the buffer allocated for the ctx tensors
 static ggml_backend_buffer_ptr check_size_matches(ggml_backend_buffer_type_t buft, ggml_context * ctx) {
@@ -1121,6 +1226,7 @@ int main() {
     run("test_buffer_size_zero", test_buffer_size_zero);
     run("test_reallocation", test_reallocation);
     run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep);
+    run("test_sched_copy_count_alloc_and_cleanup", test_sched_copy_count_alloc_and_cleanup);
     run("test_buft_alloc_buffer_n_single_buffer", test_buft_alloc_buffer_n_single_buffer);
     run("test_buft_alloc_buffer_n_multi_buffer", test_buft_alloc_buffer_n_multi_buffer);
     run("test_buft_alloc_buffer_n_zero_size", test_buft_alloc_buffer_n_zero_size);