fit: Fix memory allocation for MTP layers (#26605)

2026-08-05 13:29:45 +02:00
parent 9303cdd8d3
commit 9a688e51e6
2 changed files with 19 additions and 1 deletions
@@ -136,7 +136,10 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
        devs.push_back(llama_model_get_device(model, i));
    }

-    hp_ngl         = llama_model_n_layer(model) + llama_model_n_layer_nextn(model);
+    hp_ngl         = llama_model_n_layer(model);
+    if (mparams->load_mtp) {
+        hp_ngl    += llama_model_n_layer_nextn(model);
+    }
    hp_n_ctx_train = llama_model_n_ctx_train(model);
    hp_n_expert    = llama_model_n_expert(model);