Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3void llama_model_plm::load_arch_hparams(llama_model_loader & ml) {4 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);5 ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);6 7 switch (hparams.n_layer()) {8 case 32: type = LLM_TYPE_1_8B; break;9 default: type = LLM_TYPE_UNKNOWN;10 }11}12 13void llama_model_plm::load_arch_tensors(llama_model_loader &) {14 LLAMA_LOAD_LOCALS;15 16 const int64_t n_embd_head_qk_rope = hparams.n_rot();17 const int64_t n_embd_head_qk_nope = hparams.n_embd_head_k() - hparams.n_rot();18 const int64_t kv_lora_rank = hparams.n_lora_kv;19 20 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);21 22 // output23 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);24 // output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);25 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);26 27 for (int i = 0; i < n_layer; ++i) {28 auto & layer = layers[i];29 30 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);31 32 layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);33 layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + (n_embd_head_qk_rope)}, 0);34 layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0);35 layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v)}, 0);36 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_head * ( n_embd_head_v), n_embd}, 0);37 38 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);39 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);40 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);41 }42}43 44std::unique_ptr<llm_graph_context> llama_model_plm::build_arch_graph(const llm_graph_params & params) const {45 return std::make_unique<graph>(*this, params);46}47 48llama_model_plm::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {49 const float kq_scale = 1.0f/sqrtf(float(hparams.n_embd_head_k()));50 51 const uint32_t n_embd_head_qk_rope = hparams.n_rot();52 const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k() - hparams.n_rot();53 54 const uint32_t kv_lora_rank = hparams.n_lora_kv;55 56 ggml_tensor * cur;57 ggml_tensor * inpL;58 59 // {n_embd, n_tokens}60 inpL = build_inp_embd(model.tok_embd);61 62 // inp_pos - contains the positions63 ggml_tensor * inp_pos = build_inp_pos();64 65 auto * inp_attn = build_attn_inp_kv();66 67 ggml_tensor * inp_out_ids = build_inp_out_ids();68 69 for (int il = 0; il < n_layer; ++il) {70 ggml_tensor * inpSA = inpL;71 72 // norm73 cur = build_norm(inpL,74 model.layers[il].attn_norm, NULL,75 LLM_NORM_RMS, il);76 cb(cur, "attn_norm", il);77 78 // self_attention79 {80 ggml_tensor * q = NULL;81 q = ggml_mul_mat(ctx0, model.layers[il].wq, cur);82 cb(q, "q", il);83 84 // {n_embd_head_k, n_head, n_tokens}, RoPE is applied to the trailing dims only85 q = ggml_reshape_3d(ctx0, q, hparams.n_embd_head_k(), n_head, n_tokens);86 cb(q, "q", il);87 88 // {n_embd, kv_lora_rank + n_embd_head_qk_rope} * {n_embd, n_tokens} -> {kv_lora_rank + n_embd_head_qk_rope, n_tokens}89 ggml_tensor * kv_pe_compresseed = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);90 cb(kv_pe_compresseed, "kv_pe_compresseed", il);91 92 // split into {kv_lora_rank, n_tokens}93 ggml_tensor * kv_compressed = ggml_view_2d(ctx0, kv_pe_compresseed, kv_lora_rank, n_tokens,94 kv_pe_compresseed->nb[1],95 0);96 cb(kv_compressed, "kv_compressed", il);97 98 // and {n_embd_head_qk_rope, n_tokens}99 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_pe_compresseed, n_embd_head_qk_rope, 1, n_tokens,100 kv_pe_compresseed->nb[1],101 kv_pe_compresseed->nb[1],102 ggml_row_size(kv_pe_compresseed->type, kv_lora_rank));103 cb(k_pe, "k_pe", il);104 105 kv_compressed = build_norm(kv_compressed,106 model.layers[il].attn_kv_a_norm, NULL,107 LLM_NORM_RMS, il);108 cb(kv_compressed, "kv_compressed", il);109 110 // {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v)} * {kv_lora_rank, n_tokens} -> {n_head * (n_embd_head_qk_nope + n_embd_head_v), n_tokens}111 ggml_tensor * kv = ggml_mul_mat(ctx0, model.layers[il].wkv_b, kv_compressed);112 cb(kv, "kv", il);113 114 // split into {n_head * n_embd_head_qk_nope, n_tokens}115 ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,116 ggml_row_size(kv->type, n_embd_head_qk_nope + hparams.n_embd_head_v()),117 ggml_row_size(kv->type, n_head * (n_embd_head_qk_nope + hparams.n_embd_head_v())),118 0);119 cb(k_nope, "k_nope", il);120 121 // and {n_head * n_embd_head_v, n_tokens}122 ggml_tensor * v_states = ggml_view_3d(ctx0, kv, hparams.n_embd_head_v(), n_head, n_tokens,123 ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v())),124 ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v())*n_head),125 ggml_row_size(kv->type, (n_embd_head_qk_nope)));126 cb(v_states, "v_states", il);127 128 v_states = ggml_cont(ctx0, v_states);129 cb(v_states, "v_states", il);130 131 v_states = ggml_view_2d(ctx0, v_states, hparams.n_embd_head_v() * n_head, n_tokens,132 ggml_row_size(kv->type, hparams.n_embd_head_v() * n_head),133 0);134 cb(v_states, "v_states", il);135 136 q = ggml_rope_ext(137 ctx0, q, inp_pos, nullptr,138 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,139 ext_factor, attn_factor, beta_fast, beta_slow140 );141 q = ggml_rope_set_offset(q, n_embd_head_qk_nope);142 cb(q, "q_rope", il);143 144 // shared RoPE key145 k_pe = ggml_rope_ext(146 ctx0, k_pe, inp_pos, nullptr,147 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,148 ext_factor, attn_factor, beta_fast, beta_slow149 );150 cb(k_pe, "k_pe", il);151 152 ggml_tensor * q_states = q;153 cb(q_states, "q_states", il);154 155 ggml_tensor * k_states = ggml_concat(ctx0, k_nope,156 ggml_repeat_4d(ctx0, k_pe, n_embd_head_qk_rope, n_head, n_tokens, 1), 0);157 cb(k_states, "k_states", il);158 159 cur = build_attn(inp_attn,160 model.layers[il].wo, NULL, model.layers[il].wo_s,161 q_states, k_states, v_states, nullptr, nullptr, nullptr, kq_scale, il);162 }163 if (il == n_layer - 1 && inp_out_ids) {164 cur = ggml_get_rows(ctx0, cur, inp_out_ids);165 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);166 }167 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);168 cb(ffn_inp, "ffn_inp", il);169 170 cur = build_norm(ffn_inp,171 model.layers[il].ffn_norm, NULL,172 LLM_NORM_RMS, il);173 cb(cur, "ffn_norm", il);174 175 cur = build_ffn(cur,176 model.layers[il].ffn_up, NULL, NULL,177 NULL, NULL, NULL,178 model.layers[il].ffn_down, NULL, NULL,179 NULL,180 LLM_FFN_RELU_SQR, LLM_FFN_SEQ, il);181 cb(cur, "ffn_out", il);182 183 cur = ggml_add(ctx0, cur, ffn_inp);184 185 cur = build_cvec(cur, il);186 cb(cur, "l_out", il);187 188 // input for next layer189 inpL = cur;190 }191 cur = inpL;192 193 cur = build_norm(cur,194 model.output_norm, NULL,195 LLM_NORM_RMS, -1);196 197 cb(cur, "result_norm", -1);198 res->t_embd = cur;199 200 cur = build_lora_mm(model.output, cur, model.output_s);201 202 cb(cur, "result_output", -1);203 res->t_logits = cur;204 205 ggml_build_forward_expand(gf, cur);206}207