Per GPU fit margin (#1872)

common/common.cpp

@@ -693,6 +693,10 @@ bool gpt_params_parse_ex(int argc, char ** argv, gpt_params & params) {
     if (!params.tensor_buft_overrides.empty()) {
         params.tensor_buft_overrides.push_back({nullptr, nullptr});
     }
+    if (!params.fit_margin_array.empty()) {
+        params.fit_margin_array.push_back(-1);
+        params.fit_margin_array.push_back(0);
+    }
 
     if (!params.chat_template.empty() && !common_chat_verify_template(params.chat_template, params.use_jinja)) {
         throw std::runtime_error(string_format(
@@ -1945,6 +1949,23 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
         }
         return true;
     }
+    if (arg == "--gpu-fit-margin" || arg == "-gfm") {
+        CHECK_ARG
+        auto p = string_split_pairs<int,int>(argv[i], ',');
+        if (p.empty()) {
+            fprintf(stderr, "error: invalid GPU split margin argument: %s\n", argv[i]);
+            invalid_param = true;
+        } else {
+            auto cur_size = params.fit_margin_array.size();
+            params.fit_margin_array.resize(cur_size + 2*p.size());
+            for (auto & pair : p) {
+                params.fit_margin_array[cur_size+0] = pair.first;
+                params.fit_margin_array[cur_size+1] = pair.second;
+                cur_size += 2;
+            }
+        }
+        return true;
+    }
     if (arg == "-cuda" || arg == "--cuda-params") {
         CHECK_ARG
         params.cuda_params = argv[i];
@@ -4070,6 +4091,11 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params
     if (!mparams.flash_attn && ggml_is_quantized(mparams.type_v)) {
         throw std::runtime_error("Quantized V cache cannot be used without flash attention");
     }
+    if (!params.fit_margin_array.empty()) {
+        GGML_ASSERT(params.fit_margin_array.size() % 2 == 0 && "Fit margin array does not have even number of elements");
+        GGML_ASSERT(params.fit_margin_array[params.fit_margin_array.size()-2] == -1 && "Fit margin array is not correctly termionated");
+        mparams.fit_margin_array = params.fit_margin_array.data();
+    }
 
     return mparams;
 }

common/common.h

@@ -351,6 +351,7 @@ struct gpt_params {
     std::vector<llama_model_kv_override> kv_overrides;
     std::vector<llama_model_tensor_buft_override> tensor_buft_overrides;
     std::vector<std::pair<int,int>> offload_policy;
+    std::vector<int> fit_margin_array;
 
     bool lora_init_without_apply = false; // only load lora to memory, but do not apply it to ctx (user can manually apply lora later using llama_lora_adapter_apply)
     std::vector<llama_lora_adapter_info> lora_adapters; // lora adapter path with user defined scale