diff --git a/common/common.cpp b/common/common.cpp index 25ca838dff..d84d57ac91 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1294,11 +1294,34 @@ common_init_result::common_init_result(common_params & params, bool model_only) if (params.fit_params) { COM_TRC("%s", "fitting params to device memory ...\n"); 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); + auto cparams_dft = common_context_params_to_llama(params_dft); + if (spec_mtp) { + cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + } + cparams_dft.n_rs_seq = 0; + + const common_fit_extra_model extra = { + /*.path_model =*/ params_dft.model.path.c_str(), + /*.mparams =*/ &mparams_dft, + /*.cparams =*/ &cparams_dft, + /*.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, 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); } diff --git a/common/fit.cpp b/common/fit.cpp index 757b4937c7..c601fe405e 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -178,7 +178,7 @@ common_device_memory_data_vec common_get_device_memory_data( 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, - size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) { + size_t * margins_s, uint32_t n_ctx_min, const common_fit_extra_model * extra, enum ggml_log_level log_level) { if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) { throw common_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort"); } @@ -196,6 +196,67 @@ static void common_params_fit_impl( const uint32_t n_streams = cparams->kv_unified ? 1 : std::max(1, cparams->n_seq_max); const bool n_ctx_auto = cparams->n_ctx == 0; + dmds_t dmds_extra; // memory of the extra model, laid out on the devices of the main model + uint32_t n_ctx_extra = 0; // context that memory was measured at + + // the extra model competes for the same memory as the main model, add it to every measurement + // its memory is measured again whenever the context it follows changes + auto add_extra_memory = [&](dmds_t & dmds) { + if (extra == nullptr) { + return; + } + + 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); + } 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; + } + + 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; + } + }; + // step 1: get data for default parameters and check whether any changes are necessary in the first place LOG_TRC("%s: getting device memory data for initial parameters:\n", __func__); @@ -205,13 +266,16 @@ static void common_params_fit_impl( const uint32_t n_ctx_max = (uint32_t) std::min(uint64_t(hp_nct) * n_streams, UINT32_MAX); const uint32_t n_ctx_min_total = (uint32_t) std::min(uint64_t(n_ctx_min) * n_streams, UINT32_MAX); - // llama_context would use only hp_nct in total for n_ctx == 0, redo the estimate with the full context: - if (n_ctx_auto && n_streams > 1) { + // llama_context would use only hp_nct in total for n_ctx == 0, resolve the context before measuring anything else: + if (n_ctx_auto) { cparams->n_ctx = n_ctx_max; - LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n", - __func__, n_ctx_max, n_streams); - dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + if (n_streams > 1) { + LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n", + __func__, n_ctx_max, n_streams); + dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + } } + add_extra_memory(dmds_full); const size_t nd = devs.size(); // number of devices @@ -347,7 +411,8 @@ static void common_params_fit_impl( int64_t sum_projected_used_min_ctx = 0; cparams->n_ctx = n_ctx_min_total; - const dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + add_extra_memory(dmds_min_ctx); if (nd == 0) { sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total(); } else { @@ -527,8 +592,9 @@ static void common_params_fit_impl( llama_model_params mparams_copy = *mparams; set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy); - const dmds_t dmd_nl = common_get_device_memory_data_impl( + dmds_t dmd_nl = common_get_device_memory_data_impl( path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + add_extra_memory(dmd_nl); LOG_TRC("%s: memory for test allocation by device:\n", func_name); for (size_t id = 0; id < nd; id++) { @@ -555,8 +621,9 @@ static void common_params_fit_impl( mparams->tensor_buft_overrides = tensor_buft_overrides; LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__); - const dmds_t dmds_cpu_moe = common_get_device_memory_data_impl( + dmds_t dmds_cpu_moe = common_get_device_memory_data_impl( path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + add_extra_memory(dmds_cpu_moe); for (size_t id = 0; id < nd; id++) { global_surplus_cpu_moe += dmds_cpu_moe[id].free; @@ -816,11 +883,12 @@ enum common_params_fit_status common_fit_params( llama_model_tensor_buft_override * tensor_buft_overrides, size_t * margins, uint32_t n_ctx_min, + const common_fit_extra_model * extra, ggml_log_level log_level) { const int64_t t0_us = llama_time_us(); common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS; try { - common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, log_level); + common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, extra, log_level); LOG_TRC("%s: successfully fit params to free device memory\n", __func__); } catch (const common_params_fit_exception & e) { LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what()); diff --git a/common/fit.h b/common/fit.h index 208fc30694..824d386b07 100644 --- a/common/fit.h +++ b/common/fit.h @@ -11,6 +11,16 @@ enum common_params_fit_status { COMMON_PARAMS_FIT_STATUS_ERROR = 2, // a hard error occurred, e.g. because no model could be found at the specified path }; +// a second model that shares the devices of the main model, e.g. a draft model +// - its context follows the context of the main model, so its memory is measured again whenever that context changes +// - shares_model tells the fit that the weights are already counted in the main model, as for an MTP context +struct common_fit_extra_model { + const char * path_model; + llama_model_params * mparams; + llama_context_params * cparams; + bool shares_model; +}; + // fits mparams and cparams to free device memory (assumes system memory is unlimited) // - returns true if the parameters could be successfully modified to fit device memory // - this function is NOT thread safe because it modifies the global llama logger state @@ -24,6 +34,7 @@ common_params_fit_status common_fit_params( llama_model_tensor_buft_override * tensor_buft_overrides, // writable buffer for overrides, needs at least llama_max_tensor_buft_overrides elements size_t * margins, // margins of memory to leave per device in bytes uint32_t n_ctx_min, // minimum context size to set when trying to reduce memory use + 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 // print estimated memory to stdout diff --git a/tools/fit-params/fit-params.cpp b/tools/fit-params/fit-params.cpp index 5d897bc466..3e78c89290 100644 --- a/tools/fit-params/fit-params.cpp +++ b/tools/fit-params/fit-params.cpp @@ -33,6 +33,7 @@ int llama_fit_params(int argc, char ** argv) { if (!params.fit_params_print) { const common_params_fit_status 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, + nullptr, params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); if (status != COMMON_PARAMS_FIT_STATUS_SUCCESS) { LOG_ERR("%s: failed to fit CLI arguments to free memory, exiting...\n", __func__); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 03d59f08d1..a2da93b9a2 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -2294,6 +2294,7 @@ int llama_bench(int argc, char ** argv) { fit_overrides.data(), margins.data(), inst.fit_min_ctx, + nullptr, params.verbose ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f8976c6df3..36d9828325 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1040,74 +1040,7 @@ private: } } - // optionally reserve VRAM for the draft / MTP context before fitting the target model - if (params_base.fit_params) { - if (has_spec) { - // MTP draft context lives on the target model, only context+compute are new - bool measure_model_bytes = has_draft; - - common_params params_dft = common_base_params_to_speculative(params_base); - - auto mparams_dft = common_model_params_to_llama(params_dft); - auto cparams_dft = common_context_params_to_llama(params_dft); - if (spec_mtp) { - cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - } - cparams_dft.n_rs_seq = 0; - - std::vector devs; - uint32_t hp_ngl = 0; - uint32_t hp_nct = 0; - uint32_t hp_nex = 0; - try { - // the draft context follows the target context, measure it at the largest context the target can take - if (cparams_dft.n_ctx == 0) { - auto mparams_tgt = common_model_params_to_llama(params_base); - auto cparams_tgt = common_context_params_to_llama(params_base); - - common_get_device_memory_data( - params_base.model.path.c_str(), &mparams_tgt, &cparams_tgt, - devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR); - - cparams_dft.n_ctx = hp_nct * (params_base.kv_unified ? 1 : params_base.n_parallel); - } - - auto dmd = common_get_device_memory_data( - params_dft.model.path.c_str(), &mparams_dft, &cparams_dft, - devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR); - - GGML_ASSERT(!params_base.fit_params_target.empty()); - size_t total = 0; - - std::vector tgt_devices = params.devices; - - if (tgt_devices.empty()) { - for(size_t i = 0; i < ggml_backend_dev_count(); ++i) { - tgt_devices.push_back(ggml_backend_dev_get(i)); - } - } - - for (size_t j = 0; j < devs.size(); ++j) { - const size_t bytes = (measure_model_bytes ? dmd[j].model : 0) + dmd[j].context + dmd[j].compute; - total += bytes; - for (size_t i = 0; i < tgt_devices.size(); i++) { - if (tgt_devices[i] == devs[j]) { - SRV_DBG("[spec] adding %.2f MiB to fit_params_target for device %s\n", - bytes / (1024.0 * 1024.0), ggml_backend_dev_name(devs[j])); - params_base.fit_params_target[i] += bytes; - break; - } - } - } - SRV_TRC("[spec] estimated memory usage of %s is %.2f MiB\n", - has_draft ? "draft model" : "MTP context", - total / (1024.0 * 1024.0)); - } catch (const std::exception & e) { - SRV_WRN("[spec] failed to measure %s memory: %s\n", - has_draft ? "draft model" : "MTP context", e.what()); - } - } - } + // note: the draft / MTP context is fitted together with the target model, see common_fit_extra_model // attach a progress callback {