llamamodel: trigger CUDA OOM early so we can fall back

Signed-off-by: Jared Van Bortel <jared@nomic.ai>

gpt4all-backend/llamamodel.cpp

@@ -393,6 +393,10 @@ bool LLamaModel::loadModel(const std::string &modelPath, int n_ctx, int ngl)
             std::cerr << "warning: model was trained on only " << n_ctx_train << " context tokens ("
                       << n_ctx << " specified)\n";
         }
+
+        // GPT4All defaults to 128 tokens which is also the hardcoded maximum
+        d_ptr->ctx_params.n_batch  = LLMODEL_MAX_PROMPT_BATCH;
+        d_ptr->ctx_params.n_ubatch = LLMODEL_MAX_PROMPT_BATCH;
     }
 
     d_ptr->ctx_params.n_ctx   = n_ctx;
@@ -422,6 +426,23 @@ bool LLamaModel::loadModel(const std::string &modelPath, int n_ctx, int ngl)
         return false;
     }
 
+#ifdef GGML_USE_CUDA
+    if (d_ptr->model_params.n_gpu_layers > 0) {
+        try {
+            testModel(); // eagerly allocate memory
+        } catch (const std::runtime_error &e) {
+            std::cerr << "LLAMA ERROR: model test failed: " << e.what() << "\n";
+            llama_free(d_ptr->ctx);
+            d_ptr->ctx = nullptr;
+            llama_free_model(d_ptr->model);
+            d_ptr->model = nullptr;
+            d_ptr->device = -1;
+            d_ptr->deviceName.clear();
+            return false;
+        }
+    }
+#endif
+
     d_ptr->end_tokens = {llama_token_eos(d_ptr->model)};
 
     if (usingGPUDevice()) {
@@ -445,6 +466,26 @@ bool LLamaModel::loadModel(const std::string &modelPath, int n_ctx, int ngl)
     return true;
 }
 
+void LLamaModel::testModel() {
+    int n_ctx = llama_n_ctx(d_ptr->ctx);
+    int n_batch = LLMODEL_MAX_PROMPT_BATCH;
+    n_batch = std::min(n_batch, n_ctx);
+
+    // test with maximum batch size
+    PromptContext ctx;
+    ctx.n_batch = n_batch;
+    std::vector<int32_t> tokens(n_batch);
+
+    llama_set_skip_cpu(d_ptr->ctx, true);
+    if (!evalTokens(ctx, tokens))
+        throw std::runtime_error("llama_decode failed");
+    llama_set_skip_cpu(d_ptr->ctx, false);
+    llama_synchronize(d_ptr->ctx); // wait for GPU to finish
+
+    // clean up
+    llama_kv_cache_clear(d_ptr->ctx);
+}
+
 void LLamaModel::setThreadCount(int32_t n_threads) {
     d_ptr->n_threads = n_threads;
     llama_set_n_threads(d_ptr->ctx, n_threads, n_threads);

gpt4all-backend/llamamodel_impl.h

@@ -48,6 +48,8 @@ public:
                size_t *tokenCount = nullptr, bool doMean = true, bool atlas = false) override;
 
 private:
+    void testModel(); // used for CUDA to eagerly allocate memory
+
     std::unique_ptr<LLamaPrivate> d_ptr;
     bool m_supportsEmbedding = false;
     bool m_supportsCompletion = false;

gpt4all-backend/llmodel.h

@@ -122,7 +122,7 @@ public:
         float   top_p = 0.9f;
         float   min_p = 0.0f;
         float   temp = 0.9f;
-        int32_t n_batch = 9;
+        int32_t n_batch = 128;
         float   repeat_penalty = 1.10f;
         int32_t repeat_last_n = 64;     // last n tokens to penalize
         float   contextErase = 0.75f;   // percent of context to erase if we exceed the context window

gpt4all-chat/chatllm.cpp

@@ -374,6 +374,7 @@ bool ChatLLM::loadModel(const ModelInfo &modelInfo)
 
                 m_llModelInfo.model->setProgressCallback([this](float progress) -> bool {
                     progress = std::max(progress, std::numeric_limits<float>::min()); // keep progress above zero
+                    progress = std::min(progress, std::nextafter(1.0f, 0.0f)); // keep progress below 100% until we are actually done
                     emit modelLoadingPercentageChanged(progress);
                     return m_shouldBeLoaded;
                 });