diff --git a/common/common.cpp b/common/common.cpp index 25ca838dff0a..d84d57ac910f 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 dd1f3ef76619..c601fe405ea5 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"); } @@ -191,10 +191,92 @@ static void common_params_fit_impl( uint32_t hp_nct = 0; // hparams.n_ctx_train uint32_t hp_nex = 0; // hparams.n_expert + // with non-unified kv, we need to take into account n_streams + // for example, if memory can hold more than model's trained context size, we must extend the n_ctx to hold enough n_streams + 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__); - const dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + + // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min: + 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, resolve the context before measuring anything else: + if (n_ctx_auto) { + cparams->n_ctx = n_ctx_max; + 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 std::vector margins; // this function uses int64_t rather than size_t for memory sizes to more conveniently handle deficits @@ -307,8 +389,8 @@ static void common_params_fit_impl( "%s: cannot meet free memory targets on all devices, need to use %" PRId64 " MiB less in total\n", __func__, -global_surplus/MiB); } - if (cparams->n_ctx == 0) { - if (hp_nct > n_ctx_min) { + if (n_ctx_auto) { + if (n_ctx_max > n_ctx_min_total) { int64_t sum_used_target = sum_free; if (nd == 0) { sum_used_target -= margins[0]; @@ -328,8 +410,9 @@ static void common_params_fit_impl( } int64_t sum_projected_used_min_ctx = 0; - cparams->n_ctx = n_ctx_min; - 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); + cparams->n_ctx = n_ctx_min_total; + 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 { @@ -339,14 +422,16 @@ static void common_params_fit_impl( } if (sum_used_target > sum_projected_used_min_ctx) { // linear interpolation between minimum and maximum context size: - cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx) + cparams->n_ctx += (n_ctx_max - n_ctx_min_total) * (sum_used_target - sum_projected_used_min_ctx) / (sum_projected_used - sum_projected_used_min_ctx); - cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % 256, n_ctx_min); // round down context for CUDA backend + // round down context for CUDA backend, keep it divisible by the number of streams: + const uint32_t align = 256 * n_streams; + cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % align, n_ctx_min_total); - const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (hp_nct - n_ctx_min); - const int64_t memory_reduction = (hp_nct - cparams->n_ctx) * bytes_per_ctx; + const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (n_ctx_max - n_ctx_min_total); + const int64_t memory_reduction = (n_ctx_max - cparams->n_ctx) * bytes_per_ctx; LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", - __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); + __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB); if (nd <= 1) { LOG_TRC("%s: entire model can be fit by reducing context\n", __func__); return; @@ -355,14 +440,14 @@ static void common_params_fit_impl( } else { const int64_t memory_reduction = sum_projected_used - sum_projected_used_min_ctx; LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", - __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); + __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB); } } else { if (n_ctx_min == UINT32_MAX) { - LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, hp_nct); + LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, n_ctx_max); } else { LOG_TRC("%s: default model context size is %" PRIu32 " which is <= the min. context size of %" PRIu32 " -> no change\n", - __func__, hp_nct, n_ctx_min); + __func__, n_ctx_max, n_ctx_min_total); } } } else { @@ -507,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++) { @@ -535,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; @@ -796,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 208fc30694e0..824d386b07a1 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/common/speculative.cpp b/common/speculative.cpp index 89e9b2782c2a..12909822babc 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2385,6 +2385,9 @@ common_speculative_init_result::common_speculative_init_result( cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; } + // the draft context holds as many tokens per sequence as the target context + cparams.n_ctx = llama_n_ctx(ctx_tgt); + // note: for small models maybe we can set this to the maximum possible draft from all speculative types // the extra memory for small models is likely negligible? cparams.n_rs_seq = 0; diff --git a/tools/fit-params/fit-params.cpp b/tools/fit-params/fit-params.cpp index 5d897bc46699..3e78c89290a7 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 03d59f08d1ab..a2da93b9a282 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 1293c8640267..36d982832567 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1040,62 +1040,7 @@ struct server_context_impl { } } - // 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 { - 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 {