diff --git a/common/arg.cpp b/common/arg.cpp index 267b84874912..de070fce4cca 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 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"), diff --git a/common/common.cpp b/common/common.cpp index 6ef3b785e0e8..57014a7b7329 100644 --- a/common/common.cpp +++ b/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); diff --git a/common/common.h b/common/common.h index 70d6a671f512..f2d29ac4672f 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 --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 fit_params_target = std::vector(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); diff --git a/common/fit.cpp b/common/fit.cpp index 7a0300829508..f1963fe4beb9 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, + 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 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..5f85a413f130 100644 --- a/common/fit.h +++ b/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, diff --git a/tools/cli/README.md b/tools/cli/README.md index f1f6e6ed422e..9e7f850fd4a4 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 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) | @@ -72,6 +72,7 @@ | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(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
(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 7174a04ff984..9c0f7b55c9c4 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 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) | @@ -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')
(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
(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 56d1ea5fe885..d89dfbbd3493 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 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) | @@ -89,6 +89,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-fit, --fit [on\|off]` | whether to adjust unset arguments to fit in device memory ('on' or 'off', default: 'on')
(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
(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) |