Add support for parallel graphs to GLM MTP (#1637)
* mtp: fix split graph assert * Add mtp split graph mode * remove unused ffn function for unsupported mtp * revert cuda context syncronization
This commit is contained in:
Samuel Oliveira Alves
GitHub
parent
1163af96cf
commit
470d3a3b5b
|
|
@ -169,7 +169,7 @@ struct common_speculative_state_mtp : public common_speculative_state {
|
|||
if (ctx_mtp) {
|
||||
LOG_INF("%s: created MTP context (n_ctx=%d)\n", __func__, llama_n_ctx(ctx_mtp));
|
||||
} else {
|
||||
LOG_ERR("%s: failed to create MTP context, falling back to shared context\n", __func__);
|
||||
LOG_ERR("%s: failed to create MTP context\n", __func__);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -193,14 +193,12 @@ struct common_speculative_state_mtp : public common_speculative_state {
|
|||
int32_t n_past = (int32_t)prompt_tgt.size();
|
||||
llama_seq_id seq_id = 0;
|
||||
|
||||
if (ctx_mtp) {
|
||||
llama_pos mtp_pos_max = llama_kv_cache_seq_pos_max(ctx_mtp, seq_id);
|
||||
if (mtp_pos_max >= n_past) {
|
||||
llama_kv_cache_seq_rm(ctx_mtp, seq_id, n_past, -1);
|
||||
}
|
||||
llama_pos mtp_pos_max = llama_kv_cache_seq_pos_max(ctx_mtp, seq_id);
|
||||
if (mtp_pos_max >= n_past) {
|
||||
llama_kv_cache_seq_rm(ctx_mtp, seq_id, n_past, -1);
|
||||
}
|
||||
|
||||
llama_context * ctx = ctx_mtp ? ctx_mtp : ctx_tgt;
|
||||
llama_context * ctx = ctx_mtp;
|
||||
|
||||
result = mtp_speculative_gen_draft(
|
||||
smpl,
|
||||
|
|
@ -974,10 +972,15 @@ common_speculative * common_speculative_init(
|
|||
break;
|
||||
}
|
||||
case COMMON_SPECULATIVE_TYPE_MTP: {
|
||||
impls.push_back(std::make_unique<common_speculative_state_mtp>(config.type,
|
||||
auto mtp_state = std::make_unique<common_speculative_state_mtp>(config.type,
|
||||
/* .ctx_tgt = */ ctx_tgt,
|
||||
/* .mtp_cparams = */ params.cparams_dft
|
||||
));
|
||||
);
|
||||
if (!mtp_state->ctx_mtp) {
|
||||
LOG_ERR("%s: failed to create MTP context\n", __func__);
|
||||
return nullptr;
|
||||
}
|
||||
impls.push_back(std::move(mtp_state));
|
||||
break;
|
||||
}
|
||||
case COMMON_SPECULATIVE_TYPE_EAGLE3: {
|
||||
|
|
@ -1280,7 +1283,7 @@ void mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch, b
|
|||
return;
|
||||
}
|
||||
|
||||
llama_seq_id seq_id = batch.seq_id[0][0];
|
||||
llama_seq_id seq_id = batch.seq_id[0][0];
|
||||
llama_pos start_pos = batch.pos[0];
|
||||
|
||||
if (llama_kv_cache_seq_pos_max(ctx, seq_id) >= start_pos) {
|
||||
|
|
|
|||
|
|
@ -255,7 +255,8 @@ void server_context::init() {
|
|||
SLT_INF(slot, "%s", "speculative decoding context initialized\n");
|
||||
} else {
|
||||
if (slot.has_mtp) {
|
||||
SRV_ERR("%s", "failed to initialize MTP speculative context\n");
|
||||
SRV_ERR("%s", "failed to initialize MTP speculative context, aborting\n");
|
||||
GGML_ABORT("MTP context creation failed");
|
||||
} else {
|
||||
SLT_INF(slot, "%s", "speculative decoding context not initialized\n");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8183,59 +8183,57 @@ struct ggml_tensor * llm_build_context::build_mtp_tail(
|
|||
cb(combined, "mtp_concat", il);
|
||||
ggml_tensor* cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined);
|
||||
|
||||
struct ggml_tensor * inpSA = cur;
|
||||
|
||||
cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_norm, NULL, LLM_NORM_RMS, cb, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
|
||||
// Self-Attention
|
||||
{
|
||||
const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
|
||||
ggml_tensor * ffn_inp;
|
||||
if (rope_cache == nullptr) {
|
||||
cur = build_std_attention(gf, mtp_layer.attn_norm, cur,
|
||||
inp_pos, nullptr, nullptr,
|
||||
KQ_mask, nullptr, nullptr,
|
||||
kq_scale, 0.0f, 0, il, true, false, true, false, false, nullptr);
|
||||
ffn_inp = cur;
|
||||
} else {
|
||||
struct ggml_tensor * inpSA = cur;
|
||||
cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_norm, NULL, LLM_NORM_RMS, cb, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur,
|
||||
nullptr, nullptr, // wqkv, bqkv (not used in GLM usually?)
|
||||
nullptr, nullptr, // wqk, bqk
|
||||
nullptr, nullptr,
|
||||
nullptr, nullptr,
|
||||
mtp_layer.wq, mtp_layer.bq,
|
||||
mtp_layer.wk, mtp_layer.bk,
|
||||
mtp_layer.wv, mtp_layer.bv,
|
||||
mtp_layer.attn_q_norm, mtp_layer.attn_k_norm,
|
||||
mtp_layer.attn_q_norm, mtp_layer.attn_k_norm,
|
||||
0.f, il);
|
||||
|
||||
// RoPE
|
||||
if (rope_cache) {
|
||||
Qcur = ggml_rope_fast(ctx0, Qcur, rope_cache);
|
||||
Kcur = ggml_rope_fast(ctx0, Kcur, rope_cache);
|
||||
} else {
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
}
|
||||
|
||||
Qcur = ggml_rope_fast(ctx0, Qcur, rope_cache);
|
||||
Kcur = ggml_rope_fast(ctx0, Kcur, rope_cache);
|
||||
cb(Qcur, "Qcur", il);
|
||||
cb(Kcur, "Kcur", il);
|
||||
cb(Vcur, "Vcur", il);
|
||||
|
||||
// KV Cache & Attention
|
||||
cur = llm_build_kv(ctx0, lctx, kv_self, gf,
|
||||
model.layers[il].wo, NULL,
|
||||
mtp_layer.wo, NULL,
|
||||
Kcur, Vcur, Qcur, KQ_mask,
|
||||
n_tokens, kv_head, n_kv,
|
||||
1.0f/sqrtf(float(n_embd_head)), cb, il);
|
||||
kq_scale, cb, il);
|
||||
ffn_inp = ggml_add(ctx0, cur, inpSA);
|
||||
cb(ffn_inp, "mtp_ffn_inp", il);
|
||||
}
|
||||
|
||||
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
|
||||
cb(ffn_inp, "mtp_ffn_inp", il);
|
||||
|
||||
// FFN
|
||||
cur = llm_build_std_moe_ffn(ctx0, lctx, mtp_layer.ffn_norm, ffn_inp,
|
||||
mtp_layer.ffn_gate_inp, NULL,
|
||||
mtp_layer.ffn_up_exps, NULL,
|
||||
mtp_layer.ffn_gate_exps, NULL,
|
||||
mtp_layer.ffn_down_exps, NULL,
|
||||
mtp_layer.ffn_exp_probs_b,
|
||||
mtp_layer.ffn_up_shexp, nullptr,
|
||||
mtp_layer.ffn_gate_shexp, nullptr,
|
||||
mtp_layer.ffn_down_shexp, nullptr,
|
||||
n_expert, n_expert_used,
|
||||
LLM_FFN_SILU, hparams.expert_weights_norm, true, hparams.expert_weights_scale,
|
||||
(llm_expert_gating_func_type) hparams.expert_gating_func,
|
||||
LLM_FFN_SILU, cb, il, gf, true, mtp_layer.ffn_up_gate_exps);
|
||||
mtp_layer.ffn_gate_inp, NULL,
|
||||
mtp_layer.ffn_up_exps, NULL,
|
||||
mtp_layer.ffn_gate_exps, NULL,
|
||||
mtp_layer.ffn_down_exps, NULL,
|
||||
mtp_layer.ffn_exp_probs_b,
|
||||
mtp_layer.ffn_up_shexp, nullptr,
|
||||
mtp_layer.ffn_gate_shexp, nullptr,
|
||||
mtp_layer.ffn_down_shexp, nullptr,
|
||||
n_expert, n_expert_used,
|
||||
LLM_FFN_SILU, hparams.expert_weights_norm, true, hparams.expert_weights_scale,
|
||||
(llm_expert_gating_func_type) hparams.expert_gating_func,
|
||||
LLM_FFN_SILU, cb, il, gf, true, mtp_layer.ffn_up_gate_exps);
|
||||
|
||||
cur = lctx.cvec.apply_to(ctx0, cur, il);
|
||||
cb(cur, "ffn_out", il);
|
||||
|
||||
cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, il);
|
||||
|
|
|
|||
|
|
@ -309,6 +309,14 @@ create_tensors_helper::create_tensors_helper(llama_model_loader & _ml, llama_mod
|
|||
}
|
||||
}
|
||||
|
||||
// Split MTP layer's to graph
|
||||
if ((model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN) &&
|
||||
model.hparams.nextn_predict_layers > 0 && model.splits.size() > 1) {
|
||||
int mtp_first = n_layer - model.hparams.nextn_predict_layers;
|
||||
LLAMA_LOG_DEBUG("%s: MTP layer(s) %d-%d: split attention+FFN, nextn on per-device CUDA\n",
|
||||
__func__, mtp_first, n_layer - 1);
|
||||
}
|
||||
|
||||
auto n_tensors = ml.n_tensors;
|
||||
if (ml.merge_qkv) n_tensors += n_layer;
|
||||
if (ml.merge_up_gate_exps) n_tensors += n_layer;
|
||||
|
|
@ -2616,35 +2624,32 @@ bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) {
|
|||
layer.ffn_down = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, flags);
|
||||
layer.ffn_up = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, flags);
|
||||
}
|
||||
// --- NextN / MTP tensors (preserved but unused), on the final layer ---
|
||||
// --- NextN / MTP tensors on the final layer ---
|
||||
if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
|
||||
const int final_layer = n_layer - 1;
|
||||
// EH_PROJ: [2*embd, embd]
|
||||
layer.nextn.eh_proj = create_tensor(ctx_for_layer(final_layer),
|
||||
auto nextn_ctx = ctx_for_layer(final_layer);
|
||||
auto nextn_host_ctx = ctx_input;
|
||||
layer.nextn.eh_proj = create_tensor(nextn_ctx,
|
||||
tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", final_layer),
|
||||
{ 2*n_embd, n_embd },
|
||||
flags);
|
||||
// EMBED_TOKENS: [embd, vocab]
|
||||
layer.nextn.embed_tokens = create_tensor(ctx_for_layer(final_layer),
|
||||
layer.nextn.embed_tokens = create_tensor(nextn_host_ctx,
|
||||
tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", final_layer),
|
||||
{ n_embd, n_vocab },
|
||||
flags | llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
// ENORM, HNORM: [embd]
|
||||
layer.nextn.enorm = create_tensor(ctx_for_layer(final_layer),
|
||||
layer.nextn.enorm = create_tensor(nextn_ctx,
|
||||
tn(LLM_TENSOR_NEXTN_ENORM, "weight", final_layer),
|
||||
{ n_embd },
|
||||
flags);
|
||||
layer.nextn.hnorm = create_tensor(ctx_for_layer(final_layer),
|
||||
layer.nextn.hnorm = create_tensor(nextn_ctx,
|
||||
tn(LLM_TENSOR_NEXTN_HNORM, "weight", final_layer),
|
||||
{ n_embd },
|
||||
flags);
|
||||
// SHARED_HEAD_HEAD: [embd, vocab]
|
||||
layer.nextn.shared_head_head = create_tensor(ctx_for_layer(final_layer),
|
||||
layer.nextn.shared_head_head = create_tensor(nextn_host_ctx,
|
||||
tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", final_layer),
|
||||
{ n_embd, n_vocab },
|
||||
flags | llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
// SHARED_HEAD_NORM: [embd]
|
||||
layer.nextn.shared_head_norm = create_tensor(ctx_for_layer(final_layer),
|
||||
layer.nextn.shared_head_norm = create_tensor(nextn_ctx,
|
||||
tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", final_layer),
|
||||
{ n_embd },
|
||||
flags | llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
|
|
|
|||
Loading…
Reference in New Issue