WeeAIResearch's picture
Add pruned mixed-precision Qwen GGUF
a98d378 verified
Raw History Blame Contribute Delete
11.2 kB
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 5e61f61..9861e32 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -441,6 +441,12 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_FFN_GATE_EXP, "blk.%d.ffn_gate.%d" },
{ LLM_TENSOR_FFN_DOWN_EXP, "blk.%d.ffn_down.%d" },
{ LLM_TENSOR_FFN_UP_EXP, "blk.%d.ffn_up.%d" },
+ { LLM_TENSOR_FFN_GATE_EXPS_TIER1, "blk.%d.ffn_gate_exps_tier1" },
+ { LLM_TENSOR_FFN_UP_EXPS_TIER1, "blk.%d.ffn_up_exps_tier1" },
+ { LLM_TENSOR_FFN_DOWN_EXPS_TIER1, "blk.%d.ffn_down_exps_tier1" },
+ { LLM_TENSOR_FFN_GATE_EXPS_TIER2, "blk.%d.ffn_gate_exps_tier2" },
+ { LLM_TENSOR_FFN_UP_EXPS_TIER2, "blk.%d.ffn_up_exps_tier2" },
+ { LLM_TENSOR_FFN_DOWN_EXPS_TIER2, "blk.%d.ffn_down_exps_tier2" },
{ LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" },
{ LLM_TENSOR_FFN_GATE_UP_EXPS, "blk.%d.ffn_gate_up_exps" },
{ LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" },
@@ -889,6 +895,12 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_DEC_ATTN_REL_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
{LLM_TENSOR_ENC_ATTN_REL_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
{LLM_TENSOR_FFN_DOWN_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
+ {LLM_TENSOR_FFN_GATE_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
+ {LLM_TENSOR_FFN_UP_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
+ {LLM_TENSOR_FFN_DOWN_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
+ {LLM_TENSOR_FFN_GATE_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
+ {LLM_TENSOR_FFN_UP_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
+ {LLM_TENSOR_FFN_DOWN_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
{LLM_TENSOR_FFN_GATE_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
{LLM_TENSOR_FFN_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
{LLM_TENSOR_FFN_GATE_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}},
diff --git a/src/llama-arch.h b/src/llama-arch.h
index ca7d55a..9026666 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -469,6 +469,12 @@ enum llm_tensor {
LLM_TENSOR_FFN_DOWN_EXPS, // merged experts
LLM_TENSOR_FFN_GATE_EXPS,
LLM_TENSOR_FFN_UP_EXPS,
+ LLM_TENSOR_FFN_GATE_EXPS_TIER1,
+ LLM_TENSOR_FFN_UP_EXPS_TIER1,
+ LLM_TENSOR_FFN_DOWN_EXPS_TIER1,
+ LLM_TENSOR_FFN_GATE_EXPS_TIER2,
+ LLM_TENSOR_FFN_UP_EXPS_TIER2,
+ LLM_TENSOR_FFN_DOWN_EXPS_TIER2,
LLM_TENSOR_FFN_GATE_UP_EXPS,
LLM_TENSOR_FFN_DOWN_SHEXP,
LLM_TENSOR_FFN_GATE_SHEXP,
diff --git a/src/llama-model.h b/src/llama-model.h
index 0d7352a..024952b 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -329,6 +329,12 @@ struct llama_layer {
struct ggml_tensor * ffn_gate_exps = nullptr;
struct ggml_tensor * ffn_down_exps = nullptr;
struct ggml_tensor * ffn_up_exps = nullptr;
+ struct ggml_tensor * ffn_gate_exps_tier1 = nullptr;
+ struct ggml_tensor * ffn_up_exps_tier1 = nullptr;
+ struct ggml_tensor * ffn_down_exps_tier1 = nullptr;
+ struct ggml_tensor * ffn_gate_exps_tier2 = nullptr;
+ struct ggml_tensor * ffn_up_exps_tier2 = nullptr;
+ struct ggml_tensor * ffn_down_exps_tier2 = nullptr;
struct ggml_tensor * ffn_gate_up_exps = nullptr;
struct ggml_tensor * ffn_gate_inp_b = nullptr;
struct ggml_tensor * ffn_gate_exps_b = nullptr;
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index acfdd5b..1fcb5c7 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -198,9 +198,32 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0);
}
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0);
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0);
- create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0);
+ const auto * t1 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER1, "weight", il).str().c_str());
+ const bool mixed = t1 != nullptr;
+ const auto * router_meta = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il).str().c_str());
+ if (!router_meta) throw std::runtime_error("missing expert router tensor");
+ const int64_t layer_experts = mixed ? router_meta->ne[1] : n_expert;
+ if (layer_experts < n_expert_used || layer_experts > n_expert) throw std::runtime_error("invalid layer expert count");
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, layer_experts }, 0);
+ if (mixed) {
+ const auto * t0 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il).str().c_str());
+ const auto * t2 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER2, "weight", il).str().c_str());
+ if (!t0 || !t2 || n_ff_exp != 640) throw std::runtime_error("missing mixed-tier expert tensors");
+ const int64_t n0 = t0->ne[2], n1 = t1->ne[2], n2 = t2->ne[2];
+ if (n0 < 1 || n0 > 64 || n1 < 1 || n1 > 64 || n2 < 1 || n2 > 64 || n0+n1+n2 != layer_experts)
+ throw std::runtime_error("invalid mixed-tier bank sizes");
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n0 }, 0);
+ create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n0, 0);
+ layer.ffn_gate_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER1, "weight", il), { n_embd, n_ff_exp, n1 }, 0);
+ layer.ffn_up_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS_TIER1, "weight", il), { n_embd, n_ff_exp, n1 }, 0);
+ layer.ffn_down_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER1, "weight", il), { 768, n_embd, n1 }, 0);
+ layer.ffn_gate_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER2, "weight", il), { n_embd, n_ff_exp, n2 }, 0);
+ layer.ffn_up_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS_TIER2, "weight", il), { n_embd, n_ff_exp, n2 }, 0);
+ layer.ffn_down_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER2, "weight", il), { n_ff_exp, n_embd, n2 }, 0);
+ } else {
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0);
+ create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0);
+ }
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);
@@ -904,8 +927,52 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) {
GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr);
- ggml_tensor * moe_out =
- build_moe_ffn(cur,
+ ggml_tensor * moe_out = nullptr;
+ if (model.layers[il].ffn_gate_exps_tier1 != nullptr) {
+ const auto & layer = model.layers[il];
+ const int64_t n_tokens = cur->ne[1];
+ ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur);
+ ggml_tensor * probs = ggml_soft_max(ctx0, logits);
+ cb(probs, "ffn_moe_probs", il);
+ ggml_tensor * ids = ggml_argsort_top_k(ctx0, probs, n_expert_used);
+ cb(ids, "ffn_moe_topk", il);
+ ggml_tensor * selected = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, probs, 1, layer.ffn_gate_inp->ne[1], n_tokens), ids);
+ selected = ggml_reshape_2d(ctx0, selected, n_expert_used, n_tokens);
+ ggml_tensor * denom = ggml_clamp(ctx0, ggml_sum_rows(ctx0, selected), 6.103515625e-5f, INFINITY);
+ selected = ggml_div(ctx0, selected, denom);
+ if (hparams.expert_weights_scale != 0.0f && hparams.expert_weights_scale != 1.0f) {
+ selected = ggml_scale(ctx0, selected, hparams.expert_weights_scale);
+ }
+ ggml_tensor * ids_f = ggml_cast(ctx0, ids, GGML_TYPE_F32);
+ ggml_tensor * inp = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, n_tokens);
+ ggml_tensor * weighted = nullptr;
+ ggml_tensor * ups[] = {layer.ffn_up_exps, layer.ffn_up_exps_tier1, layer.ffn_up_exps_tier2};
+ ggml_tensor * gates[] = {layer.ffn_gate_exps, layer.ffn_gate_exps_tier1, layer.ffn_gate_exps_tier2};
+ ggml_tensor * downs[] = {layer.ffn_down_exps, layer.ffn_down_exps_tier1, layer.ffn_down_exps_tier2};
+ int64_t tier_start = 0;
+ for (int tier = 0; tier < 3; ++tier) {
+ const int64_t tier_size = ups[tier]->ne[2];
+ const float base = float(tier_start);
+ ggml_tensor * local_f = ggml_scale_bias(ctx0, ids_f, 1.0f, -base);
+ ggml_tensor * local = ggml_cast(ctx0, ggml_clamp(ctx0, local_f, 0.0f, float(tier_size - 1)), GGML_TYPE_I32);
+ ggml_tensor * low = ggml_step(ctx0, ggml_scale_bias(ctx0, local_f, 1.0f, 0.5f));
+ ggml_tensor * high = ggml_step(ctx0, ggml_scale_bias(ctx0, local_f, -1.0f, float(tier_size) - 0.5f));
+ tier_start += tier_size;
+ ggml_tensor * mask = ggml_mul(ctx0, low, high);
+ ggml_tensor * weights = ggml_reshape_3d(ctx0, ggml_mul(ctx0, selected, mask), 1, n_expert_used, n_tokens);
+ ggml_tensor * up = build_lora_mm_id(ups[tier], inp, local);
+ ggml_tensor * gate = build_lora_mm_id(gates[tier], inp, local);
+ ggml_tensor * act = ggml_swiglu_split(ctx0, gate, up);
+ if (tier == 1) act = ggml_pad(ctx0, act, 128, 0, 0, 0);
+ ggml_tensor * result = ggml_mul(ctx0, build_lora_mm_id(downs[tier], act, local), weights);
+ weighted = weighted ? ggml_add(ctx0, weighted, result) : result;
+ }
+ for (int i = 0; i < n_expert_used; ++i) {
+ ggml_tensor * piece = ggml_view_2d(ctx0, weighted, cur->ne[0], n_tokens, weighted->nb[2], i*weighted->nb[1]);
+ moe_out = moe_out ? ggml_add(ctx0, moe_out, piece) : piece;
+ }
+ } else {
+ moe_out = build_moe_ffn(cur,
model.layers[il].ffn_gate_inp,
model.layers[il].ffn_up_exps,
model.layers[il].ffn_gate_exps,
@@ -919,6 +986,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, co
model.layers[il].ffn_up_exps_s,
model.layers[il].ffn_gate_exps_s,
model.layers[il].ffn_down_exps_s);
+ }
cb(moe_out, "ffn_moe_out", il);
// shared experts, as in the Qwen3Next reference