diff --git a/common/arg.cpp b/common/arg.cpp index b75f4f0..0140dcc 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1713,6 +1713,21 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str())); } }).set_env("LLAMA_ARG_FLASH_ATTN")); + add_opt(common_arg({ "--pipeline-parallel" }, "[on|off|auto]", + string_format("set pipeline parallelism ('on', 'off', or 'auto', default: '%s')", + llama_pipeline_parallel_type_name(params.pipeline_parallel_type)), + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED; + } else if (is_falsey(value)) { + params.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED; + } else if (is_autoy(value)) { + params.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_AUTO; + } else { + throw std::runtime_error( + string_format("error: unknown value for --pipeline-parallel: '%s'\n", value.c_str())); + } + }).set_env("LLAMA_ARG_PIPELINE_PARALLEL")); add_opt(common_arg( {"-p", "--prompt"}, "PROMPT", "prompt to start generation with; for system message, use -sys", diff --git a/common/common.cpp b/common/common.cpp index d9ce575..2a45d38 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1667,6 +1667,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.pooling_type = params.pooling_type; cparams.attention_type = params.attention_type; cparams.flash_attn_type = params.flash_attn_type; + cparams.pipeline_parallel_type = params.pipeline_parallel_type; cparams.cb_eval = params.cb_eval; cparams.cb_eval_user_data = params.cb_eval_user_data; cparams.offload_kqv = !params.no_kv_offload; diff --git a/common/common.h b/common/common.h index 78d0877..3718997 100644 --- a/common/common.h +++ b/common/common.h @@ -494,6 +494,7 @@ struct common_params { enum llama_pooling_type pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED; // pooling type for embeddings enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention + enum llama_pipeline_parallel_type pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_AUTO; // whether to use pipeline parallelism struct common_params_sampling sampling; struct common_params_speculative speculative; 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..5c22d87 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2321,6 +2321,10 @@ common_speculative_init_result::common_speculative_init_result( if (spec_mtp) { cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED; + // 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/include/llama.h b/include/llama.h index fb2ca38..9ab521d 100644 --- a/include/llama.h +++ b/include/llama.h @@ -195,6 +195,14 @@ extern "C" { LLAMA_API const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_type); + enum llama_pipeline_parallel_type { + LLAMA_PIPELINE_PARALLEL_TYPE_AUTO = -1, + LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED = 0, + LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED = 1, + }; + + LLAMA_API const char * llama_pipeline_parallel_type_name(enum llama_pipeline_parallel_type pipeline_parallel_type); + enum llama_split_mode { LLAMA_SPLIT_MODE_NONE = 0, // single GPU LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs @@ -362,6 +370,7 @@ extern "C" { enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id enum llama_attention_type attention_type; // attention type to use for embeddings enum llama_flash_attn_type flash_attn_type; // when to enable Flash Attention + enum llama_pipeline_parallel_type pipeline_parallel_type; // when to enable pipeline parallelism // ref: https://github.com/ggml-org/llama.cpp/pull/2054 float rope_freq_base; // RoPE base frequency, 0 = from model diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 19cca7d..3a17395 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -307,6 +307,7 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); + LLAMA_LOG_INFO("%s: pipeline mode = %s\n", __func__, llama_pipeline_parallel_type_name(params.pipeline_parallel_type)); LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); @@ -422,6 +423,8 @@ 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 && + params.pipeline_parallel_type != LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED && model.n_devices() > 1 && model.n_gpu_layers() > model.hparams.n_layer_all && model.split_mode() == LLAMA_SPLIT_MODE_LAYER && @@ -448,6 +451,11 @@ llama_context::llama_context( } } + if (cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP && + params.pipeline_parallel_type == LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED && !pipeline_parallel) { + LLAMA_LOG_WARN("%s: pipeline parallelism requested but unavailable\n", __func__); + } + cparams.pipeline_parallel = pipeline_parallel; if (cparams.pipeline_parallel) { @@ -3495,6 +3503,7 @@ llama_context_params llama_context_default_params() { /*.pooling_type =*/ LLAMA_POOLING_TYPE_UNSPECIFIED, /*.attention_type =*/ LLAMA_ATTENTION_TYPE_UNSPECIFIED, /*.flash_attn_type =*/ LLAMA_FLASH_ATTN_TYPE_AUTO, + /*.pipeline_parallel_type =*/ LLAMA_PIPELINE_PARALLEL_TYPE_AUTO, /*.rope_freq_base =*/ 0.0f, /*.rope_freq_scale =*/ 0.0f, /*.yarn_ext_factor =*/ -1.0f, diff --git a/src/llama.cpp b/src/llama.cpp --- a/src/llama.cpp +++ b/src/llama.cpp @@ -46,6 +46,18 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty GGML_ABORT("fatal error"); } +const char * llama_pipeline_parallel_type_name(enum llama_pipeline_parallel_type pipeline_parallel_type) { + switch (pipeline_parallel_type) { + case LLAMA_PIPELINE_PARALLEL_TYPE_AUTO: + return "auto"; + case LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED: + return "disabled"; + case LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED: + return "enabled"; + } + GGML_ABORT("fatal error"); +} + const char * llama_load_mode_name(enum llama_load_mode load_mode) { switch (load_mode) { case LLAMA_LOAD_MODE_NONE: diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5d2798c..d2a0966 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1133,6 +1133,20 @@ private: auto cparams_dft = common_context_params_to_llama(params_dft); if (spec_mtp) { cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams_dft.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED; + // 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 +1170,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 +1179,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 +1187,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());