add arg --cuda-mem

2025-01-16 09:15:34 +04:00
parent dab6b2e1c2
commit 46e99218b4
5 changed files with 19 additions and 3 deletions
@@ -730,6 +730,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
            params.unload = true;
        }
    ).set_env("LLAMA_ARG_UNLOAD"));
+    add_opt(llama_arg(
+        {"-cm", "--cuda-mem"}, "N", 
+        format("maximum cuda memory to use (default: %d)", params.cuda_mem),
+        [](gpt_params & params, int value) {
+            params.cuda_mem = value; // in GiB
+        }
+    ).set_env("LLAMA_ARG_CUDA_MEM"));
    add_opt(llama_arg(
        {"-n", "--predict", "--n-predict"}, "N",
        format("number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)", params.n_predict),