From 34765b338ee68e2882517640d5b3eb9310934339 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Thu, 22 Oct 2026 18:00:59 -0500 Subject: [PATCH 1/5] common : MoE models default n_ubatch to n_batch A MoE ubatch reads the weights of nearly every expert, and uploads them when the experts are in host memory, so prompt processing pays one pass over all experts per ubatch. Without -ub, MoE models now use the logical batch (up to 4096) instead of 512. -ub still overrides it. With a draft or MTP model the larger default applies only when the tensor overrides keep the experts in host memory. The draft context follows the main context's ubatch, so both compute buffers grow, and with the whole model on the GPU that can overrun VRAM. --- common/arg.cpp | 5 ++-- common/common.cpp | 53 +++++++++++++++++++++++++++++++++++--- common/common.h | 2 ++ tools/cli/README.md | 2 +- tools/completion/README.md | 2 +- tools/server/README.md | 2 +- 6 files changed, 57 insertions(+), 9 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 267b84874912..31c46391f0e1 100644 --- a/common/arg.cpp +++ b/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: n_batch up to 4096)", 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( diff --git a/common/common.cpp b/common/common.cpp index 6ef3b785e0e8..06a6a77e8c91 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1238,11 +1238,60 @@ 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; } +// whether the tensor overrides (--cpu-moe, --n-cpu-moe, -ot) put the experts of some layer in a host buffer +// the loader applies the first override whose pattern matches the tensor name, so match a representative expert name +static bool common_overrides_put_experts_on_host(const std::vector & overrides) { + std::vector patterns; + std::vector on_host; + for (const auto & ovr : overrides) { + if (ovr.pattern == nullptr) { + break; + } + try { + patterns.emplace_back(ovr.pattern); + } catch (const std::regex_error &) { + continue; // the loader rejects an invalid pattern anyway + } + on_host.push_back(ovr.buft != nullptr && ggml_backend_buft_is_host(ovr.buft)); + } + for (int il = 0; il < 512; ++il) { + const std::string name = "blk." + std::to_string(il) + ".ffn_down_exps.weight"; + for (size_t i = 0; i < patterns.size(); ++i) { + if (std::regex_search(name, patterns[i])) { + if (on_host[i]) { + return true; + } + break; + } + } + } + return false; +} + common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { + const bool has_draft = params.speculative.has_dft(); + const bool spec_mtp = std::find(params.speculative.types.begin(), params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); + + // 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, MoE models use the logical batch + // except with a draft/MTP model and the experts on the GPU: the draft context follows this ubatch, so its compute buffers grow too + if (!params.n_ubatch_set && params.n_ubatch < params.n_batch && + (!(has_draft || spec_mtp) || common_overrides_put_experts_on_host(params.tensor_buft_overrides)) && + common_get_gguf_info(params.model.path).n_expert > 0) { + params.n_ubatch = std::min(params.n_batch, 4096); + COM_INF("MoE model, n_ubatch = %d (set -ub to override)\n", params.n_ubatch); + } + auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); @@ -1251,10 +1300,6 @@ common_init_result::common_init_result(common_params & params, bool model_only) COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n"); // the draft context is created from the same base params and follows the main context, fit both together - const bool has_draft = params.speculative.has_dft(); - const bool spec_mtp = std::find(params.speculative.types.begin(), params.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); - common_params params_dft = common_base_params_to_speculative(params); auto mparams_dft = common_model_params_to_llama(params_dft); diff --git a/common/common.h b/common/common.h index 70d6a671f512..2e5ff3a7029b 100644 --- a/common/common.h +++ b/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 MoE models use n_batch (up to 4096) 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 @@ -986,6 +987,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); diff --git a/tools/cli/README.md b/tools/cli/README.md index f1f6e6ed422e..eadc7280a215 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -29,7 +29,7 @@ | `-c, --ctx-size N` | size of the prompt context (default: 0, 0 = loaded from model)
(env: LLAMA_ARG_CTX_SIZE) | | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity)
(env: LLAMA_ARG_N_PREDICT) | | `-b, --batch-size N` | logical maximum batch size (default: 2048)
(env: LLAMA_ARG_BATCH) | -| `-ub, --ubatch-size N` | physical maximum batch size (default: 512)
(env: LLAMA_ARG_UBATCH) | +| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models: n_batch up to 4096)
(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)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | diff --git a/tools/completion/README.md b/tools/completion/README.md index 7174a04ff984..c55c03a5a32b 100644 --- a/tools/completion/README.md +++ b/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)
(env: LLAMA_ARG_CTX_SIZE) | | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity, -2 = until context filled)
(env: LLAMA_ARG_N_PREDICT) | | `-b, --batch-size N` | logical maximum batch size (default: 2048)
(env: LLAMA_ARG_BATCH) | -| `-ub, --ubatch-size N` | physical maximum batch size (default: 512)
(env: LLAMA_ARG_UBATCH) | +| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models: n_batch up to 4096)
(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)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | diff --git a/tools/server/README.md b/tools/server/README.md index 56d1ea5fe885..ffcbb185d72f 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -50,7 +50,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-c, --ctx-size N` | size of the prompt context (default: 0, 0 = loaded from model)
(env: LLAMA_ARG_CTX_SIZE) | | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity)
(env: LLAMA_ARG_N_PREDICT) | | `-b, --batch-size N` | logical maximum batch size (default: 2048)
(env: LLAMA_ARG_BATCH) | -| `-ub, --ubatch-size N` | physical maximum batch size (default: 512)
(env: LLAMA_ARG_UBATCH) | +| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models: n_batch up to 4096)
(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)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | From 27c496ba9d6e42546e1e04d9aeadf4167966f0e4 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Sat, 24 Oct 2026 12:59:08 -0500 Subject: [PATCH 2/5] common : --fit grows the ubatch of MoE models in free memory Replaces the fixed MoE default from the previous commit. Without -ub, --fit now places the weights and the context at the default ubatch as before, then doubles n_ubatch towards n_batch (up to 4096) while the larger compute buffers still fit in the memory left over on every device. The fit never gives up context or device weights for a bigger ubatch, and a draft or MTP context that follows the main ubatch is measured with it. -ub still overrides it, and -fit off keeps the default. --- common/arg.cpp | 2 +- common/common.cpp | 70 +++++++------------- common/common.h | 2 +- common/fit.cpp | 127 +++++++++++++++++++++++++++++-------- common/fit.h | 11 ++++ tools/cli/README.md | 2 +- tools/completion/README.md | 2 +- tools/server/README.md | 2 +- 8 files changed, 139 insertions(+), 79 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 31c46391f0e1..3f2c35231494 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1677,7 +1677,7 @@ 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, MoE models: n_batch up to 4096)", 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_set = true; diff --git a/common/common.cpp b/common/common.cpp index 06a6a77e8c91..2d525eae2205 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1246,52 +1246,8 @@ common_gguf_info common_get_gguf_info(const std::string & fname) { return info; } -// whether the tensor overrides (--cpu-moe, --n-cpu-moe, -ot) put the experts of some layer in a host buffer -// the loader applies the first override whose pattern matches the tensor name, so match a representative expert name -static bool common_overrides_put_experts_on_host(const std::vector & overrides) { - std::vector patterns; - std::vector on_host; - for (const auto & ovr : overrides) { - if (ovr.pattern == nullptr) { - break; - } - try { - patterns.emplace_back(ovr.pattern); - } catch (const std::regex_error &) { - continue; // the loader rejects an invalid pattern anyway - } - on_host.push_back(ovr.buft != nullptr && ggml_backend_buft_is_host(ovr.buft)); - } - for (int il = 0; il < 512; ++il) { - const std::string name = "blk." + std::to_string(il) + ".ffn_down_exps.weight"; - for (size_t i = 0; i < patterns.size(); ++i) { - if (std::regex_search(name, patterns[i])) { - if (on_host[i]) { - return true; - } - break; - } - } - } - return false; -} - common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { - const bool has_draft = params.speculative.has_dft(); - const bool spec_mtp = std::find(params.speculative.types.begin(), params.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); - - // 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, MoE models use the logical batch - // except with a draft/MTP model and the experts on the GPU: the draft context follows this ubatch, so its compute buffers grow too - if (!params.n_ubatch_set && params.n_ubatch < params.n_batch && - (!(has_draft || spec_mtp) || common_overrides_put_experts_on_host(params.tensor_buft_overrides)) && - common_get_gguf_info(params.model.path).n_expert > 0) { - params.n_ubatch = std::min(params.n_batch, 4096); - COM_INF("MoE model, n_ubatch = %d (set -ub to override)\n", params.n_ubatch); - } - auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); @@ -1300,6 +1256,10 @@ common_init_result::common_init_result(common_params & params, bool model_only) COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n"); // the draft context is created from the same base params and follows the main context, fit both together + const bool has_draft = params.speculative.has_dft(); + const bool spec_mtp = std::find(params.speculative.types.begin(), params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); + common_params params_dft = common_base_params_to_speculative(params); auto mparams_dft = common_model_params_to_llama(params_dft); @@ -1316,13 +1276,31 @@ 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.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, + 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 to override)\n", params.n_ubatch); + } + } } llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); diff --git a/common/common.h b/common/common.h index 2e5ff3a7029b..1de83d5bee82 100644 --- a/common/common.h +++ b/common/common.h @@ -457,7 +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 MoE models use n_batch (up to 4096) + 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 diff --git a/common/fit.cpp b/common/fit.cpp index 7a0300829508..5372581be81b 100644 --- a/common/fit.cpp +++ b/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 common_get_extra_memory_data( + const common_fit_extra_model * extra, const std::vector & devs, ggml_log_level log_level) { + std::vector devs_extra; + uint32_t ngl_extra = 0; + uint32_t nct_extra = 0; + uint32_t nex_extra = 0; + + const std::vector 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 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 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, + ggml_log_level log_level) { + constexpr int64_t MiB = 1024*1024; + + // a draft context follows the ubatch of the main context unless it was given its own + const bool extra_follows = extra != nullptr && extra->cparams->n_ubatch == cparams->n_ubatch; + + 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 devs; + uint32_t hp_ngl = 0; + uint32_t hp_nct = 0; + uint32_t hp_nex = 0; + std::vector 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 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(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); diff --git a/common/fit.h b/common/fit.h index 824d386b07a1..bf77639f6140 100644 --- a/common/fit.h +++ b/common/fit.h @@ -37,6 +37,17 @@ 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 + ggml_log_level log_level); + // print estimated memory to stdout void common_fit_print( const char * path_model, diff --git a/tools/cli/README.md b/tools/cli/README.md index eadc7280a215..dea08b2dd647 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -29,7 +29,7 @@ | `-c, --ctx-size N` | size of the prompt context (default: 0, 0 = loaded from model)
(env: LLAMA_ARG_CTX_SIZE) | | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity)
(env: LLAMA_ARG_N_PREDICT) | | `-b, --batch-size N` | logical maximum batch size (default: 2048)
(env: LLAMA_ARG_BATCH) | -| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models: n_batch up to 4096)
(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)
(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)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | diff --git a/tools/completion/README.md b/tools/completion/README.md index c55c03a5a32b..98e6e8288ce3 100644 --- a/tools/completion/README.md +++ b/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)
(env: LLAMA_ARG_CTX_SIZE) | | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity, -2 = until context filled)
(env: LLAMA_ARG_N_PREDICT) | | `-b, --batch-size N` | logical maximum batch size (default: 2048)
(env: LLAMA_ARG_BATCH) | -| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models: n_batch up to 4096)
(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)
(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)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | diff --git a/tools/server/README.md b/tools/server/README.md index ffcbb185d72f..3aa9a1dde46b 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -50,7 +50,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-c, --ctx-size N` | size of the prompt context (default: 0, 0 = loaded from model)
(env: LLAMA_ARG_CTX_SIZE) | | `-n, --predict, --n-predict N` | number of tokens to predict (default: -1, -1 = infinity)
(env: LLAMA_ARG_N_PREDICT) | | `-b, --batch-size N` | logical maximum batch size (default: 2048)
(env: LLAMA_ARG_BATCH) | -| `-ub, --ubatch-size N` | physical maximum batch size (default: 512, MoE models: n_batch up to 4096)
(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)
(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)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | From 658939a1617063eae638de0b70b704f29a297ccb Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Sat, 24 Oct 2026 15:12:30 -0500 Subject: [PATCH 3/5] common : add --fit-ubatch to turn off the MoE ubatch growth The growth stays on by default. -fitub off (or LLAMA_ARG_FIT_UBATCH=off) keeps the ubatch at its default while --fit still places the weights and the context, so the behaviour can be switched off without pinning -ub 512. --- common/arg.cpp | 14 ++++++++++++++ common/common.cpp | 2 +- common/common.h | 1 + tools/cli/README.md | 1 + tools/completion/README.md | 1 + tools/server/README.md | 1 + 6 files changed, 19 insertions(+), 1 deletion(-) diff --git a/common/arg.cpp b/common/arg.cpp index 3f2c35231494..7c45e38d3e46 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2893,6 +2893,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_FIT")); + add_opt(common_arg( + { "-fitub", "--fit-ubatch" }, "[on|off]", + string_format("whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('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( { "-fitp", "--fit-print" }, "[on|off]", string_format("print the estimated required memory ('on' or 'off', default: '%s')", params.fit_params_print ? "on" : "off"), diff --git a/common/common.cpp b/common/common.cpp index 2d525eae2205..a9203f79eb2f 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1289,7 +1289,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) // 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.n_ubatch_set && params.n_ubatch < params.n_batch && + 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(), diff --git a/common/common.h b/common/common.h index 1de83d5bee82..c9b851c701f3 100644 --- a/common/common.h +++ b/common/common.h @@ -482,6 +482,7 @@ struct common_params { int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs bool fit_params = true; // whether to fit unset model/context parameters to free device memory + bool fit_ubatch = true; // whether --fit grows the ubatch of MoE models in 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 diff --git a/tools/cli/README.md b/tools/cli/README.md index dea08b2dd647..f95740a61490 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -70,6 +70,7 @@ | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | | `-mg, --main-gpu INDEX` | the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: 0)
(env: LLAMA_ARG_MAIN_GPU) | | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT) | +| `-fitub, --fit-ubatch [on\|off]` | whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT_UBATCH) | | `-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
(env: LLAMA_ARG_FIT_TARGET) | | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096
(env: LLAMA_ARG_FIT_CTX) | | `--check-tensors` | check model tensor data for invalid values (default: false) | diff --git a/tools/completion/README.md b/tools/completion/README.md index 98e6e8288ce3..5f7a7bfc40ab 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -153,6 +153,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | | `-mg, --main-gpu INDEX` | the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: 0)
(env: LLAMA_ARG_MAIN_GPU) | | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT) | +| `-fitub, --fit-ubatch [on\|off]` | whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT_UBATCH) | | `-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
(env: LLAMA_ARG_FIT_TARGET) | | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096
(env: LLAMA_ARG_FIT_CTX) | | `--check-tensors` | check model tensor data for invalid values (default: false) | diff --git a/tools/server/README.md b/tools/server/README.md index 3aa9a1dde46b..fa865156bf2d 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -87,6 +87,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | | `-mg, --main-gpu INDEX` | the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: 0)
(env: LLAMA_ARG_MAIN_GPU) | | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT) | +| `-fitub, --fit-ubatch [on\|off]` | whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT_UBATCH) | | `-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
(env: LLAMA_ARG_FIT_TARGET) | | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096
(env: LLAMA_ARG_FIT_CTX) | | `--check-tensors` | check model tensor data for invalid values (default: false) | From ae1a7faafcf1b2a78679c88b1d7417e742e75f34 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Sat, 24 Oct 2026 16:26:46 -0500 Subject: [PATCH 4/5] common : name the MoE ubatch switch -fitu like the other --fit options -fitub becomes -fitu (as -fitp, -fitt, -fitc), the option moves next to the other --fit options, and the log line names it as the way to turn the growth off. --- common/arg.cpp | 28 ++++++++++++++-------------- common/common.cpp | 2 +- common/common.h | 2 +- tools/cli/README.md | 2 +- tools/completion/README.md | 2 +- tools/server/README.md | 2 +- 6 files changed, 19 insertions(+), 19 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 7c45e38d3e46..de070fce4cca 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2893,20 +2893,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_FIT")); - add_opt(common_arg( - { "-fitub", "--fit-ubatch" }, "[on|off]", - string_format("whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('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( { "-fitp", "--fit-print" }, "[on|off]", string_format("print the estimated required memory ('on' or 'off', default: '%s')", params.fit_params_print ? "on" : "off"), @@ -2953,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"), diff --git a/common/common.cpp b/common/common.cpp index a9203f79eb2f..9869291474bb 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1298,7 +1298,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) 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 to override)\n", 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); } } } diff --git a/common/common.h b/common/common.h index c9b851c701f3..f2d29ac4672f 100644 --- a/common/common.h +++ b/common/common.h @@ -482,9 +482,9 @@ struct common_params { int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs bool fit_params = true; // whether to fit unset model/context parameters to free device memory - bool fit_ubatch = true; // whether --fit grows the ubatch of MoE models in 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 fit_params_target = std::vector(llama_max_devices(), 1024 * 1024*1024); diff --git a/tools/cli/README.md b/tools/cli/README.md index f95740a61490..9e7f850fd4a4 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -70,9 +70,9 @@ | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | | `-mg, --main-gpu INDEX` | the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: 0)
(env: LLAMA_ARG_MAIN_GPU) | | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT) | -| `-fitub, --fit-ubatch [on\|off]` | whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT_UBATCH) | | `-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
(env: LLAMA_ARG_FIT_TARGET) | | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096
(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')
(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.
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) | diff --git a/tools/completion/README.md b/tools/completion/README.md index 5f7a7bfc40ab..9c0f7b55c9c4 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -153,9 +153,9 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | | `-mg, --main-gpu INDEX` | the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: 0)
(env: LLAMA_ARG_MAIN_GPU) | | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT) | -| `-fitub, --fit-ubatch [on\|off]` | whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT_UBATCH) | | `-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
(env: LLAMA_ARG_FIT_TARGET) | | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096
(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')
(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.
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) | diff --git a/tools/server/README.md b/tools/server/README.md index fa865156bf2d..d89dfbbd3493 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -87,9 +87,9 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | | `-mg, --main-gpu INDEX` | the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: 0)
(env: LLAMA_ARG_MAIN_GPU) | | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT) | -| `-fitub, --fit-ubatch [on\|off]` | whether --fit grows the ubatch of MoE models into free device memory when -ub is not set ('on' or 'off', default: 'on')
(env: LLAMA_ARG_FIT_UBATCH) | | `-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
(env: LLAMA_ARG_FIT_TARGET) | | `-fitc, --fit-ctx N` | minimum ctx size that can be set by --fit option, default: 4096
(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')
(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.
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) | From 6074bd5f6ab8784b6366c99d8dd3ec2cbaf833a7 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Fri, 30 Oct 2026 08:38:25 -0500 Subject: [PATCH 5/5] common : keep an explicit --spec-draft-n-ubatch when growing the MoE ubatch The growth inferred that the draft follows the target from equal ubatch values, so --spec-draft-n-ubatch set to the target's ubatch was measured at the grown size while the draft context is created at the user's value. Pass whether the draft follows from the caller instead. --- common/common.cpp | 1 + common/fit.cpp | 4 ++-- common/fit.h | 1 + 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 9869291474bb..57014a7b7329 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1295,6 +1295,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) 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 diff --git a/common/fit.cpp b/common/fit.cpp index 5372581be81b..f1963fe4beb9 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -914,11 +914,11 @@ uint32_t common_fit_n_ubatch( 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; - // a draft context follows the ubatch of the main context unless it was given its own - const bool extra_follows = extra != nullptr && extra->cparams->n_ubatch == cparams->n_ubatch; + extra_follows = extra_follows && extra != nullptr; while (cparams->n_ubatch < n_ubatch_max) { const uint32_t n_ubatch_prev = cparams->n_ubatch; diff --git a/common/fit.h b/common/fit.h index bf77639f6140..5f85a413f130 100644 --- a/common/fit.h +++ b/common/fit.h @@ -46,6 +46,7 @@ uint32_t common_fit_n_ubatch( 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