Better placement of MoE tensors with -ncmoe and 1 GPU

This commit is contained in:
Kawrakow
2026-08-06 06:52:57 +00:00
parent cf1aa57e1a
commit fb5fb74880
+11 -4
View File
@@ -264,10 +264,17 @@ create_tensors_helper::create_tensors_helper(llama_model_loader & _ml, llama_mod
if (ml.ncmoe > 0) {
auto buft = llama_default_buffer_type_cpu(true);
if (model.split_mode == LLAMA_SPLIT_MODE_ATTN || model.split_mode == LLAMA_SPLIT_MODE_GRAPH || ml.ncmoe >= n_layer || model.devices.size() < 2) {
int nmax = std::min(ml.ncmoe, n_layer);
for (int i = 0; i < nmax; ++i) {
std::string pattern = "blk\\." + std::to_string(i) + "\\.ffn_(up|down|gate|gate_up)_exps\\.(weight|scale)";
this->overrides.emplace_back(std::make_pair(std::regex(pattern), buft));
const auto tn = LLM_TN(model.arch);
int last_layer = n_layer - model.hparams.nextn_predict_layers;
int ndone = 0;
for (int i = last_layer-1; i >= 0; --i) {
auto d_name = tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i);
auto d_meta = ml.get_tensor_meta(d_name.c_str());
if (d_meta) {
std::string pattern = "blk\\." + std::to_string(i) + "\\.ffn_(up|down|gate|gate_up)_exps\\.(weight|scale)";
this->overrides.emplace_back(std::make_pair(std::regex(pattern), buft));
if (++ndone == ml.ncmoe) break;
}
}
}
else if (model.split_mode == LLAMA_SPLIT_MODE_LAYER) {