diff --git a/common/fit.cpp b/common/fit.cpp index c82d066ad4..dd1f3ef766 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -136,7 +136,10 @@ static std::vector 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); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index dda311c47b..4cc1c0a1c2 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2890,6 +2890,21 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid, int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_, int flags) { const int64_t n_embd_qkv = n_embd_q_ + n_embd_k_ + n_embd_v_; + + if (flags & TENSOR_SKIP) { + const int skip = TENSOR_NOT_REQUIRED | TENSOR_SKIP; + + create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, skip | TENSOR_SKIP_IF_VIRTUAL); + create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, skip | TENSOR_SKIP_IF_VIRTUAL); + create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", bid), {n_embd_, n_embd_q_}, skip); + create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", bid), {n_embd_, n_embd_k_}, skip); + create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", bid), {n_embd_, n_embd_v_}, skip); + create_tensor(tn(LLM_TENSOR_ATTN_Q, "bias", bid), {n_embd_q_}, skip); + create_tensor(tn(LLM_TENSOR_ATTN_K, "bias", bid), {n_embd_k_}, skip); + create_tensor(tn(LLM_TENSOR_ATTN_V, "bias", bid), {n_embd_v_}, skip); + return; + } + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); if (layer.wqkv) { layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);