diff --git a/common/fit.cpp b/common/fit.cpp index c82d066..bb1a547 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -182,6 +182,7 @@ static void common_params_fit_impl( constexpr int64_t MiB = 1024*1024; typedef std::vector dmds_t; const llama_model_params default_mparams = llama_model_default_params(); + const bool fixed_gpu_layers = mparams->n_gpu_layers != default_mparams.n_gpu_layers; std::vector devs; uint32_t hp_ngl = 0; // hparams.n_gpu_layers @@ -314,7 +315,7 @@ static void common_params_fit_impl( sum_used_target -= margins[id]; } } - if (nd > 1) { + if (nd > 1 && !fixed_gpu_layers) { // for multiple devices we need to be more conservative in terms of how much context we think can fit: // - for dense models only whole layers can be assigned to devices // - for MoE models only whole tensors can be assigned to devices, which we estimate to be <= 1/3 of a layer @@ -334,6 +335,30 @@ static void common_params_fit_impl( sum_projected_used_min_ctx += dmds_min_ctx[id].mb.total(); } } + if (nd > 1 && fixed_gpu_layers) { + uint32_t n_ctx_fit = hp_nct; + bool fits = true; + for (size_t id = 0; id < nd; id++) { + const int64_t used_full = dmds_full[id].mb.total(); + const int64_t used_min = dmds_min_ctx[id].mb.total(); + const int64_t target = dmds_full[id].free - margins[id]; + if (used_full <= target) { + continue; + } + uint32_t n_ctx_device = n_ctx_min; + fits = fits && target >= used_min; + if (target > used_min && used_full > used_min) { + n_ctx_device += (hp_nct - n_ctx_min) * (target - used_min) / (used_full - used_min); + } + n_ctx_fit = std::min(n_ctx_fit, n_ctx_device); + } + cparams->n_ctx = std::max(n_ctx_fit - n_ctx_fit % 256, n_ctx_min); + LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " using per-device limits\n", + __func__, hp_nct, cparams->n_ctx); + if (fits) { + return; + } + } 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) @@ -371,7 +396,7 @@ static void common_params_fit_impl( throw common_params_fit_exception("was unable to fit model into system memory by reducing context, abort"); } - if (mparams->n_gpu_layers != default_mparams.n_gpu_layers) { + if (fixed_gpu_layers) { throw common_params_fit_exception("n_gpu_layers already set by user to " + std::to_string(mparams->n_gpu_layers) + ", abort"); } if (nd > 1) { diff --git a/common/speculative.cpp b/common/speculative.cpp index 70dc0ac..f7980e0 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2321,6 +2321,9 @@ common_speculative_init_result::common_speculative_init_result( if (spec_mtp) { cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + // use the same (possibly fit-reduced) context size as the target + // avoids MTP draft reserving flash attention workspace for full n_ctx + 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 diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 19cca7d..a095c1b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -422,6 +422,7 @@ llama_context::llama_context( // TODO: move these checks to ggml_backend_sched // enabling pipeline parallelism in the scheduler increases memory usage, so it is only done when necessary bool pipeline_parallel = + cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP && model.n_devices() > 1 && model.n_gpu_layers() > model.hparams.n_layer_all && model.split_mode() == LLAMA_SPLIT_MODE_LAYER && diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5d2798c..2ff4abc 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1133,6 +1133,19 @@ private: auto cparams_dft = common_context_params_to_llama(params_dft); if (spec_mtp) { cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + // MTP draft uses the target's (possibly fit-reduced) n_ctx at runtime. + // Estimate with the same n_ctx to avoid over-reserving fit_params_target. + { + auto mparams_tgt = common_model_params_to_llama(params_base); + auto cparams_tgt = common_context_params_to_llama(params_base); + common_fit_params(params_dft.model.path.c_str(), &mparams_tgt, &cparams_tgt, + params_base.tensor_split, + params_base.tensor_buft_overrides.data(), + params_base.fit_params_target.data(), + params_base.fit_params_min_ctx, + GGML_LOG_LEVEL_ERROR); + cparams_dft.n_ctx = cparams_tgt.n_ctx; + } } cparams_dft.n_rs_seq = 0; @@ -1156,6 +1169,7 @@ private: } } + std::vector measured_bytes(tgt_devices.size(), 0); 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; @@ -1164,6 +1178,7 @@ private: 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; + measured_bytes[i] += bytes; break; } } @@ -1171,6 +1186,47 @@ private: SRV_TRC("[spec] estimated memory usage of %s is %.2f MiB\n", has_draft ? "draft model" : "MTP context", total / (1024.0 * 1024.0)); + + if (spec_mtp) { + for (int pass = 0; pass < 2; ++pass) { + auto mparams_tgt = common_model_params_to_llama(params_base); + auto cparams_tgt = common_context_params_to_llama(params_base); + common_fit_params(params_dft.model.path.c_str(), &mparams_tgt, &cparams_tgt, + params_base.tensor_split, + params_base.tensor_buft_overrides.data(), + params_base.fit_params_target.data(), + params_base.fit_params_min_ctx, + GGML_LOG_LEVEL_ERROR); + if (cparams_dft.n_ctx == cparams_tgt.n_ctx) { + break; + } + + cparams_dft.n_ctx = cparams_tgt.n_ctx; + auto dmd_refined = 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); + + std::vector refined_bytes(tgt_devices.size(), 0); + total = 0; + for (size_t j = 0; j < devs.size(); ++j) { + const size_t bytes = (measure_model_bytes ? dmd_refined[j].model : 0) + dmd_refined[j].context + dmd_refined[j].compute; + total += bytes; + for (size_t i = 0; i < tgt_devices.size(); ++i) { + if (tgt_devices[i] == devs[j]) { + refined_bytes[i] += bytes; + break; + } + } + } + for (size_t i = 0; i < tgt_devices.size(); ++i) { + params_base.fit_params_target[i] -= measured_bytes[i]; + params_base.fit_params_target[i] += refined_bytes[i]; + } + measured_bytes = std::move(refined_bytes); + SRV_TRC("[spec] refined MTP memory estimate at n_ctx=%" PRIu32 " is %.2f MiB\n", + cparams_dft.n_ctx, 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());