char-fork/llama.cpp

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

Yaniss Amazouzmodel: add NVIDIA Nemotron-3-Puzzle-75B-A9B (NemotronHPuzzle) support (#25444)c61b98b87

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
11.1 KiB274 linesraw
1#include "models.h"
2
3void llama_model_llama4::load_arch_hparams(llama_model_loader & ml) {
4    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
5    ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
6    ml.get_key(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,   hparams.n_moe_layer_step);
7
8    const bool found_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
9    if (found_swa && hparams.n_swa == 0) {
10        hparams.swa_type             = LLAMA_SWA_TYPE_NONE;
11        hparams.n_no_rope_layer_step = hparams.n_layer(); // always use rope
12    } else {
13        hparams.swa_type                = LLAMA_SWA_TYPE_CHUNKED;
14        hparams.n_swa                   = 8192;
15        hparams.n_attn_temp_floor_scale = 8192;
16        hparams.f_attn_temp_scale       = 0.1f;
17        hparams.f_attn_temp_offset      = 1.0f;
18
19        uint32_t swa_period = 4; // pattern: 3 chunked - 1 full
20        ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false);
21        hparams.set_swa_pattern(swa_period);
22
23        hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;
24        hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
25        ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
26    }
27
28    switch (hparams.n_expert) {
29        case 0: {
30            // MobileLLM (no MoE)
31            switch (hparams.n_embd) {
32                case 2048: type = LLM_TYPE_140M; break;
33                case 4096: type = LLM_TYPE_360M; break;
34                case 6144: type = LLM_TYPE_950M; break;
35                default:   type = LLM_TYPE_UNKNOWN;
36            }
37        } break;
38        case 16:  type = LLM_TYPE_17B_16E; break;
39        case 128: type = LLM_TYPE_17B_128E; break;
40        default:  type = LLM_TYPE_UNKNOWN;
41    }
42
43    hparams.use_kq_norm = type != LLM_TYPE_17B_128E;
44}
45
46void llama_model_llama4::load_arch_tensors(llama_model_loader &) {
47    LLAMA_LOAD_LOCALS;
48
49    if (n_expert == 0) {
50        throw std::runtime_error(arch_name() + " model cannot have zero experts");
51    }
52    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
53
54    // output
55    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
56    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
57
58    // if output is NULL, init from the input tok embed
59    if (output == NULL) {
60        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
61    }
62
63    for (int i = 0; i < n_layer; ++i) {
64        const bool is_moe_layer = hparams.n_moe_layer_step > 0 && (i + 1) % hparams.n_moe_layer_step == 0;
65
66        auto & layer = layers[i];
67
68        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
69
70        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
71        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
72
73        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
74
75        layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
76
77        if (is_moe_layer) {
78            const int64_t n_ff_exp = hparams.n_ff_exp();
79
80            layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), {n_embd, n_expert}, 0);
81            layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd,   n_ff_exp, n_expert}, 0);
82            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {  n_ff_exp, n_embd, n_expert}, 0);
83            layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd,   n_ff_exp, n_expert}, 0);
84
85            // Shared expert
86            const int64_t n_ff_shexp = n_ff_exp;
87            layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {    n_embd, n_ff_shexp}, 0);
88            layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd    }, 0);
89            layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {    n_embd, n_ff_shexp}, 0);
90        } else {
91            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);
92            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);
93            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
94        }
95    }
96}
97
98std::unique_ptr<llm_graph_context> llama_model_llama4::build_arch_graph(const llm_graph_params & params) const {
99    if (hparams.swa_type == LLAMA_SWA_TYPE_NONE) {
100        return std::make_unique<graph<false>>(*this, params);
101    } else {
102        return std::make_unique<graph<true>>(*this, params);
103    }
104}
105
106template <bool iswa>
107llama_model_llama4::graph<iswa>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
108    const int64_t n_embd_head = hparams.n_embd_head_v();
109
110    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
111    GGML_ASSERT(n_embd_head == n_rot);
112
113    ggml_tensor * cur;
114    ggml_tensor * inpL;
115
116    inpL = build_inp_embd(model.tok_embd);
117
118    // inp_pos - contains the positions
119    ggml_tensor * inp_pos = build_inp_pos();
120
121    // temperature tuning
122    ggml_tensor * inp_attn_scale = nullptr;
123    inp_attn_scale = build_inp_attn_scale();
124
125    using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;
126    inp_attn_type * inp_attn = nullptr;
127
128    if constexpr (iswa) {
129        inp_attn = build_attn_inp_kv_iswa();
130    } else {
131        inp_attn = build_attn_inp_kv();
132    }
133
134    const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
135
136    ggml_tensor * inp_out_ids = build_inp_out_ids();
137
138    for (int il = 0; il < n_layer; ++il) {
139        const float freq_base_l  = model.get_rope_freq_base (cparams, il);
140        const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
141
142        ggml_tensor * inpSA = inpL;
143
144        // This overlaps with SWA layers in current models, so get_rope_freq_base/scale may be superfluous
145        const bool use_rope = hparams.n_no_rope_layer_step > 0 &&
146                              (il + 1) % hparams.n_no_rope_layer_step != 0;
147
148        // norm
149        cur = build_norm(inpL,
150                model.layers[il].attn_norm, NULL,
151                LLM_NORM_RMS, il);
152        cb(cur, "attn_norm", il);
153
154        // self-attention
155        {
156            // rope freq factors for llama3; may return nullptr for llama2 and other models
157            ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
158
159            // compute Q and K and RoPE them
160            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
161                    n_embd_head, n_head, n_head_kv, il);
162
163            if (use_rope) {
164                Qcur = ggml_rope_ext(
165                        ctx0, Qcur, inp_pos, rope_factors,
166                        n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
167                        ext_factor, attn_factor, beta_fast, beta_slow
168                        );
169
170                Kcur = ggml_rope_ext(
171                        ctx0, Kcur, inp_pos, rope_factors,
172                        n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
173                        ext_factor, attn_factor, beta_fast, beta_slow
174                        );
175            } else if (inp_attn_scale) {
176                Qcur = ggml_mul(ctx0, Qcur, inp_attn_scale);
177            }
178            cb(Qcur, "Qcur", il);
179            cb(Kcur, "Kcur", il);
180            cb(Vcur, "Vcur", il);
181
182            if (use_rope && hparams.use_kq_norm) {
183                // Llama4TextL2Norm
184                Qcur = ggml_rms_norm(ctx0, Qcur, hparams.f_norm_rms_eps);
185                Kcur = ggml_rms_norm(ctx0, Kcur, hparams.f_norm_rms_eps);
186                cb(Qcur, "Qcur_normed", il);
187                cb(Kcur, "Kcur_normed", il);
188            }
189            cur = build_attn(inp_attn,
190                    model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,
191                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
192            cb(cur, "attn_out", il);
193        }
194        if (il == n_layer - 1 && inp_out_ids) {
195            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
196            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
197        }
198        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
199        cb(ffn_inp, "ffn_inp", il);
200
201        // feed-forward network (non-MoE)
202        if (model.layers[il].ffn_gate_inp == nullptr) {
203            cur = build_norm(ffn_inp,
204                    model.layers[il].ffn_norm, NULL,
205                    LLM_NORM_RMS, il);
206            cb(cur, "ffn_norm", il);
207
208            cur = build_ffn(cur,
209                    model.layers[il].ffn_up,   model.layers[il].ffn_up_b,   NULL,
210                    model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL,
211                    model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL,
212                    NULL,
213                    LLM_FFN_SILU, LLM_FFN_PAR, il);
214            cb(cur, "ffn_out", il);
215        } else {
216            ggml_tensor * ffn_inp_normed = build_norm(ffn_inp,
217                    model.layers[il].ffn_norm, NULL,
218                    LLM_NORM_RMS, il);
219            cb(cur, "ffn_norm", il);
220
221            ggml_tensor * moe_out = build_moe_ffn(ffn_inp_normed,
222                    model.layers[il].ffn_gate_inp,
223                    model.layers[il].ffn_up_exps,
224                    model.layers[il].ffn_gate_exps,
225                    model.layers[il].ffn_down_exps,
226                    nullptr,
227                    n_expert, n_expert_used,
228                    LLM_FFN_SILU, false,
229                    hparams.expert_weights_scale,
230                    LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,
231                    il);
232
233            // Shared experts
234            ggml_tensor * shexp_out = build_ffn(ffn_inp_normed,
235                model.layers[il].ffn_up_shexp,   NULL, NULL,
236                model.layers[il].ffn_gate_shexp, NULL, NULL,
237                model.layers[il].ffn_down_shexp, NULL, NULL,
238                NULL,
239                LLM_FFN_SILU, LLM_FFN_PAR, il);
240            cb(shexp_out, "ffn_moe_shexp", il);
241
242            cur = ggml_add(ctx0, moe_out, shexp_out);
243            cb(cur, "ffn_moe_out_merged", il);
244        }
245        cur = ggml_add(ctx0, cur, ffn_inp);
246        cb(cur, "ffn_out", il);
247
248        cur = build_cvec(cur, il);
249        cb(cur, "l_out", il);
250
251        // input for next layer
252        inpL = cur;
253    }
254    cur = inpL;
255
256    cur = build_norm(cur,
257            model.output_norm, NULL,
258            LLM_NORM_RMS, -1);
259
260    cb(cur, "result_norm", -1);
261    res->t_embd = cur;
262
263    // lm_head
264    cur = build_lora_mm(model.output, cur, model.output_s);
265
266    cb(cur, "result_output", -1);
267    res->t_logits = cur;
268
269    ggml_build_forward_expand(gf, cur);
270}
271
272// Explicit template instantiations
273template struct llama_model_llama4::graph<false>;
274template struct llama_model_llama4::graph<true>;