mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-25 23:47:23 -05:00
Pascal
542348a35c
* llama: read the SWA pattern as a period or a per-layer array Add llama_model_base::load_swa_pattern(), which reads sliding_window_pattern either as one flag per layer or as a period expanded by set_swa_pattern(), and use it in every loader that reads the key as a period. These loaders silently ignored an array and applied their default period, although the converters of olmo2, gemma3n and exaone4 write arrays. The published GGUFs match the defaults, so their outputs do not change. The loaders that already accepted both forms lose their duplicated scalar-then-array block, and use their declared default period when the key is absent. * model-saver: write the SWA pattern and the MLA SWA geometry Write sliding_window_pattern as one flag per layer, nextn layers included, for every model using SWA. The array is never collapsed to a scalar, since the loaders read a scalar as a period. Also write the MLA key/value lengths and KV LoRA rank of the SWA layers, required by dots3note. This enables the saver for plamo3, gemma3, cohere2, cohere2moe, olmo2, exaone-moe, afmoe, mimo2, spark2_5, muse-glimmer, mellum, laguna, granite_swa, dots3note and maple, all passing the bit-exact roundtrip of test-llama-archs.
284 lines
12 KiB
C++
284 lines
12 KiB
C++
#include "models.h"
|
|
|
|
void llama_model_afmoe::load_arch_hparams(llama_model_loader & ml) {
|
|
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
|
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
|
|
ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
|
|
ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
|
|
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);
|
|
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
|
|
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
|
|
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
|
|
|
|
// Set up interleaved sliding window attention (ISWA)
|
|
// Pattern: 3 sliding - 1 full (global_attn_every_n_layers = 4)
|
|
if (hparams.n_swa > 0) {
|
|
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
|
|
load_swa_pattern(ml, 4);
|
|
|
|
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
|
|
hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
|
|
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
|
|
} else {
|
|
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
|
|
}
|
|
|
|
// Default to sigmoid if not set
|
|
if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
|
|
hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
|
|
}
|
|
|
|
switch (hparams.n_layer()) {
|
|
case 56: type = LLM_TYPE_6B; break;
|
|
case 32: type = LLM_TYPE_26B; break;
|
|
default: type = LLM_TYPE_UNKNOWN;
|
|
}
|
|
}
|
|
|
|
void llama_model_afmoe::load_arch_tensors(llama_model_loader &) {
|
|
LLAMA_LOAD_LOCALS;
|
|
const int64_t n_expert_shared = hparams.n_expert_shared;
|
|
|
|
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
|
|
|
|
// output
|
|
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
|
|
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
|
|
|
|
// if output is NULL, init from the input tok embed
|
|
if (output == NULL) {
|
|
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
|
|
}
|
|
|
|
const int64_t n_ff_exp = hparams.n_ff_exp();
|
|
|
|
for (int i = 0; i < n_layer; ++i) {
|
|
auto & layer = layers[i];
|
|
|
|
// dual attention normalization
|
|
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
|
|
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);
|
|
|
|
// attention projections
|
|
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
|
|
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
|
|
|
|
// Q/K normalization
|
|
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
|
|
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
|
|
|
|
// attention gating
|
|
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
|
|
|
|
// dual ffn normalization
|
|
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
|
|
layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
|
|
|
|
if (static_cast<uint32_t>(i) >= hparams.n_layer_dense_lead) {
|
|
// MoE layers
|
|
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
|
|
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
|
|
|
|
// grouped expert weights
|
|
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
|
|
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
|
|
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
|
|
|
|
// shared expert
|
|
if (n_expert_shared > 0) {
|
|
const int64_t n_ff_shexp = n_ff_exp * n_expert_shared;
|
|
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
|
|
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);
|
|
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
|
|
}
|
|
} else {
|
|
// Dense layers
|
|
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
|
|
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
|
|
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
|
|
}
|
|
}
|
|
}
|
|
|
|
std::unique_ptr<llm_graph_context> llama_model_afmoe::build_arch_graph(const llm_graph_params & params) const {
|
|
return std::make_unique<graph>(*this, params);
|
|
}
|
|
|
|
llama_model_afmoe::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
|
|
const int64_t n_embd_head = hparams.n_embd_head_v();
|
|
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
|
|
|
|
ggml_tensor * cur;
|
|
ggml_tensor * inpL;
|
|
|
|
inpL = build_inp_embd(model.tok_embd);
|
|
|
|
// MuP scaling: embeddings * sqrt(hidden_size)
|
|
// mup_enabled = true, hidden_size = 1024, scale = 32.0
|
|
inpL = ggml_scale(ctx0, inpL, sqrtf(float(n_embd)));
|
|
cb(inpL, "inp_embd_scaled", -1);
|
|
|
|
// inp_pos - contains the positions
|
|
ggml_tensor * inp_pos = build_inp_pos();
|
|
auto * inp_attn = build_attn_inp_kv_iswa();
|
|
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
|
|
|
const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
|
|
|
|
for (int il = 0; il < n_layer; ++il) {
|
|
const float freq_base_l = model.get_rope_freq_base (cparams, il);
|
|
const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
|
|
|
|
ggml_tensor * inpSA = inpL;
|
|
|
|
// This overlaps with SWA layers in current models, so get_rope_freq_base/scale may be superfluous
|
|
const bool use_rope = hparams.n_no_rope_layer_step > 0 &&
|
|
(il + 1) % hparams.n_no_rope_layer_step != 0;
|
|
|
|
// dual attention normalization (pre)
|
|
cur = build_norm(inpL,
|
|
model.layers[il].attn_norm, NULL,
|
|
LLM_NORM_RMS, il);
|
|
cb(cur, "attn_norm", il);
|
|
|
|
// self-attention
|
|
{
|
|
ggml_tensor * attn_inp = cur; // save input for gate computation
|
|
|
|
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
|
|
n_embd_head, n_head, n_head_kv, il);
|
|
|
|
// compute gate from input
|
|
ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);
|
|
cb(gate, "attn_gate_proj", il);
|
|
|
|
// Q/K normalization
|
|
Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
|
|
Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
|
|
cb(Qcur, "Qcur_normed", il);
|
|
cb(Kcur, "Kcur_normed", il);
|
|
|
|
if (use_rope) {
|
|
Qcur = ggml_rope_ext(
|
|
ctx0, Qcur, inp_pos, nullptr,
|
|
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
|
|
ext_factor, attn_factor, beta_fast, beta_slow);
|
|
cb(Qcur, "Qcur_rope", il);
|
|
|
|
Kcur = ggml_rope_ext(
|
|
ctx0, Kcur, inp_pos, nullptr,
|
|
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
|
|
ext_factor, attn_factor, beta_fast, beta_slow);
|
|
cb(Kcur, "Kcur_rope", il);
|
|
}
|
|
|
|
cur = build_attn(inp_attn,
|
|
NULL, NULL, NULL, // wo will be applied after gating
|
|
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
|
|
cb(cur, "attn_out", il);
|
|
|
|
// attention gating: attn_out * sigmoid(gate) BEFORE o_proj
|
|
gate = ggml_sigmoid(ctx0, gate);
|
|
cb(gate, "attn_gate_sig", il);
|
|
cur = ggml_mul(ctx0, cur, gate);
|
|
cb(cur, "attn_gated", il);
|
|
|
|
// now apply output projection
|
|
cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
|
|
cb(cur, "attn_o_proj", il);
|
|
}
|
|
|
|
// dual attention normalization (post)
|
|
cur = build_norm(cur,
|
|
model.layers[il].attn_post_norm, NULL,
|
|
LLM_NORM_RMS, il);
|
|
cb(cur, "attn_post_norm", il);
|
|
|
|
if (il == n_layer - 1 && inp_out_ids) {
|
|
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
|
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
|
}
|
|
|
|
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
|
|
cb(ffn_inp, "ffn_inp", il);
|
|
|
|
// dual ffn normalization (pre)
|
|
cur = build_norm(ffn_inp,
|
|
model.layers[il].ffn_norm, NULL,
|
|
LLM_NORM_RMS, il);
|
|
cb(cur, "ffn_norm", il);
|
|
|
|
// MoE or dense FFN
|
|
if ((uint32_t)il >= hparams.n_layer_dense_lead) {
|
|
// MoE layer with sigmoid routing, normalization, and scaling
|
|
ggml_tensor * moe_out = build_moe_ffn(cur,
|
|
model.layers[il].ffn_gate_inp,
|
|
model.layers[il].ffn_up_exps,
|
|
model.layers[il].ffn_gate_exps,
|
|
model.layers[il].ffn_down_exps,
|
|
model.layers[il].ffn_exp_probs_b,
|
|
n_expert, n_expert_used,
|
|
LLM_FFN_SILU,
|
|
hparams.expert_weights_norm, // norm_w (route_norm=True)
|
|
hparams.expert_weights_scale, // w_scale (route_scale=2.826)
|
|
(llama_expert_gating_func_type) hparams.expert_gating_func,
|
|
il);
|
|
cb(moe_out, "ffn_moe_out", il);
|
|
|
|
// shared expert
|
|
if (hparams.n_expert_shared > 0) {
|
|
ggml_tensor * ffn_shexp = build_ffn(cur,
|
|
model.layers[il].ffn_up_shexp, NULL, NULL,
|
|
model.layers[il].ffn_gate_shexp, NULL, NULL,
|
|
model.layers[il].ffn_down_shexp, NULL, NULL,
|
|
NULL,
|
|
LLM_FFN_SILU, LLM_FFN_PAR, il);
|
|
cb(ffn_shexp, "ffn_shexp", il);
|
|
|
|
cur = ggml_add(ctx0, moe_out, ffn_shexp);
|
|
cb(cur, "ffn_out", il);
|
|
} else {
|
|
cur = moe_out;
|
|
}
|
|
} else {
|
|
// dense layer
|
|
cur = build_ffn(cur,
|
|
model.layers[il].ffn_up, NULL, NULL,
|
|
model.layers[il].ffn_gate, NULL, NULL,
|
|
model.layers[il].ffn_down, NULL, NULL,
|
|
NULL,
|
|
LLM_FFN_SILU, LLM_FFN_PAR, il);
|
|
cb(cur, "ffn_out", il);
|
|
}
|
|
|
|
// dual ffn normalization (post)
|
|
cur = build_norm(cur,
|
|
model.layers[il].ffn_post_norm, NULL,
|
|
LLM_NORM_RMS, il);
|
|
cb(cur, "ffn_post_norm", il);
|
|
|
|
cur = ggml_add(ctx0, cur, ffn_inp);
|
|
cur = build_cvec(cur, il);
|
|
cb(cur, "l_out", il);
|
|
|
|
// input for next layer
|
|
inpL = cur;
|
|
}
|
|
|
|
cur = inpL;
|
|
|
|
cur = build_norm(cur,
|
|
model.output_norm, NULL,
|
|
LLM_NORM_RMS, -1);
|
|
cb(cur, "result_norm", -1);
|
|
|
|
res->t_embd = cur;
|
|
|
|
// lm_head
|
|
cur = build_lora_mm(model.output, cur, model.output_s);
|
|
cb(cur, "result_output", -1);
|
|
res->t_logits = cur;
|
|
|
|
ggml_build_forward_expand(gf, cur);
|
|
}
|