diff --git a/common/common.cpp b/common/common.cpp index 73245b87..0427f21c 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1517,6 +1517,11 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa } return true; } + if (arg == "-wgt" || arg == "--worst-graph-tokens") { + CHECK_ARG; + params.worst_graph_tokens = std::stoi(argv[i]); + return true; + } if (arg == "--no-mmap") { params.use_mmap = false; return true; @@ -2502,6 +2507,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", " --cpu-moe", "keep all MoE weights in CPU memory"}); options.push_back({ "*", " --n-cpu-moe N", "keep MoE weights of the first N layers in CPU memory"}); options.push_back({ "*", " --fit-margin N", "safety margin in MiB when auto-fitting model offloading"}); + options.push_back({ "*", "-wgt, --worst-graph-tokens N", "number of tokens to use for worst-case graph"}); options.push_back({ "*", " --fit", "automatically determine which tensors to offload to the GPU(s)"}); options.push_back({ "*", " --numa TYPE", "attempt optimizations that help on some NUMA systems\n" " - distribute: spread execution evenly over all nodes\n" @@ -3337,6 +3343,7 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params mparams.ncmoe = params.ncmoe; mparams.fit = params.fit; mparams.fit_margin = params.fit_margin; + mparams.worst_graph_tokens = params.worst_graph_tokens; mparams.type_k = kv_cache_type_from_str(params.cache_type_k); mparams.type_v = kv_cache_type_from_str(params.cache_type_v); mparams.max_ctx_size = params.n_ctx; @@ -3404,6 +3411,7 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa cparams.seed = params.seed; cparams.logits_all = params.logits_all; cparams.embeddings = params.embedding; + cparams.worst_case_tokens = params.worst_graph_tokens; cparams.rope_scaling_type = params.rope_scaling_type; cparams.rope_freq_base = params.rope_freq_base; cparams.rope_freq_scale = params.rope_freq_scale; @@ -4386,6 +4394,7 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l fprintf(stream, "ncmoe: %d # default: 0\n", params.ncmoe); fprintf(stream, "fit: %d # default: false\n", params.fit); fprintf(stream, "fit_margin: %d # default: 0\n", params.fit_margin); + fprintf(stream, "worst_graph_tokens: %d # default: 0\n", params.worst_graph_tokens); fprintf(stream, "min_keep: %d # default: 0 (disabled)\n", sparams.min_keep); fprintf(stream, "mirostat: %d # default: 0 (disabled)\n", sparams.mirostat); fprintf(stream, "mirostat_ent: %f # default: 5.0\n", sparams.mirostat_tau); diff --git a/common/common.h b/common/common.h index 31cca0d3..1a5ab59a 100644 --- a/common/common.h +++ b/common/common.h @@ -228,6 +228,7 @@ struct gpt_params { int32_t ncmoe = 0; // number of layers in which MoE tensors are left in VRAM int32_t fit_margin = 0; // safety margin for auto-fit in MiB bool fit = false; // automatically fit model (for now just using MoE tensor overrides) + int32_t worst_graph_tokens = 0; // number of tokens to use when reserving the worst graph float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs int32_t grp_attn_n = 1; // group-attention factor int32_t grp_attn_w = 512; // group-attention width diff --git a/include/llama.h b/include/llama.h index c99e21b0..a2f04409 100644 --- a/include/llama.h +++ b/include/llama.h @@ -380,6 +380,7 @@ extern "C" { int32_t amb; int32_t fit_margin; bool fit; + int32_t worst_graph_tokens; // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; @@ -426,6 +427,7 @@ extern "C" { uint32_t n_threads; // number of threads to use for generation uint32_t n_threads_batch; // number of threads to use for batch processing int32_t max_extra_alloc; // Max. additional VRAM the scheduler is allowed to allocate + int32_t worst_case_tokens; // number of tokens to use when reserving worst case graphs enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type` enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 67955280..6338a391 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -34,7 +34,8 @@ llm_build_context::llm_build_context( const llama_batch & batch, const llm_build_cb & cb, bool worst_case, - bool warmup) : + bool warmup, + int n_outputs_) : model (lctx.model), lctx (lctx), hparams (model.hparams), @@ -63,7 +64,7 @@ llm_build_context::llm_build_context( norm_rms_eps (hparams.f_norm_rms_eps), n_tokens (batch.n_tokens), n_kv (worst_case ? kv_self.size : kv_self.n), - n_outputs (worst_case ? n_tokens : lctx.n_outputs), + n_outputs (worst_case ? n_outputs_ > 0 ? n_outputs_ : n_tokens : lctx.n_outputs), n_outputs_enc (worst_case ? n_tokens : lctx.embd_enc.size() / hparams.n_embd), kv_head (worst_case ? (kv_self.recurrent ? 0 : kv_self.size - n_tokens) : kv_self.head), n_ctx_orig (cparams.n_ctx_orig_yarn), @@ -308,17 +309,17 @@ struct ggml_tensor * llm_build_context::build_inp_embd_mtp(struct ggml_tensor * if (batch.token) { lctx.inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch.n_tokens); - + cb(lctx.inp_tokens, "inp_tokens", -1); ggml_set_input(lctx.inp_tokens); cur = ggml_get_rows(ctx0, mtp_tok_embd, lctx.inp_tokens); } else { - return nullptr; + return nullptr; } cb(cur, "inp_embd", -1); - + return cur; } @@ -9597,7 +9598,8 @@ struct ggml_cgraph * llm_build_context::llama_build_graph_s_copy(llama_context & ggml_cgraph * llm_build_context::llama_build_graph( llama_context & lctx, const llama_batch & batch, - bool worst_case) { + bool worst_case, + int n_outputs) { const auto & model = lctx.model; #if IK_PRINT_TIMING @@ -9654,7 +9656,7 @@ ggml_cgraph * llm_build_context::llama_build_graph( llama_token bos = vocab->token_bos(); llama_token eos = vocab->token_eos(); bool is_warming_up = lctx.n_eval == 0 && (batch.n_tokens == 1 && (batch.token[0] == ((bos != -1) ? bos : eos))); - struct llm_build_context llm(lctx, batch, cb, worst_case, is_warming_up); + struct llm_build_context llm(lctx, batch, cb, worst_case, is_warming_up, n_outputs); llm.init(); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index cdaa7422..29edf2d7 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -102,7 +102,8 @@ struct llm_build_context { const llama_batch & batch, const llm_build_cb & cb, bool worst_case, - bool warmup); + bool warmup, + int n_outputs = 0); void init(); @@ -428,7 +429,7 @@ llm_expert_gating_func_type gating_op, static ggml_cgraph * llama_build_graph_s_copy(llama_context & lctx); - static ggml_cgraph * llama_build_graph(llama_context & lctx, const llama_batch & batch, bool worst_case); + static ggml_cgraph * llama_build_graph(llama_context & lctx, const llama_batch & batch, bool worst_case, int n_outputs = 0); ggml_tensor * build_std_attention(ggml_cgraph * gf, ggml_tensor * attn_norm, ggml_tensor * cur, ggml_tensor * inp_pos, ggml_tensor * inp_out_ids, ggml_tensor * rope_factors, @@ -439,8 +440,8 @@ llm_expert_gating_func_type gating_op, static uint32_t llama_kv_qnext_state_slots(const llama_kv_cache & kv_self); struct ggml_tensor * build_mtp_tail( - const struct llama_layer & mtp_layer, - struct ggml_tensor * prev_embeddings, + const struct llama_layer & mtp_layer, + struct ggml_tensor * prev_embeddings, int64_t n_embd_head, struct ggml_cgraph * gf, struct ggml_tensor * inp_pos, diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 49e8cbea..adc25099 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -47,6 +47,7 @@ struct llama_cparams { int min_experts; float thresh_experts; bool mtp; + int worst_graph_tokens; enum ggml_type reduce_type; enum llama_pooling_type pooling_type; diff --git a/src/llama.cpp b/src/llama.cpp index 01797641..ee92f55a 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -2007,7 +2007,8 @@ struct expert_tensors { }; static std::pair, double> get_layer_sizes(const llama_model_loader & ml, const llama_model & model, - ggml_type cache_type_k, ggml_type cache_type_v, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, int amb, bool flash_attn, + ggml_type cache_type_k, ggml_type cache_type_v, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, + int amb, int worst_case_tokens, bool flash_attn, std::vector & experts) { int n_layer = model.hparams.n_layer; std::vector result(n_layer+1, 0); @@ -2181,7 +2182,8 @@ static std::pair, double> get_layer_sizes(const llama_model_ tot_cache += kv_size; result[il] += kv_size; } - size_t output_size = model.hparams.n_vocab * n_ubatch * sizeof(float); + int n_output = worst_case_tokens > 0 ? worst_case_tokens : n_ubatch; + size_t output_size = model.hparams.n_vocab * n_output * sizeof(float); if (output_size < max_compute) output_size = max_compute; output_size -= max_compute; LLAMA_LOG_INFO("Layer %2d: %9.2f, %9.2f, %9.2f MiB (output layer)\n", n_layer, result[n_layer]/1024./1024., output_size/1024./1024., (result[n_layer] + output_size)/1024./1024.); @@ -2209,6 +2211,7 @@ static bool llm_load_tensors( int n_ubatch, int amb, int fit_margin, + int worst_case_tokens, bool flash_attn, bool use_mlock, bool validate_quants, @@ -2331,7 +2334,8 @@ static bool llm_load_tensors( std::vector overrides; if (device_count > 0) { std::vector experts; - auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, n_ubatch, amb, flash_attn, experts); + auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, n_ubatch, + amb, worst_case_tokens, flash_attn, experts); size_t required_mem = 0; for (int i = 0; i <= n_layer; ++i) { required_mem += layer_sizes[i]; @@ -2894,7 +2898,8 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam if (!llm_load_tensors( ml, model, params.n_gpu_layers, params.mla, params.split_mode, params.main_gpu, params.max_gpu, params.tensor_split, - params.type_k, params.type_v, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.fit_margin, params.flash_attn, + params.type_k, params.type_v, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.fit_margin, + params.worst_graph_tokens, params.flash_attn, params.use_mlock, params.validate_quants, params.mtp, params.fit, params.dry_run, params.progress_callback, params.progress_callback_user_data )) { @@ -4651,7 +4656,7 @@ static int32_t llama_kv_cache_update_internal(struct llama_context & lctx) { int n_tokens = (int)std::min(lctx.cparams.n_ctx, lctx.cparams.n_ubatch); int n_past = lctx.cparams.n_ctx - n_tokens; llama_token token = llama_token_bos(&lctx.model); // not actually used by llama_build_graph, but required to choose between token and embedding inputs graph - ggml_cgraph * gf = llm_build_context::llama_build_graph(lctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true); + ggml_cgraph * gf = llm_build_context::llama_build_graph(lctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true, lctx.cparams.worst_graph_tokens); // initialize scheduler with the worst-case graph lctx.reset_scheduler(); @@ -4892,6 +4897,7 @@ struct llama_model_params llama_model_default_params() { /*.amb =*/ 0, /*.fit_margin =*/ 0, /*.fit =*/ false, + /*.worst_graph_tokens =*/ 0, /*.tensor_split =*/ nullptr, /*.rpc_servers =*/ nullptr, /*.progress_callback =*/ nullptr, @@ -4930,6 +4936,7 @@ struct llama_context_params llama_context_default_params() { /*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default /*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS, /*.max_extra_alloc =*/ 256, + /*.worst_case_tokens =*/ 0, /*.rope_scaling_type =*/ LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED, /*.pooling_type =*/ LLAMA_POOLING_TYPE_UNSPECIFIED, /*.attention_type =*/ LLAMA_ATTENTION_TYPE_UNSPECIFIED, @@ -5350,6 +5357,7 @@ struct llama_context * llama_init_from_model( cparams.thresh_experts = params.thresh_experts; cparams.cuda_params = params.cuda_params; cparams.mtp = params.mtp; + cparams.worst_graph_tokens = params.worst_case_tokens; cparams.reduce_type = params.type_reduce; cparams.pooling_type = params.pooling_type; @@ -5749,7 +5757,7 @@ struct llama_context * llama_init_from_model( // build worst-case graph int n_past = cparams.n_ctx - n_tokens; llama_token token = llama_token_bos(&ctx->model); // not actually used by llama_build_graph, but required to choose between token and embedding inputs graph - ggml_cgraph * gf = llm_build_context::llama_build_graph(*ctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true); + ggml_cgraph * gf = llm_build_context::llama_build_graph(*ctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true, cparams.worst_graph_tokens); // initialize scheduler with the worst-case graph bool gf_success = ggml_backend_sched_reserve(ctx->sched, gf);