Skip to content

Commit 5446ccc

Browse files
ubergarmubergarm
authored
add hunyuan moe support for 561 (#565)
* add hunyuan moe * Don't reshape Vcur * Apply chat template fix from mainline PR14584
1 parent 97c34f4 commit 5446ccc

File tree

3 files changed

+258
-0
lines changed

3 files changed

+258
-0
lines changed

‎include/llama.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ extern "C" {
111111
LLAMA_VOCAB_PRE_TYPE_FALCON_3 = 34,
112112
LLAMA_VOCAB_PRE_TYPE_FALCON_E = 35,
113113
LLAMA_VOCAB_PRE_TYPE_SEED_CODER = 36, //llama.cpp lists this as 35
114+
LLAMA_VOCAB_PRE_TYPE_HUNYUAN = 37, //llama.cpp lists this as 36
114115
};
115116

116117
// note: these values should be synchronized with ggml_rope

‎src/llama-vocab.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,7 @@ struct llm_tokenizer_bpe {
427427
break;
428428
case LLAMA_VOCAB_PRE_TYPE_STABLELM2:
429429
case LLAMA_VOCAB_PRE_TYPE_QWEN2:
430+
case LLAMA_VOCAB_PRE_TYPE_HUNYUAN:
430431
regex_exprs = {
431432
// original regex from tokenizer.json
432433
// "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"

‎src/llama.cpp‎

Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,7 @@ enum llm_arch {
235235
LLM_ARCH_GRANITE,
236236
LLM_ARCH_GRANITE_MOE,
237237
LLM_ARCH_COHERE2,
238+
LLM_ARCH_HUNYUAN_MOE,
238239
LLM_ARCH_UNKNOWN,
239240
};
240241

@@ -291,6 +292,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
291292
{ LLM_ARCH_GRANITE, "granite" },
292293
{ LLM_ARCH_GRANITE_MOE, "granitemoe" },
293294
{ LLM_ARCH_COHERE2, "cohere2" },
295+
{ LLM_ARCH_HUNYUAN_MOE, "hunyuan-moe" },
294296
{ LLM_ARCH_UNKNOWN, "(unknown)" },
295297
};
296298

@@ -1595,6 +1597,29 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
15951597
{ LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" },
15961598
},
15971599
},
1600+
{
1601+
LLM_ARCH_HUNYUAN_MOE,
1602+
{
1603+
{ LLM_TENSOR_TOKEN_EMBD, "token_embd" },
1604+
{ LLM_TENSOR_OUTPUT_NORM, "output_norm" },
1605+
{ LLM_TENSOR_OUTPUT, "output" },
1606+
{ LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" },
1607+
{ LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" },
1608+
{ LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" },
1609+
{ LLM_TENSOR_ATTN_K, "blk.%d.attn_k" },
1610+
{ LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" },
1611+
{ LLM_TENSOR_ATTN_V, "blk.%d.attn_v" },
1612+
{ LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" },
1613+
{ LLM_TENSOR_FFN_GATE_INP, "blk.%d.ffn_gate_inp" },
1614+
{ LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" },
1615+
{ LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" },
1616+
{ LLM_TENSOR_FFN_DOWN_SHEXP, "blk.%d.ffn_down_shexp" },
1617+
{ LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" },
1618+
{ LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" },
1619+
{ LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" },
1620+
{ LLM_TENSOR_FFN_UP_EXPS, "blk.%d.ffn_up_exps" },
1621+
},
1622+
},
15981623
{
15991624
LLM_ARCH_UNKNOWN,
16001625
{
@@ -1638,6 +1663,7 @@ enum llm_chat_template {
16381663
LLM_CHAT_TEMPLATE_MEGREZ,
16391664
LLM_CHAT_TEMPLATE_LLAMA4,
16401665
LLM_CHAT_TEMPLATE_BITNET,
1666+
LLM_CHAT_TEMPLATE_HUNYUAN_MOE,
16411667
LLM_CHAT_TEMPLATE_UNKNOWN,
16421668
};
16431669

@@ -1675,6 +1701,7 @@ static const std::map<std::string, llm_chat_template> LLM_CHAT_TEMPLATES = {
16751701
{ "gigachat", LLM_CHAT_TEMPLATE_GIGACHAT },
16761702
{ "megrez", LLM_CHAT_TEMPLATE_MEGREZ },
16771703
{ "llama4", LLM_CHAT_TEMPLATE_LLAMA4 },
1704+
{ "hunyuan-moe", LLM_CHAT_TEMPLATE_HUNYUAN_MOE },
16781705
{ "bitnet", LLM_CHAT_TEMPLATE_BITNET },
16791706
};
16801707

@@ -2570,6 +2597,7 @@ enum e_model {
25702597
MODEL_27B,
25712598
MODEL_17B_16E,
25722599
MODEL_17B_128E,
2600+
MODEL_80B_A13B,
25732601
};
25742602

25752603
static const size_t kiB = 1024;
@@ -5203,6 +5231,7 @@ static const char * llama_model_type_name(e_model type) {
52035231
case MODEL_27B: return "27B";
52045232
case MODEL_17B_16E: return "17Bx16E (Scout)";
52055233
case MODEL_17B_128E: return "17Bx128E (Maverick)";
5234+
case MODEL_80B_A13B: return "80B.A13B";
52065235
default: return "?B";
52075236
}
52085237
}
@@ -6037,6 +6066,17 @@ static void llm_load_hparams(
60376066
default: model.type = e_model::MODEL_UNKNOWN;
60386067
}
60396068
} break;
6069+
case LLM_ARCH_HUNYUAN_MOE:
6070+
{
6071+
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
6072+
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
6073+
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
6074+
6075+
switch (hparams.n_layer) {
6076+
case 32: model.type = e_model::MODEL_80B_A13B; break;
6077+
default: model.type = e_model::MODEL_UNKNOWN;
6078+
}
6079+
} break;
60406080
default: (void)0;
60416081
}
60426082

@@ -6306,6 +6346,10 @@ static void llm_load_vocab(
63066346
tokenizer_pre == "seed-coder") {
63076347
vocab.type_pre = LLAMA_VOCAB_PRE_TYPE_SEED_CODER;
63086348
vocab.tokenizer_clean_spaces = false;
6349+
} else if (
6350+
tokenizer_pre == "hunyuan") {
6351+
vocab.type_pre = LLAMA_VOCAB_PRE_TYPE_HUNYUAN;
6352+
vocab.tokenizer_clean_spaces = false;
63096353
} else {
63106354
throw std::runtime_error(format("unknown pre-tokenizer type: '%s'", tokenizer_pre.c_str()));
63116355
}
@@ -9164,6 +9208,47 @@ static bool llm_load_tensors(
91649208
layer.ffn_post_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
91659209
}
91669210
} break;
9211+
case LLM_ARCH_HUNYUAN_MOE:
9212+
{
9213+
model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
9214+
9215+
// output
9216+
model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
9217+
model.output = create_tensor(ctx_output_split, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
9218+
9219+
// if output is NULL, init from the input tok embed
9220+
if (model.output == NULL) {
9221+
model.output = create_tensor(ctx_output_split, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED);
9222+
}
9223+
9224+
for (int i = 0; i < n_layer; ++i) {
9225+
ggml_context * ctx_layer = ctx_for_layer(i);
9226+
ggml_context * ctx_split = ctx_for_layer_split(i);
9227+
9228+
auto & layer = model.layers[i];
9229+
9230+
layer.attn_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
9231+
9232+
layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
9233+
layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0);
9234+
layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0);
9235+
layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
9236+
9237+
layer.attn_k_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
9238+
layer.attn_q_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
9239+
9240+
layer.ffn_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
9241+
9242+
layer.ffn_gate_inp = create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
9243+
layer.ffn_gate_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0);
9244+
layer.ffn_down_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff, n_embd, n_expert}, 0);
9245+
layer.ffn_up_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0);
9246+
9247+
layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, 0);
9248+
layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, 0);
9249+
layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, 0);
9250+
}
9251+
} break;
91679252
default:
91689253
throw std::runtime_error("unknown architecture");
91699254
}
@@ -16862,6 +16947,158 @@ struct llm_build_context {
1686216947

1686316948
return gf;
1686416949
}
16950+
16951+
struct ggml_cgraph * build_hunyuan_moe() {
16952+
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false);
16953+
16954+
const int64_t n_embd_head = hparams.n_embd_head_v;
16955+
16956+
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
16957+
GGML_ASSERT(n_embd_head == hparams.n_rot);
16958+
16959+
ggml_tensor * cur;
16960+
ggml_tensor * inpL;
16961+
16962+
inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
16963+
16964+
// inp_pos - contains the positions
16965+
ggml_tensor * inp_pos = build_inp_pos();
16966+
16967+
struct ggml_tensor * KQ_mask = build_inp_KQ_mask();
16968+
16969+
const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
16970+
16971+
ggml_tensor * inp_out_ids = build_inp_out_ids();
16972+
16973+
for (int il = 0; il < n_layer; ++il) {
16974+
ggml_tensor * inpSA = inpL;
16975+
16976+
// norm
16977+
cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il);
16978+
cb(cur, "attn_norm", il);
16979+
16980+
// self-attention
16981+
{
16982+
// rope freq factors for llama3; may return nullptr for llama2 and other models
16983+
struct ggml_tensor * rope_factors = build_rope_factors(il);
16984+
16985+
// compute Q and K and RoPE them
16986+
ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur);
16987+
cb(Qcur, "Qcur", il);
16988+
if (model.layers[il].bq) {
16989+
Qcur = ggml_add(ctx0, Qcur, model.layers[il].bq);
16990+
cb(Qcur, "Qcur", il);
16991+
}
16992+
16993+
ggml_tensor * Kcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur);
16994+
cb(Kcur, "Kcur", il);
16995+
if (model.layers[il].bk) {
16996+
Kcur = ggml_add(ctx0, Kcur, model.layers[il].bk);
16997+
cb(Kcur, "Kcur", il);
16998+
}
16999+
17000+
ggml_tensor * Vcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, cur);
17001+
cb(Vcur, "Vcur", il);
17002+
if (model.layers[il].bv) {
17003+
Vcur = ggml_add(ctx0, Vcur, model.layers[il].bv);
17004+
cb(Vcur, "Vcur", il);
17005+
}
17006+
17007+
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
17008+
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
17009+
17010+
Qcur = ggml_rope_ext(
17011+
ctx0, Qcur, inp_pos, rope_factors,
17012+
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
17013+
ext_factor, attn_factor, beta_fast, beta_slow
17014+
);
17015+
17016+
cb(Qcur, "Qcur", il);
17017+
cb(Kcur, "Kcur", il);
17018+
cb(Vcur, "Vcur", il);
17019+
17020+
Kcur = ggml_rope_ext(
17021+
ctx0, Kcur, inp_pos, rope_factors,
17022+
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
17023+
ext_factor, attn_factor, beta_fast, beta_slow
17024+
);
17025+
17026+
Kcur = llm_build_norm(ctx0, Kcur, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il);
17027+
cb(Kcur, "Kcur_norm", il);
17028+
17029+
Qcur = llm_build_norm(ctx0, Qcur, hparams, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, cb, il);
17030+
cb(Qcur, "Qcur_norm", il);
17031+
17032+
cur = llm_build_kv(ctx0, lctx, kv_self, gf, model.layers[il].wo, model.layers[il].bo, Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il);
17033+
cb(cur, "attn_out", il);
17034+
}
17035+
17036+
if (il == n_layer - 1 && inp_out_ids) {
17037+
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
17038+
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
17039+
}
17040+
17041+
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
17042+
cb(ffn_inp, "ffn_inp", il);
17043+
17044+
cur = llm_build_norm(ctx0,ffn_inp, hparams, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, cb, il);
17045+
cb(cur, "ffn_norm", il);
17046+
17047+
// feed-forward network (non-MoE)
17048+
ggml_tensor * cur_mlp = llm_build_ffn(ctx0, lctx, cur,
17049+
model.layers[il].ffn_up_shexp, NULL, NULL,
17050+
model.layers[il].ffn_gate_shexp, NULL, NULL,
17051+
model.layers[il].ffn_down_shexp, NULL, NULL,
17052+
NULL,
17053+
LLM_FFN_SILU, LLM_FFN_PAR, cb, il);
17054+
cb(cur_mlp, "ffn_mlp", il);
17055+
17056+
// MoE branch
17057+
ggml_tensor * cur_moe = llm_build_moe_ffn(ctx0, lctx, cur,
17058+
model.layers[il].ffn_gate_inp,
17059+
model.layers[il].ffn_up_exps,
17060+
model.layers[il].ffn_gate_exps,
17061+
model.layers[il].ffn_down_exps,
17062+
nullptr,
17063+
n_expert, n_expert_used,
17064+
LLM_FFN_SILU,
17065+
true, // norm_topk_prob
17066+
false,
17067+
0.0,
17068+
LLM_EXPERT_GATING_FUNC_SOFTMAX,
17069+
cb,
17070+
il);
17071+
cb(cur_moe, "ffn_moe_out", il);
17072+
17073+
ggml_tensor * ffn_out = ggml_add(ctx0, cur_moe, cur_mlp);
17074+
cb(ffn_out, "ffn_out", il);
17075+
17076+
cur = ggml_add(ctx0, ffn_out, ffn_inp);
17077+
17078+
cur = lctx.cvec.apply_to(ctx0, cur, il);
17079+
cb(cur, "l_out", il);
17080+
17081+
// input for next layer
17082+
inpL = cur;
17083+
}
17084+
17085+
cur = inpL;
17086+
17087+
cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM_RMS, cb, -1);
17088+
17089+
cb(cur, "result_norm", -1);
17090+
//res->t_embd = cur;
17091+
17092+
// lm_head
17093+
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
17094+
17095+
cb(cur, "result_output", -1);
17096+
//res->t_logits = cur;
17097+
17098+
ggml_build_forward_expand(gf, cur);
17099+
17100+
return gf;
17101+
}
1686517102
};
1686617103

