Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3void llama_model_afmoe::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_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);6 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);7 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);8 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);9 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);10 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);11 ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);12 13 // Set up interleaved sliding window attention (ISWA)14 // Pattern: 3 sliding - 1 full (global_attn_every_n_layers = 4)15 if (hparams.n_swa > 0) {16 hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;17 load_swa_pattern(ml, 4);18 19 hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;20 hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;21 ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);22 } else {23 hparams.swa_type = LLAMA_SWA_TYPE_NONE;24 }25 26 // Default to sigmoid if not set27 if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {28 hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;29 }30 31 switch (hparams.n_layer()) {32 case 56: type = LLM_TYPE_6B; break;33 case 32: type = LLM_TYPE_26B; break;34 default: type = LLM_TYPE_UNKNOWN;35 }36}37 38void llama_model_afmoe::load_arch_tensors(llama_model_loader &) {39 LLAMA_LOAD_LOCALS;40 const int64_t n_expert_shared = hparams.n_expert_shared;41 42 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);43 44 // output45 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);46 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);47 48 // if output is NULL, init from the input tok embed49 if (output == NULL) {50 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);51 }52 53 const int64_t n_ff_exp = hparams.n_ff_exp();54 55 for (int i = 0; i < n_layer; ++i) {56 auto & layer = layers[i];57 58 // dual attention normalization59 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);60 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);61 62 // attention projections63 create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);64 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);65 66 // Q/K normalization67 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);68 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);69 70 // attention gating71 layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);72 73 // dual ffn normalization74 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);75 layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);76 77 if (static_cast<uint32_t>(i) >= hparams.n_layer_dense_lead) {78 // MoE layers79 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);80 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);81 82 // grouped expert weights83 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);84 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);85 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);86 87 // shared expert88 if (n_expert_shared > 0) {89 const int64_t n_ff_shexp = n_ff_exp * n_expert_shared;90 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);91 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);92 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);93 }94 } else {95 // Dense layers96 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);97 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);98 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);99 }100 }101}102 103std::unique_ptr<llm_graph_context> llama_model_afmoe::build_arch_graph(const llm_graph_params & params) const {104 return std::make_unique<graph>(*this, params);105}106 107llama_model_afmoe::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {108 const int64_t n_embd_head = hparams.n_embd_head_v();109 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());110 111 ggml_tensor * cur;112 ggml_tensor * inpL;113 114 inpL = build_inp_embd(model.tok_embd);115 116 // MuP scaling: embeddings * sqrt(hidden_size)117 // mup_enabled = true, hidden_size = 1024, scale = 32.0118 inpL = ggml_scale(ctx0, inpL, sqrtf(float(n_embd)));119 cb(inpL, "inp_embd_scaled", -1);120 121 // inp_pos - contains the positions122 ggml_tensor * inp_pos = build_inp_pos();123 auto * inp_attn = build_attn_inp_kv_iswa();124 ggml_tensor * inp_out_ids = build_inp_out_ids();125 126 const float kq_scale = 1.0f/sqrtf(float(n_embd_head));127 128 for (int il = 0; il < n_layer; ++il) {129 const float freq_base_l = model.get_rope_freq_base (cparams, il);130 const float freq_scale_l = model.get_rope_freq_scale(cparams, il);131 132 ggml_tensor * inpSA = inpL;133 134 // This overlaps with SWA layers in current models, so get_rope_freq_base/scale may be superfluous135 const bool use_rope = hparams.n_no_rope_layer_step > 0 &&136 (il + 1) % hparams.n_no_rope_layer_step != 0;137 138 // dual attention normalization (pre)139 cur = build_norm(inpL,140 model.layers[il].attn_norm, NULL,141 LLM_NORM_RMS, il);142 cb(cur, "attn_norm", il);143 144 // self-attention145 {146 ggml_tensor * attn_inp = cur; // save input for gate computation147 148 auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,149 n_embd_head, n_head, n_head_kv, il);150 151 // compute gate from input152 ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);153 cb(gate, "attn_gate_proj", il);154 155 // Q/K normalization156 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);157 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);158 cb(Qcur, "Qcur_normed", il);159 cb(Kcur, "Kcur_normed", il);160 161 if (use_rope) {162 Qcur = ggml_rope_ext(163 ctx0, Qcur, inp_pos, nullptr,164 n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,165 ext_factor, attn_factor, beta_fast, beta_slow);166 cb(Qcur, "Qcur_rope", il);167 168 Kcur = ggml_rope_ext(169 ctx0, Kcur, inp_pos, nullptr,170 n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,171 ext_factor, attn_factor, beta_fast, beta_slow);172 cb(Kcur, "Kcur_rope", il);173 }174 175 cur = build_attn(inp_attn,176 NULL, NULL, NULL, // wo will be applied after gating177 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);178 cb(cur, "attn_out", il);179 180 // attention gating: attn_out * sigmoid(gate) BEFORE o_proj181 gate = ggml_sigmoid(ctx0, gate);182 cb(gate, "attn_gate_sig", il);183 cur = ggml_mul(ctx0, cur, gate);184 cb(cur, "attn_gated", il);185 186 // now apply output projection187 cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);188 cb(cur, "attn_o_proj", il);189 }190 191 // dual attention normalization (post)192 cur = build_norm(cur,193 model.layers[il].attn_post_norm, NULL,194 LLM_NORM_RMS, il);195 cb(cur, "attn_post_norm", il);196 197 if (il == n_layer - 1 && inp_out_ids) {198 cur = ggml_get_rows(ctx0, cur, inp_out_ids);199 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);200 }201 202 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);203 cb(ffn_inp, "ffn_inp", il);204 205 // dual ffn normalization (pre)206 cur = build_norm(ffn_inp,207 model.layers[il].ffn_norm, NULL,208 LLM_NORM_RMS, il);209 cb(cur, "ffn_norm", il);210 211 // MoE or dense FFN212 if ((uint32_t)il >= hparams.n_layer_dense_lead) {213 // MoE layer with sigmoid routing, normalization, and scaling214 ggml_tensor * moe_out = build_moe_ffn(cur,215 model.layers[il].ffn_gate_inp,216 model.layers[il].ffn_up_exps,217 model.layers[il].ffn_gate_exps,218 model.layers[il].ffn_down_exps,219 model.layers[il].ffn_exp_probs_b,220 n_expert, n_expert_used,221 LLM_FFN_SILU,222 hparams.expert_weights_norm, // norm_w (route_norm=True)223 hparams.expert_weights_scale, // w_scale (route_scale=2.826)224 (llama_expert_gating_func_type) hparams.expert_gating_func,225 il);226 cb(moe_out, "ffn_moe_out", il);227 228 // shared expert229 if (hparams.n_expert_shared > 0) {230 ggml_tensor * ffn_shexp = build_ffn(cur,231 model.layers[il].ffn_up_shexp, NULL, NULL,232 model.layers[il].ffn_gate_shexp, NULL, NULL,233 model.layers[il].ffn_down_shexp, NULL, NULL,234 NULL,235 LLM_FFN_SILU, LLM_FFN_PAR, il);236 cb(ffn_shexp, "ffn_shexp", il);237 238 cur = ggml_add(ctx0, moe_out, ffn_shexp);239 cb(cur, "ffn_out", il);240 } else {241 cur = moe_out;242 }243 } else {244 // dense layer245 cur = build_ffn(cur,246 model.layers[il].ffn_up, NULL, NULL,247 model.layers[il].ffn_gate, NULL, NULL,248 model.layers[il].ffn_down, NULL, NULL,249 NULL,250 LLM_FFN_SILU, LLM_FFN_PAR, il);251 cb(cur, "ffn_out", il);252 }253 254 // dual ffn normalization (post)255 cur = build_norm(cur,256 model.layers[il].ffn_post_norm, NULL,257 LLM_NORM_RMS, il);258 cb(cur, "ffn_post_norm", il);259 260 cur = ggml_add(ctx0, cur, ffn_inp);261 cur = build_cvec(cur, il);262 cb(cur, "l_out", il);263 264 // input for next layer265 inpL = cur;266 }267 268 cur = inpL;269 270 cur = build_norm(cur,271 model.output_norm, NULL,272 LLM_NORM_RMS, -1);273 cb(cur, "result_norm", -1);274 275 res->t_embd = cur;276 277 // lm_head278 cur = build_lora_mm(model.output, cur, model.output_s);279 cb(cur, "result_output", -1);280 res->t_logits = cur;281 282 ggml_build_forward_expand(gf, cur);283}284