common : --fit grows the ubatch of MoE models into spare VRAM, for up to 2.5x faster prompt processing

common/arg.cpp

@@ -1677,9 +1677,10 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_env("LLAMA_ARG_BATCH"));
     add_opt(common_arg(
         {"-ub", "--ubatch-size"}, "N",
-        string_format("physical maximum batch size (default: %d)", params.n_ubatch),
+        string_format("physical maximum batch size (default: %d, MoE models with --fit: up to n_batch (max 4096) in free memory)", params.n_ubatch),
         [](common_params & params, int value) {
-            params.n_ubatch = value;
+            params.n_ubatch     = value;
+            params.n_ubatch_set = true;
         }
     ).set_env("LLAMA_ARG_UBATCH"));
     add_opt(common_arg(
@@ -2938,6 +2939,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.fit_params_min_ctx = value;
         }
     ).set_env("LLAMA_ARG_FIT_CTX"));
+    add_opt(common_arg(
+        { "-fitu", "--fit-ubatch" }, "[on|off]",
+        string_format("whether --fit may raise the ubatch size of MoE models into leftover device memory ('on' or 'off', default: '%s')", params.fit_ubatch ? "on" : "off"),
+        [](common_params & params, const std::string & value) {
+            if (is_truthy(value)) {
+                params.fit_ubatch = true;
+            } else if (is_falsey(value)) {
+                params.fit_ubatch = false;
+            } else {
+                throw std::runtime_error(
+                    string_format("error: unknown value for --fit-ubatch: '%s'\n", value.c_str()));
+            }
+        }
+    ).set_env("LLAMA_ARG_FIT_UBATCH"));
     add_opt(common_arg(
         {"--check-tensors"},
         string_format("check model tensor data for invalid values (default: %s)", params.check_tensors ? "true" : "false"),

common/common.cpp

@@ -1238,6 +1238,11 @@ common_gguf_info common_get_gguf_info(const std::string & fname) {
     // same key as llama_model_chat_template(model, nullptr)
     info.has_chat_tmpl = gguf_find_key(gguf_ctx.get(), "tokenizer.chat_template") >= 0;
 
+    const int64_t exp_id = gguf_find_key(gguf_ctx.get(), (arch + ".expert_count").c_str());
+    if (exp_id >= 0 && gguf_get_kv_type(gguf_ctx.get(), exp_id) == GGUF_TYPE_UINT32) {
+        info.n_expert = gguf_get_val_u32(gguf_ctx.get(), exp_id);
+    }
+
     return info;
 }
 
@@ -1271,13 +1276,32 @@ common_init_result::common_init_result(common_params & params, bool model_only)
             /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model
         };
 
-        common_fit_params(params.model.path.c_str(), &mparams, &cparams,
+        const ggml_log_level fit_log_level = params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR;
+
+        const common_params_fit_status fit_status = common_fit_params(params.model.path.c_str(), &mparams, &cparams,
             params.tensor_split,
             params.tensor_buft_overrides.data(),
             params.fit_params_target.data(),
             params.fit_params_min_ctx,
             has_draft || spec_mtp ? &extra : nullptr,
-            params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
+            fit_log_level);
+
+        // a MoE ubatch reads (and, for experts in host memory, uploads) the weights of nearly every expert, so a prompt
+        // costs one pass over all experts per ubatch: without -ub, grow the ubatch of a MoE model towards the logical batch
+        // in the memory the fit left over, so that it never costs context or weights on the device
+        if (fit_status == COMMON_PARAMS_FIT_STATUS_SUCCESS && params.fit_ubatch && !params.n_ubatch_set && params.n_ubatch < params.n_batch &&
+            common_get_gguf_info(params.model.path).n_expert > 0) {
+            common_fit_n_ubatch(params.model.path.c_str(), &mparams, &cparams,
+                params.fit_params_target.data(),
+                std::min(params.n_batch, 4096),
+                has_draft || spec_mtp ? &extra : nullptr,
+                params.speculative.draft.n_ubatch <= 0, // otherwise the draft keeps its own --spec-draft-n-ubatch
+                fit_log_level);
+            if ((int32_t) cparams.n_ubatch != params.n_ubatch) {
+                params.n_ubatch = cparams.n_ubatch; // the draft context follows params.n_ubatch
+                COM_INF("MoE model, n_ubatch = %d fits in free memory (set -ub or -fitu off to override)\n", params.n_ubatch);
+            }
+        }
     }
 
     llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams);

common/common.h

