ggml : report allocation failure from the meta buffer type

ggml/src/ggml-backend-meta.cpp

@@ -1700,9 +1700,15 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_bac
     std::vector<ggml_backend_buffer_t> bufs;
     bufs.reserve(n_simple_bufts);
     for (size_t i = 0; i < n_simple_bufts; i++) {
-        bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size));
-        GGML_ASSERT(bufs.back() != nullptr);
-        max_size = std::max(max_size, ggml_backend_buffer_get_size(bufs.back()));
+        ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size);
+        if (buf == nullptr) {
+            for (ggml_backend_buffer_t b : bufs) {
+                ggml_backend_buffer_free(b);
+            }
+            return nullptr;
+        }
+        bufs.push_back(buf);
+        max_size = std::max(max_size, ggml_backend_buffer_get_size(buf));
     }
     ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs);
 
@@ -1737,27 +1743,35 @@ struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struc
         t->data = (void *) 0x2000000000000000; // FIXME
     }
     for (size_t i = 0; i < n_simple_bufts; i++) {
-        ggml_context * ctx = meta_buf_ctx->stc_static.ctxs[i].get();
+        ggml_context * ctx_simple = meta_buf_ctx->stc_static.ctxs[i].get();
         ggml_backend_buffer_type_t simple_buft = ggml_backend_meta_buft_simple_buft(buft, i);
 
         // If a ggml_context only has zero-sized tensors, ggml_backend_alloc_ctx_tensors_from_buft returns NULL.
         // For those edge cases, allocate a dummy buffer instead.
         bool any_nonzero_slice = false;
-        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx_simple); t != nullptr; t = ggml_get_next_tensor(ctx_simple, t)) {
             if (ggml_nelements(t) != 0) {
                 any_nonzero_slice = true;
                 break;
             }
         }
         if (any_nonzero_slice) {
-            meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx, simple_buft));
+            meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_simple, simple_buft));
         } else {
             meta_buf_ctx->bufs[i].reset(ggml_backend_buft_alloc_buffer(simple_buft, 0));
-            for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx_simple); t != nullptr; t = ggml_get_next_tensor(ctx_simple, t)) {
                 t->buffer = meta_buf_ctx->bufs[i].get();
             }
         }
-        GGML_ASSERT(meta_buf_ctx->bufs[i]);
+        if (meta_buf_ctx->bufs[i] == nullptr) {
+            // the tensors of ctx were pointed at meta_buf, which is freed here
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+                t->buffer = nullptr;
+                t->data   = nullptr;
+            }
+            ggml_backend_buffer_free(meta_buf);
+            return nullptr;
+        }
         meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->bufs[i].get()));
     }
     return meta_buf;

tests/CMakeLists.txt

@@ -214,6 +214,13 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
         ARGS --test-live-context-workspace
     )
 
+    llama_test(
+        test-llama-archs
+        NAME test-meta-alloc-failure
+        LABEL main
+        ARGS --test-meta-alloc-failure
+    )
+
     set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/")
     file(MAKE_DIRECTORY "${MODEL_DIR}")
 

tests/test-llama-archs.cpp

