From c5fc7e34885ba31217e330809437afa993d27745 Mon Sep 17 00:00:00 2001 From: Jonas J <111707981+John-194@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:26:42 +0300 Subject: [PATCH] llama : add --n-cpu-ffn option (#26622) * common : dedupe --n-cpu-moe / --spec-draft-n-cpu-moe override loops * common : add --n-cpu-ffn to CPU-offload dense FFN weights of first N layers * common : generalize llm_ffn_block_regex over the FFN regex, drop TODO --- common/arg.cpp | 24 +++++++++++++----------- common/common.h | 18 +++++++++++++++--- tools/llama-bench/llama-bench.cpp | 2 +- 3 files changed, 29 insertions(+), 15 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a5..aad8266ed 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2750,14 +2750,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex if (value < 0) { throw std::invalid_argument("invalid value"); } - for (int i = 0; i < value; ++i) { - // keep strings alive and avoid leaking memory by storing them in a static vector - static std::list buft_overrides; - buft_overrides.push_back(llm_ffn_exps_block_regex(i)); - params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()}); - } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides); } ).set_env("LLAMA_ARG_N_CPU_MOE")); + add_opt(common_arg( + {"-ncffn", "--n-cpu-ffn"}, "N", + "keep the dense FFN weights of the first N layers in the CPU\n" + "(dense models; for MoE expert weights use --n-cpu-moe)", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_DENSE_REGEX, params.tensor_buft_overrides); + } + ).set_env("LLAMA_ARG_N_CPU_FFN")); GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 add_opt(common_arg( {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N", @@ -4084,11 +4090,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex if (value < 0) { throw std::invalid_argument("invalid value"); } - for (int i = 0; i < value; ++i) { - static std::list buft_overrides_draft; - buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i)); - params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()}); - } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.speculative.draft.tensor_buft_overrides); } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE")); diff --git a/common/common.h b/common/common.h index de49dac9f..51518f343 100644 --- a/common/common.h +++ b/common/common.h @@ -8,6 +8,7 @@ #include "ggml.h" #include "llama.h" +#include #include #include #include @@ -1108,19 +1109,30 @@ const char * const LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"; } // -// MoE utils +// FFN offload utils // const char * const LLM_FFN_EXPS_REGEX = "\\.ffn_(up|down|gate|gate_up)_(ch|)exps"; -inline std::string llm_ffn_exps_block_regex(int idx) { - return string_format("blk\\.%d%s", idx, LLM_FFN_EXPS_REGEX); +const char * const LLM_FFN_DENSE_REGEX = "\\.ffn_(up|down|gate)\\."; + +inline std::string llm_ffn_block_regex(int idx, const char * ffn_regex) { + return string_format("blk\\.%d%s", idx, ffn_regex); } inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; } +inline void llm_add_n_cpu_ffn_overrides(int n, const char * ffn_regex, std::vector & overrides) { + // keep strings alive and avoid leaking memory by storing them in a static list + static std::list buft_override_strings; + for (int i = 0; i < n; ++i) { + buft_override_strings.push_back(llm_ffn_block_regex(i, ffn_regex)); + overrides.push_back({buft_override_strings.back().c_str(), ggml_backend_cpu_buffer_type()}); + } +} + // // training utils // diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index a2da93b9a..bc14d15c7 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -1254,7 +1254,7 @@ struct cmd_params_instance { merged.reserve(merged.size() + (size_t) n_cpu_moe + 1); for (int i = 0; i < n_cpu_moe; ++i) { - patterns.push_back(llm_ffn_exps_block_regex(i)); + patterns.push_back(llm_ffn_block_regex(i, LLM_FFN_EXPS_REGEX)); merged.push_back({ patterns.back().c_str(), ggml_backend_cpu_buffer_type() }); }