diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 7256bc7c..baa1f6f3 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -8452,89 +8452,107 @@ ggml_cgraph * llm_build_context::build_bailingmoe2() { const int n_transformer_layers = n_layer - hparams.nextn_predict_layers; - auto rope_cache = cparams.rope_cache && (rope_type == LLAMA_ROPE_TYPE_NEOX || rope_type == LLAMA_ROPE_TYPE_NORM) ? - ggml_rope_cache(ctx0, inp_pos, nullptr, n_embd_head, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow) : nullptr; - for (int il = 0; il < n_transformer_layers; ++il) { - ggml_tensor * inpSA = inpL; + //ggml_tensor * inpSA = inpL; - // norm - cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il); - cb(cur, "attn_norm", il); + //auto wqkv = model.split_mode == LLAMA_SPLIT_MODE_GRAPH ? nullptr : + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos, nullptr, KQ_mask, + nullptr, nullptr, kq_scale, 0.0f, 0, il, true, false, true); - // self_attention - { - auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, model.layers[il].bqkv, - nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, - model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, 0.0f, il); + //// norm + //cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il); + //cb(cur, "attn_norm", il); - 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); - } + //// self_attention + //{ + // //auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wqkv, model.layers[il].bqkv, + // // nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + // // model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, 0.0f, il); + // auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, nullptr, nullptr, nullptr, nullptr, + // model.layers[il].wq, nullptr, model.layers[il].wk, nullptr, model.layers[il].wv, nullptr, + // model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, 0.0f, il); - cb(Qcur, "Qcur", il); - cb(Kcur, "Kcur", il); - cb(Vcur, "Vcur", il); + // 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); + // } - 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); - } + // cb(Qcur, "Qcur", il); + // cb(Kcur, "Kcur", il); + // cb(Vcur, "Vcur", il); + + // 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); + //} if (il == n_transformer_layers - 1 && inp_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); - inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + //inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } - ggml_tensor * sa_out = ggml_add(ctx0, cur, inpSA); - cb(sa_out, "sa_out", il); + //ggml_tensor * sa_out = ggml_add(ctx0, cur, inpSA); + //cb(sa_out, "sa_out", il); // MoE branch - cur = llm_build_norm(ctx0, sa_out, hparams, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, cb, il); - cb(cur, "ffn_norm", il); + //cur = llm_build_norm(ctx0, sa_out, hparams, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, cb, il); + //cb(cur, "ffn_norm", il); if (static_cast(il) < hparams.n_layer_dense_lead) { - cur = llm_build_ffn(ctx0, lctx, nullptr, cur, + cur = llm_build_ffn(ctx0, lctx, model.layers[il].ffn_norm, 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, cb, il); + LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, true); cb(cur, "ffn_out", il); } else { - ggml_tensor * moe_out = - llm_build_moe_ffn(ctx0, lctx, 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, - true, hparams.expert_weights_scale, - (llm_expert_gating_func_type) hparams.expert_gating_func, - cb, il, gf); - cb(moe_out, "ffn_moe_out", il); + cur = llm_build_std_moe_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur, + model.layers[il].ffn_gate_inp, nullptr, + model.layers[il].ffn_up_exps, nullptr, + model.layers[il].ffn_gate_exps, nullptr, + model.layers[il].ffn_down_exps, nullptr, + model.layers[il].ffn_exp_probs_b, + model.layers[il].ffn_up_shexp, nullptr, + model.layers[il].ffn_gate_shexp, nullptr, + model.layers[il].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); - ggml_tensor * ffn_shexp = llm_build_ffn(ctx0, lctx, nullptr, 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, cb, il); - cb(ffn_shexp, "ffn_shexp", il); + //ggml_tensor * moe_out = + // llm_build_moe_ffn(ctx0, lctx, 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, + // true, hparams.expert_weights_scale, + // (llm_expert_gating_func_type) hparams.expert_gating_func, + // cb, il, gf); + //cb(moe_out, "ffn_moe_out", il); - cur = ggml_add(ctx0, moe_out, ffn_shexp); - cb(cur, "ffn_out", il); + //ggml_tensor * ffn_shexp = llm_build_ffn(ctx0, lctx, nullptr, 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, cb, il); + //cb(ffn_shexp, "ffn_shexp", il); + + //cur = ggml_add(ctx0, moe_out, ffn_shexp); + //cb(cur, "ffn_out", il); } - cur = ggml_add(ctx0, cur, sa_out); + //cur = ggml_add(ctx0, cur, sa_out); cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "l_out", il); @@ -8542,17 +8560,21 @@ ggml_cgraph * llm_build_context::build_bailingmoe2() { // input for next layer inpL = cur; } - cur = inpL; - - cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM_RMS, cb, -1); - - cb(cur, "result_norm", -1); - - // lm_head - cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); + cur = build_output(lctx, ctx0, inpL, model.output, model.output_norm, cb); cb(cur, "result_output", -1); + //cur = inpL; + + //cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM_RMS, cb, -1); + + //cb(cur, "result_norm", -1); + + //// lm_head + //cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); + + //cb(cur, "result_output", -1); + ggml_build_forward_expand(gf, cur); return gf; } @@ -9160,7 +9182,9 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens float freq_base_l = n_swa > 0 ? hparams.rope_freq_base_train_swa : cparams.rope_freq_base; float freq_scale_l = n_swa > 0 ? hparams.rope_freq_scale_train_swa : hparams.rope_freq_scale_train; - if (!model.layers[il].wqkv && !model.layers[il].wqk && cparams.flash_attn && + auto wqkv = lctx.model.split_mode == LLAMA_SPLIT_MODE_GRAPH ? nullptr : model.layers[il].wqkv; + + if (!wqkv && !model.layers[il].wqk && cparams.flash_attn && model.layers[il].wq->extra && model.layers[il].wk->extra && model.layers[il].wv->extra && model.layers[il].wo->extra) { if (kv_self.k_l[il]->extra && kv_self.v_l[il]->extra) { ggml_split_tensor_t * attn_norm = the_attn_norm ? (ggml_split_tensor_t *)the_attn_norm->extra : nullptr; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 722c245e..ee22def6 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -2428,14 +2428,13 @@ bool create_tensors_helper::create_bailingmoe2_tensors(const LLM_TN & tn) { // output model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); - model.output = create_tensor(ctx_output_split, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); GGML_ASSERT(n_expert > 0 && "n_expert must be > 0 for bailingmoe2"); GGML_ASSERT(n_expert_used > 0 && "n_expert_used must be > 0 for bailingmoe2"); for (int i = 0; i < n_layer; ++i) { auto & layer = model.layers[i]; - ggml_context * ctx_layer = ctx_for_layer(i); ggml_context * ctx_split = ctx_for_layer_split(i); int flags = 0; @@ -2444,21 +2443,32 @@ bool create_tensors_helper::create_bailingmoe2_tensors(const LLM_TN & tn) { flags |= llama_model_loader::TENSOR_SKIP; } - layer.attn_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); layer.wqkv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd + 2*n_embd_gqa}, flags); layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, flags); - layer.attn_q_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, flags); - layer.attn_k_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, flags); + if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) { + layer.wq = ggml_view_2d(ctx_split, layer.wqkv, layer.wqkv->ne[0], n_embd_head_k*n_head, layer.wqkv->nb[1], 0); + ggml_set_name(layer.wq, tn(LLM_TENSOR_ATTN_Q, "weight", i).c_str()); + layer.wk = ggml_view_2d(ctx_split, layer.wqkv, layer.wqkv->ne[0], n_embd_head_k*n_head_kv, layer.wqkv->nb[1], + n_embd_head_k*n_head*layer.wqkv->nb[1]); + ggml_set_name(layer.wk, tn(LLM_TENSOR_ATTN_K, "weight", i).c_str()); + layer.wv = ggml_view_2d(ctx_split, layer.wqkv, layer.wqkv->ne[0], n_embd_head_k*n_head_kv, layer.wqkv->nb[1], + n_embd_head_k*(n_head + n_head_kv)*layer.wqkv->nb[1]); + ggml_set_name(layer.wv, tn(LLM_TENSOR_ATTN_V, "weight", i).c_str()); + } - layer.ffn_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, flags); + layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, flags); + + layer.ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); if (static_cast(i) >= hparams.n_layer_dense_lead) { // MoE layers const int64_t n_ff_shexp = (hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff_exp) * n_expert_shared; - layer.ffn_gate_inp = create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); - layer.ffn_exp_probs_b = create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, + layer.ffn_gate_inp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, llama_model_loader::TENSOR_NOT_REQUIRED | flags); layer.ffn_gate_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); @@ -2479,11 +2489,11 @@ bool create_tensors_helper::create_bailingmoe2_tensors(const LLM_TN & tn) { layer.nextn.eh_proj = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags); layer.nextn.embed_tokens = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, llama_model_loader::TENSOR_NOT_REQUIRED | flags); - layer.nextn.enorm = create_tensor(ctx_layer, tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags); - layer.nextn.hnorm = create_tensor(ctx_layer, tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags); + layer.nextn.enorm = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags); layer.nextn.shared_head_head = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, llama_model_loader::TENSOR_NOT_REQUIRED | flags); - layer.nextn.shared_head_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, llama_model_loader::TENSOR_NOT_REQUIRED | flags); - layer.layer_out_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, flags); + layer.nextn.shared_head_norm = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, llama_model_loader::TENSOR_NOT_REQUIRED | flags); + layer.layer_out_norm = create_tensor(ctx_split, tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, flags); } } return use_mmap_buffer; diff --git a/src/llama.cpp b/src/llama.cpp index 10db6d58..3b653602 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1736,6 +1736,7 @@ static bool is_model_split_supported(const llama_model & model) { LLM_ARCH_HUNYUAN_MOE, LLM_ARCH_OPENAI_MOE, LLM_ARCH_ERNIE4_5_MOE, + LLM_ARCH_BAILINGMOE2, }; auto it = k_supported.find(model.arch); return it != k_supported.end();