1686717104
static struct ggml_cgraph * llama_build_graph_defrag(llama_context & lctx, const std::vector<uint32_t> & ids) {
@@ -17157,6 +17394,10 @@ static struct ggml_cgraph * llama_build_graph(
1715717394
{
1715817395
result = llm.build_jais();
1715917396
} break;
17397+
case LLM_ARCH_HUNYUAN_MOE:
17398+
{
17399+
result = llm.build_hunyuan_moe();
17400+
} break;
1716017401
default:
1716117402
GGML_ABORT("fatal error");
1716217403
}
@@ -20929,6 +21170,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) {
2092921170
case LLM_ARCH_OPENELM:
2093021171
case LLM_ARCH_GPTNEOX:
2093121172
case LLM_ARCH_CODESHELL:
21173+
case LLM_ARCH_HUNYUAN_MOE:
2093221174
return LLAMA_ROPE_TYPE_NEOX;
2093321175

2093421176
// all model arches should be listed explicitly here
@@ -22742,6 +22984,8 @@ static llm_chat_template llama_chat_detect_template(const std::string & tmpl) {
2274222984
return LLM_CHAT_TEMPLATE_MEGREZ;
2274322985
} else if (tmpl_contains("<|header_start|>") && tmpl_contains("<|header_end|>")) {
2274422986
return LLM_CHAT_TEMPLATE_LLAMA4;
22987+
} else if (tmpl_contains("<|startoftext|>") && tmpl_contains("<|extra_4|>")) {
22988+
return LLM_CHAT_TEMPLATE_HUNYUAN_MOE;
2274522989
}
2274622990
return LLM_CHAT_TEMPLATE_UNKNOWN;
2274722991
}
@@ -23160,6 +23404,18 @@ static int32_t llama_chat_apply_template_internal(
2316023404
ss << message->content;
2316123405
}
2316223406
}
23407+
} else if (tmpl == LLM_CHAT_TEMPLATE_HUNYUAN_MOE) {
23408+
// tencent/Hunyuan-A13B-Instruct
23409+
for (auto message : chat) {
23410+
std::string role(message->role);
23411+
if (role == "system") {
23412+
ss << "<|startoftext|>" << message->content << "<|extra_4|>";
23413+
} else if (role == "assistant") {
23414+
ss << "<|startoftext|>" << message->content << "<|eos|>";
23415+
} else {
23416+
ss << "<|startoftext|>" << message->content << "<|extra_0|>";
23417+
}
23418+
}
2316323419
} else {
2316423420
// template not supported
2316523421
return -1;

0 commit comments

Comments
 (0)