diff --git a/common/arg.cpp b/common/arg.cpp index da40874740ce..1aaa62394216 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2671,14 +2671,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", @@ -3982,11 +3988,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 2e15ec3f815f..cde7916ea8cd 100644 --- a/common/common.h +++ b/common/common.h @@ -8,6 +8,7 @@ #include "ggml.h" #include "llama.h" +#include #include #include #include @@ -1083,19 +1084,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 c17a27b54019..76644b262e94 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -1252,7 +1252,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() }); }