@@ -69,7 +69,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {
 }
 
 static void usage(char ** argv) {
-    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace]\n", argv[0]);
+    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace] [--test-meta-alloc-failure]\n", argv[0]);
 }
 
 static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){
@@ -1052,6 +1052,115 @@ static void test_phase_workspace_mismatched_placement(size_t seed) {
     GGML_ASSERT(llama_contexts_share_workspace(target.get(), draft.get()) == (status == 1));
 }
 
+static size_t meta_alloc_n_devices = 2;
+
+static ggml_backend_meta_split_state meta_alloc_get_split_state(const ggml_tensor * tensor, void * userdata) {
+    ggml_backend_meta_split_state ss;
+    memset(&ss, 0, sizeof(ss));
+    ss.axis = GGML_BACKEND_SPLIT_AXIS_1;
+    int64_t low = 0;
+    for (size_t i = 0; i < meta_alloc_n_devices; i++) {
+        const int64_t high = tensor->ne[1]*(i + 1)/meta_alloc_n_devices;
+        ss.ne[i] = high - low;
+        low = high;
+    }
+    ss.nr[0] = 1;
+    ss.n_segments = 1;
+    return ss;
+    GGML_UNUSED(userdata);
+}
+
+static int meta_alloc_calls        = 0;
+static int meta_alloc_fail_from    = -1;
+static int meta_alloc_live_buffers = 0;
+static ggml_backend_buffer_t (*meta_alloc_buffer_orig)(ggml_backend_buffer_type_t, size_t) = nullptr;
+static void (*meta_alloc_free_buffer_orig)(ggml_backend_buffer_t) = nullptr;
+
+static void meta_alloc_count_free_buffer(ggml_backend_buffer_t buffer) {
+    meta_alloc_live_buffers--;
+    meta_alloc_free_buffer_orig(buffer);
+}
+
+static ggml_backend_buffer_t meta_alloc_test_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
+    const bool fail = meta_alloc_fail_from >= 0 && meta_alloc_calls >= meta_alloc_fail_from;
+    meta_alloc_calls++;
+    if (fail) {
+        return nullptr;
+    }
+    ggml_backend_buffer_t buf = meta_alloc_buffer_orig(buft, size);
+    if (buf != nullptr && buf->iface.free_buffer != nullptr) {
+        meta_alloc_free_buffer_orig = buf->iface.free_buffer;
+        buf->iface.free_buffer      = meta_alloc_count_free_buffer;
+        meta_alloc_live_buffers++;
+    }
+    return buf;
+}
+
+// a meta buffer type must report an allocation failure instead of aborting, and must leave nothing behind
+static void test_meta_alloc_failure() {
+    ggml_backend_dev_t dev_cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+    GGML_ASSERT(dev_cpu != nullptr);
+
+    ggml_backend_dev_t devs[] = { dev_cpu, dev_cpu };
+    meta_alloc_n_devices = sizeof(devs)/sizeof(devs[0]);
+    ggml_backend_dev_t dev_meta = ggml_backend_meta_device(devs, meta_alloc_n_devices, meta_alloc_get_split_state, nullptr);
+    GGML_ASSERT(dev_meta != nullptr);
+    ggml_backend_buffer_type_t buft = ggml_backend_dev_buffer_type(dev_meta);
+
+    ggml_backend_buffer_type_t buft_cpu = ggml_backend_dev_buffer_type(dev_cpu);
+    meta_alloc_buffer_orig       = buft_cpu->iface.alloc_buffer;
+    buft_cpu->iface.alloc_buffer = meta_alloc_test_alloc_buffer;
+
+    // plain allocation, the second device fails
+    meta_alloc_calls     = 0;
+    meta_alloc_fail_from = 1;
+    GGML_ASSERT(ggml_backend_buft_alloc_buffer(buft, 1024) == nullptr);
+    GGML_ASSERT(meta_alloc_calls == 2);
+    GGML_ASSERT(meta_alloc_live_buffers == 0);
+
+    meta_alloc_fail_from = -1;
+    ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(buft, 1024);
+    GGML_ASSERT(buf != nullptr);
+    ggml_backend_buffer_free(buf);
+    GGML_ASSERT(meta_alloc_live_buffers == 0);
+
+    // allocation of a whole context, the second device fails
+    const ggml_init_params params = {
+        /*.mem_size   =*/ 4*ggml_tensor_overhead(),
+        /*.mem_buffer =*/ nullptr,
+        /*.no_alloc   =*/ true,
+    };
+    ggml_context_ptr ctx(ggml_init(params));
+    GGML_ASSERT(ctx);
+    std::vector<ggml_tensor *> tensors;
+    for (int i = 0; i < 2; i++) {
+        ggml_tensor * t = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, 32, 32);
+        ggml_set_name(t, ("t" + std::to_string(i)).c_str());
+        tensors.push_back(t);
+    }
+
+    meta_alloc_calls     = 0;
+    meta_alloc_fail_from = 1;
+    GGML_ASSERT(ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft) == nullptr);
+    GGML_ASSERT(meta_alloc_calls == 2);
+    GGML_ASSERT(meta_alloc_live_buffers == 0);
+    for (ggml_tensor * t : tensors) {
+        GGML_ASSERT(t->buffer == nullptr);
+        GGML_ASSERT(t->data   == nullptr);
+    }
+
+    meta_alloc_fail_from = -1;
+    ggml_backend_buffer_t buf_ctx = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);
+    GGML_ASSERT(buf_ctx != nullptr);
+    for (ggml_tensor * t : tensors) {
+        GGML_ASSERT(t->buffer == buf_ctx);
+    }
+    ggml_backend_buffer_free(buf_ctx);
+    GGML_ASSERT(meta_alloc_live_buffers == 0);
+
+    buft_cpu->iface.alloc_buffer = meta_alloc_buffer_orig;
+}
+
 static std::vector<float> get_logits(
         llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens, bool encode = false) {
     const uint32_t n_vocab  = llama_vocab_n_tokens(llama_model_get_vocab(model));
@@ -1541,6 +1650,7 @@ int main(int argc, char ** argv) {
     std::string out;
     bool test_phase_workspace = false;
     bool test_live_context_workspace = false;
+    bool test_meta_alloc = false;
 
     int verbosity = LOG_LEVEL_ERROR;
 
@@ -1594,6 +1704,10 @@ int main(int argc, char ** argv) {
             test_live_context_workspace = true;
             continue;
         }
+        if (strcmp(argv[i], "--test-meta-alloc-failure") == 0) {
+            test_meta_alloc = true;
+            continue;
+        }
     }
     printf("%s: using seed %zu\n", __func__, seed);
 
@@ -1611,6 +1725,10 @@ int main(int argc, char ** argv) {
             test_live_context_workspace_unsupported(seed);
             return 0;
         }
+        if (test_meta_alloc) {
+            test_meta_alloc_failure();
+            return 0;
+        }
         if (!out.empty()) {
             return save_models(arch, seed, verbosity, out);
         }