@@ -457,6 +457,7 @@ struct common_params {
     int32_t n_ctx                 =     0; // context size, 0 == context the model was trained with
     int32_t n_batch               =  2048; // logical batch size for prompt processing (must be >=32 to use BLAS)
     int32_t n_ubatch              =   512; // physical batch size for prompt processing (must be >=32 to use BLAS)
+    bool    n_ubatch_set          = false; // -ub was given; without it --fit grows the ubatch of MoE models in free memory
     int32_t n_keep                =     0; // number of tokens to keep from initial prompt
     int32_t n_chunks              =    -1; // max number of chunks to process (-1 = unlimited)
     int32_t n_parallel            =     1; // number of parallel sequences to decode
@@ -483,6 +484,7 @@ struct common_params {
     bool    fit_params         = true;  // whether to fit unset model/context parameters to free device memory
     bool    fit_params_print   = false; // print the estimated required memory to run the model
     int32_t fit_params_min_ctx = 4096;  // minimum context size to set when trying to reduce memory use
+    bool    fit_ubatch         = true;  // whether --fit may raise n_ubatch of MoE models into leftover device memory
 
     // margin per device in bytes for fitting parameters to free memory:
     std::vector<size_t> fit_params_target = std::vector<size_t>(llama_max_devices(), 1024 * 1024*1024);
@@ -986,6 +988,7 @@ struct common_gguf_info {
     common_decision_type decision_type = COMMON_DECISION_TYPE_UNKNOWN; // UNKNOWN if the file is missing, unreadable, or invalid
     uint32_t             n_ctx_train   = 0;                            // 0 if unknown
     bool                 has_chat_tmpl = false;                        // has the default tokenizer.chat_template
+    uint32_t             n_expert      = 0;                            // 0 if unknown or not a MoE model
 };
 
 common_gguf_info common_get_gguf_info(const std::string & fname);

common/fit.cpp

@@ -176,6 +176,37 @@ common_device_memory_data_vec common_get_device_memory_data(
     return ret;
 }
 
+// memory of the extra model, laid out on the devices of the main model
+static std::vector<llama_device_memory_data> common_get_extra_memory_data(
+        const common_fit_extra_model * extra, const std::vector<ggml_backend_dev_t> & devs, ggml_log_level log_level) {
+    std::vector<ggml_backend_dev_t> devs_extra;
+    uint32_t ngl_extra = 0;
+    uint32_t nct_extra = 0;
+    uint32_t nex_extra = 0;
+
+    const std::vector<llama_device_memory_data> measured = common_get_device_memory_data_impl(
+        extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
+
+    std::vector<llama_device_memory_data> ret(devs.size() + 1);
+    ret.back().mb = measured.back().mb;
+    for (size_t je = 0; je < devs_extra.size(); je++) {
+        for (size_t id = 0; id < devs.size(); id++) {
+            if (devs_extra[je] == devs[id]) {
+                ret[id].mb.model   += measured[je].mb.model;
+                ret[id].mb.context += measured[je].mb.context;
+                ret[id].mb.compute += measured[je].mb.compute;
+                break;
+            }
+        }
+    }
+    if (extra->shares_model) {
+        for (llama_device_memory_data & dmd : ret) {
+            dmd.mb.model = 0;
+        }
+    }
+    return ret;
+}
+
 static void common_params_fit_impl(
         const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
         float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides,
@@ -208,44 +239,17 @@ static void common_params_fit_impl(
         }
 
         if (dmds_extra.empty() || n_ctx_extra != cparams->n_ctx) {
-            std::vector<ggml_backend_dev_t> devs_extra;
-            uint32_t ngl_extra = 0;
-            uint32_t nct_extra = 0;
-            uint32_t nex_extra = 0;
-
             extra->cparams->n_ctx = cparams->n_ctx;
 
             LOG_TRC("%s: getting device memory data for the extra model at a context size of %" PRIu32 ":\n",
                 __func__, cparams->n_ctx);
 
-            dmds_t measured;
             try {
-                measured = common_get_device_memory_data_impl(
-                    extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
+                dmds_extra = common_get_extra_memory_data(extra, devs, log_level);
             } catch (const std::runtime_error & e) {
                 // the extra model is optional, fit the main model alone rather than giving up
                 LOG_WRN("%s: failed to measure the memory of the extra model, fitting without it: %s\n", __func__, e.what());
                 dmds_extra = dmds_t(devs.size() + 1);
-                n_ctx_extra = cparams->n_ctx;
-                return;
-            }
-
-            dmds_extra = dmds_t(devs.size() + 1);
-            dmds_extra.back().mb = measured.back().mb;
-            for (size_t je = 0; je < devs_extra.size(); je++) {
-                for (size_t id = 0; id < devs.size(); id++) {
-                    if (devs_extra[je] == devs[id]) {
-                        dmds_extra[id].mb.model   += measured[je].mb.model;
-                        dmds_extra[id].mb.context += measured[je].mb.context;
-                        dmds_extra[id].mb.compute += measured[je].mb.compute;
-                        break;
-                    }
-                }
-            }
-            if (extra->shares_model) {
-                for (llama_device_memory_data & dmd : dmds_extra) {
-                    dmd.mb.model = 0;
-                }
             }
 
             n_ctx_extra = cparams->n_ctx;
@@ -903,6 +907,73 @@ enum common_params_fit_status common_fit_params(
     return status;
 }
 
+uint32_t common_fit_n_ubatch(
+        const char * path_model,
+        const llama_model_params * mparams,
+        llama_context_params * cparams,
+        const size_t * margins,
+        uint32_t n_ubatch_max,
+        const common_fit_extra_model * extra,
+        bool extra_follows,
+        ggml_log_level log_level) {
+    constexpr int64_t MiB = 1024*1024;
+
+    extra_follows = extra_follows && extra != nullptr;
+
+    while (cparams->n_ubatch < n_ubatch_max) {
+        const uint32_t n_ubatch_prev = cparams->n_ubatch;
+        cparams->n_ubatch = std::min(2*n_ubatch_prev, n_ubatch_max);
+        if (extra_follows) {
+            extra->cparams->n_ubatch = cparams->n_ubatch;
+        }
+
+        // the weights and the context are already placed, only the compute buffers grow
+        bool fits = true;
+        try {
+            std::vector<ggml_backend_dev_t> devs;
+            uint32_t hp_ngl = 0;
+            uint32_t hp_nct = 0;
+            uint32_t hp_nex = 0;
+            std::vector<llama_device_memory_data> dmds = common_get_device_memory_data_impl(
+                path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+            if (extra != nullptr) {
+                extra->cparams->n_ctx = cparams->n_ctx;
+                const std::vector<llama_device_memory_data> dmds_extra = common_get_extra_memory_data(extra, devs, log_level);
+                for (size_t id = 0; id < dmds.size(); id++) {
+                    dmds[id].mb.model   += dmds_extra[id].mb.model;
+                    dmds[id].mb.context += dmds_extra[id].mb.context;
+                    dmds[id].mb.compute += dmds_extra[id].mb.compute;
+                }
+            }
+
+            // without devices the fit targets system memory
+            const size_t nd = std::max<size_t>(devs.size(), 1);
+            for (size_t id = 0; id < nd; id++) {
+                const llama_device_memory_data & dmd = devs.empty() ? dmds.back() : dmds[id];
+                const int64_t projected_free = (devs.empty() ? dmd.total : dmd.free) - int64_t(dmd.mb.total());
+                LOG_TRC("%s: n_ubatch=%" PRIu32 ", id=%zu: %" PRId64 " MiB free vs. target of %" PRId64 " MiB\n",
+                    __func__, cparams->n_ubatch, id, projected_free/MiB, int64_t(margins[id])/MiB);
+                if (projected_free < int64_t(margins[id])) {
+                    fits = false;
+                }
+            }
+        } catch (const std::runtime_error & e) {
+            LOG_WRN("%s: failed to measure memory at n_ubatch = %" PRIu32 ": %s\n", __func__, cparams->n_ubatch, e.what());
+            fits = false;
+        }
+
+        if (!fits) {
+            cparams->n_ubatch = n_ubatch_prev;
+            if (extra_follows) {
+                extra->cparams->n_ubatch = n_ubatch_prev;
+            }
+            break;
+        }
+    }
+
+    return cparams->n_ubatch;
+}
+
 void common_memory_breakdown_print(const struct llama_context * ctx) {
     //const auto & devices = ctx->get_model().devices;
     const auto * model = llama_get_model(ctx);

common/fit.h

@@ -37,6 +37,18 @@ common_params_fit_status common_fit_params(
       const common_fit_extra_model * extra,                  // model to fit alongside the main one, nullptr if there is none
                      ggml_log_level   log_level);            // minimum log level to print during fitting, lower levels go to debug log
 
+// to call after common_fit_params: grows cparams->n_ubatch in powers of 2 up to n_ubatch_max while the projected memory use
+// stays within the margins on every device, without changing the weights or the context; returns the new n_ubatch
+uint32_t common_fit_n_ubatch(
+                         const char * path_model,
+           const llama_model_params * mparams,
+               llama_context_params * cparams,
+                       const size_t * margins,      // margins of memory to leave per device in bytes
+                           uint32_t   n_ubatch_max, // largest n_ubatch to try
+      const common_fit_extra_model * extra,         // model to fit alongside the main one, nullptr if there is none
+                               bool   extra_follows, // the extra context has no n_ubatch of its own and grows with the main one
+                     ggml_log_level   log_level);
+
 // print estimated memory to stdout
 void common_fit_print(
                          const char * path_model,

tools/cli/README.md

@@ -29,7 +29,7 @@
 | `-c, --ctx-size N` | size of the prompt context (default: 0, 0 = loaded from model)<br/>(env: LLAMA_ARG_CTX_SIZE) |
 | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity)<br/>(env: LLAMA_ARG_N_PREDICT) |
 | `-b, --batch-size N` | logical maximum batch size (default: 2048)<br/>(env: LLAMA_ARG_BATCH) |
-| `-ub, --ubatch-size N` | physical maximum batch size (default: 512)<br/>(env: LLAMA_ARG_UBATCH) |
+| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models with --fit: up to n_batch (max 4096) in free memory)<br/>(env: LLAMA_ARG_UBATCH) |
 | `--keep N` | number of tokens to keep from the initial prompt (default: 0, -1 = all) |
 | `--swa-full` | use full-size SWA cache (default: false)<br/>[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)<br/>(env: LLAMA_ARG_SWA_FULL) |
 | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')<br/>(env: LLAMA_ARG_FLASH_ATTN) |
@@ -72,6 +72,7 @@
 | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')<br/>(env: LLAMA_ARG_FIT) |
 | `-fitt, --fit-target MiB0,MiB1,MiB2,...` | target margin per device for --fit, comma-separated list of values, single value is broadcast across all devices, default: 1024<br/>(env: LLAMA_ARG_FIT_TARGET) |
 | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096<br/>(env: LLAMA_ARG_FIT_CTX) |
+| `-fitu, --fit-ubatch [on\|off]` | whether --fit may raise the ubatch size of MoE models into leftover device memory ('on' or 'off', default: 'on')<br/>(env: LLAMA_ARG_FIT_UBATCH) |
 | `--check-tensors` | check model tensor data for invalid values (default: false) |
 | `--override-kv KEY=TYPE:VALUE,...` | advanced option to override model metadata by key. to specify multiple overrides, either use comma-separated values.<br/>types: int, float, bool, str. example: --override-kv tokenizer.ggml.add_bos_token=bool:false,tokenizer.ggml.add_eos_token=bool:false |
 | `--op-offload, --no-op-offload` | whether to offload host tensor operations to device (default: true) |

tools/completion/README.md

@@ -112,7 +112,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
 | `-c, --ctx-size N` | size of the prompt context (default: 0, 0 = loaded from model)<br/>(env: LLAMA_ARG_CTX_SIZE) |
 | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity, -2 = until context filled)<br/>(env: LLAMA_ARG_N_PREDICT) |
 | `-b, --batch-size N` | logical maximum batch size (default: 2048)<br/>(env: LLAMA_ARG_BATCH) |
-| `-ub, --ubatch-size N` | physical maximum batch size (default: 512)<br/>(env: LLAMA_ARG_UBATCH) |
+| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models with --fit: up to n_batch (max 4096) in free memory)<br/>(env: LLAMA_ARG_UBATCH) |
 | `--keep N` | number of tokens to keep from the initial prompt (default: 0, -1 = all) |
 | `--swa-full` | use full-size SWA cache (default: false)<br/>[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)<br/>(env: LLAMA_ARG_SWA_FULL) |
 | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')<br/>(env: LLAMA_ARG_FLASH_ATTN) |
@@ -155,6 +155,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
 | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')<br/>(env: LLAMA_ARG_FIT) |
 | `-fitt, --fit-target MiB0,MiB1,MiB2,...` | target margin per device for --fit, comma-separated list of values, single value is broadcast across all devices, default: 1024<br/>(env: LLAMA_ARG_FIT_TARGET) |
 | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096<br/>(env: LLAMA_ARG_FIT_CTX) |
+| `-fitu, --fit-ubatch [on\|off]` | whether --fit may raise the ubatch size of MoE models into leftover device memory ('on' or 'off', default: 'on')<br/>(env: LLAMA_ARG_FIT_UBATCH) |
 | `--check-tensors` | check model tensor data for invalid values (default: false) |
 | `--override-kv KEY=TYPE:VALUE,...` | advanced option to override model metadata by key. to specify multiple overrides, either use comma-separated values.<br/>types: int, float, bool, str. example: --override-kv tokenizer.ggml.add_bos_token=bool:false,tokenizer.ggml.add_eos_token=bool:false |
 | `--op-offload, --no-op-offload` | whether to offload host tensor operations to device (default: true) |