char-fork/llama.cpp

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

Xuan-Son Nguyenfix: check gguf array type before reading (#27075)27df9199d

char/aliceai-t5-moe
35.4 KiB959 linesraw
1#include "ggml.h"
2#include "gguf.h"
3
4#include "llama.h"
5#include "common.h"
6#include "log.h"
7
8#include <algorithm>
9#include <cassert>
10#include <cinttypes>
11#include <climits>
12#include <clocale>
13#include <cstdarg>
14#include <cstring>
15#include <ctime>
16#include <unordered_map>
17#include <vector>
18#include <random>
19#include <stdexcept>
20#include <sstream>
21#include <algorithm>
22#include <string>
23
24// GGUF keys & tensor names.
25
26#define KV_GENERAL_ARCHITECTURE          "general.architecture"
27#define KV_GENERAL_NAME                  "general.name"
28
29#define KV_TOKENIZER_MODEL               "tokenizer.ggml.model"
30#define KV_TOKENIZER_LIST                "tokenizer.ggml.tokens"
31#define KV_TOKENIZER_TOKEN_TYPE          "tokenizer.ggml.token_type"
32#define KV_TOKENIZER_SCORES              "tokenizer.ggml.scores"
33#define KV_TOKENIZER_BOS_ID              "tokenizer.ggml.bos_token_id"
34#define KV_TOKENIZER_EOS_ID              "tokenizer.ggml.eos_token_id"
35#define KV_TOKENIZER_UNK_ID              "tokenizer.ggml.unknown_token_id"
36#define KV_TOKENIZER_SEP_ID              "tokenizer.ggml.seperator_token_id"
37#define KV_TOKENIZER_PAD_ID              "tokenizer.ggml.padding_token_id"
38#define KV_TOKENIZER_HF_JSON             "tokenizer.huggingface.json"
39
40#define KV_CONTEXT_LENGTH                "llama.context_length"
41#define KV_EMBEDDING_LENGTH              "llama.embedding_length"
42#define KV_BLOCK_COUNT                   "llama.block_count"
43#define KV_FEED_FORWARD_LENGTH           "llama.feed_forward_length"
44#define KV_ATTENTION_HEAD_COUNT          "llama.attention.head_count"
45#define KV_ATTENTION_HEAD_COUNT_KV       "llama.attention.head_count_kv"
46#define KV_ATTENTION_LAYERNORM_RMS_EPS   "llama.attention.layer_norm_rms_epsilon"
47#define KV_ROPE_DIMENSION_COUNT          "llama.rope.dimension_count"
48
49#define TN_TOKEN_EMBD  "token_embd.weight"
50#define TN_OUTPUT_NORM "output_norm.weight"
51#define TN_OUTPUT      "output.weight"
52#define TN_ATTN_NORM   "blk.%d.attn_norm.weight"
53#define TN_ATTN_Q      "blk.%d.attn_q.weight"
54#define TN_ATTN_K      "blk.%d.attn_k.weight"
55#define TN_ATTN_V      "blk.%d.attn_v.weight"
56#define TN_ATTN_OUTPUT "blk.%d.attn_output.weight"
57#define TN_FFN_NORM    "blk.%d.ffn_norm.weight"
58#define TN_FFN_GATE    "blk.%d.ffn_gate.weight"
59#define TN_FFN_DOWN    "blk.%d.ffn_down.weight"
60#define TN_FFN_UP      "blk.%d.ffn_up.weight"
61
62#if defined(_MSC_VER)
63#pragma warning(disable: 4244 4267) // possible loss of data
64#endif
65
66#define LLAMA_FILE_MAGIC_GGJT        0x67676a74u // 'ggjt'
67#define LLAMA_FILE_VERSION_GGJT_V3   3
68
69#define TOKENIZER_NAME "llama"
70#define UNKNOWN_TOKEN_ID 0
71#define BOS_TOKEN_ID 1
72#define EOS_TOKEN_ID 2
73
74//////////////////////////////////////// llama2.c model structs and functions to load models, alloc memory etc.
75typedef struct {
76    int dim; // transformer dimension
77    int hidden_dim; // for ffn layers
78    int n_layers; // number of layers
79    int n_heads; // number of query heads
80    int n_kv_heads; // number of key/value heads (can be < query heads because of multiquery)
81    int vocab_size; // vocabulary size, usually 256 (byte-level)
82    int seq_len; // max sequence length
83} Config;
84
85struct TransformerWeights {
86    // token embedding table
87    std::vector<float> token_embedding_table;    // (vocab_size, dim)
88    // weights for rmsnorms
89    std::vector<float> rms_att_weight; // (layer, dim) rmsnorm weights
90    std::vector<float> rms_ffn_weight; // (layer, dim)
91    // weights for matmuls
92    std::vector<float> wq; // (layer, dim, dim)
93    std::vector<float> wk; // (layer, dim, dim)
94    std::vector<float> wv; // (layer, dim, dim)
95    std::vector<float> wo; // (layer, dim, dim)
96    // weights for ffn
97    std::vector<float> w1; // (layer, hidden_dim, dim)
98    std::vector<float> w2; // (layer, dim, hidden_dim)
99    std::vector<float> w3; // (layer, hidden_dim, dim)
100    // final rmsnorm
101    std::vector<float> rms_final_weight; // (dim,)
102    // freq_cis for RoPE relatively positional embeddings
103    // std::vector<float> freq_cis_real; // (seq_len, dim/2)
104    // std::vector<float> freq_cis_imag; // (seq_len, dim/2)
105    // (optional) classifier weights for the logits, on the last layer
106    std::vector<float> wcls;
107};
108
109static void alloc_weights(TransformerWeights * w, const Config * p, bool shared_weights) {
110    const int n_multiqueries = p->n_kv_heads <= 0 || p->n_kv_heads >= p->n_heads ? 1 : p->n_heads / p->n_kv_heads;
111    try {
112        w->token_embedding_table.resize(p->vocab_size * p->dim);
113        LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->token_embedding_table\n",__func__,p->vocab_size , p->dim, p->vocab_size * p->dim);
114
115        w->rms_att_weight.resize(p->n_layers * p->dim);
116        LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->rms_att_weight\n",__func__,p->n_layers, p->dim, p->n_layers * p->dim);
117
118        w->rms_ffn_weight.resize(p->n_layers * p->dim);
119        LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->rms_ffn_weight\n",__func__,p->n_layers , p->dim, p->n_layers * p->dim);
120
121        w->wq.resize(p->n_layers * p->dim * p->dim);
122        LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wq\n",__func__,p->n_layers, p->dim, p->dim, p->n_layers * p->dim * p->dim);
123
124        w->wk.resize(p->n_layers * p->dim * p->dim / n_multiqueries);
125        LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wk\n",__func__,p->n_layers, p->dim, p->dim / n_multiqueries, p->n_layers * p->dim * p->dim / n_multiqueries);
126
127        w->wv.resize(p->n_layers * p->dim * p->dim / n_multiqueries);
128        LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wv\n",__func__, p->n_layers, p->dim, p->dim / n_multiqueries, p->n_layers * p->dim * p->dim / n_multiqueries);
129
130        w->wo.resize(p->n_layers * p->dim * p->dim);
131        LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wo\n",__func__,p->n_layers, p->dim, p->dim, p->n_layers * p->dim * p->dim);
132
133        w->w1.resize(p->n_layers * p->hidden_dim * p->dim);
134        LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->w1\n",__func__,p->n_layers, p->hidden_dim, p->dim, p->n_layers * p->hidden_dim * p->dim);
135
136        w->w2.resize(p->n_layers * p->hidden_dim * p->dim);
137        LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->w2\n",__func__,p->n_layers, p->dim, p->hidden_dim, p->n_layers * p->hidden_dim * p->dim);
138
139        w->w3.resize(p->n_layers * p->hidden_dim * p->dim);
140        LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->w3\n",__func__,p->n_layers, p->hidden_dim, p->dim, p->n_layers * p->hidden_dim * p->dim);
141
142        w->rms_final_weight.resize(p->dim);
143        LOG_INF("%s: Allocating [%d] float space for w->rms_final_weight\n",__func__,p->dim);
144
145        if (shared_weights) {
146            w->wcls = {};
147        } else {
148            w->wcls.resize(p->vocab_size * p->dim);
149            LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->wcls\n",__func__,p->vocab_size , p->dim, p->vocab_size * p->dim);
150        }
151    }
152    catch (std::length_error &) {
153        die("Invalid configuration. Failed to allocate memory for weights");
154    }
155}
156
157static int checkpoint_init_weights(TransformerWeights * w, const Config * p, FILE * f, bool shared_weights) {
158    if (fread(w->token_embedding_table.data(), sizeof(float), w->token_embedding_table.size(), f) != w->token_embedding_table.size()) return 1;
159    if (fread(w->rms_att_weight.data(), sizeof(float), w->rms_att_weight.size(), f) != w->rms_att_weight.size()) return 1;
160    if (fread(w->wq.data(), sizeof(float), w->wq.size(), f) != w->wq.size()) return 1;
161    if (fread(w->wk.data(), sizeof(float), w->wk.size(), f) != w->wk.size()) return 1;
162    if (fread(w->wv.data(), sizeof(float), w->wv.size(), f) != w->wv.size()) return 1;
163    if (fread(w->wo.data(), sizeof(float), w->wo.size(), f) != w->wo.size()) return 1;
164    if (fread(w->rms_ffn_weight.data(), sizeof(float), w->rms_ffn_weight.size(), f) != w->rms_ffn_weight.size()) return 1;
165    if (fread(w->w1.data(), sizeof(float), w->w1.size(), f) != w->w1.size()) return 1;
166    if (fread(w->w2.data(), sizeof(float), w->w2.size(), f) != w->w2.size()) return 1;
167    if (fread(w->w3.data(), sizeof(float), w->w3.size(), f) != w->w3.size()) return 1;
168    if (fread(w->rms_final_weight.data(), sizeof(float), w->rms_final_weight.size(), f) != w->rms_final_weight.size()) return 1;
169
170    // Skip freq_cis_real & freq_cis_imag
171    int head_size = p->dim / p->n_heads;
172    fseek(f, p->seq_len * head_size * sizeof(float), SEEK_CUR);
173
174    if (!shared_weights && fread(w->wcls.data(), sizeof(float), w->wcls.size(), f) != w->wcls.size()) return 1;
175
176    // Check we didn't forget to read anything
177    auto curr = ftell(f);
178    fseek(f, 0, SEEK_END);
179    auto end = ftell(f);
180    if (curr != end) {
181        LOG_ERR("%s: Error: failed to read the checkpoint file to the end (curr = %ld, end =  %ld)\n", __func__, curr, end);
182        return 1;
183    }
184
185    return 0;
186}
187
188static void print_sample_weights(TransformerWeights *w){
189    LOG_INF("----- Quick print of first of the weight vales of all the variables\n");
190    LOG_INF("%f\n", w->token_embedding_table[0]);
191    LOG_INF("%f\n", w->rms_att_weight[0]);
192    LOG_INF("%f\n", w->rms_ffn_weight[0]);
193
194    LOG_INF("%f\n", w->wq[0]);
195    LOG_INF("%f\n", w->wk[0]);
196    LOG_INF("%f\n", w->wv[0]);
197    LOG_INF("%f\n", w->wo[0]);
198    LOG_INF("%f\n", w->w1[0]);
199    LOG_INF("%f\n", w->w2[0]);
200    LOG_INF("%f\n", w->w3[0]);
201    LOG_INF("%f\n", w->rms_att_weight[0]);
202    if (!w->wcls.empty()) LOG_INF("%f\n", w->wcls[0]);
203}
204////////////////////////////////////////////////////////////////////////////////////////////////////////////
205
206//////////////////////////////////////// ggml structs and functions required to load models, configs and save the model.
207
208struct my_llama_vocab {
209    using id    = int32_t;
210    using token = std::string;
211    using ttype = llama_token_type;
212
213    struct token_data {
214        token text;
215        float score;
216        ttype type;
217    };
218
219    std::unordered_map<token, id> token_to_id;
220    std::vector<token_data> id_to_token;
221};
222
223struct my_llama_hparams {
224    uint32_t n_vocab   = 32000;
225    uint32_t n_ctx     = 512;   // this is provided as user input?
226    uint32_t n_embd    = 4096;
227    uint32_t n_ff      = 11008;
228    uint32_t n_mult    = 4;
229    uint32_t n_head    = 32;
230    uint32_t n_head_kv = 32;
231    uint32_t n_layer   = 32;
232    uint32_t n_rot     = 64;
233
234    bool operator!=(const my_llama_hparams& other) const {
235        return memcmp(this, &other, sizeof(my_llama_hparams));
236    }
237};
238
239struct my_llama_layer {
240    // normalization
241    struct ggml_tensor * attention_norm;
242
243    // attention
244    struct ggml_tensor * wq;
245    struct ggml_tensor * wk;
246    struct ggml_tensor * wv;
247    struct ggml_tensor * wo;
248
249    // normalization
250    struct ggml_tensor * ffn_norm;
251
252    // ff
253    struct ggml_tensor * w1;
254    struct ggml_tensor * w2;
255    struct ggml_tensor * w3;
256};
257
258struct my_llama_model {
259    struct ggml_context * ctx = NULL;
260
261    std::string name;
262
263    my_llama_hparams hparams;
264
265    struct ggml_tensor * tok_embeddings;
266
267    struct ggml_tensor * norm;
268    struct ggml_tensor * output;
269
270    std::vector<my_llama_layer> layers;
271
272    uint32_t train_its = 0;
273    uint32_t train_samples = 0;
274    uint32_t train_tokens = 0;
275};
276
277struct train_params {
278    const char * fn_vocab_model;
279    const char * fn_llama2c_model;
280    const char * fn_llama2c_output_model;
281    const char * fn_train_data;
282    const char * fn_checkpoint_in;
283    const char * fn_checkpoint_out;
284    const char * fn_model_out;
285
286    uint32_t seed;
287
288    int n_ctx;
289    int n_embd;
290    int n_mult;
291    int n_head;
292    int n_layer;
293    int n_rotmax;
294
295    int n_threads;
296    int n_batch;
297    int n_examples;
298    int n_predict;
299
300    int print_info_interval;
301    int print_details_interval;
302
303    bool samples_start_after_nl;
304    bool use_adam;
305    bool use_flash;
306    bool use_scratch;
307
308    // only adam
309    int   warmup;
310    int   cos_decay_steps;
311    float cos_decay_restart;
312    float cos_decay_alpha;
313
314    int   lbfgs_n_iter;
315    int   adam_n_iter;
316    float adam_alpha;
317    float adam_decay;
318
319    int mem_model_gb;
320    int mem_compute_gb;
321    int mem_compute0_gb;
322    int mem_compute1_gb;
323};
324
325static void print_params(struct my_llama_hparams * params) {
326    LOG_INF("%s: n_vocab:   %u\n", __func__, params->n_vocab);
327    LOG_INF("%s: n_ctx:     %u\n", __func__, params->n_ctx);
328    LOG_INF("%s: n_embd:    %u\n", __func__, params->n_embd);
329    LOG_INF("%s: n_mult:    %u\n", __func__, params->n_mult);
330    LOG_INF("%s: n_head:    %u\n", __func__, params->n_head);
331    LOG_INF("%s: n_head_kv: %u\n", __func__, params->n_head_kv);
332    LOG_INF("%s: n_ff:      %u\n", __func__, params->n_ff);
333    LOG_INF("%s: n_layer:   %u\n", __func__, params->n_layer);
334    LOG_INF("%s: n_rot:     %u\n", __func__, params->n_rot);
335}
336
337static void print_tensor_info(const struct ggml_context * ctx) {
338    for (auto * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
339        LOG_INF("%s: Allocating ", __func__);
340        int64_t total = 1;
341        int i = 0;
342        for (; i < ggml_n_dims(t); ++i) {
343            if (i > 0) { LOG_INF("x "); }
344            LOG_INF("[%" PRId64 "] ", t->ne[i]);
345            total *= t->ne[i];
346        }
347        if (i > 1) { LOG_INF("= [%" PRId64 "] ", total); }
348        LOG_INF("float space for %s\n", ggml_get_name(t));
349    }
350}
351
352static void init_model(struct my_llama_model * model) {
353    const auto & hparams = model->hparams;
354
355    const uint32_t n_embd  = hparams.n_embd;
356    const uint32_t n_layer = hparams.n_layer;
357    const uint32_t n_vocab = hparams.n_vocab;
358
359    const uint32_t n_multiqueries = hparams.n_head_kv <= 0 || hparams.n_head_kv >= hparams.n_head ? 1 : hparams.n_head / hparams.n_head_kv;
360
361    const uint32_t n_ff = hparams.n_ff;
362    struct ggml_context * ctx = model->ctx;
363
364    model->train_its = 0;
365    model->train_samples = 0;
366    model->train_tokens = 0;
367
368    model->tok_embeddings = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_vocab);
369    model->norm           = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
370    model->output         = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_vocab);
371
372    ggml_set_name(model->tok_embeddings, "tok_embeddings.weight");
373    ggml_set_name(model->norm,           "norm.weight");
374    ggml_set_name(model->output,         "output.weight");
375
376    model->layers.resize(n_layer);
377    for (uint32_t i = 0; i < n_layer; ++i) {
378        auto & layer = model->layers[i];
379
380        std::string layers_i = "layers." + std::to_string(i);
381
382        layer.attention_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
383
384        layer.wq = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd);
385        layer.wk = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd / n_multiqueries);
386        layer.wv = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd / n_multiqueries);
387        layer.wo = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd);
388
389        layer.ffn_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
390
391        layer.w1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_ff);
392        layer.w2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_ff, n_embd);
393        layer.w3 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_ff);
394
395        ggml_set_name(layer.attention_norm, (layers_i + ".attention_norm.weight").c_str());
396
397        ggml_set_name(layer.wq, (layers_i + ".attention.wq.weight").c_str());
398        ggml_set_name(layer.wk, (layers_i + ".attention.wk.weight").c_str());
399        ggml_set_name(layer.wv, (layers_i + ".attention.wv.weight").c_str());
400        ggml_set_name(layer.wo, (layers_i + ".attention.wo.weight").c_str());
401
402        ggml_set_name(layer.ffn_norm, (layers_i + ".ffn_norm.weight").c_str());
403
404        ggml_format_name(layer.w1, "%s.feed_forward.w1.weight", layers_i.c_str());
405        ggml_format_name(layer.w2, "%s.feed_forward.w2.weight", layers_i.c_str());
406        ggml_format_name(layer.w3, "%s.feed_forward.w3.weight", layers_i.c_str());
407    }
408
409    print_tensor_info(ctx);
410}
411
412static float get_f32_2d(struct ggml_tensor * tensor, int64_t i0, int64_t i1) {
413    float * ptr = (float *) ((char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1]);
414    return *ptr;
415}
416
417static int32_t get_i32_2d(struct ggml_tensor * tensor, int64_t i0, int64_t i1) {
418    int32_t * ptr = (int32_t *) ((char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1]);
419    return *ptr;
420}
421
422static void print_row(struct ggml_tensor * probs, int i) {
423    for (int k = 0; k < probs->ne[0]; ++k) {
424        float p = get_f32_2d(probs, k, i);
425        LOG(" %f", p);
426    }
427    LOG("\n");
428}
429
430static void print_matrix(struct ggml_tensor * probs) {
431    assert(ggml_is_matrix(probs));
432    for (int i = 0; i < probs->ne[1]; ++i) {
433        for (int k = 0; k < probs->ne[0]; ++k) {
434            float p = get_f32_2d(probs, k, i);
435            LOG(" %.2f", p);
436        }
437        LOG("\n");
438    }
439}
440
441struct my_llama_file {
442    // use FILE * so we don't have to re-open the file to mmap
443    FILE * fp;
444    size_t size;
445
446    my_llama_file(const char * fname, const char * mode) {
447        fp = std::fopen(fname, mode);
448        if (fp == NULL) {
449            size = 0;
450        } else {
451            seek(0, SEEK_END);
452            size = tell();
453            seek(0, SEEK_SET);
454        }
455    }
456
457    size_t tell() const {
458#ifdef _WIN32
459        __int64 ret = _ftelli64(fp);
460#else
461        long ret = std::ftell(fp);
462#endif
463        GGML_ASSERT(ret != -1); // this really shouldn't fail
464        return (size_t) ret;
465    }
466
467    void seek(size_t offset, int whence) {
468#ifdef _WIN32
469        int ret = _fseeki64(fp, (__int64) offset, whence);
470#else
471        int ret = std::fseek(fp, (long) offset, whence);
472#endif
473        GGML_ASSERT(ret == 0); // same
474    }
475
476    void read_raw(void * ptr, size_t size) {
477        if (size == 0) {
478            return;
479        }
480        errno = 0;
481        std::size_t ret = std::fread(ptr, size, 1, fp);
482        if (ferror(fp)) {
483            die_fmt("fread failed: %s", strerror(errno));
484        }
485        if (ret != 1) {
486            die("unexpectedly reached end of file");
487        }
488    }
489
490    std::uint32_t read_u32() {
491        std::uint32_t ret;
492        read_raw(&ret, sizeof(ret));
493        return ret;
494    }
495    std::float_t read_f32() {
496        std::float_t ret;
497        read_raw(&ret, sizeof(ret));
498        return ret;
499    }
500
501    std::string read_string(std::uint32_t len) {
502        std::vector<char> chars(len);
503        read_raw(chars.data(), len);
504        return std::string(chars.data(), len);
505    }
506
507    ~my_llama_file() {
508        if (fp) {
509            std::fclose(fp);
510        }
511    }
512};
513
514static bool is_ggml_file(const char * filename) {
515    my_llama_file file(filename, "rb");
516    if (file.size < 4) {
517        return false;
518    }
519    std::string magic = file.read_string(4);
520    return magic == GGUF_MAGIC;
521}
522
523static std::string llama_escape_whitespaces(const std::string & text) {
524    std::ostringstream out;
525    for (char c : text) {
526        if (c == ' ') out << "\xe2\x96\x81";
527        else out << c;
528    }
529    return out.str();
530}
531
532static void load_vocab(const char * filename, const Config * config, struct my_llama_vocab * vocab) {
533    if (is_ggml_file(filename)) {
534        LOG_INF("%s: Loading vocabulary from gguf file %s\n", __func__, filename);
535        struct ggml_context * ctx_data = NULL;
536
537        struct gguf_init_params params = {
538            /*.no_alloc = */ false,
539            /*.ctx      = */ &ctx_data,
540        };
541
542        struct gguf_context * ctx = gguf_init_from_file(filename, params);
543        GGML_ASSERT(ctx != NULL);
544
545        const int model_idx = gguf_find_key(ctx, KV_TOKENIZER_MODEL);
546        GGML_ASSERT(model_idx >= 0);
547        std::string tokenizer_name = gguf_get_val_str(ctx, model_idx);
548        GGML_ASSERT(tokenizer_name == TOKENIZER_NAME);
549
550        const int token_idx = gguf_find_key(ctx, KV_TOKENIZER_LIST);
551        GGML_ASSERT(token_idx >= 0);
552        if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
553            gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
554            die_fmt("invalid gguf type for %s", KV_TOKENIZER_LIST);
555        }
556
557        const uint32_t n_vocab = gguf_get_arr_n(ctx, token_idx);
558        if (n_vocab != static_cast<uint32_t>(config->vocab_size)) {
559            die_fmt("vocab size mismatch: (gguf) %u != (llama2c) %d", n_vocab, config->vocab_size);
560        }
561
562        const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES);
563        GGML_ASSERT(score_idx >= 0);
564        if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
565            gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32 ||
566            gguf_get_arr_n(ctx, score_idx) < n_vocab) {
567            die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_SCORES);
568        }
569        const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
570
571        const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE);
572        GGML_ASSERT(toktype_idx >= 0);
573        if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
574            gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32 ||
575            gguf_get_arr_n(ctx, toktype_idx) < n_vocab) {
576            die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_TOKEN_TYPE);
577        }
578        const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx);
579
580        vocab->id_to_token.resize(n_vocab);
581
582        for (uint32_t i = 0; i < n_vocab; i++) {
583            std::string word = gguf_get_arr_str(ctx, token_idx, i);
584
585            vocab->token_to_id[word] = i;
586
587            auto & token_data = vocab->id_to_token[i];
588            token_data.text  = std::move(word);
589            token_data.score = scores[i];
590            token_data.type  = (llama_token_type) toktypes[i];
591        }
592        ggml_free(ctx_data);
593        gguf_free(ctx);
594    } else {
595        // assume llama2.c vocabulary
596        LOG_INF("%s: Assuming llama2.c vocabulary since %s is not a gguf file\n", __func__, filename);
597        my_llama_file file(filename, "rb");
598        if (!file.fp) {
599            die_fmt("%s: %s", strerror(errno), filename);
600        }
601        const int  n_vocab = config->vocab_size;
602        /* uint32_t max_token_length =  */ file.read_u32(); // unused
603        vocab->id_to_token.resize(n_vocab);
604        for (my_llama_vocab::id id=0; id<n_vocab; ++id) {
605            float_t score = file.read_f32();
606            uint32_t len = file.read_u32();
607            std::string text = file.read_string(len);
608
609            unsigned char byte_val;
610            my_llama_vocab::ttype type = LLAMA_TOKEN_TYPE_NORMAL;
611            if (id == UNKNOWN_TOKEN_ID) {
612                text = "<unk>";
613                type = LLAMA_TOKEN_TYPE_UNKNOWN;
614            } else if (id == BOS_TOKEN_ID) {
615                text = "<s>";
616                type = LLAMA_TOKEN_TYPE_CONTROL;
617            } else if (id == EOS_TOKEN_ID) {
618                text = "</s>";
619                type = LLAMA_TOKEN_TYPE_CONTROL;
620            } else if (text.empty()) {
621                type = LLAMA_TOKEN_TYPE_CONTROL;
622            } else if (sscanf(text.c_str(), "<0x%02hhX>", &byte_val) == 1) {
623                // Text of byte tokens is already in the expected format.
624                type = LLAMA_TOKEN_TYPE_BYTE;
625            } else {
626                type = LLAMA_TOKEN_TYPE_NORMAL;
627            }
628            text = llama_escape_whitespaces(text);
629
630            vocab->id_to_token[id].text = text;
631            vocab->id_to_token[id].score = score;
632            vocab->id_to_token[id].type = type;
633            vocab->token_to_id.emplace(text, id);
634        }
635    }
636}
637
638static void convert_weights_ak_to_gg(struct ggml_tensor * gg_weights, const float * karpathy_weights) {
639    int size = 1;
640    for (int dim = 0; dim < ggml_n_dims(gg_weights); ++dim) {
641        size *= gg_weights->ne[dim];
642    }
643    for (int ct = 0; ct < size; ++ct) {
644        int64_t i0 = 0; int64_t i1 = 0;
645        int64_t i2 = 0; int64_t i3 = 0;
646        ggml_unravel_index(gg_weights, ct, &i0, &i1, &i2, &i3);
647        ggml_set_f32_nd(gg_weights, i0, i1, i2, i3, karpathy_weights[ct]);
648    }
649}
650
651static void save_as_llama_model(
652    struct my_llama_vocab * vocab, struct my_llama_model * model, TransformerWeights* w, const char * filename
653) {
654    // convert AK weights into GG weights one by one.
655    // w->token_embedding_table -> model->tok_embeddings
656    // float*                   -> struct ggml_tensor
657    convert_weights_ak_to_gg(model->tok_embeddings, w->token_embedding_table.data());
658    convert_weights_ak_to_gg(model->output, !w->wcls.empty() ? w->wcls.data() : w->token_embedding_table.data());
659
660    convert_weights_ak_to_gg(model->norm, w->rms_final_weight.data());
661    //print_row(model->norm, 0);
662
663    // for rms-att-weight
664    int row_length = model->hparams.n_embd;
665    int n_ff = model->hparams.n_ff;
666
667    const uint32_t n_multiqueries = model->hparams.n_head_kv <= 0 || model->hparams.n_head_kv >= model->hparams.n_head ? 1 : model->hparams.n_head / model->hparams.n_head_kv;
668
669    for (uint32_t i = 0; i < model->hparams.n_layer; ++i){
670        auto & layer = model->layers[i];
671        // 1d
672        convert_weights_ak_to_gg(layer.attention_norm, &w->rms_att_weight[i*row_length]);
673        convert_weights_ak_to_gg(layer.ffn_norm      , &w->rms_ffn_weight[i*row_length]);
674
675        // from 3d matrix layer x dim x dim to 2d matrix dim x dim
676        convert_weights_ak_to_gg(layer.wq            , &w->wq[i*row_length*row_length]);
677        convert_weights_ak_to_gg(layer.wo            , &w->wo[i*row_length*row_length]);
678        // from 3d matrix layer x dim x dim to 2d matrix dim x dim / n_multiqueries
679        convert_weights_ak_to_gg(layer.wk            , &w->wk[i*row_length*row_length/n_multiqueries]);
680        convert_weights_ak_to_gg(layer.wv            , &w->wv[i*row_length*row_length/n_multiqueries]);
681
682        convert_weights_ak_to_gg(layer.w1            , &w->w1[i*row_length*n_ff]);
683        convert_weights_ak_to_gg(layer.w2            , &w->w2[i*n_ff*row_length]);
684        convert_weights_ak_to_gg(layer.w3            , &w->w3[i*row_length*n_ff]);
685    }
686
687    struct gguf_context * ctx = gguf_init_empty();
688
689    std::vector<const char*> tokens;
690    std::vector<float> scores;
691    std::vector<llama_token_type> token_types;
692    for (const my_llama_vocab::token_data & token_data : vocab->id_to_token) {
693        tokens.push_back(token_data.text.c_str());
694        scores.push_back(token_data.score);
695        token_types.push_back(token_data.type);
696    }
697    gguf_set_arr_str(ctx, KV_TOKENIZER_LIST, tokens.data(), tokens.size());
698    gguf_set_arr_data(ctx, KV_TOKENIZER_SCORES, GGUF_TYPE_FLOAT32, scores.data(), scores.size());
699    gguf_set_arr_data(ctx, KV_TOKENIZER_TOKEN_TYPE, GGUF_TYPE_INT32, token_types.data(), token_types.size());
700
701    gguf_set_val_str(ctx, KV_TOKENIZER_MODEL, TOKENIZER_NAME);
702
703    gguf_set_val_str(ctx, KV_GENERAL_ARCHITECTURE, "llama");
704    gguf_set_val_str(ctx, KV_GENERAL_NAME, "llama");
705
706    // special tokens
707    gguf_set_val_u32(ctx, KV_TOKENIZER_UNK_ID, UNKNOWN_TOKEN_ID);
708    gguf_set_val_u32(ctx, KV_TOKENIZER_BOS_ID, BOS_TOKEN_ID);
709    gguf_set_val_u32(ctx, KV_TOKENIZER_EOS_ID, EOS_TOKEN_ID);
710    gguf_set_val_u32(ctx, KV_TOKENIZER_SEP_ID, LLAMA_TOKEN_NULL);
711    gguf_set_val_u32(ctx, KV_TOKENIZER_PAD_ID, LLAMA_TOKEN_NULL);
712
713    gguf_set_val_u32(ctx, KV_CONTEXT_LENGTH, model->hparams.n_ctx);
714    gguf_set_val_u32(ctx, KV_EMBEDDING_LENGTH, model->hparams.n_embd);
715    gguf_set_val_u32(ctx, KV_FEED_FORWARD_LENGTH, model->hparams.n_ff);
716    gguf_set_val_u32(ctx, KV_ATTENTION_HEAD_COUNT, model->hparams.n_head);
717    gguf_set_val_u32(ctx, KV_ATTENTION_HEAD_COUNT, model->hparams.n_head);
718    gguf_set_val_u32(ctx, KV_ATTENTION_HEAD_COUNT_KV, model->hparams.n_head_kv);
719    gguf_set_val_u32(ctx, KV_BLOCK_COUNT, model->hparams.n_layer);
720    gguf_set_val_u32(ctx, KV_ROPE_DIMENSION_COUNT, model->hparams.n_rot);
721    gguf_set_val_f32(ctx, KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f);
722
723    // write tensors
724    ggml_set_name(model->tok_embeddings, TN_TOKEN_EMBD);
725    gguf_add_tensor(ctx, model->tok_embeddings);
726
727    ggml_set_name(model->norm, TN_OUTPUT_NORM);
728    gguf_add_tensor(ctx, model->norm);
729
730    ggml_set_name(model->output, TN_OUTPUT);
731    gguf_add_tensor(ctx, model->output);
732
733    for (uint32_t i = 0; i < model->hparams.n_layer; ++i) {
734        auto & layer = model->layers[i];
735
736        ggml_format_name(layer.wq, TN_ATTN_Q, i);
737        gguf_add_tensor(ctx, layer.wq);
738
739        ggml_format_name(layer.wk, TN_ATTN_K, i);
740        gguf_add_tensor(ctx, layer.wk);
741
742        ggml_format_name(layer.wv, TN_ATTN_V, i);
743        gguf_add_tensor(ctx, layer.wv);
744
745        ggml_format_name(layer.wo, TN_ATTN_OUTPUT, i);
746        gguf_add_tensor(ctx, layer.wo);
747
748        ggml_format_name(layer.attention_norm, TN_ATTN_NORM, i);
749        gguf_add_tensor(ctx, layer.attention_norm);
750
751        ggml_format_name(layer.w1, TN_FFN_GATE, i);
752        gguf_add_tensor(ctx, layer.w1);
753
754        ggml_format_name(layer.w2, TN_FFN_DOWN, i);
755        gguf_add_tensor(ctx, layer.w2);
756
757        ggml_format_name(layer.w3, TN_FFN_UP, i);
758        gguf_add_tensor(ctx, layer.w3);
759
760        ggml_format_name(layer.ffn_norm, TN_FFN_NORM, i);
761        gguf_add_tensor(ctx, layer.ffn_norm);
762    }
763
764    gguf_write_to_file(ctx, filename, false);
765    gguf_free(ctx);
766}
767
768static struct train_params get_default_train_params() {
769    struct train_params params;
770    params.fn_vocab_model          = "models/7B/ggml-model-f16.gguf";
771    params.fn_llama2c_output_model = "ak_llama_model.bin";
772    params.fn_train_data           = "shakespeare.txt";
773    params.fn_checkpoint_in        = "checkpoint.bin";
774    params.fn_checkpoint_out       = "checkpoint.bin";
775    params.fn_model_out            = "ggml-checkpoint-f32.bin";
776
777    params.seed       =   -1;
778
779    params.n_ctx      =  128;
780    params.n_embd     =  256;
781    params.n_mult     =  256;
782    params.n_head     =    8;
783    params.n_layer    =   16;
784    params.n_rotmax   =   64;
785
786    params.n_threads  =    6;
787    params.n_batch    =    8;
788    params.n_examples =    8;
789    params.n_predict  = 1024;
790
791    params.print_info_interval    = 1;
792    params.print_details_interval = 2;
793
794    params.samples_start_after_nl = false;
795    params.use_adam               = true;
796    params.use_flash              = false;
797    params.use_scratch            = true;
798
799    // only adam
800    params.warmup            =  100;
801    params.cos_decay_steps   = 1000;
802    params.cos_decay_restart = 1.1f;
803    params.cos_decay_alpha   = 0.0f;
804
805    params.lbfgs_n_iter      = 16;
806    params.adam_n_iter       = 16;
807    params.adam_alpha        = 1e-3f;
808    params.adam_decay        = 1e-3f;
809
810    params.mem_model_gb    = 2;
811    params.mem_compute_gb  = 24;
812    params.mem_compute0_gb = 8;
813    params.mem_compute1_gb = 2;
814
815    return params;
816}
817
818static void print_usage(int /*argc*/, char ** argv, const struct train_params * params) {
819    fprintf(stderr, "usage: %s [options]\n", argv[0]);
820    fprintf(stderr, "\n");
821    fprintf(stderr, "options:\n");
822    fprintf(stderr, "  -h, --help                       show this help message and exit\n");
823    fprintf(stderr, "  --copy-vocab-from-model FNAME    path of gguf llama model or llama2.c vocabulary from which to copy vocab (default '%s')\n", params->fn_vocab_model);
824    fprintf(stderr, "  --llama2c-model FNAME            [REQUIRED] model path from which to load Karpathy's llama2.c model\n");
825    fprintf(stderr, "  --llama2c-output-model FNAME     model path to save the converted llama2.c model (default %s')\n", params->fn_llama2c_output_model);
826    fprintf(stderr, "\n");
827}
828
829static bool params_parse(int argc, char ** argv, struct train_params * params) {
830    bool invalid_param = false;
831    bool reqd_param_found = false;
832    std::string arg;
833    struct train_params default_params = get_default_train_params();
834    const std::string arg_prefix = "--";
835
836    for (int i = 1; i < argc; i++) {
837        arg = argv[i];
838        if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) {
839            std::replace(arg.begin(), arg.end(), '_', '-');
840        }
841
842        if (arg == "--copy-vocab-from-model") {
843            if (++i >= argc) {
844                invalid_param = true;
845                break;
846            }
847            params->fn_vocab_model = argv[i];
848        } else if (arg == "--llama2c-model") {
849            if (++i >= argc) {
850                invalid_param = true;
851                break;
852            }
853            reqd_param_found = true;
854            params->fn_llama2c_model = argv[i];
855        } else if (arg == "--llama2c-output-model") {
856            if (++i >= argc) {
857                invalid_param = true;
858                break;
859            }
860            params->fn_llama2c_output_model = argv[i];
861        } else if (arg == "-h" || arg == "--help") {
862            print_usage(argc, argv, &default_params);
863            exit(0);
864        } else {
865            fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());
866            print_usage(argc, argv, &default_params);
867            exit(1);
868        }
869    }
870    if (invalid_param) {
871        fprintf(stderr, "error: invalid parameter for argument: %s\n", arg.c_str());
872        print_usage(argc, argv, &default_params);
873        exit(1);
874    }
875    if (!reqd_param_found){
876        fprintf(stderr, "error: please specify a llama2.c .bin file to be converted with argument --llama2c-model\n");
877        print_usage(argc, argv, &default_params);
878        exit(1);
879    }
880
881    return true;
882}
883
884static std::string basename(const std::string &path) {
885    size_t pos = path.find_last_of("/\\");
886    if (pos == std::string::npos) {
887        return path;
888    }
889    return path.substr(pos + 1);
890}
891
892int main(int argc, char ** argv) {
893    std::setlocale(LC_NUMERIC, "C");
894
895    common_init();
896
897    struct train_params params = get_default_train_params();
898    if (!params_parse(argc, argv, &params)) {
899        return 1;
900    }
901
902    Config config;
903    TransformerWeights weights = {};
904    {
905        LOG_INF("%s: Loading llama2c model from %s\n", __func__, params.fn_llama2c_model);
906        FILE * file = fopen(params.fn_llama2c_model, "rb");
907        if (!file) {
908            LOG_ERR("%s: Unable to open the checkpoint file %s!\n", __func__, params.fn_llama2c_model);
909            return 1;
910        }
911        // read in the config header
912        if (fread(&config, sizeof(Config), 1, file) != 1) {
913            LOG_ERR("%s: Unable to read llama2c config from %s!\n",__func__,params.fn_llama2c_model);
914            return 1;
915        }
916        auto shared_weights = config.vocab_size > 0;
917        config.vocab_size = abs(config.vocab_size);
918
919        // read in the Transformer weights
920        alloc_weights(&weights, &config, shared_weights);
921        if (checkpoint_init_weights(&weights, &config, file, shared_weights)) {
922            LOG_ERR("%s: Unable to initialize transformer weights from %s!",__func__,params.fn_llama2c_model);
923            return 1;
924        }
925        fclose(file);
926    }
927
928    struct my_llama_vocab vocab;
929    load_vocab(params.fn_vocab_model, &config, &vocab);
930
931    struct my_llama_model model;
932    model.hparams.n_vocab   = config.vocab_size; //llama_vocab_n_vocab(lctx);
933    model.hparams.n_ctx     = params.n_ctx;
934    model.hparams.n_embd    = config.dim; //params.n_embd;
935    model.hparams.n_ff      = config.hidden_dim;
936    model.hparams.n_mult    = 32;//params.n_mult;
937    model.hparams.n_head    = config.n_heads; //params.n_head;
938    model.hparams.n_head_kv = config.n_kv_heads;
939    model.hparams.n_layer   = config.n_layers; //params.n_layer;
940    model.hparams.n_rot     = std::min((uint32_t)params.n_rotmax, model.hparams.n_embd / model.hparams.n_head);
941
942    print_params(&model.hparams);
943
944    struct ggml_init_params lcparams;
945    lcparams.mem_size   = 1024ll*1024ll*1024ll*((size_t) params.mem_model_gb);
946    lcparams.mem_buffer = NULL;
947    lcparams.no_alloc   = false;
948
949    model.ctx = ggml_init(lcparams);
950
951    init_model(&model);
952    model.name = basename(params.fn_llama2c_model);
953    save_as_llama_model(&vocab, &model, &weights, params.fn_llama2c_output_model);
954
955    LOG_INF("%s: Saving llama.c model file %s in ggml format at %s\n", __func__, params.fn_llama2c_model, params.fn_llama2c_output_model);
956
957    ggml_free(model.ctx);
958    return 0;
959}