echodict/llama.cpp

1#include "models.h"2 3llm_build_bert::llm_build_bert(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {4    const int64_t n_embd_head = hparams.n_embd_head_v();5 6    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());7 8    ggml_tensor * cur;9    ggml_tensor * inpL;10    ggml_tensor * inp_pos = nullptr;11 12    if (model.arch != LLM_ARCH_JINA_BERT_V2) {13        inp_pos = build_inp_pos();14    }15 16    // construct input embeddings (token, type, position)17    inpL = build_inp_embd(model.tok_embd);18 19    // token types are hardcoded to zero ("Sentence A")20    if (model.type_embd) {21        ggml_tensor * type_row0 = ggml_view_1d(ctx0, model.type_embd, n_embd, 0);22        inpL                    = ggml_add(ctx0, inpL, type_row0);23    }24    if (model.arch == LLM_ARCH_BERT) {25        inpL = ggml_add(ctx0, ggml_get_rows(ctx0, model.pos_embd, inp_pos), inpL);26    }27    cb(inpL, "inp_embd", -1);28 29    // embed layer norm30    inpL = build_norm(inpL, model.tok_norm, model.tok_norm_b, LLM_NORM, 0);31    cb(inpL, "inp_norm", 0);32 33    auto * inp_attn = build_attn_inp_no_cache();34 35    ggml_tensor * inp_out_ids = build_inp_out_ids();36 37    for (int il = 0; il < n_layer; ++il) {38        ggml_tensor * cur = inpL;39 40        {41            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,42                    n_embd_head, n_head, n_head_kv, il);43 44            if (model.layers[il].attn_q_norm) {45                Qcur = ggml_reshape_2d(ctx0, Qcur, n_embd_head * n_head, n_tokens);46 47                Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, model.layers[il].attn_q_norm_b, LLM_NORM, il);48 49                Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);50            }51 52            if (model.layers[il].attn_k_norm) {53                Kcur = ggml_reshape_2d(ctx0, Kcur, n_embd_head * n_head_kv, n_tokens);54 55                Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, model.layers[il].attn_k_norm_b, LLM_NORM, il);56 57                Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);58            }59 60            // RoPE61            if (model.arch == LLM_ARCH_NOMIC_BERT || model.arch == LLM_ARCH_NOMIC_BERT_MOE ||62                model.arch == LLM_ARCH_JINA_BERT_V3) {63                Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,64                                     ext_factor, attn_factor, beta_fast, beta_slow);65 66                Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,67                                     ext_factor, attn_factor, beta_fast, beta_slow);68            }69 70            cb(Qcur, "Qcur", il);71            cb(Kcur, "Kcur", il);72            cb(Vcur, "Vcur", il);73 74            cur = build_attn(inp_attn,75                    model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,76                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);77            cb(cur, "kqv_out", il);78        }79 80        if (il == n_layer - 1 && inp_out_ids) {81            cur  = ggml_get_rows(ctx0, cur, inp_out_ids);82            inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);83        }84 85        // re-add the layer input86        cur = ggml_add(ctx0, cur, inpL);87 88        // attention layer norm89        cur = build_norm(cur, model.layers[il].attn_out_norm, model.layers[il].attn_out_norm_b, LLM_NORM, il);90 91        if (model.layers[il].attn_norm_2 != nullptr) {92            cur = ggml_add(ctx0, cur, inpL);  // re-add the layer input93            cur = build_norm(cur, model.layers[il].attn_norm_2, model.layers[il].attn_norm_2_b, LLM_NORM, il);94        }95 96        ggml_tensor * ffn_inp = cur;97        cb(ffn_inp, "ffn_inp", il);98 99        // feed-forward network100        if (hparams.moe_every_n_layers > 0 && il % hparams.moe_every_n_layers == 1) {101            // MoE branch102            cur = build_moe_ffn(cur,103                    model.layers[il].ffn_gate_inp,104                    model.layers[il].ffn_up_exps,105                    nullptr,106                    model.layers[il].ffn_down_exps,107                    nullptr,108                    hparams.n_expert, hparams.n_expert_used,109                    LLM_FFN_GELU, false,110                    hparams.expert_weights_scale,111                    LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,112                    il);113            cb(cur, "ffn_moe_out", il);114        } else if (model.arch == LLM_ARCH_BERT || model.arch == LLM_ARCH_NOMIC_BERT_MOE ||115                   model.arch == LLM_ARCH_JINA_BERT_V3) {116            cur = build_ffn(cur,117                    model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL,118                    NULL, NULL, NULL,119                    model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, NULL,120                    LLM_FFN_GELU, LLM_FFN_SEQ, il);121            cb(cur, "ffn_out", il);122        } else if (model.arch == LLM_ARCH_JINA_BERT_V2) {123            const bool up_contains_gate = !model.layers[il].ffn_gate && model.layers[il].ffn_up->ne[1] != hparams.n_ff();124            auto type_op = up_contains_gate ? LLM_FFN_GEGLU : LLM_FFN_GELU;125            cur = build_ffn(cur,126                    model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL,127                    model.layers[il].ffn_gate, NULL, NULL,128                    model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, NULL,129                    type_op, LLM_FFN_PAR, il);130            cb(cur, "ffn_out", il);131        } else {132            cur = build_ffn(cur,133                model.layers[il].ffn_up, NULL, NULL,134                model.layers[il].ffn_gate, NULL, NULL,135                model.layers[il].ffn_down, NULL, NULL,136                NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);137            cb(cur, "ffn_out", il);138        }139 140        // attentions bypass the intermediate layer141        cur = ggml_add(ctx0, cur, ffn_inp);142 143        // output layer norm144        cur = build_norm(cur, model.layers[il].layer_out_norm, model.layers[il].layer_out_norm_b, LLM_NORM, il);145 146        // input for next layer147        inpL = cur;148    }149 150    cur = inpL;151 152    cb(cur, "result_embd", -1);153    res->t_embd = cur;154 155    ggml_build_forward_expand(gf, cur);156}157