common : fractional -ncmoe for tensor-granularity expert placement
Extends --n-cpu-moe to accept a fractional layer count. The integer part offloads whole layers as before; the fractional part offloads a subset of the boundary layer expert tensors (gate, then up), keeping down_proj resident. This realizes ATSInfer tensor-granularity static placement, giving sub-layer control over expert VRAM residency. Placement-only, lossless. Assisted-by: Claude
This commit is contained in:
+7
-5
@@ -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<std::string> 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()});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<std::string> llm_ffn_exps_cpu_block_regexes(double n_cpu_moe) {
|
||||
std::vector<std::string> 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
|
||||
//
|
||||
|
||||
Reference in New Issue
Block a user