Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3void llama_model_glm4_moe::load_arch_hparams(llama_model_loader & ml) {4 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);5 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);6 ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);7 8 // MoE parameters9 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);10 ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);11 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);12 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);13 14 // Expert gating function (GLM-4.5 uses sigmoid)15 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);16 if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {17 hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;18 }19 20 switch (hparams.n_layer()) {21 case 46: type = LLM_TYPE_106B_A12B; break; // GLM-4.5-Air22 case 48: type = LLM_TYPE_102B_A12B; break; // Solar Open23 case 92: type = LLM_TYPE_355B_A32B; break; // GLM-4.524 default: type = LLM_TYPE_UNKNOWN;25 }26}27 28void llama_model_glm4_moe::load_arch_tensors(llama_model_loader & ml) {29 LLAMA_LOAD_LOCALS;30 const int64_t n_expert_shared = hparams.n_expert_shared;31 32 const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);33 const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";34 const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);35 const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;36 int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;37 38 if (!ml.load_mtp) {39 mtp_flags |= TENSOR_SKIP;40 }41 42 GGML_ASSERT(hparams.n_expert > 0 && "n_expert must be > 0 for GLM4_MOE MoE layers");43 GGML_ASSERT(hparams.n_expert_used() > 0 && "n_expert_used must be > 0 for GLM4_MOE MoE layers");44 45 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);46 47 // output48 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0);49 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);50 // if output is NULL, init from the input tok embed51 if (output == NULL) {52 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);53 }54 55 for (int i = 0; i < n_layer_all; ++i) {56 auto & layer = layers[i];57 const int flags = i < n_layer ? trunk_flags : mtp_flags;58 59 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, flags);60 61 // GLM-style attention with bias terms62 create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, flags);63 64 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, flags);65 66 // K/Q norm tensors (optional for GLM-4.5 355B variant)67 layer.attn_q_norm = create_tensor(68 tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, TENSOR_NOT_REQUIRED | flags);69 layer.attn_k_norm = create_tensor(70 tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, TENSOR_NOT_REQUIRED | flags);71 72 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, flags);73 74 // Check if this layer uses MoE or dense FFN based on n_layer_dense_lead75 // GLM 4.5 uses hybrid architecture: layer 0 is dense, layers 1+ are MoE76 const bool use_moe = (static_cast<uint32_t>(i) >= hparams.n_layer_dense_lead);77 78 if (use_moe) {79 // MoE layers80 layer.ffn_gate_inp =81 create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, flags);82 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), { n_expert }, flags);83 84 // MoE branch85 const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff / n_expert_used;86 87 layer.ffn_gate_exps = create_tensor(88 tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert }, flags);89 layer.ffn_down_exps = create_tensor(90 tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, flags);91 layer.ffn_up_exps = create_tensor(92 tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert }, flags);93 94 // Shared expert95 if (n_expert_shared > 0) {96 const int64_t n_ff_shexp = n_ff_exp * n_expert_shared;97 layer.ffn_gate_shexp = create_tensor(98 tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, flags);99 layer.ffn_down_shexp = create_tensor(100 tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, flags);101 layer.ffn_up_shexp = create_tensor(102 tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, flags);103 }104 } else {105 // Dense layers (first k layers) - GLM uses separate gate/up projections106 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, flags);107 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, flags);108 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, flags);109 }110 111 // NextN/MTP tensors112 if (i >= n_layer) {113 layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);114 layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);115 layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);116 117 // Optional tensors118 layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED | flags);119 layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED | flags);120 layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED | flags);121 }122 }123}124 125std::unique_ptr<llm_graph_context> llama_model_glm4_moe::build_arch_graph(const llm_graph_params & params) const {126 if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {127 return std::make_unique<graph_mtp>(*this, params);128 }129 return std::make_unique<graph>(*this, params);130}131 132llama_model_glm4_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)133 : llm_graph_context(params) {134 GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM4_MOE MTP requires n_layer_nextn > 0");135 GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM4_MOE MTP currently only supports a single MTP block");136 137 const int64_t n_embd_head = hparams.n_embd_head_v();138 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());139 140 const int il = hparams.n_layer() + cparams.nextn_layer_offset;141 GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&142 cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&143 "nextn_layer_offset out of range [0, n_layer_nextn)");144 145 const auto & layer = model.layers[il];146 147 GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");148 GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");149 GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");150 GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");151 152 auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);153 154 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);155 ggml_set_input(inp->tokens);156 157 inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);158 ggml_set_input(inp->embd);159 160 ggml_tensor * tok_embd;161 if (ubatch.token) {162 ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;163 tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);164 } else {165 tok_embd = inp->embd;166 }167 cb(tok_embd, "mtp_tok_embd", il);168 169 inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);170 ggml_set_input(inp->h);171 ggml_set_name(inp->h, "mtp_h_input");172 173 ggml_tensor * h_embd = inp->h;174 175 res->add_input(std::move(inp));176 177 ggml_tensor * inp_pos = build_inp_pos();178 ggml_tensor * inp_out_ids = build_inp_out_ids();179 180 auto * inp_attn = build_attn_inp_kv();181 182 ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);183 cb(h_norm, "mtp_hnorm", il);184 185 ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);186 cb(e_norm, "mtp_enorm", il);187 188 ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);189 cb(concat, "mtp_concat", il);190 191 ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);192 cb(cur, "mtp_eh_proj", il);193 194 ggml_tensor * inpSA = cur;195 196 cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);197 cb(cur, "mtp_attn_norm", il);198 199 auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur,200 n_embd_head, n_head, n_head_kv, il);201 202 if (layer.attn_q_norm) {203 Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);204 cb(Qcur, "mtp_Qcur_normed", il);205 }206 if (layer.attn_k_norm) {207 Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);208 cb(Kcur, "mtp_Kcur_normed", il);209 }210 211 Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot,212 rope_type, n_ctx_orig, freq_base, freq_scale,213 ext_factor, attn_factor, beta_fast, beta_slow);214 215 Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot,216 rope_type, n_ctx_orig, freq_base, freq_scale,217 ext_factor, attn_factor, beta_fast, beta_slow);218 219 cb(Qcur, "mtp_Qcur", il);220 cb(Kcur, "mtp_Kcur", il);221 cb(Vcur, "mtp_Vcur", il);222 223 cur = build_attn(inp_attn,224 layer.wo, nullptr, layer.wo_s,225 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,226 1.0f / sqrtf(float(n_embd_head)), il);227 cb(cur, "mtp_attn_out", il);228 229 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);230 cb(ffn_inp, "mtp_ffn_inp", il);231 232 cur = build_norm(ffn_inp, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);233 cb(cur, "mtp_post_attn_norm", il);234 235 ggml_tensor * routed_out = build_moe_ffn(cur,236 layer.ffn_gate_inp,237 layer.ffn_up_exps,238 layer.ffn_gate_exps,239 layer.ffn_down_exps,240 layer.ffn_exp_probs_b,241 n_expert, n_expert_used,242 LLM_FFN_SILU, hparams.expert_weights_norm,243 hparams.expert_weights_scale,244 (llama_expert_gating_func_type) hparams.expert_gating_func,245 il);246 cb(routed_out, "mtp_ffn_moe_out", il);247 248 ggml_tensor * shared_out = build_ffn(cur,249 layer.ffn_up_shexp, nullptr, nullptr,250 layer.ffn_gate_shexp, nullptr, nullptr,251 layer.ffn_down_shexp, nullptr, nullptr,252 nullptr,253 LLM_FFN_SILU, LLM_FFN_PAR, il);254 cb(shared_out, "mtp_ffn_shexp_out", il);255 256 cur = ggml_add(ctx0, routed_out, shared_out);257 cb(cur, "mtp_ffn_out", il);258 259 cur = ggml_add(ctx0, cur, ffn_inp);260 cb(cur, "mtp_post_ffn", il);261 262 ggml_tensor * head_norm_w = layer.nextn.shared_head_norm263 ? layer.nextn.shared_head_norm264 : model.output_norm;265 GGML_ASSERT(head_norm_w && "GLM4_MOE MTP: missing both nextn.shared_head_norm and output_norm");266 267 cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);268 cb(cur, "h_nextn", -1);269 res->t_h_nextn = cur;270 271 if (inp_out_ids) {272 cur = ggml_get_rows(ctx0, cur, inp_out_ids);273 }274 cb(cur, "mtp_shared_head_norm", -1);275 276 ggml_tensor * head_w = layer.nextn.shared_head_head277 ? layer.nextn.shared_head_head278 : model.output;279 ggml_tensor * head_s = layer.nextn.shared_head_head280 ? layer.nextn.shared_head_head_s281 : model.output_s;282 GGML_ASSERT(head_w && "GLM4_MOE MTP: missing LM head (nextn.shared_head_head or model.output)");283 284 cur = build_lora_mm(head_w, cur, head_s);285 cb(cur, "result_output", -1);286 287 res->t_logits = cur;288 ggml_build_forward_expand(gf, cur);289}290 291llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {292 const int64_t n_embd_head = hparams.n_embd_head_v();293 294 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());295 296 int sections[4];297 std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);298 299 ggml_tensor * cur;300 ggml_tensor * inpL;301 302 inpL = build_inp_embd(model.tok_embd);303 304 bool use_mrope = hparams.use_mrope();305 if (ubatch.embd && !use_mrope) {306 // unfortunately, we need to forcefully stop here, to avoid users complaining about wrong results307 GGML_ABORT("This GGUF does not support multimodal. Please reconvert it.");308 }309 310 // inp_pos - contains the positions311 ggml_tensor * inp_pos = build_inp_pos();312 313 auto * inp_attn = build_attn_inp_kv();314 315 ggml_tensor * inp_out_ids = build_inp_out_ids();316 317 // NextN layers are processed by graph_mtp.318 for (int il = 0; il < n_layer; ++il) {319 ggml_tensor * inpSA = inpL;320 321 // Pre-attention norm322 cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);323 cb(cur, "attn_norm", il);324 325 // self-attention326 {327 auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,328 n_embd_head, n_head, n_head_kv, il);329 330 // Apply Q/K norm if available (GLM-4.5 355B variant)331 if (model.layers[il].attn_q_norm) {332 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);333 cb(Qcur, "Qcur_normed", il);334 }335 if (model.layers[il].attn_k_norm) {336 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);337 cb(Kcur, "Kcur_normed", il);338 }339 340 if (use_mrope) {341 Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr,342 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,343 ext_factor, attn_factor, beta_fast, beta_slow);344 345 Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr,346 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,347 ext_factor, attn_factor, beta_fast, beta_slow);348 } else {349 // Normal RoPE350 Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot,351 rope_type, n_ctx_orig, freq_base, freq_scale,352 ext_factor, attn_factor, beta_fast, beta_slow);353 354 Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot,355 rope_type, n_ctx_orig, freq_base, freq_scale,356 ext_factor, attn_factor, beta_fast, beta_slow);357 }358 359 cb(Qcur, "Qcur", il);360 cb(Kcur, "Kcur", il);361 cb(Vcur, "Vcur", il);362 363 cur = build_attn(inp_attn,364 model.layers[il].wo, NULL, model.layers[il].wo_s,365 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);366 }367 if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {368 cur = ggml_get_rows(ctx0, cur, inp_out_ids);369 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);370 }371 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);372 cb(ffn_inp, "ffn_inp", il);373 374 // Post-attention norm375 cur = build_norm(ffn_inp, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, il);376 cb(cur, "post_attn_norm", il);377 378 // Check if this is a dense layer (n_layer_dense_lead=1, so layer 0 is dense)379 if (static_cast<uint32_t>(il) < hparams.n_layer_dense_lead) {380 // Dense FFN layer381 cur = build_ffn(cur,382 model.layers[il].ffn_up, NULL, NULL,383 model.layers[il].ffn_gate, NULL, NULL,384 model.layers[il].ffn_down, NULL, NULL,385 NULL,386 LLM_FFN_SILU, LLM_FFN_PAR, il);387 cb(cur, "ffn_out", il);388 } else {389 // Process routed experts using existing MoE infrastructure390 ggml_tensor * routed_out = build_moe_ffn(cur,391 model.layers[il].ffn_gate_inp,392 model.layers[il].ffn_up_exps,393 model.layers[il].ffn_gate_exps,394 model.layers[il].ffn_down_exps,395 model.layers[il].ffn_exp_probs_b,396 n_expert, n_expert_used,397 LLM_FFN_SILU, hparams.expert_weights_norm,398 hparams.expert_weights_scale,399 (llama_expert_gating_func_type) hparams.expert_gating_func,400 il);401 cb(routed_out, "ffn_moe_out", il);402 403 // Process shared expert on original input404 ggml_tensor * shared_out = build_ffn(cur,405 model.layers[il].ffn_up_shexp, NULL, NULL,406 model.layers[il].ffn_gate_shexp, NULL, NULL,407 model.layers[il].ffn_down_shexp, NULL, NULL,408 NULL,409 LLM_FFN_SILU, LLM_FFN_PAR, il);410 cb(shared_out, "ffn_shexp_out", il);411 412 // Final output: routed_output + shared_output413 cur = ggml_add(ctx0, routed_out, shared_out);414 cb(cur, "ffn_out", il);415 }416 cur = ggml_add(ctx0, cur, ffn_inp);417 418 cur = build_cvec(cur, il);419 cb(cur, "l_out", il);420 421 // input for next layer422 inpL = cur;423 }424 cur = inpL;425 cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);426 427 cb(cur, "h_nextn", -1);428 res->t_h_nextn = cur;429 430 if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {431 cur = ggml_get_rows(ctx0, cur, inp_out_ids);432 }433 434 cb(cur, "result_norm", -1);435 res->t_embd = cur;436 437 // lm_head438 cur = build_lora_mm(model.output, cur, model.output_s);439 440 cb(cur, "result_output", -1);441 res->t_logits = cur;442 443 ggml_build_forward_expand(gf, cur);444}445