fit: Fix memory allocation for MTP layers (#26605)
This commit is contained in:
+4
-1
@@ -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));
|
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_ctx_train = llama_model_n_ctx_train(model);
|
||||||
hp_n_expert = llama_model_n_expert(model);
|
hp_n_expert = llama_model_n_expert(model);
|
||||||
|
|
||||||
|
|||||||
@@ -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_,
|
int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_,
|
||||||
int flags) {
|
int flags) {
|
||||||
const int64_t n_embd_qkv = n_embd_q_ + n_embd_k_ + n_embd_v_;
|
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);
|
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) {
|
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);
|
layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
|
||||||
|
|||||||
Reference in New Issue
Block a user