char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Charlotte Somsupport aliceai_t5_moe98a01afd9

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
11.1 KiB208 linesraw
1#include "models.h"
2
3#include <type_traits>
4
5void llama_model_aliceai_t5_moe::load_arch_hparams(llama_model_loader & ml) {
6    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
7    ml.get_key(LLM_KV_DECODER_BLOCK_COUNT, hparams.dec_n_layer);
8    uint32_t decoder_start_token_id;
9    ml.get_key(LLM_KV_DECODER_START_TOKEN_ID, decoder_start_token_id);
10    hparams.dec_start_token_id = decoder_start_token_id;
11    ml.get_key(LLM_KV_ENCODER_ATTENTION_HEAD_COUNT_KV, n_head_kv_enc);
12    ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
13    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);
14    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm);
15    ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
16
17    if (hparams.dec_n_layer == 0 || hparams.dec_n_layer > hparams.n_layer()) {
18        throw std::runtime_error("AliceAI requires 0 < decoder_block_count <= block_count");
19    }
20    if (hparams.n_head() == 0 || n_head_kv_enc == 0 || hparams.n_head() % n_head_kv_enc != 0 ||
21            hparams.n_head_kv() == 0 || hparams.n_head() % hparams.n_head_kv() != 0) {
22        throw std::runtime_error("AliceAI KV head counts must divide the query head count");
23    }
24    if (hparams.n_expert == 0 || hparams.n_ff_exp() == 0 ||
25            hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID) {
26        throw std::runtime_error("AliceAI requires sigmoid routing and nonempty experts");
27    }
28
29    if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_YARN) {
30        ml.get_key(LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR, hparams.rope_attn_factor);
31        // GGML applies the default YaRN magnitude inside RoPE; metadata gives the total magnitude.
32        hparams.rope_attn_factor /= 1.0f + 0.1f * logf(1.0f / hparams.rope_freq_scale_train);
33    }
34
35    // The published causal flag describes the encoder. llama_encode disables causality separately.
36    hparams.causal_attn = true;
37    type = LLM_TYPE_UNKNOWN;
38}
39
40void llama_model_aliceai_t5_moe::load_arch_tensors(llama_model_loader &) {
41    LLAMA_LOAD_LOCALS;
42
43    if (hparams.dec_start_token_id < 0 || hparams.dec_start_token_id >= n_vocab) {
44        throw std::runtime_error("AliceAI decoder start token is outside the vocabulary");
45    }
46    if (n_embd_head_k != n_embd_head_v || n_embd_head_k == 0) {
47        throw std::runtime_error("AliceAI requires equal, nonzero key and value head dimensions");
48    }
49
50    const int64_t n_ff_exp = hparams.n_ff_exp();
51    const int64_t n_embd_q = n_embd_head_k * n_head;
52    const int64_t n_embd_k_enc = n_embd_head_k * n_head_kv_enc;
53    const int64_t n_embd_v_enc = n_embd_head_v * n_head_kv_enc;
54
55    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
56    output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), {n_embd}, 0);
57    output_norm     = create_tensor(tn(LLM_TENSOR_DEC_OUTPUT_NORM, "weight"), {n_embd}, 0);
58    output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
59    if (output == nullptr) {
60        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
61    }
62
63    layers_enc.resize(n_layer);
64    for (int i = 0; i < n_layer; ++i) {
65        auto & layer = layers_enc[i];
66
67        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ENC_ATTN_NORM, "weight", i), {n_embd}, 0);
68        layer.wq = create_tensor(tn(LLM_TENSOR_ENC_ATTN_Q, "weight", i), {n_embd, n_embd_q}, 0);
69        layer.wk = create_tensor(tn(LLM_TENSOR_ENC_ATTN_K, "weight", i), {n_embd, n_embd_k_enc}, 0);
70        layer.wv = create_tensor(tn(LLM_TENSOR_ENC_ATTN_V, "weight", i), {n_embd, n_embd_v_enc}, 0);
71        layer.wo = create_tensor(tn(LLM_TENSOR_ENC_ATTN_OUT, "weight", i), {n_embd_head_v * n_head, n_embd}, 0);
72
73        layer.ffn_norm      = create_tensor(tn(LLM_TENSOR_ENC_FFN_NORM, "weight", i), {n_embd}, 0);
74        layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_ENC_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
75        layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_ENC_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
76        layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_ENC_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
77        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_ENC_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
78        layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_ENC_EXP_PROBS_B, i), {n_expert}, 0);
79        layer.ffn_down_b      = create_tensor(tn(LLM_TENSOR_ENC_FFN_OUTPUT_B, i), {n_embd}, 0);
80    }
81
82    for (uint32_t i = 0; i < hparams.dec_n_layer; ++i) {
83        auto & layer = layers[i];
84
85        layer.attn_norm = create_tensor(tn(LLM_TENSOR_DEC_ATTN_NORM, "weight", i), {n_embd}, 0);
86        layer.wq = create_tensor(tn(LLM_TENSOR_DEC_ATTN_Q, "weight", i), {n_embd, n_embd_q}, 0);
87        layer.wk = create_tensor(tn(LLM_TENSOR_DEC_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0);
88        layer.wv = create_tensor(tn(LLM_TENSOR_DEC_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0);
89        layer.wo = create_tensor(tn(LLM_TENSOR_DEC_ATTN_OUT, "weight", i), {n_embd_head_v * n_head, n_embd}, 0);
90
91        layer.attn_norm_cross = create_tensor(tn(LLM_TENSOR_DEC_CROSS_ATTN_NORM, "weight", i), {n_embd}, 0);
92        layer.wq_cross = create_tensor(tn(LLM_TENSOR_DEC_CROSS_ATTN_Q, "weight", i), {n_embd, n_embd_q}, 0);
93        layer.wk_cross = create_tensor(tn(LLM_TENSOR_DEC_CROSS_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0);
94        layer.wv_cross = create_tensor(tn(LLM_TENSOR_DEC_CROSS_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0);
95        layer.wo_cross = create_tensor(tn(LLM_TENSOR_DEC_CROSS_ATTN_OUT, "weight", i), {n_embd_head_v * n_head, n_embd}, 0);
96
97        layer.ffn_norm      = create_tensor(tn(LLM_TENSOR_DEC_FFN_NORM, "weight", i), {n_embd}, 0);
98        layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_DEC_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
99        layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_DEC_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
100        layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_DEC_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
101        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_DEC_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
102        layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_DEC_EXP_PROBS_B, i), {n_expert}, 0);
103        layer.ffn_down_b      = create_tensor(tn(LLM_TENSOR_DEC_FFN_OUTPUT_B, i), {n_embd}, 0);
104    }
105}
106
107std::unique_ptr<llm_graph_context> llama_model_aliceai_t5_moe::build_arch_graph(const llm_graph_params & params) const {
108    switch (params.gtype) {
109        case LLM_GRAPH_TYPE_ENCODER:
110            return std::make_unique<graph<true>>(*this, params);
111        case LLM_GRAPH_TYPE_DEFAULT:
112        case LLM_GRAPH_TYPE_DECODER:
113            return std::make_unique<graph<false>>(*this, params);
114        default:
115            GGML_ABORT("invalid graph type");
116    }
117}
118
119template <bool is_enc>
120llama_model_aliceai_t5_moe::graph<is_enc>::graph(const llama_model_aliceai_t5_moe & model, const llm_graph_params & params) : llm_graph_context(params) {
121    const int64_t n_embd_head = hparams.n_embd_head_k();
122    const int64_t n_head_kv_self = is_enc ? model.n_head_kv_enc : n_head_kv;
123    const int n_layers = is_enc ? n_layer : hparams.dec_n_layer;
124    const auto & layers = is_enc ? model.layers_enc : model.layers;
125    const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
126
127    ggml_tensor * inpL = build_inp_embd(model.tok_embd);
128    ggml_tensor * inp_pos = build_inp_pos();
129    ggml_tensor * inp_out_ids = is_enc ? nullptr : build_inp_out_ids();
130    ggml_tensor * embd_enc = nullptr;
131    llm_graph_input_attn_cross * inp_cross = nullptr;
132
133    using inp_attn_type = std::conditional_t<is_enc, llm_graph_input_attn_no_cache, llm_graph_input_attn_kv>;
134    inp_attn_type * inp_attn;
135    if constexpr (is_enc) {
136        inp_attn = build_attn_inp_no_cache();
137    } else {
138        inp_attn = build_attn_inp_kv();
139        embd_enc = build_inp_cross_embd();
140        inp_cross = build_attn_inp_cross();
141    }
142
143    for (int il = 0; il < n_layers; ++il) {
144        const auto & layer = layers[il];
145        ggml_tensor * residual = inpL;
146        ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
147        cb(cur, "attn_norm", il);
148
149        auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv_self, il);
150        Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
151                ext_factor, attn_factor, beta_fast, beta_slow);
152        Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
153                ext_factor, attn_factor, beta_fast, beta_slow);
154        cb(Qcur, "Qcur", il);
155        cb(Kcur, "Kcur", il);
156
157        cur = build_attn(inp_attn, layer.wo, nullptr, nullptr,
158                Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
159        cur = ggml_add(ctx0, cur, residual);
160        cb(cur, "attn_out", il);
161
162        if constexpr (!is_enc) {
163            residual = cur;
164            cur = build_norm(cur, layer.attn_norm_cross, nullptr, LLM_NORM_RMS, il);
165            cb(cur, "attn_norm_cross", il);
166
167            Qcur = build_lora_mm(layer.wq_cross, cur);
168            Kcur = build_lora_mm(layer.wk_cross, embd_enc);
169            Vcur = build_lora_mm(layer.wv_cross, embd_enc);
170            Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
171            Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, embd_enc->ne[1]);
172            Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, embd_enc->ne[1]);
173            cur = build_attn(inp_cross, layer.wo_cross, nullptr, nullptr,
174                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
175            cur = ggml_add(ctx0, cur, residual);
176            cb(cur, "cross_attn_out", il);
177        }
178
179        if (il == n_layers - 1 && inp_out_ids) {
180            cur = ggml_get_rows(ctx0, cur, inp_out_ids);
181        }
182        residual = cur;
183        cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
184        cb(cur, "ffn_norm", il);
185        cur = build_moe_ffn(cur, layer.ffn_gate_inp, layer.ffn_up_exps, layer.ffn_gate_exps, layer.ffn_down_exps,
186                layer.ffn_exp_probs_b, n_expert, hparams.n_expert_used(il), LLM_FFN_SILU,
187                hparams.expert_weights_norm, hparams.expert_weights_scale, LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID, il);
188        cur = ggml_add(ctx0, cur, layer.ffn_down_b);
189        cb(cur, "ffn_out", il);
190        cur = ggml_add(ctx0, cur, residual);
191        inpL = build_cvec(cur, il);
192        cb(inpL, "l_out", il);
193    }
194
195    ggml_tensor * cur = build_norm(inpL, is_enc ? model.output_norm_enc : model.output_norm, nullptr, LLM_NORM_RMS, -1);
196    cb(cur, "result_norm", -1);
197    res->t_embd = cur;
198
199    if constexpr (!is_enc) {
200        cur = build_lora_mm(model.output, cur);
201        cb(cur, "result_output", -1);
202        res->t_logits = cur;
203    }
204    ggml_build_forward_expand(gf, cur);
205}
206
207template struct llama_model_aliceai_t5_moe::graph<true>;
208template struct llama_model_aliceai_t5_moe::graph<false>;