nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Take into account layer sizes for setting GPU layers (count 2) (#1498) 

* Also take into account KV cache

* Take into account attn_wkv_b and mla = 3 compute buffers

* WIP

* Minor
This commit is contained in:
Kawrakow 2026-03-24 08:18:28 +01:00 committed by GitHub
parent f6c8b5f2cb
commit 233225db8f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 74 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/llama-arch.h
View File

@ -340,6 +340,8 @@ enum llm_tensor {
LLM_TENSOR_INDEXER_PROJ,
LLM_TENSOR_INDEXER_ATTN_K,
LLM_TENSOR_INDEXER_ATTN_Q_B,
LLM_TENSOR_UNKNOWN,
};
llm_arch llm_arch_from_string(const std::string & name);
@ -348,3 +350,5 @@ const char * llama_model_arch_name(llm_arch arch);
bool llm_arch_is_recurrent(const llm_arch & arch);
bool llm_arch_is_hybrid(const llm_arch & arch);
llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int il);

23
src/llama-model.cpp
View File

@ -1837,6 +1837,29 @@ bool llama_model_has_recurrent(const llama_model * model) {
return llm_arch_is_hybrid(model->arch) || llm_arch_is_recurrent(model->arch);
}
llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int il) {
auto it = LLM_TENSOR_NAMES.find(arch);
if (it == LLM_TENSOR_NAMES.end()) {
printf("%s: Oops, did not find arch\n", __func__);
return LLM_TENSOR_UNKNOWN;
}
if (il < 0) {
for (auto & entry : it->second) {
if (tensor_name.find(entry.second) == 0) {
return entry.first;
}
}
return LLM_TENSOR_UNKNOWN;
}
for (auto & entry : it->second) {
auto this_name = ::format(entry.second.c_str(), il);
if (tensor_name.find(this_name) == 0) {
return entry.first;
}
}
return LLM_TENSOR_UNKNOWN;
}
size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const {
if (il < 0 || il >= hparams.n_layer) return 0;
if (hparams.recurrent_layer_arr[il]) {

47
src/llama.cpp
View File

@ -2013,6 +2013,9 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
if (has_mla) {
mla_tensors.resize(n_layer);
}
std::vector<size_t> ffn_exps(n_layer, 0), ffn_shexp(n_layer, 0), ffn_dense(n_layer, 0);
std::vector<size_t> attn(n_layer, 0);
std::vector<bool> has_layer_norm(n_layer, false);
size_t ow_size = 0;
size_t embd_size = 0;
for (int i = 0; i < ml.n_tensors; ++i) {
@ -2027,6 +2030,9 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
ow_size = size;
continue;
}
if (name == "output_norm.weight") {
continue;
}
auto pos = name.find("blk.");
if (pos != 0) {
printf("Oops: tensor with strange name %s\n", name.c_str());
@ -2049,17 +2055,58 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
continue;
}
result[il] += size;
if (auto pos = name.rfind(".bias"); pos < name.size() && name.size() - pos == 4) {
// bias, we don't need to account for those
continue;
}
bool is_mla = false;
if (has_mla) {
if (name.find("attn_k_b.weight") != std::string::npos) {
mla_tensors[il].wk_b = t;
is_mla = true;
}
else if (name.find("attn_v_b.weight") != std::string::npos) {
mla_tensors[il].wv_b = t;
is_mla = true;
}
else if (name.find("attn_kv_b.weight") != std::string::npos) {
mla_tensors[il].wkv_b = t;
is_mla = true;
}
}
if (!is_mla) {
auto ttype = llm_tensor_type(model.arch, name, il);
if (ttype == LLM_TENSOR_UNKNOWN) printf("Oops: got unknows for tensor %s\n", name.c_str());
if (ttype == LLM_TENSOR_FFN_GATE_UP_EXPS || ttype == LLM_TENSOR_FFN_GATE_EXPS || ttype == LLM_TENSOR_FFN_UP_EXPS || ttype == LLM_TENSOR_FFN_DOWN_EXPS) {
auto size = t->ne[1] * n_ubatch * model.hparams.n_expert_used * sizeof(float);
ffn_exps[il] += size;
}
else if (ttype == LLM_TENSOR_FFN_UP_SHEXP || ttype == LLM_TENSOR_FFN_GATE_SHEXP || ttype == LLM_TENSOR_FFN_DOWN_SHEXP) {
ffn_shexp[il] += t->ne[1] * n_ubatch * sizeof(float);
}
else if (ttype == LLM_TENSOR_FFN_UP || ttype == LLM_TENSOR_FFN_GATE || ttype == LLM_TENSOR_FFN_DOWN) {
ffn_dense[il] += t->ne[1] * n_ubatch * sizeof(float);
}
else if (ttype == LLM_TENSOR_ATTN_Q || ttype == LLM_TENSOR_ATTN_K || ttype == LLM_TENSOR_ATTN_V) {
attn[il] += t->ne[1] * n_ubatch * sizeof(float);
}
else if (ttype == LLM_TENSOR_ATTN_OUT) {
// This assume FA. With FA, we need to have at the same time Q, K, V and FA(Q, K, V).
// FA(Q, K, V) has the dimension t->ne[0] x n_ubatch
attn[il] += t->ne[0] * n_ubatch * sizeof(float);
}
else if (ttype == LLM_TENSOR_ATTN_NORM || ttype == LLM_TENSOR_ATTN_NORM_2 || ttype == LLM_TENSOR_ATTN_OUT_NORM || ttype == LLM_TENSOR_ATTN_POST_NORM ||
ttype == LLM_TENSOR_FFN_NORM || ttype == LLM_TENSOR_FFN_POST_NORM || ttype == LLM_TENSOR_LAYER_OUT_NORM) {
has_layer_norm[il] = true;
}
}
}
for (int il = 0; il < n_layer; ++il) {
auto ffn = std::max(std::max(ffn_exps[il], ffn_shexp[il]), ffn_dense[il]);
auto inp = model.hparams.n_embd * n_ubatch * sizeof(float);
if (has_layer_norm[il]) inp *= 2;
ffn += inp;
compute[il] = std::max<double>(compute[il], ffn);
}
if (has_mla) {
for (int il = 0; il < n_layer; ++il) {