mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-10-03 11:35:46 +00:00
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.
272 lines
11 KiB
C++
272 lines
11 KiB
C++
#include "models.h"
|
|
|
|
void llama_model_llama4::load_arch_hparams(llama_model_loader & ml) {
|
|
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
|
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_INTERLEAVE_MOE_LAYER_STEP, hparams.n_moe_layer_step);
|
|
|
|
const bool found_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
|
|
if (found_swa && hparams.n_swa == 0) {
|
|
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
|
|
hparams.n_no_rope_layer_step = hparams.n_layer(); // always use rope
|
|
} else {
|
|
hparams.swa_type = LLAMA_SWA_TYPE_CHUNKED;
|
|
hparams.n_swa = 8192;
|
|
hparams.n_attn_temp_floor_scale = 8192;
|
|
hparams.f_attn_temp_scale = 0.1f;
|
|
hparams.f_attn_temp_offset = 1.0f;
|
|
|
|
load_swa_pattern(ml, 4); // pattern: 3 chunked - 1 full
|
|
|
|
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);
|
|
}
|
|
|
|
switch (hparams.n_expert) {
|
|
case 0: {
|
|
// MobileLLM (no MoE)
|
|
switch (hparams.n_embd) {
|
|
case 2048: type = LLM_TYPE_140M; break;
|
|
case 4096: type = LLM_TYPE_360M; break;
|
|
case 6144: type = LLM_TYPE_950M; break;
|
|
default: type = LLM_TYPE_UNKNOWN;
|
|
}
|
|
} break;
|
|
case 16: type = LLM_TYPE_17B_16E; break;
|
|
case 128: type = LLM_TYPE_17B_128E; break;
|
|
default: type = LLM_TYPE_UNKNOWN;
|
|
}
|
|
|
|
hparams.use_kq_norm = type != LLM_TYPE_17B_128E;
|
|
}
|
|
|
|
void llama_model_llama4::load_arch_tensors(llama_model_loader &) {
|
|
LLAMA_LOAD_LOCALS;
|
|
|
|
if (n_expert == 0) {
|
|
throw std::runtime_error(arch_name() + " model cannot have zero experts");
|
|
}
|
|
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);
|
|
}
|
|
|
|
for (int i = 0; i < n_layer; ++i) {
|
|
const bool is_moe_layer = hparams.n_moe_layer_step > 0 && (i + 1) % hparams.n_moe_layer_step == 0;
|
|
|
|
auto & layer = layers[i];
|
|
|
|
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
|
|
|
|
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);
|
|
|
|
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
|
|
|
|
layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
|
|
|
|
if (is_moe_layer) {
|
|
const int64_t n_ff_exp = hparams.n_ff_exp();
|
|
|
|
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
|
|
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
|
|
const int64_t n_ff_shexp = n_ff_exp;
|
|
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 {
|
|
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_llama4::build_arch_graph(const llm_graph_params & params) const {
|
|
if (hparams.swa_type == LLAMA_SWA_TYPE_NONE) {
|
|
return std::make_unique<graph<false>>(*this, params);
|
|
} else {
|
|
return std::make_unique<graph<true>>(*this, params);
|
|
}
|
|
}
|
|
|
|
template <bool iswa>
|
|
llama_model_llama4::graph<iswa>::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_ASSERT(n_embd_head == n_rot);
|
|
|
|
ggml_tensor * cur;
|
|
ggml_tensor * inpL;
|
|
|
|
inpL = build_inp_embd(model.tok_embd);
|
|
|
|
// inp_pos - contains the positions
|
|
ggml_tensor * inp_pos = build_inp_pos();
|
|
|
|
// temperature tuning
|
|
ggml_tensor * inp_attn_scale = nullptr;
|
|
inp_attn_scale = build_inp_attn_scale();
|
|
|
|
using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;
|
|
inp_attn_type * inp_attn = nullptr;
|
|
|
|
if constexpr (iswa) {
|
|
inp_attn = build_attn_inp_kv_iswa();
|
|
} else {
|
|
inp_attn = build_attn_inp_kv();
|
|
}
|
|
|
|
const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
|
|
|
|
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
|
|
|
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;
|
|
|
|
// norm
|
|
cur = build_norm(inpL,
|
|
model.layers[il].attn_norm, NULL,
|
|
LLM_NORM_RMS, il);
|
|
cb(cur, "attn_norm", il);
|
|
|
|
// self-attention
|
|
{
|
|
// rope freq factors for llama3; may return nullptr for llama2 and other models
|
|
ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
|
|
|
|
// compute Q and K and RoPE them
|
|
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
|
|
n_embd_head, n_head, n_head_kv, il);
|
|
|
|
if (use_rope) {
|
|
Qcur = ggml_rope_ext(
|
|
ctx0, Qcur, inp_pos, rope_factors,
|
|
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
|
|
ext_factor, attn_factor, beta_fast, beta_slow
|
|
);
|
|
|
|
Kcur = ggml_rope_ext(
|
|
ctx0, Kcur, inp_pos, rope_factors,
|
|
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
|
|
ext_factor, attn_factor, beta_fast, beta_slow
|
|
);
|
|
} else if (inp_attn_scale) {
|
|
Qcur = ggml_mul(ctx0, Qcur, inp_attn_scale);
|
|
}
|
|
cb(Qcur, "Qcur", il);
|
|
cb(Kcur, "Kcur", il);
|
|
cb(Vcur, "Vcur", il);
|
|
|
|
if (use_rope && hparams.use_kq_norm) {
|
|
// Llama4TextL2Norm
|
|
Qcur = ggml_rms_norm(ctx0, Qcur, hparams.f_norm_rms_eps);
|
|
Kcur = ggml_rms_norm(ctx0, Kcur, hparams.f_norm_rms_eps);
|
|
cb(Qcur, "Qcur_normed", il);
|
|
cb(Kcur, "Kcur_normed", il);
|
|
}
|
|
cur = build_attn(inp_attn,
|
|
model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,
|
|
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
|
|
cb(cur, "attn_out", 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);
|
|
|
|
// feed-forward network (non-MoE)
|
|
if (model.layers[il].ffn_gate_inp == nullptr) {
|
|
cur = build_norm(ffn_inp,
|
|
model.layers[il].ffn_norm, NULL,
|
|
LLM_NORM_RMS, il);
|
|
cb(cur, "ffn_norm", il);
|
|
|
|
cur = build_ffn(cur,
|
|
model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL,
|
|
model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL,
|
|
model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL,
|
|
NULL,
|
|
LLM_FFN_SILU, LLM_FFN_PAR, il);
|
|
cb(cur, "ffn_out", il);
|
|
} else {
|
|
ggml_tensor * ffn_inp_normed = build_norm(ffn_inp,
|
|
model.layers[il].ffn_norm, NULL,
|
|
LLM_NORM_RMS, il);
|
|
cb(cur, "ffn_norm", il);
|
|
|
|
ggml_tensor * moe_out = build_moe_ffn(ffn_inp_normed,
|
|
model.layers[il].ffn_gate_inp,
|
|
model.layers[il].ffn_up_exps,
|
|
model.layers[il].ffn_gate_exps,
|
|
model.layers[il].ffn_down_exps,
|
|
nullptr,
|
|
n_expert, n_expert_used,
|
|
LLM_FFN_SILU, false,
|
|
hparams.expert_weights_scale,
|
|
LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,
|
|
il);
|
|
|
|
// Shared experts
|
|
ggml_tensor * shexp_out = build_ffn(ffn_inp_normed,
|
|
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(shexp_out, "ffn_moe_shexp", il);
|
|
|
|
cur = ggml_add(ctx0, moe_out, shexp_out);
|
|
cb(cur, "ffn_moe_out_merged", il);
|
|
}
|
|
cur = ggml_add(ctx0, cur, ffn_inp);
|
|
cb(cur, "ffn_out", il);
|
|
|
|
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);
|
|
}
|
|
|
|
// Explicit template instantiations
|
|
template struct llama_model_llama4::graph<false>;
|
|
template struct llama_model_llama4::graph<true>;
|