diff --git a/common/arg.cpp b/common/arg.cpp index 975344131..0dbfa88fc 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2606,15 +2606,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_CPU_MOE")); add_opt(common_arg( {"-ncmoe", "--n-cpu-moe"}, "N", - "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU", - [](common_params & params, int value) { - if (value < 0) { + "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU; " + "fractional N offloads part of the boundary layer at tensor granularity", + [](common_params & params, const std::string & value) { + const double n = std::stod(value); + if (n < 0) { throw std::invalid_argument("invalid value"); } - for (int i = 0; i < value; ++i) { + for (const std::string & re : llm_ffn_exps_cpu_block_regexes(n)) { // 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)); + buft_overrides.push_back(re); params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()}); } } diff --git a/common/common.h b/common/common.h index 279252183..7e63cb39f 100644 --- a/common/common.h +++ b/common/common.h @@ -1083,6 +1083,28 @@ inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; } +// ATSInfer-style tensor-granularity static placement of MoE expert weights. +// Offloads the expert weights of the first floor(n) layers to the CPU, plus a +// subset of the boundary layer's three expert tensors for the fractional part. +// Expert tensors are dropped from the GPU in ascending performance-density +// order (gate, then up), keeping the higher-value down_proj resident longest. +inline std::vector llm_ffn_exps_cpu_block_regexes(double n_cpu_moe) { + std::vector regexes; + const int n_full = n_cpu_moe > 0 ? (int) n_cpu_moe : 0; + for (int i = 0; i < n_full; ++i) { + regexes.push_back(llm_ffn_exps_block_regex(i)); + } + const int k = (int) ((n_cpu_moe - n_full) * 3.0 + 0.5); + if (k >= 3) { + regexes.push_back(llm_ffn_exps_block_regex(n_full)); + } else if (k == 2) { + regexes.push_back(string_format("blk\\.%d\\.ffn_(gate|up)_(ch|)exps", n_full)); + } else if (k == 1) { + regexes.push_back(string_format("blk\\.%d\\.ffn_gate_(ch|)exps", n_full)); + } + return regexes; +} + // // training utils //