From 8f02f1b46f66d7a04d4177f733228c235f5ab2ac Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Wed, 15 Mar 2023 20:33:43 -0700 Subject: [PATCH 01/18] Add llama.cpp and llama.h --- Makefile | 7 +- llama.cpp | 682 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ llama.h | 66 ++++++ 3 files changed, 753 insertions(+), 2 deletions(-) create mode 100644 llama.cpp create mode 100644 llama.h diff --git a/Makefile b/Makefile index 1601079a48685..a74929c82fd7e 100644 --- a/Makefile +++ b/Makefile @@ -188,11 +188,14 @@ ggml.o: ggml.c ggml.h utils.o: utils.cpp utils.h $(CXX) $(CXXFLAGS) -c utils.cpp -o utils.o +llama.o: llama.cpp llama.h + $(CXX) $(CXXFLAGS) -c llama.cpp -o llama.o + clean: rm -f *.o main quantize -main: main.cpp ggml.o utils.o - $(CXX) $(CXXFLAGS) main.cpp ggml.o utils.o -o main $(LDFLAGS) +main: main.cpp ggml.o utils.o llama.o + $(CXX) $(CXXFLAGS) main.cpp ggml.o llama.o utils.o -o main $(LDFLAGS) ./main -h quantize: quantize.cpp ggml.o utils.o diff --git a/llama.cpp b/llama.cpp new file mode 100644 index 0000000000000..034666327401c --- /dev/null +++ b/llama.cpp @@ -0,0 +1,682 @@ +#include "llama.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +// determine number of model parts based on the dimension +static const std::map LLAMA_N_PARTS = { + { 4096, 1 }, + { 5120, 2 }, + { 6656, 4 }, + { 8192, 8 }, +}; + + +// load the model's weights from a file +bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx) { + fprintf(stderr, "%s: loading model from '%s' - please wait ...\n", __func__, fname.c_str()); + + std::vector f_buf(1024*1024); + + auto fin = std::ifstream(fname, std::ios::binary); + fin.rdbuf()->pubsetbuf(f_buf.data(), f_buf.size()); + if (!fin) { + fprintf(stderr, "%s: failed to open '%s'\n", __func__, fname.c_str()); + return false; + } + + // verify magic + { + uint32_t magic; + fin.read((char *) &magic, sizeof(magic)); + if (magic != 0x67676d6c) { + fprintf(stderr, "%s: invalid model file '%s' (bad magic)\n", __func__, fname.c_str()); + return false; + } + } + + int n_ff = 0; + int n_parts = 0; + + // load hparams + { + auto & hparams = model.hparams; + + fin.read((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); + //fin.read((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); + fin.read((char *) &hparams.n_embd, sizeof(hparams.n_embd)); + fin.read((char *) &hparams.n_mult, sizeof(hparams.n_mult)); + fin.read((char *) &hparams.n_head, sizeof(hparams.n_head)); + fin.read((char *) &hparams.n_layer, sizeof(hparams.n_layer)); + fin.read((char *) &hparams.n_rot, sizeof(hparams.n_rot)); + fin.read((char *) &hparams.f16, sizeof(hparams.f16)); + + hparams.n_ctx = n_ctx; + + n_ff = ((2*(4*hparams.n_embd)/3 + hparams.n_mult - 1)/hparams.n_mult)*hparams.n_mult; + n_parts = LLAMA_N_PARTS.at(hparams.n_embd); + + fprintf(stderr, "%s: n_vocab = %d\n", __func__, hparams.n_vocab); + fprintf(stderr, "%s: n_ctx = %d\n", __func__, hparams.n_ctx); + fprintf(stderr, "%s: n_embd = %d\n", __func__, hparams.n_embd); + fprintf(stderr, "%s: n_mult = %d\n", __func__, hparams.n_mult); + fprintf(stderr, "%s: n_head = %d\n", __func__, hparams.n_head); + fprintf(stderr, "%s: n_layer = %d\n", __func__, hparams.n_layer); + fprintf(stderr, "%s: n_rot = %d\n", __func__, hparams.n_rot); + fprintf(stderr, "%s: f16 = %d\n", __func__, hparams.f16); + fprintf(stderr, "%s: n_ff = %d\n", __func__, n_ff); + fprintf(stderr, "%s: n_parts = %d\n", __func__, n_parts); + } + + // load vocab + { + std::string word; + for (int i = 0; i < model.hparams.n_vocab; i++) { + uint32_t len; + fin.read((char *) &len, sizeof(len)); + + word.resize(len); + fin.read((char *) word.data(), len); + + vocab.token_to_id[word] = i; + vocab.id_to_token[i] = word; + + //if (i < 30000) { + // fprintf(stderr, "%s: vocab[%d] = '%s'\n", __func__, i, word.c_str()); + //} + } + } + + // for the big tensors, we have the option to store the data in 16-bit floats or quantized + // in order to save memory and also to speed up the computation + ggml_type wtype = GGML_TYPE_COUNT; + switch (model.hparams.f16) { + case 0: wtype = GGML_TYPE_F32; break; + case 1: wtype = GGML_TYPE_F16; break; + case 2: wtype = GGML_TYPE_Q4_0; break; + case 3: wtype = GGML_TYPE_Q4_1; break; + default: + { + fprintf(stderr, "%s: invalid model file '%s' (bad f16 value %d)\n", + __func__, fname.c_str(), model.hparams.f16); + return false; + } + } + + const ggml_type wtype2 = GGML_TYPE_F32; + + auto & ctx = model.ctx; + + size_t ctx_size = 0; + + { + const auto & hparams = model.hparams; + + const int n_embd = hparams.n_embd; + const int n_layer = hparams.n_layer; + const int n_ctx = hparams.n_ctx; + const int n_vocab = hparams.n_vocab; + + ctx_size += n_embd*n_vocab*ggml_type_sizef(wtype); // tok_embeddings + + ctx_size += n_embd*ggml_type_sizef(GGML_TYPE_F32); // norm + + ctx_size += n_embd*n_vocab*ggml_type_sizef(wtype); // output + + ctx_size += n_layer*(n_embd*ggml_type_sizef(GGML_TYPE_F32)); // attention_norm + + ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wq + ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wk + ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wv + ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wo + + ctx_size += n_layer*(n_embd*ggml_type_sizef(GGML_TYPE_F32)); // ffn_norm + + ctx_size += n_layer*(n_ff*n_embd*ggml_type_sizef(wtype)); // w1 + ctx_size += n_layer*(n_ff*n_embd*ggml_type_sizef(wtype)); // w2 + ctx_size += n_layer*(n_ff*n_embd*ggml_type_sizef(wtype)); // w3 + + ctx_size += n_ctx*n_layer*n_embd*ggml_type_sizef(GGML_TYPE_F32); // memory_k + ctx_size += n_ctx*n_layer*n_embd*ggml_type_sizef(GGML_TYPE_F32); // memory_v + + ctx_size += (5 + 10*n_layer)*256; // object overhead + + fprintf(stderr, "%s: ggml ctx size = %6.2f MB\n", __func__, ctx_size/(1024.0*1024.0)); + } + + // create the ggml context + { + struct ggml_init_params params = { + /*.mem_size =*/ ctx_size, + /*.mem_buffer =*/ NULL, + }; + + model.ctx = ggml_init(params); + if (!model.ctx) { + fprintf(stderr, "%s: ggml_init() failed\n", __func__); + return false; + } + } + + // prepare memory for the weights + { + const auto & hparams = model.hparams; + + const int n_embd = hparams.n_embd; + const int n_layer = hparams.n_layer; + const int n_ctx = hparams.n_ctx; + const int n_vocab = hparams.n_vocab; + + model.layers.resize(n_layer); + + model.tok_embeddings = ggml_new_tensor_2d(ctx, wtype, n_embd, n_vocab); + + model.norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); + model.output = ggml_new_tensor_2d(ctx, wtype, n_embd, n_vocab); + + // map by name + model.tensors["tok_embeddings.weight"] = model.tok_embeddings; + + model.tensors["norm.weight"] = model.norm; + model.tensors["output.weight"] = model.output; + + for (int i = 0; i < n_layer; ++i) { + auto & layer = model.layers[i]; + + layer.attention_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); + + layer.wq = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); + layer.wk = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); + layer.wv = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); + layer.wo = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); + + layer.ffn_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); + + layer.w1 = ggml_new_tensor_2d(ctx, wtype, n_embd, n_ff); + layer.w2 = ggml_new_tensor_2d(ctx, wtype, n_ff, n_embd); + layer.w3 = ggml_new_tensor_2d(ctx, wtype, n_embd, n_ff); + + // map by name + model.tensors["layers." + std::to_string(i) + ".attention_norm.weight"] = layer.attention_norm; + + model.tensors["layers." + std::to_string(i) + ".attention.wq.weight"] = layer.wq; + model.tensors["layers." + std::to_string(i) + ".attention.wk.weight"] = layer.wk; + model.tensors["layers." + std::to_string(i) + ".attention.wv.weight"] = layer.wv; + model.tensors["layers." + std::to_string(i) + ".attention.wo.weight"] = layer.wo; + + model.tensors["layers." + std::to_string(i) + ".ffn_norm.weight"] = layer.ffn_norm; + + model.tensors["layers." + std::to_string(i) + ".feed_forward.w1.weight"] = layer.w1; + model.tensors["layers." + std::to_string(i) + ".feed_forward.w2.weight"] = layer.w2; + model.tensors["layers." + std::to_string(i) + ".feed_forward.w3.weight"] = layer.w3; + } + } + + // key + value memory + { + const auto & hparams = model.hparams; + + const int n_embd = hparams.n_embd; + const int n_layer = hparams.n_layer; + const int n_ctx = hparams.n_ctx; + + const int n_mem = n_layer*n_ctx; + const int n_elements = n_embd*n_mem; + + model.memory_k = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_elements); + model.memory_v = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_elements); + + const size_t memory_size = ggml_nbytes(model.memory_k) + ggml_nbytes(model.memory_v); + + fprintf(stderr, "%s: memory_size = %8.2f MB, n_mem = %d\n", __func__, memory_size/1024.0/1024.0, n_mem); + } + + const size_t file_offset = fin.tellg(); + + fin.close(); + + std::vector tmp; + + for (int i = 0; i < n_parts; ++i) { + const int part_id = i; + //const int part_id = n_parts - i - 1; + + std::string fname_part = fname; + if (i > 0) { + fname_part += "." + std::to_string(i); + } + + fprintf(stderr, "%s: loading model part %d/%d from '%s'\n", __func__, i+1, n_parts, fname_part.c_str()); + + fin = std::ifstream(fname_part, std::ios::binary); + fin.rdbuf()->pubsetbuf(f_buf.data(), f_buf.size()); + fin.seekg(file_offset); + + // load weights + { + int n_tensors = 0; + size_t total_size = 0; + + fprintf(stderr, "%s: ", __func__); + + while (true) { + int32_t n_dims; + int32_t length; + int32_t ftype; + + fin.read(reinterpret_cast(&n_dims), sizeof(n_dims)); + fin.read(reinterpret_cast(&length), sizeof(length)); + fin.read(reinterpret_cast(&ftype), sizeof(ftype)); + + if (fin.eof()) { + break; + } + + int32_t nelements = 1; + int32_t ne[2] = { 1, 1 }; + for (int i = 0; i < n_dims; ++i) { + fin.read(reinterpret_cast(&ne[i]), sizeof(ne[i])); + nelements *= ne[i]; + } + + std::string name(length, 0); + fin.read(&name[0], length); + + if (model.tensors.find(name.data()) == model.tensors.end()) { + fprintf(stderr, "%s: unknown tensor '%s' in model file\n", __func__, name.data()); + return false; + } + + // split_type = 0: split by columns + // split_type = 1: split by rows + int split_type = 0; + + // split_type = 0: + // regex: + // - tok_embeddings.* + // - layers.*.attention.wo.weight + // - layers.*.feed_forward.w2.weight + + // split_type = 1: + // regex: + // - output.* + // - layers.*.attention.wq.weight + // - layers.*.attention.wk.weight + // - layers.*.attention.wv.weight + // - layers.*.feed_forward.w1.weight + // - layers.*.feed_forward.w3.weight + if (name.find("tok_embeddings") != std::string::npos) { + split_type = 0; + } else if (name.find("layers") != std::string::npos) { + if (name.find("attention.wo.weight") != std::string::npos) { + split_type = 0; + } else if (name.find("feed_forward.w2.weight") != std::string::npos) { + split_type = 0; + } else { + split_type = 1; + } + } else if (name.find("output") != std::string::npos) { + split_type = 1; + } + + auto tensor = model.tensors[name.data()]; + + if (n_dims == 1) { + if (ggml_nelements(tensor) != nelements) { + fprintf(stderr, "%s: tensor '%s' has wrong size in model file\n", __func__, name.data()); + return false; + } + } else { + if (ggml_nelements(tensor)/n_parts != nelements) { + fprintf(stderr, "%s: tensor '%s' has wrong size in model file\n", __func__, name.data()); + return false; + } + } + + if (n_dims == 1) { + if (tensor->ne[0] != ne[0] || tensor->ne[1] != ne[1]) { + fprintf(stderr, "%s: tensor '%s' has wrong shape in model file: got [%d, %d], expected [%d, %d]\n", + __func__, name.data(), tensor->ne[0], tensor->ne[1], ne[0], ne[1]); + return false; + } + } else { + if (split_type == 0) { + if (tensor->ne[0]/n_parts != ne[0] || tensor->ne[1] != ne[1]) { + fprintf(stderr, "%s: tensor '%s' has wrong shape in model file: got [%d, %d], expected [%d, %d]\n", + __func__, name.data(), tensor->ne[0]/n_parts, tensor->ne[1], ne[0], ne[1]); + return false; + } + } else { + if (tensor->ne[0] != ne[0] || tensor->ne[1]/n_parts != ne[1]) { + fprintf(stderr, "%s: tensor '%s' has wrong shape in model file: got [%d, %d], expected [%d, %d]\n", + __func__, name.data(), tensor->ne[0], tensor->ne[1]/n_parts, ne[0], ne[1]); + return false; + } + } + } + + if (0) { + static const char * ftype_str[] = { "f32", "f16", "q4_0", "q4_1", }; + fprintf(stderr, "%24s - [%5d, %5d], type = %6s, split = %d\n", name.data(), ne[0], ne[1], ftype_str[ftype], split_type); + } + + size_t bpe = 0; + + switch (ftype) { + case 0: bpe = ggml_type_size(GGML_TYPE_F32); break; + case 1: bpe = ggml_type_size(GGML_TYPE_F16); break; + case 2: bpe = ggml_type_size(GGML_TYPE_Q4_0); assert(ne[0] % 64 == 0); break; + case 3: bpe = ggml_type_size(GGML_TYPE_Q4_1); assert(ne[0] % 64 == 0); break; + default: + { + fprintf(stderr, "%s: unknown ftype %d in model file\n", __func__, ftype); + return false; + } + }; + + if (n_dims == 1 || n_parts == 1) { + if ((nelements*bpe)/ggml_blck_size(tensor->type) != ggml_nbytes(tensor)) { + fprintf(stderr, "%s: tensor '%s' has wrong size in model file: got %zu, expected %zu\n", + __func__, name.data(), ggml_nbytes(tensor), nelements*bpe); + return false; + } + + if (part_id == 0) { + fin.read(reinterpret_cast(tensor->data), ggml_nbytes(tensor)); + } else { + fin.seekg(ggml_nbytes(tensor), std::ios::cur); + } + + total_size += ggml_nbytes(tensor); + } else { + if ((nelements*bpe)/ggml_blck_size(tensor->type) != ggml_nbytes(tensor)/n_parts) { + fprintf(stderr, "%s: tensor '%s' has wrong size in model file: got %zu, expected %zu\n", + __func__, name.data(), ggml_nbytes(tensor)/n_parts, nelements*bpe); + return false; + } + + if (split_type == 0) { + const int np0 = ne[0]; + + const size_t row_size = (tensor->ne[0]/ggml_blck_size(tensor->type))*ggml_type_size(tensor->type); + assert(row_size == tensor->nb[1]); + + for (int i1 = 0; i1 < ne[1]; ++i1) { + const size_t offset_row = i1*row_size; + const size_t offset = offset_row + ((part_id*np0)/ggml_blck_size(tensor->type))*ggml_type_size(tensor->type); + fin.read(reinterpret_cast(tensor->data) + offset, row_size/n_parts); + } + } else { + const int np1 = ne[1]; + + const size_t row_size = (tensor->ne[0]/ggml_blck_size(tensor->type))*ggml_type_size(tensor->type); + + for (int i1 = 0; i1 < ne[1]; ++i1) { + const size_t offset_row = (i1 + part_id*np1)*row_size; + fin.read(reinterpret_cast(tensor->data) + offset_row, row_size); + } + } + + total_size += ggml_nbytes(tensor)/n_parts; + } + + //fprintf(stderr, "%42s - [%5d, %5d], type = %6s, %6.2f MB\n", name.data(), ne[0], ne[1], ftype == 0 ? "float" : "f16", ggml_nbytes(tensor)/1024.0/1024.0); + if (++n_tensors % 8 == 0) { + fprintf(stderr, "."); + fflush(stderr); + } + } + + fprintf(stderr, " done\n"); + + fprintf(stderr, "%s: model size = %8.2f MB / num tensors = %d\n", __func__, total_size/1024.0/1024.0, n_tensors); + } + + fin.close(); + } + + return true; +} + +// evaluate the transformer +// +// - model: the model +// - n_threads: number of threads to use +// - n_past: the context size so far +// - embd_inp: the embeddings of the tokens in the context +// - embd_w: the predicted logits for the next token +// +// The GPT-J model requires about 16MB of memory per input token. +// +bool llama_eval( + const llama_model & model, + const int n_threads, + const int n_past, + const std::vector & embd_inp, + std::vector & embd_w, + size_t & mem_per_token) { + const int N = embd_inp.size(); + + const auto & hparams = model.hparams; + + const int n_embd = hparams.n_embd; + const int n_layer = hparams.n_layer; + const int n_ctx = hparams.n_ctx; + const int n_head = hparams.n_head; + const int n_vocab = hparams.n_vocab; + const int n_rot = hparams.n_embd/hparams.n_head; + + const int d_key = n_embd/n_head; + + // TODO: check if this size scales with n_ctx linearly and remove constant. somehow I feel it wasn't the case + // static size_t buf_size = hparams.n_ctx*1024*1024; + static size_t buf_size = 512u*1024*1024; + static void * buf = malloc(buf_size); + + if (mem_per_token > 0 && mem_per_token*N > buf_size) { + const size_t buf_size_new = 1.1*(mem_per_token*N); // add 10% to account for ggml object overhead + //fprintf(stderr, "\n%s: reallocating buffer from %zu to %zu bytes\n", __func__, buf_size, buf_size_new); + + // reallocate + buf_size = buf_size_new; + buf = realloc(buf, buf_size); + if (buf == nullptr) { + fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size); + return false; + } + } + + struct ggml_init_params params = { + /*.mem_size =*/ buf_size, + /*.mem_buffer =*/ buf, + }; + + struct ggml_context * ctx0 = ggml_init(params); + ggml_cgraph gf = {}; + gf.n_threads = n_threads; + + struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N); + memcpy(embd->data, embd_inp.data(), N*ggml_element_size(embd)); + + struct ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embeddings, embd); + + for (int il = 0; il < n_layer; ++il) { + struct ggml_tensor * inpSA = inpL; + + struct ggml_tensor * cur; + + // norm + { + cur = ggml_rms_norm(ctx0, inpL); + + // cur = attention_norm*cur + cur = ggml_mul(ctx0, + ggml_repeat(ctx0, model.layers[il].attention_norm, cur), + cur); + } + + // self-attention + { + struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, model.layers[il].wq, cur); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, model.layers[il].wv, cur); + + // store key and value to memory + if (N >= 1) { + struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_k, N*n_embd, (ggml_element_size(model.memory_k)*n_embd)*(il*n_ctx + n_past)); + struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_v, N*n_embd, (ggml_element_size(model.memory_v)*n_embd)*(il*n_ctx + n_past)); + + ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Kcur, k)); + ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcur, v)); + } + + // Q = Qcur.contiguous().view(n_embd/n_head, n_head, N).permute(0, 2, 1, 3) + struct ggml_tensor * Q = + ggml_permute(ctx0, + ggml_rope(ctx0, + ggml_cpy(ctx0, + Qcur, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd/n_head, n_head, N)), + n_past, n_rot, 0), + 0, 2, 1, 3); + + // K = Kmem.view(n_embd/n_head, n_head, n_past + N).permute(0, 2, 1, 3) + struct ggml_tensor * K = + ggml_permute(ctx0, + ggml_rope(ctx0, + ggml_reshape_3d(ctx0, + ggml_view_1d(ctx0, model.memory_k, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_k)*n_embd), + n_embd/n_head, n_head, n_past + N), + n_past, n_rot, 1), + 0, 2, 1, 3); + + // K * Q + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + + // KQ_scaled = KQ / sqrt(n_embd/n_head) + struct ggml_tensor * KQ_scaled = + ggml_scale(ctx0, + KQ, + ggml_new_f32(ctx0, 1.0f/sqrt(float(n_embd)/n_head)) + ); + + // KQ_masked = mask_past(KQ_scaled) + struct ggml_tensor * KQ_masked = ggml_diag_mask_inf(ctx0, KQ_scaled, n_past); + + // KQ = soft_max(KQ_masked) + struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctx0, KQ_masked); + + // V_trans = Vmem.view(n_embd/n_head, n_head, n_past + N).permute(1, 2, 0, 3).contiguous() + struct ggml_tensor * V_trans = + ggml_permute(ctx0, + ggml_reshape_3d(ctx0, + ggml_view_1d(ctx0, model.memory_v, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_v)*n_embd), + n_embd/n_head, n_head, n_past + N), + 1, 2, 0, 3); + + // KQV = transpose(V) * KQ_soft_max + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V_trans, KQ_soft_max); + + // KQV_merged = KQV.permute(0, 2, 1, 3) + struct ggml_tensor * KQV_merged = ggml_permute(ctx0, KQV, 0, 2, 1, 3); + + // cur = KQV_merged.contiguous().view(n_embd, N) + cur = ggml_cpy(ctx0, + KQV_merged, + ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, N)); + + // projection (no bias) + cur = ggml_mul_mat(ctx0, + model.layers[il].wo, + cur); + } + + struct ggml_tensor * inpFF = ggml_add(ctx0, cur, inpSA); + + // feed-forward network + { + // norm + { + cur = ggml_rms_norm(ctx0, inpFF); + + // cur = ffn_norm*cur + cur = ggml_mul(ctx0, + ggml_repeat(ctx0, model.layers[il].ffn_norm, cur), + cur); + } + + struct ggml_tensor * tmp = ggml_mul_mat(ctx0, + model.layers[il].w3, + cur); + + + cur = ggml_mul_mat(ctx0, + model.layers[il].w1, + cur); + + // SILU activation + cur = ggml_silu(ctx0, cur); + + cur = ggml_mul(ctx0, cur, tmp); + + cur = ggml_mul_mat(ctx0, + model.layers[il].w2, + cur); + } + + cur = ggml_add(ctx0, cur, inpFF); + + // input for next layer + inpL = cur; + } + + // norm + { + inpL = ggml_rms_norm(ctx0, inpL); + + // inpL = norm*inpL + inpL = ggml_mul(ctx0, + ggml_repeat(ctx0, model.norm, inpL), + inpL); + } + + // lm_head + { + inpL = ggml_mul_mat(ctx0, model.output, inpL); + } + + // logits -> probs + //inpL = ggml_soft_max(ctx0, inpL); + + // run the computation + ggml_build_forward_expand(&gf, inpL); + ggml_graph_compute (ctx0, &gf); + + //if (n_past%100 == 0) { + // ggml_graph_print (&gf); + // ggml_graph_dump_dot(&gf, NULL, "gpt-2.dot"); + //} + + //embd_w.resize(n_vocab*N); + //memcpy(embd_w.data(), ggml_get_data(inpL), sizeof(float)*n_vocab*N); + + // return result for just the last token + embd_w.resize(n_vocab); + memcpy(embd_w.data(), (float *) ggml_get_data(inpL) + (n_vocab*(N-1)), sizeof(float)*n_vocab); + + if (mem_per_token == 0) { + mem_per_token = ggml_used_mem(ctx0)/N; + } + //fprintf(stderr, "used_mem = %zu\n", ggml_used_mem(ctx0)); + + ggml_free(ctx0); + + return true; +} diff --git a/llama.h b/llama.h new file mode 100644 index 0000000000000..0ceab22bbbef0 --- /dev/null +++ b/llama.h @@ -0,0 +1,66 @@ +#pragma once + +#include +#include +#include + +#include "utils.h" + +// default hparams (LLaMA 7B) +struct llama_hparams { + int32_t n_vocab = 32000; + int32_t n_ctx = 512; // this is provided as user input? + int32_t n_embd = 4096; + int32_t n_mult = 256; + int32_t n_head = 32; + int32_t n_layer = 32; + int32_t n_rot = 64; + int32_t f16 = 1; +}; + +struct llama_layer { + // normalization + struct ggml_tensor * attention_norm; + + // attention + struct ggml_tensor * wq; + struct ggml_tensor * wk; + struct ggml_tensor * wv; + struct ggml_tensor * wo; + + // normalization + struct ggml_tensor * ffn_norm; + + // ff + struct ggml_tensor * w1; + struct ggml_tensor * w2; + struct ggml_tensor * w3; +}; + +struct llama_model { + llama_hparams hparams; + + struct ggml_tensor * tok_embeddings; + + struct ggml_tensor * norm; + struct ggml_tensor * output; + + std::vector layers; + + // key + value memory + struct ggml_tensor * memory_k; + struct ggml_tensor * memory_v; + + // + struct ggml_context * ctx; + std::map tensors; +}; + +bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx); +bool llama_eval( + const llama_model & model, + const int n_threads, + const int n_past, + const std::vector & embd_inp, + std::vector & embd_w, + size_t & mem_per_token); From 9c2109e6fc5fe3cabaaa3f18fa8d0d2d7be6532f Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Mon, 13 Mar 2023 21:07:49 -0700 Subject: [PATCH 02/18] Refactor out library code from main.cpp --- main.cpp | 710 +------------------------------------------------------ 1 file changed, 2 insertions(+), 708 deletions(-) diff --git a/main.cpp b/main.cpp index c7186e0dff273..9e09a6a8579ee 100644 --- a/main.cpp +++ b/main.cpp @@ -1,5 +1,5 @@ #include "ggml.h" - +#include "llama.h" #include "utils.h" #include @@ -38,712 +38,6 @@ static const std::map LLAMA_N_PARTS = { { 8192, 8 }, }; -// default hparams (LLaMA 7B) -struct llama_hparams { - int32_t n_vocab = 32000; - int32_t n_ctx = 512; // this is provided as user input? - int32_t n_embd = 4096; - int32_t n_mult = 256; - int32_t n_head = 32; - int32_t n_layer = 32; - int32_t n_rot = 64; - int32_t f16 = 1; -}; - -struct llama_layer { - // normalization - struct ggml_tensor * attention_norm; - - // attention - struct ggml_tensor * wq; - struct ggml_tensor * wk; - struct ggml_tensor * wv; - struct ggml_tensor * wo; - - // normalization - struct ggml_tensor * ffn_norm; - - // ff - struct ggml_tensor * w1; - struct ggml_tensor * w2; - struct ggml_tensor * w3; -}; - -struct llama_model { - llama_hparams hparams; - - struct ggml_tensor * tok_embeddings; - - struct ggml_tensor * norm; - struct ggml_tensor * output; - - std::vector layers; - - // key + value memory - struct ggml_tensor * memory_k; - struct ggml_tensor * memory_v; - - // - struct ggml_context * ctx; - std::map tensors; -}; - -// load the model's weights from a file -bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx, ggml_type memory_type = GGML_TYPE_F32) { - fprintf(stderr, "%s: loading model from '%s' - please wait ...\n", __func__, fname.c_str()); - - std::vector f_buf(1024*1024); - - auto fin = std::ifstream(fname, std::ios::binary); - fin.rdbuf()->pubsetbuf(f_buf.data(), f_buf.size()); - if (!fin) { - fprintf(stderr, "%s: failed to open '%s'\n", __func__, fname.c_str()); - return false; - } - - // verify magic - { - uint32_t magic; - fin.read((char *) &magic, sizeof(magic)); - if (magic != 0x67676d6c) { - fprintf(stderr, "%s: invalid model file '%s' (bad magic)\n", __func__, fname.c_str()); - return false; - } - } - - int n_ff = 0; - int n_parts = 0; - - // load hparams - { - auto & hparams = model.hparams; - - fin.read((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); - //fin.read((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); - fin.read((char *) &hparams.n_embd, sizeof(hparams.n_embd)); - fin.read((char *) &hparams.n_mult, sizeof(hparams.n_mult)); - fin.read((char *) &hparams.n_head, sizeof(hparams.n_head)); - fin.read((char *) &hparams.n_layer, sizeof(hparams.n_layer)); - fin.read((char *) &hparams.n_rot, sizeof(hparams.n_rot)); - fin.read((char *) &hparams.f16, sizeof(hparams.f16)); - - hparams.n_ctx = n_ctx; - - n_ff = ((2*(4*hparams.n_embd)/3 + hparams.n_mult - 1)/hparams.n_mult)*hparams.n_mult; - n_parts = LLAMA_N_PARTS.at(hparams.n_embd); - - fprintf(stderr, "%s: n_vocab = %d\n", __func__, hparams.n_vocab); - fprintf(stderr, "%s: n_ctx = %d\n", __func__, hparams.n_ctx); - fprintf(stderr, "%s: n_embd = %d\n", __func__, hparams.n_embd); - fprintf(stderr, "%s: n_mult = %d\n", __func__, hparams.n_mult); - fprintf(stderr, "%s: n_head = %d\n", __func__, hparams.n_head); - fprintf(stderr, "%s: n_layer = %d\n", __func__, hparams.n_layer); - fprintf(stderr, "%s: n_rot = %d\n", __func__, hparams.n_rot); - fprintf(stderr, "%s: f16 = %d\n", __func__, hparams.f16); - fprintf(stderr, "%s: n_ff = %d\n", __func__, n_ff); - fprintf(stderr, "%s: n_parts = %d\n", __func__, n_parts); - } - - // load vocab - { - std::string word; - for (int i = 0; i < model.hparams.n_vocab; i++) { - uint32_t len; - fin.read((char *) &len, sizeof(len)); - - word.resize(len); - fin.read((char *) word.data(), len); - - vocab.token_to_id[word] = i; - vocab.id_to_token[i] = word; - - //if (i < 30000) { - // fprintf(stderr, "%s: vocab[%d] = '%s'\n", __func__, i, word.c_str()); - //} - } - } - - // for the big tensors, we have the option to store the data in 16-bit floats or quantized - // in order to save memory and also to speed up the computation - ggml_type wtype = GGML_TYPE_COUNT; - switch (model.hparams.f16) { - case 0: wtype = GGML_TYPE_F32; break; - case 1: wtype = GGML_TYPE_F16; break; - case 2: wtype = GGML_TYPE_Q4_0; break; - case 3: wtype = GGML_TYPE_Q4_1; break; - default: - { - fprintf(stderr, "%s: invalid model file '%s' (bad f16 value %d)\n", - __func__, fname.c_str(), model.hparams.f16); - return false; - } - } - - auto & ctx = model.ctx; - - size_t ctx_size = 0; - - { - const auto & hparams = model.hparams; - - const int n_embd = hparams.n_embd; - const int n_layer = hparams.n_layer; - const int n_ctx = hparams.n_ctx; - const int n_vocab = hparams.n_vocab; - - ctx_size += n_embd*n_vocab*ggml_type_sizef(wtype); // tok_embeddings - - ctx_size += n_embd*ggml_type_sizef(GGML_TYPE_F32); // norm - - ctx_size += n_embd*n_vocab*ggml_type_sizef(wtype); // output - - ctx_size += n_layer*(n_embd*ggml_type_sizef(GGML_TYPE_F32)); // attention_norm - - ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wq - ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wk - ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wv - ctx_size += n_layer*(n_embd*n_embd*ggml_type_sizef(wtype)); // wo - - ctx_size += n_layer*(n_embd*ggml_type_sizef(GGML_TYPE_F32)); // ffn_norm - - ctx_size += n_layer*(n_ff*n_embd*ggml_type_sizef(wtype)); // w1 - ctx_size += n_layer*(n_ff*n_embd*ggml_type_sizef(wtype)); // w2 - ctx_size += n_layer*(n_ff*n_embd*ggml_type_sizef(wtype)); // w3 - - ctx_size += n_ctx*n_layer*n_embd*ggml_type_sizef(memory_type); // memory_k - ctx_size += n_ctx*n_layer*n_embd*ggml_type_sizef(memory_type); // memory_v - - ctx_size += (5 + 10*n_layer)*256; // object overhead - - fprintf(stderr, "%s: ggml ctx size = %6.2f MB\n", __func__, ctx_size/(1024.0*1024.0)); - } - - // create the ggml context - { - struct ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ NULL, - }; - - model.ctx = ggml_init(params); - if (!model.ctx) { - fprintf(stderr, "%s: ggml_init() failed\n", __func__); - return false; - } - } - - // prepare memory for the weights - { - const auto & hparams = model.hparams; - - const int n_embd = hparams.n_embd; - const int n_layer = hparams.n_layer; - const int n_vocab = hparams.n_vocab; - - model.layers.resize(n_layer); - - model.tok_embeddings = ggml_new_tensor_2d(ctx, wtype, n_embd, n_vocab); - - model.norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); - model.output = ggml_new_tensor_2d(ctx, wtype, n_embd, n_vocab); - - // map by name - model.tensors["tok_embeddings.weight"] = model.tok_embeddings; - - model.tensors["norm.weight"] = model.norm; - model.tensors["output.weight"] = model.output; - - for (int i = 0; i < n_layer; ++i) { - auto & layer = model.layers[i]; - - layer.attention_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); - - layer.wq = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); - layer.wk = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); - layer.wv = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); - layer.wo = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd); - - layer.ffn_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); - - layer.w1 = ggml_new_tensor_2d(ctx, wtype, n_embd, n_ff); - layer.w2 = ggml_new_tensor_2d(ctx, wtype, n_ff, n_embd); - layer.w3 = ggml_new_tensor_2d(ctx, wtype, n_embd, n_ff); - - // map by name - model.tensors["layers." + std::to_string(i) + ".attention_norm.weight"] = layer.attention_norm; - - model.tensors["layers." + std::to_string(i) + ".attention.wq.weight"] = layer.wq; - model.tensors["layers." + std::to_string(i) + ".attention.wk.weight"] = layer.wk; - model.tensors["layers." + std::to_string(i) + ".attention.wv.weight"] = layer.wv; - model.tensors["layers." + std::to_string(i) + ".attention.wo.weight"] = layer.wo; - - model.tensors["layers." + std::to_string(i) + ".ffn_norm.weight"] = layer.ffn_norm; - - model.tensors["layers." + std::to_string(i) + ".feed_forward.w1.weight"] = layer.w1; - model.tensors["layers." + std::to_string(i) + ".feed_forward.w2.weight"] = layer.w2; - model.tensors["layers." + std::to_string(i) + ".feed_forward.w3.weight"] = layer.w3; - } - } - - // key + value memory - { - const auto & hparams = model.hparams; - - const int n_embd = hparams.n_embd; - const int n_layer = hparams.n_layer; - const int n_ctx = hparams.n_ctx; - - const int n_mem = n_layer*n_ctx; - const int n_elements = n_embd*n_mem; - - model.memory_k = ggml_new_tensor_1d(ctx, memory_type, n_elements); - model.memory_v = ggml_new_tensor_1d(ctx, memory_type, n_elements); - - const size_t memory_size = ggml_nbytes(model.memory_k) + ggml_nbytes(model.memory_v); - - fprintf(stderr, "%s: memory_size = %8.2f MB, n_mem = %d\n", __func__, memory_size/1024.0/1024.0, n_mem); - } - - const size_t file_offset = fin.tellg(); - - fin.close(); - - std::vector tmp; - - for (int i = 0; i < n_parts; ++i) { - const int part_id = i; - //const int part_id = n_parts - i - 1; - - std::string fname_part = fname; - if (i > 0) { - fname_part += "." + std::to_string(i); - } - - fprintf(stderr, "%s: loading model part %d/%d from '%s'\n", __func__, i+1, n_parts, fname_part.c_str()); - - fin = std::ifstream(fname_part, std::ios::binary); - fin.rdbuf()->pubsetbuf(f_buf.data(), f_buf.size()); - fin.seekg(file_offset); - - // load weights - { - int n_tensors = 0; - size_t total_size = 0; - - fprintf(stderr, "%s: ", __func__); - - while (true) { - int32_t n_dims; - int32_t length; - int32_t ftype; - - fin.read(reinterpret_cast(&n_dims), sizeof(n_dims)); - fin.read(reinterpret_cast(&length), sizeof(length)); - fin.read(reinterpret_cast(&ftype), sizeof(ftype)); - - if (fin.eof()) { - break; - } - - int32_t nelements = 1; - int32_t ne[2] = { 1, 1 }; - for (int i = 0; i < n_dims; ++i) { - fin.read(reinterpret_cast(&ne[i]), sizeof(ne[i])); - nelements *= ne[i]; - } - - std::string name(length, 0); - fin.read(&name[0], length); - - if (model.tensors.find(name.data()) == model.tensors.end()) { - fprintf(stderr, "%s: unknown tensor '%s' in model file\n", __func__, name.data()); - return false; - } - - // split_type = 0: split by columns - // split_type = 1: split by rows - int split_type = 0; - - // split_type = 0: - // regex: - // - tok_embeddings.* - // - layers.*.attention.wo.weight - // - layers.*.feed_forward.w2.weight - - // split_type = 1: - // regex: - // - output.* - // - layers.*.attention.wq.weight - // - layers.*.attention.wk.weight - // - layers.*.attention.wv.weight - // - layers.*.feed_forward.w1.weight - // - layers.*.feed_forward.w3.weight - if (name.find("tok_embeddings") != std::string::npos) { - split_type = 0; - } else if (name.find("layers") != std::string::npos) { - if (name.find("attention.wo.weight") != std::string::npos) { - split_type = 0; - } else if (name.find("feed_forward.w2.weight") != std::string::npos) { - split_type = 0; - } else { - split_type = 1; - } - } else if (name.find("output") != std::string::npos) { - split_type = 1; - } - - auto tensor = model.tensors[name.data()]; - - if (n_dims == 1) { - if (ggml_nelements(tensor) != nelements) { - fprintf(stderr, "%s: tensor '%s' has wrong size in model file\n", __func__, name.data()); - return false; - } - } else { - if (ggml_nelements(tensor)/n_parts != nelements) { - fprintf(stderr, "%s: tensor '%s' has wrong size in model file\n", __func__, name.data()); - return false; - } - } - - if (n_dims == 1) { - if (tensor->ne[0] != ne[0] || tensor->ne[1] != ne[1]) { - fprintf(stderr, "%s: tensor '%s' has wrong shape in model file: got [%d, %d], expected [%d, %d]\n", - __func__, name.data(), tensor->ne[0], tensor->ne[1], ne[0], ne[1]); - return false; - } - } else { - if (split_type == 0) { - if (tensor->ne[0]/n_parts != ne[0] || tensor->ne[1] != ne[1]) { - fprintf(stderr, "%s: tensor '%s' has wrong shape in model file: got [%d, %d], expected [%d, %d]\n", - __func__, name.data(), tensor->ne[0]/n_parts, tensor->ne[1], ne[0], ne[1]); - return false; - } - } else { - if (tensor->ne[0] != ne[0] || tensor->ne[1]/n_parts != ne[1]) { - fprintf(stderr, "%s: tensor '%s' has wrong shape in model file: got [%d, %d], expected [%d, %d]\n", - __func__, name.data(), tensor->ne[0], tensor->ne[1]/n_parts, ne[0], ne[1]); - return false; - } - } - } - - if (0) { - static const char * ftype_str[] = { "f32", "f16", "q4_0", "q4_1", }; - fprintf(stderr, "%24s - [%5d, %5d], type = %6s, split = %d\n", name.data(), ne[0], ne[1], ftype_str[ftype], split_type); - } - - size_t bpe = 0; - - switch (ftype) { - case 0: bpe = ggml_type_size(GGML_TYPE_F32); break; - case 1: bpe = ggml_type_size(GGML_TYPE_F16); break; - case 2: bpe = ggml_type_size(GGML_TYPE_Q4_0); assert(ne[0] % 64 == 0); break; - case 3: bpe = ggml_type_size(GGML_TYPE_Q4_1); assert(ne[0] % 64 == 0); break; - default: - { - fprintf(stderr, "%s: unknown ftype %d in model file\n", __func__, ftype); - return false; - } - }; - - if (n_dims == 1 || n_parts == 1) { - if ((nelements*bpe)/ggml_blck_size(tensor->type) != ggml_nbytes(tensor)) { - fprintf(stderr, "%s: tensor '%s' has wrong size in model file: got %zu, expected %zu\n", - __func__, name.data(), ggml_nbytes(tensor), nelements*bpe); - return false; - } - - if (part_id == 0) { - fin.read(reinterpret_cast(tensor->data), ggml_nbytes(tensor)); - } else { - fin.seekg(ggml_nbytes(tensor), std::ios::cur); - } - - total_size += ggml_nbytes(tensor); - } else { - if ((nelements*bpe)/ggml_blck_size(tensor->type) != ggml_nbytes(tensor)/n_parts) { - fprintf(stderr, "%s: tensor '%s' has wrong size in model file: got %zu, expected %zu\n", - __func__, name.data(), ggml_nbytes(tensor)/n_parts, nelements*bpe); - return false; - } - - if (split_type == 0) { - const int np0 = ne[0]; - - const size_t row_size = (tensor->ne[0]/ggml_blck_size(tensor->type))*ggml_type_size(tensor->type); - assert(row_size == tensor->nb[1]); - - for (int i1 = 0; i1 < ne[1]; ++i1) { - const size_t offset_row = i1*row_size; - const size_t offset = offset_row + ((part_id*np0)/ggml_blck_size(tensor->type))*ggml_type_size(tensor->type); - fin.read(reinterpret_cast(tensor->data) + offset, row_size/n_parts); - } - } else { - const int np1 = ne[1]; - - const size_t row_size = (tensor->ne[0]/ggml_blck_size(tensor->type))*ggml_type_size(tensor->type); - - for (int i1 = 0; i1 < ne[1]; ++i1) { - const size_t offset_row = (i1 + part_id*np1)*row_size; - fin.read(reinterpret_cast(tensor->data) + offset_row, row_size); - } - } - - total_size += ggml_nbytes(tensor)/n_parts; - } - - //fprintf(stderr, "%42s - [%5d, %5d], type = %6s, %6.2f MB\n", name.data(), ne[0], ne[1], ftype == 0 ? "float" : "f16", ggml_nbytes(tensor)/1024.0/1024.0); - if (++n_tensors % 8 == 0) { - fprintf(stderr, "."); - fflush(stderr); - } - } - - fprintf(stderr, " done\n"); - - fprintf(stderr, "%s: model size = %8.2f MB / num tensors = %d\n", __func__, total_size/1024.0/1024.0, n_tensors); - } - - fin.close(); - } - - return true; -} - -// evaluate the transformer -// -// - model: the model -// - n_threads: number of threads to use -// - n_past: the context size so far -// - embd_inp: the embeddings of the tokens in the context -// - embd_w: the predicted logits for the next token -// -// The GPT-J model requires about 16MB of memory per input token. -// -bool llama_eval( - const llama_model & model, - const int n_threads, - const int n_past, - const std::vector & embd_inp, - std::vector & embd_w, - size_t & mem_per_token) { - const int N = embd_inp.size(); - - const auto & hparams = model.hparams; - - const int n_embd = hparams.n_embd; - const int n_layer = hparams.n_layer; - const int n_ctx = hparams.n_ctx; - const int n_head = hparams.n_head; - const int n_vocab = hparams.n_vocab; - const int n_rot = hparams.n_embd/hparams.n_head; - - // TODO: check if this size scales with n_ctx linearly and remove constant. somehow I feel it wasn't the case - // static size_t buf_size = hparams.n_ctx*1024*1024; - static size_t buf_size = 512u*1024*1024; - static void * buf = malloc(buf_size); - - if (mem_per_token > 0 && mem_per_token*N > buf_size) { - const size_t buf_size_new = 1.1*(mem_per_token*N); // add 10% to account for ggml object overhead - //fprintf(stderr, "\n%s: reallocating buffer from %zu to %zu bytes\n", __func__, buf_size, buf_size_new); - - // reallocate - buf_size = buf_size_new; - buf = realloc(buf, buf_size); - if (buf == nullptr) { - fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size); - return false; - } - } - - struct ggml_init_params params = { - /*.mem_size =*/ buf_size, - /*.mem_buffer =*/ buf, - }; - - struct ggml_context * ctx0 = ggml_init(params); - ggml_cgraph gf = {}; - gf.n_threads = n_threads; - - struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N); - memcpy(embd->data, embd_inp.data(), N*ggml_element_size(embd)); - - struct ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embeddings, embd); - - for (int il = 0; il < n_layer; ++il) { - struct ggml_tensor * inpSA = inpL; - - struct ggml_tensor * cur; - - // norm - { - cur = ggml_rms_norm(ctx0, inpL); - - // cur = attention_norm*cur - cur = ggml_mul(ctx0, - ggml_repeat(ctx0, model.layers[il].attention_norm, cur), - cur); - } - - // self-attention - { - struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, model.layers[il].wq, cur); - struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur); - struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, model.layers[il].wv, cur); - - // store key and value to memory - if (N >= 1) { - struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_k, N*n_embd, (ggml_element_size(model.memory_k)*n_embd)*(il*n_ctx + n_past)); - struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_v, N*n_embd, (ggml_element_size(model.memory_v)*n_embd)*(il*n_ctx + n_past)); - - ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Kcur, k)); - ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcur, v)); - } - - // Q = Qcur.contiguous().view(n_embd/n_head, n_head, N).permute(0, 2, 1, 3) - struct ggml_tensor * Q = - ggml_permute(ctx0, - ggml_rope(ctx0, - ggml_cpy(ctx0, - Qcur, - ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd/n_head, n_head, N)), - n_past, n_rot, 0), - 0, 2, 1, 3); - - // K = Kmem.view(n_embd/n_head, n_head, n_past + N).permute(0, 2, 1, 3) - struct ggml_tensor * K = - ggml_permute(ctx0, - ggml_rope(ctx0, - ggml_reshape_3d(ctx0, - ggml_view_1d(ctx0, model.memory_k, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_k)*n_embd), - n_embd/n_head, n_head, n_past + N), - n_past, n_rot, 1), - 0, 2, 1, 3); - - // K * Q - struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); - - // KQ_scaled = KQ / sqrt(n_embd/n_head) - struct ggml_tensor * KQ_scaled = - ggml_scale(ctx0, - KQ, - ggml_new_f32(ctx0, 1.0f/sqrt(float(n_embd)/n_head)) - ); - - // KQ_masked = mask_past(KQ_scaled) - struct ggml_tensor * KQ_masked = ggml_diag_mask_inf(ctx0, KQ_scaled, n_past); - - // KQ = soft_max(KQ_masked) - struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctx0, KQ_masked); - - // V_trans = Vmem.view(n_embd/n_head, n_head, n_past + N).permute(1, 2, 0, 3).contiguous() - struct ggml_tensor * V_trans = - ggml_permute(ctx0, - ggml_reshape_3d(ctx0, - ggml_view_1d(ctx0, model.memory_v, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_v)*n_embd), - n_embd/n_head, n_head, n_past + N), - 1, 2, 0, 3); - - // KQV = transpose(V) * KQ_soft_max - struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V_trans, KQ_soft_max); - - // KQV_merged = KQV.permute(0, 2, 1, 3) - struct ggml_tensor * KQV_merged = ggml_permute(ctx0, KQV, 0, 2, 1, 3); - - // cur = KQV_merged.contiguous().view(n_embd, N) - cur = ggml_cpy(ctx0, - KQV_merged, - ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, N)); - - // projection (no bias) - cur = ggml_mul_mat(ctx0, - model.layers[il].wo, - cur); - } - - struct ggml_tensor * inpFF = ggml_add(ctx0, cur, inpSA); - - // feed-forward network - { - // norm - { - cur = ggml_rms_norm(ctx0, inpFF); - - // cur = ffn_norm*cur - cur = ggml_mul(ctx0, - ggml_repeat(ctx0, model.layers[il].ffn_norm, cur), - cur); - } - - struct ggml_tensor * tmp = ggml_mul_mat(ctx0, - model.layers[il].w3, - cur); - - - cur = ggml_mul_mat(ctx0, - model.layers[il].w1, - cur); - - // SILU activation - cur = ggml_silu(ctx0, cur); - - cur = ggml_mul(ctx0, cur, tmp); - - cur = ggml_mul_mat(ctx0, - model.layers[il].w2, - cur); - } - - cur = ggml_add(ctx0, cur, inpFF); - - // input for next layer - inpL = cur; - } - - // norm - { - inpL = ggml_rms_norm(ctx0, inpL); - - // inpL = norm*inpL - inpL = ggml_mul(ctx0, - ggml_repeat(ctx0, model.norm, inpL), - inpL); - } - - // lm_head - { - inpL = ggml_mul_mat(ctx0, model.output, inpL); - } - - // logits -> probs - //inpL = ggml_soft_max(ctx0, inpL); - - // run the computation - ggml_build_forward_expand(&gf, inpL); - ggml_graph_compute (ctx0, &gf); - - //if (n_past%100 == 0) { - // ggml_graph_print (&gf); - // ggml_graph_dump_dot(&gf, NULL, "gpt-2.dot"); - //} - - //embd_w.resize(n_vocab*N); - //memcpy(embd_w.data(), ggml_get_data(inpL), sizeof(float)*n_vocab*N); - - // return result for just the last token - embd_w.resize(n_vocab); - memcpy(embd_w.data(), (float *) ggml_get_data(inpL) + (n_vocab*(N-1)), sizeof(float)*n_vocab); - - if (mem_per_token == 0) { - mem_per_token = ggml_used_mem(ctx0)/N; - } - //fprintf(stderr, "used_mem = %zu\n", ggml_used_mem(ctx0)); - - ggml_free(ctx0); - - return true; -} static bool is_interacting = false; @@ -786,7 +80,7 @@ int main(int argc, char ** argv) { const int64_t t_main_start_us = ggml_time_us(); gpt_params params; - params.model = "models/llama-7B/ggml-model.bin"; + params.model = "models/7B/ggml-model-q4_0.bin"; if (gpt_params_parse(argc, argv, params) == false) { return 1; From ebfc622197ed779c47b4c34b7fabca20f515dc88 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Sun, 12 Mar 2023 18:48:48 -0700 Subject: [PATCH 03/18] Move llama_model_quantize() into llama.cpp --- Makefile | 4 +- llama.cpp | 260 ++++++++++++++++++++++++++++++++++++++++++++++++ llama.h | 1 + quantize.cpp | 271 +-------------------------------------------------- 4 files changed, 264 insertions(+), 272 deletions(-) diff --git a/Makefile b/Makefile index a74929c82fd7e..20b200a2dff4f 100644 --- a/Makefile +++ b/Makefile @@ -198,8 +198,8 @@ main: main.cpp ggml.o utils.o llama.o $(CXX) $(CXXFLAGS) main.cpp ggml.o llama.o utils.o -o main $(LDFLAGS) ./main -h -quantize: quantize.cpp ggml.o utils.o - $(CXX) $(CXXFLAGS) quantize.cpp ggml.o utils.o -o quantize $(LDFLAGS) +quantize: quantize.cpp ggml.o utils.o llama.o + $(CXX) $(CXXFLAGS) quantize.cpp ggml.o llama.o utils.o -o quantize $(LDFLAGS) # # Tests diff --git a/llama.cpp b/llama.cpp index 034666327401c..772600b4c01d3 100644 --- a/llama.cpp +++ b/llama.cpp @@ -9,6 +9,11 @@ #include #include #include +#include + +// TODO: move somewhere else +#define QK 32 + // determine number of model parts based on the dimension static const std::map LLAMA_N_PARTS = { @@ -680,3 +685,258 @@ bool llama_eval( return true; } +bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype) { + ggml_type type = GGML_TYPE_Q4_1; + + switch (itype) { + case 2: type = GGML_TYPE_Q4_0; break; + case 3: type = GGML_TYPE_Q4_1; break; + default: fprintf(stderr, "%s: invalid quantization type %d\n", __func__, itype); return 1; + }; + + if (type != GGML_TYPE_Q4_0 && type != GGML_TYPE_Q4_1) { + fprintf(stderr, "%s: invalid quantization type %d\n", __func__, type); + return false; + } + + gpt_vocab vocab; + + printf("%s: loading model from '%s'\n", __func__, fname_inp.c_str()); + + auto finp = std::ifstream(fname_inp, std::ios::binary); + if (!finp) { + fprintf(stderr, "%s: failed to open '%s' for reading\n", __func__, fname_inp.c_str()); + return false; + } + + auto fout = std::ofstream(fname_out, std::ios::binary); + if (!fout) { + fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_out.c_str()); + return false; + } + + // verify magic + { + uint32_t magic; + finp.read((char *) &magic, sizeof(magic)); + if (magic != 0x67676d6c) { + fprintf(stderr, "%s: invalid model file '%s' (bad magic)\n", __func__, fname_inp.c_str()); + return false; + } + + fout.write((char *) &magic, sizeof(magic)); + } + + llama_hparams hparams; + + // load hparams + { + finp.read((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); + //finp.read((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); + finp.read((char *) &hparams.n_embd, sizeof(hparams.n_embd)); + finp.read((char *) &hparams.n_mult, sizeof(hparams.n_mult)); + finp.read((char *) &hparams.n_head, sizeof(hparams.n_head)); + finp.read((char *) &hparams.n_layer, sizeof(hparams.n_layer)); + finp.read((char *) &hparams.n_rot, sizeof(hparams.n_rot)); + finp.read((char *) &hparams.f16, sizeof(hparams.f16)); + + printf("%s: n_vocab = %d\n", __func__, hparams.n_vocab); + printf("%s: n_ctx = %d\n", __func__, hparams.n_ctx); + printf("%s: n_embd = %d\n", __func__, hparams.n_embd); + printf("%s: n_mult = %d\n", __func__, hparams.n_mult); + printf("%s: n_head = %d\n", __func__, hparams.n_head); + printf("%s: n_layer = %d\n", __func__, hparams.n_layer); + printf("%s: f16 = %d\n", __func__, hparams.f16); + + fout.write((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); + //fout.write((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); + fout.write((char *) &hparams.n_embd, sizeof(hparams.n_embd)); + fout.write((char *) &hparams.n_mult, sizeof(hparams.n_mult)); + fout.write((char *) &hparams.n_head, sizeof(hparams.n_head)); + fout.write((char *) &hparams.n_layer, sizeof(hparams.n_layer)); + fout.write((char *) &hparams.n_rot, sizeof(hparams.n_rot)); + fout.write((char *) &itype, sizeof(hparams.f16)); + } + + // load vocab + { + const int32_t n_vocab = hparams.n_vocab; + + if (n_vocab != hparams.n_vocab) { + fprintf(stderr, "%s: invalid model file '%s' (bad vocab size %d != %d)\n", + __func__, fname_inp.c_str(), n_vocab, hparams.n_vocab); + return false; + } + + std::string word; + for (int i = 0; i < n_vocab; i++) { + uint32_t len; + finp.read ((char *) &len, sizeof(len)); + fout.write((char *) &len, sizeof(len)); + + word.resize(len); + finp.read ((char *) word.data(), len); + fout.write((char *) word.data(), len); + + vocab.token_to_id[word] = i; + vocab.id_to_token[i] = word; + } + } + + // load weights + { + size_t total_size_org = 0; + size_t total_size_new = 0; + + std::vector work; + + std::vector data_u8; + std::vector data_f16; + std::vector data_f32; + + std::vector hist_all(1 << 4, 0); + + while (true) { + int32_t n_dims; + int32_t length; + int32_t ftype; + + finp.read(reinterpret_cast(&n_dims), sizeof(n_dims)); + finp.read(reinterpret_cast(&length), sizeof(length)); + finp.read(reinterpret_cast(&ftype), sizeof(ftype)); + + if (finp.eof()) { + break; + } + + int32_t nelements = 1; + int32_t ne[2] = { 1, 1 }; + for (int i = 0; i < n_dims; ++i) { + finp.read (reinterpret_cast(&ne[i]), sizeof(ne[i])); + nelements *= ne[i]; + } + + std::string name(length, 0); + finp.read (&name[0], length); + + { + static const char * ftype_str[] = { "f32", "f16", "q4_0", "q4_1", }; + printf("%48s - [%5d, %5d], type = %6s ", name.data(), ne[0], ne[1], ftype_str[ftype]); + } + + // regexes of tensor names to be quantized + const std::vector k_names = { + ".*weight", + }; + + bool quantize = false; + for (const auto & s : k_names) { + if (std::regex_match(name, std::regex(s))) { + quantize = true; + break; + } + } + + // quantize only 2D tensors + quantize &= (n_dims == 2); + + if (quantize) { + if (ftype != 0 && ftype != 1) { + fprintf(stderr, "%s: unsupported ftype %d for integer quantization\n", __func__, ftype); + return false; + } + + if (ftype == 1) { + data_f16.resize(nelements); + finp.read(reinterpret_cast(data_f16.data()), nelements * sizeof(ggml_fp16_t)); + data_f32.resize(nelements); + for (int i = 0; i < nelements; ++i) { + data_f32[i] = ggml_fp16_to_fp32(data_f16[i]); + } + } else { + data_f32.resize(nelements); + finp.read(reinterpret_cast(data_f32.data()), nelements * sizeof(float)); + } + + ftype = itype; + } else { + const int bpe = (ftype == 0) ? sizeof(float) : sizeof(uint16_t); + + data_u8.resize(nelements*bpe); + finp.read(reinterpret_cast(data_u8.data()), nelements * bpe); + } + + fout.write(reinterpret_cast(&n_dims), sizeof(n_dims)); + fout.write(reinterpret_cast(&length), sizeof(length)); + fout.write(reinterpret_cast(&ftype), sizeof(ftype)); + for (int i = 0; i < n_dims; ++i) { + fout.write(reinterpret_cast(&ne[i]), sizeof(ne[i])); + } + fout.write(&name[0], length); + + if (quantize) { + printf("quantizing .. "); + work.resize(nelements); // for quantization + + size_t cur_size = 0; + std::vector hist_cur(1 << 4, 0); + + switch (type) { + case GGML_TYPE_Q4_0: + { + cur_size = ggml_quantize_q4_0(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); + } break; + case GGML_TYPE_Q4_1: + { + cur_size = ggml_quantize_q4_1(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); + } break; + default: + { + fprintf(stderr, "%s: unsupported quantization type %d\n", __func__, type); + return false; + } + } + + fout.write(reinterpret_cast(work.data()), cur_size); + total_size_new += cur_size; + + printf("size = %8.2f MB -> %8.2f MB | hist: ", nelements * sizeof(float)/1024.0/1024.0, cur_size/1024.0/1024.0); + for (int i = 0; i < hist_cur.size(); ++i) { + hist_all[i] += hist_cur[i]; + } + + for (int i = 0; i < hist_cur.size(); ++i) { + printf("%5.3f ", hist_cur[i] / (float)nelements); + } + printf("\n"); + } else { + printf("size = %8.3f MB\n", data_u8.size()/1024.0/1024.0); + fout.write(reinterpret_cast(data_u8.data()), data_u8.size()); + total_size_new += data_u8.size(); + } + + total_size_org += nelements * sizeof(float); + } + + printf("%s: model size = %8.2f MB\n", __func__, total_size_org/1024.0/1024.0); + printf("%s: quant size = %8.2f MB\n", __func__, total_size_new/1024.0/1024.0); + + { + int64_t sum_all = 0; + for (int i = 0; i < hist_all.size(); ++i) { + sum_all += hist_all[i]; + } + + printf("%s: hist: ", __func__); + for (int i = 0; i < hist_all.size(); ++i) { + printf("%5.3f ", hist_all[i] / (float)sum_all); + } + printf("\n"); + } + } + + finp.close(); + fout.close(); + + return true; +} diff --git a/llama.h b/llama.h index 0ceab22bbbef0..a902273048ce6 100644 --- a/llama.h +++ b/llama.h @@ -64,3 +64,4 @@ bool llama_eval( const std::vector & embd_inp, std::vector & embd_w, size_t & mem_per_token); +bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype); diff --git a/quantize.cpp b/quantize.cpp index 14c7b277a4024..b4ff997a191c1 100644 --- a/quantize.cpp +++ b/quantize.cpp @@ -1,6 +1,7 @@ #include "ggml.h" #include "utils.h" +#include "llama.h" #include #include @@ -10,280 +11,10 @@ #include #include #include -#include - -// TODO: move somewhere else -#define QK 32 - -// default hparams (LLaMA76B) -struct llama_hparams { - int32_t n_vocab = 32000; - int32_t n_ctx = 512; // this is provided as user input? - int32_t n_embd = 4096; - int32_t n_mult = 256; - int32_t n_head = 32; - int32_t n_layer = 32; - int32_t n_rot = 64; - int32_t f16 = 1; -}; // quantize a model -bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype) { - ggml_type type = GGML_TYPE_Q4_1; - - switch (itype) { - case 2: type = GGML_TYPE_Q4_0; break; - case 3: type = GGML_TYPE_Q4_1; break; - default: fprintf(stderr, "%s: invalid quantization type %d\n", __func__, itype); return 1; - }; - - if (type != GGML_TYPE_Q4_0 && type != GGML_TYPE_Q4_1) { - fprintf(stderr, "%s: invalid quantization type %d\n", __func__, type); - return false; - } - - gpt_vocab vocab; - - printf("%s: loading model from '%s'\n", __func__, fname_inp.c_str()); - - auto finp = std::ifstream(fname_inp, std::ios::binary); - if (!finp) { - fprintf(stderr, "%s: failed to open '%s' for reading\n", __func__, fname_inp.c_str()); - return false; - } - - auto fout = std::ofstream(fname_out, std::ios::binary); - if (!fout) { - fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_out.c_str()); - return false; - } - - // verify magic - { - uint32_t magic; - finp.read((char *) &magic, sizeof(magic)); - if (magic != 0x67676d6c) { - fprintf(stderr, "%s: invalid model file '%s' (bad magic)\n", __func__, fname_inp.c_str()); - return false; - } - - fout.write((char *) &magic, sizeof(magic)); - } - - llama_hparams hparams; - - // load hparams - { - finp.read((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); - //finp.read((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); - finp.read((char *) &hparams.n_embd, sizeof(hparams.n_embd)); - finp.read((char *) &hparams.n_mult, sizeof(hparams.n_mult)); - finp.read((char *) &hparams.n_head, sizeof(hparams.n_head)); - finp.read((char *) &hparams.n_layer, sizeof(hparams.n_layer)); - finp.read((char *) &hparams.n_rot, sizeof(hparams.n_rot)); - finp.read((char *) &hparams.f16, sizeof(hparams.f16)); - - printf("%s: n_vocab = %d\n", __func__, hparams.n_vocab); - printf("%s: n_ctx = %d\n", __func__, hparams.n_ctx); - printf("%s: n_embd = %d\n", __func__, hparams.n_embd); - printf("%s: n_mult = %d\n", __func__, hparams.n_mult); - printf("%s: n_head = %d\n", __func__, hparams.n_head); - printf("%s: n_layer = %d\n", __func__, hparams.n_layer); - printf("%s: f16 = %d\n", __func__, hparams.f16); - - fout.write((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); - //fout.write((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); - fout.write((char *) &hparams.n_embd, sizeof(hparams.n_embd)); - fout.write((char *) &hparams.n_mult, sizeof(hparams.n_mult)); - fout.write((char *) &hparams.n_head, sizeof(hparams.n_head)); - fout.write((char *) &hparams.n_layer, sizeof(hparams.n_layer)); - fout.write((char *) &hparams.n_rot, sizeof(hparams.n_rot)); - fout.write((char *) &itype, sizeof(hparams.f16)); - } - - // load vocab - { - const int32_t n_vocab = hparams.n_vocab; - - if (n_vocab != hparams.n_vocab) { - fprintf(stderr, "%s: invalid model file '%s' (bad vocab size %d != %d)\n", - __func__, fname_inp.c_str(), n_vocab, hparams.n_vocab); - return false; - } - - std::string word; - for (int i = 0; i < n_vocab; i++) { - uint32_t len; - finp.read ((char *) &len, sizeof(len)); - fout.write((char *) &len, sizeof(len)); - - word.resize(len); - finp.read ((char *) word.data(), len); - fout.write((char *) word.data(), len); - - vocab.token_to_id[word] = i; - vocab.id_to_token[i] = word; - } - } - - // load weights - { - size_t total_size_org = 0; - size_t total_size_new = 0; - std::vector work; - - std::vector data_u8; - std::vector data_f16; - std::vector data_f32; - - std::vector hist_all(1 << 4, 0); - - while (true) { - int32_t n_dims; - int32_t length; - int32_t ftype; - - finp.read(reinterpret_cast(&n_dims), sizeof(n_dims)); - finp.read(reinterpret_cast(&length), sizeof(length)); - finp.read(reinterpret_cast(&ftype), sizeof(ftype)); - - if (finp.eof()) { - break; - } - - int32_t nelements = 1; - int32_t ne[2] = { 1, 1 }; - for (int i = 0; i < n_dims; ++i) { - finp.read (reinterpret_cast(&ne[i]), sizeof(ne[i])); - nelements *= ne[i]; - } - - std::string name(length, 0); - finp.read (&name[0], length); - - { - static const char * ftype_str[] = { "f32", "f16", "q4_0", "q4_1", }; - printf("%48s - [%5d, %5d], type = %6s ", name.data(), ne[0], ne[1], ftype_str[ftype]); - } - - // regexes of tensor names to be quantized - const std::vector k_names = { - ".*weight", - }; - - bool quantize = false; - for (const auto & s : k_names) { - if (std::regex_match(name, std::regex(s))) { - quantize = true; - break; - } - } - - // quantize only 2D tensors - quantize &= (n_dims == 2); - - if (quantize) { - if (ftype != 0 && ftype != 1) { - fprintf(stderr, "%s: unsupported ftype %d for integer quantization\n", __func__, ftype); - return false; - } - - if (ftype == 1) { - data_f16.resize(nelements); - finp.read(reinterpret_cast(data_f16.data()), nelements * sizeof(ggml_fp16_t)); - data_f32.resize(nelements); - for (int i = 0; i < nelements; ++i) { - data_f32[i] = ggml_fp16_to_fp32(data_f16[i]); - } - } else { - data_f32.resize(nelements); - finp.read(reinterpret_cast(data_f32.data()), nelements * sizeof(float)); - } - - ftype = itype; - } else { - const int bpe = (ftype == 0) ? sizeof(float) : sizeof(uint16_t); - - data_u8.resize(nelements*bpe); - finp.read(reinterpret_cast(data_u8.data()), nelements * bpe); - } - - fout.write(reinterpret_cast(&n_dims), sizeof(n_dims)); - fout.write(reinterpret_cast(&length), sizeof(length)); - fout.write(reinterpret_cast(&ftype), sizeof(ftype)); - for (int i = 0; i < n_dims; ++i) { - fout.write(reinterpret_cast(&ne[i]), sizeof(ne[i])); - } - fout.write(&name[0], length); - - if (quantize) { - printf("quantizing .. "); - work.resize(nelements); // for quantization - - size_t cur_size = 0; - std::vector hist_cur(1 << 4, 0); - - switch (type) { - case GGML_TYPE_Q4_0: - { - cur_size = ggml_quantize_q4_0(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); - } break; - case GGML_TYPE_Q4_1: - { - cur_size = ggml_quantize_q4_1(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); - } break; - default: - { - fprintf(stderr, "%s: unsupported quantization type %d\n", __func__, type); - return false; - } - } - - fout.write(reinterpret_cast(work.data()), cur_size); - total_size_new += cur_size; - - printf("size = %8.2f MB -> %8.2f MB | hist: ", nelements * sizeof(float)/1024.0/1024.0, cur_size/1024.0/1024.0); - for (int i = 0; i < hist_cur.size(); ++i) { - hist_all[i] += hist_cur[i]; - } - - for (int i = 0; i < hist_cur.size(); ++i) { - printf("%5.3f ", hist_cur[i] / (float)nelements); - } - printf("\n"); - } else { - printf("size = %8.3f MB\n", data_u8.size()/1024.0/1024.0); - fout.write(reinterpret_cast(data_u8.data()), data_u8.size()); - total_size_new += data_u8.size(); - } - - total_size_org += nelements * sizeof(float); - } - - printf("%s: model size = %8.2f MB\n", __func__, total_size_org/1024.0/1024.0); - printf("%s: quant size = %8.2f MB\n", __func__, total_size_new/1024.0/1024.0); - - { - int64_t sum_all = 0; - for (int i = 0; i < hist_all.size(); ++i) { - sum_all += hist_all[i]; - } - - printf("%s: hist: ", __func__); - for (int i = 0; i < hist_all.size(); ++i) { - printf("%5.3f ", hist_all[i] / (float)sum_all); - } - printf("\n"); - } - } - - finp.close(); - fout.close(); - - return true; -} // usage: // ./llama-quantize models/llama/ggml-model.bin models/llama/ggml-model-quant.bin type From 0995df5c9e66b23242960538b8682b3ca086a64c Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Mon, 13 Mar 2023 23:10:22 -0700 Subject: [PATCH 04/18] Add to CMakeLists.txt --- CMakeLists.txt | 65 +++++++++++++++++++++++++++++--------------------- 1 file changed, 38 insertions(+), 27 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 38e7266dca630..ded54b4b6c351 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -7,56 +7,57 @@ set(CMAKE_C_STANDARD 11) set(THREADS_PREFER_PTHREAD_FLAG ON) find_package(Threads REQUIRED) -if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE) +if(NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE) set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo") endif() -option(LLAMA_ALL_WARNINGS "llama: enable all compiler warnings" ON) -option(LLAMA_ALL_WARNINGS_3RD_PARTY "llama: enable all compiler warnings in 3rd party libs" OFF) +option(LLAMA_ALL_WARNINGS "llama: enable all compiler warnings" ON) +option(LLAMA_ALL_WARNINGS_3RD_PARTY "llama: enable all compiler warnings in 3rd party libs" OFF) -option(LLAMA_SANITIZE_THREAD "llama: enable thread sanitizer" OFF) -option(LLAMA_SANITIZE_ADDRESS "llama: enable address sanitizer" OFF) -option(LLAMA_SANITIZE_UNDEFINED "llama: enable undefined sanitizer" OFF) +option(LLAMA_SANITIZE_THREAD "llama: enable thread sanitizer" OFF) +option(LLAMA_SANITIZE_ADDRESS "llama: enable address sanitizer" OFF) +option(LLAMA_SANITIZE_UNDEFINED "llama: enable undefined sanitizer" OFF) -if (APPLE) - option(LLAMA_NO_ACCELERATE "llama: disable Accelerate framework" OFF) - option(LLAMA_NO_AVX "llama: disable AVX" OFF) - option(LLAMA_NO_AVX2 "llama: disable AVX2" OFF) - option(LLAMA_NO_FMA "llama: disable FMA" OFF) +if(APPLE) + option(LLAMA_NO_ACCELERATE "llama: disable Accelerate framework" OFF) + option(LLAMA_NO_AVX "llama: disable AVX" OFF) + option(LLAMA_NO_AVX2 "llama: disable AVX2" OFF) + option(LLAMA_NO_FMA "llama: disable FMA" OFF) endif() -if (NOT MSVC) - if (LLAMA_SANITIZE_THREAD) - set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=thread") +if(NOT MSVC) + if(LLAMA_SANITIZE_THREAD) + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=thread") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=thread") endif() - if (LLAMA_SANITIZE_ADDRESS) + if(LLAMA_SANITIZE_ADDRESS) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=address -fno-omit-frame-pointer") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=address -fno-omit-frame-pointer") endif() - if (LLAMA_SANITIZE_UNDEFINED) + if(LLAMA_SANITIZE_UNDEFINED) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=undefined") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=undefined") endif() endif() -if (APPLE AND NOT LLAMA_NO_ACCELERATE) +if(APPLE AND NOT LLAMA_NO_ACCELERATE) find_library(ACCELERATE_FRAMEWORK Accelerate) - if (ACCELERATE_FRAMEWORK) + + if(ACCELERATE_FRAMEWORK) message(STATUS "Accelerate framework found") - set(LLAMA_EXTRA_LIBS ${LLAMA_EXTRA_LIBS} ${ACCELERATE_FRAMEWORK}) + set(LLAMA_EXTRA_LIBS ${LLAMA_EXTRA_LIBS} ${ACCELERATE_FRAMEWORK}) set(LLAMA_EXTRA_FLAGS ${LLAMA_EXTRA_FLAGS} -DGGML_USE_ACCELERATE) else() message(WARNING "Accelerate framework not found") endif() endif() -if (LLAMA_ALL_WARNINGS) - if (NOT MSVC) +if(LLAMA_ALL_WARNINGS) + if(NOT MSVC) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} \ -Wall \ -Wextra \ @@ -80,11 +81,12 @@ endif() message(STATUS "CMAKE_SYSTEM_PROCESSOR: ${CMAKE_SYSTEM_PROCESSOR}") -if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64") +if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64") message(STATUS "ARM detected") else() message(STATUS "x86 detected") - if (MSVC) + + if(MSVC) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /arch:AVX2") set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /arch:AVX2") set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /arch:AVX2") @@ -92,22 +94,25 @@ else() if(NOT LLAMA_NO_AVX) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx") endif() + if(NOT LLAMA_NO_AVX2) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx2") endif() + if(NOT LLAMA_NO_FMA) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mfma") endif() + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mf16c") endif() endif() # if (LLAMA_PERF) -# set(LLAMA_EXTRA_FLAGS ${LLAMA_EXTRA_FLAGS} -DGGML_PERF) +# set(LLAMA_EXTRA_FLAGS ${LLAMA_EXTRA_FLAGS} -DGGML_PERF) # endif() - add_executable(llama main.cpp + llama.cpp utils.cpp utils.h) @@ -116,6 +121,9 @@ add_executable(quantize utils.cpp utils.h) +add_library(llamalib + llama.cpp) + add_library(ggml ggml.c ggml.h) @@ -125,7 +133,10 @@ target_compile_definitions(llama PUBLIC ${LLAMA_EXTRA_FLAGS}) target_compile_definitions(quantize PUBLIC ${LLAMA_EXTRA_FLAGS}) target_link_libraries(ggml PRIVATE ${LLAMA_EXTRA_LIBS}) +target_link_libraries(llamalib PRIVATE ggml ${LLAMA_EXTRA_LIBS}) + target_include_directories(ggml PUBLIC .) -target_link_libraries(quantize PRIVATE ggml) -target_link_libraries(llama PRIVATE ggml) + +target_link_libraries(quantize PRIVATE ggml llamalib) +target_link_libraries(llama PRIVATE ggml llamalib) target_link_libraries(ggml PRIVATE Threads::Threads) From b14486e1c00c66471ba3baee248cb14c379b37fc Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Sun, 12 Mar 2023 18:48:48 -0700 Subject: [PATCH 05/18] Refactor code structure in llama.cpp and llama.h --- llama.cpp | 1014 +++++++++++++++++++++++++++++++++++------------------ llama.h | 75 ++-- 2 files changed, 716 insertions(+), 373 deletions(-) diff --git a/llama.cpp b/llama.cpp index 772600b4c01d3..70f763fd58196 100644 --- a/llama.cpp +++ b/llama.cpp @@ -10,11 +10,11 @@ #include #include #include +#include // TODO: move somewhere else #define QK 32 - // determine number of model parts based on the dimension static const std::map LLAMA_N_PARTS = { { 4096, 1 }, @@ -23,6 +23,96 @@ static const std::map LLAMA_N_PARTS = { { 8192, 8 }, }; +struct llama_layer { + // normalization + struct ggml_tensor * attention_norm; + + // attention + struct ggml_tensor * wq; + struct ggml_tensor * wk; + struct ggml_tensor * wv; + struct ggml_tensor * wo; + + // normalization + struct ggml_tensor * ffn_norm; + + // ff + struct ggml_tensor * w1; + struct ggml_tensor * w2; + struct ggml_tensor * w3; +}; +struct llama_model { + llama_hparams hparams; + + struct ggml_tensor * tok_embeddings; + + struct ggml_tensor * norm; + struct ggml_tensor * output; + + std::vector layers; + + // key + value memory + struct ggml_tensor * memory_k; + struct ggml_tensor * memory_v; + + // + struct ggml_context * ctx; + std::map tensors; +}; +struct llama_state +{ + int64_t t_sample_us = 0; + int64_t t_predict_us = 0; + + std::vector logits; + + mutable std::mt19937 rng; + + std::vector embd{}; + + int input_consumed = 0; + std::vector embd_inp; + std::vector last_n_tokens; + int remaining_tokens = 0; + int n_past = 0; + size_t mem_per_token = 0; + bool is_initialized = false; + llama_state() {} + + bool has_more_input() const { + return input_consumed < embd_inp.size(); + } +}; +struct llama_context +{ + int64_t t_load_us = 0; + int64_t t_start_us = 0; + + ggml_type wtype = ggml_type::GGML_TYPE_F16; // weight type (FP32 or FP16) + + llama_model model{}; + gpt_vocab vocab{}; + gpt_params params{}; + + std::unique_ptr state = nullptr; + + // default constructor + llama_context() = default; + // constructor + llama_context(llama_model&& model, gpt_vocab&& vocab, const gpt_params& params): + t_load_us(0), + t_start_us(0), + wtype(ggml_type::GGML_TYPE_F16), + model(std::move(model)), + vocab(std::move(vocab)), + params(params), + state(std::make_unique()) + { + } + ~llama_context(){ + ggml_free(model.ctx); + } +}; // load the model's weights from a file bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx) { @@ -450,147 +540,489 @@ bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab return true; } -// evaluate the transformer -// -// - model: the model -// - n_threads: number of threads to use -// - n_past: the context size so far -// - embd_inp: the embeddings of the tokens in the context -// - embd_w: the predicted logits for the next token -// -// The GPT-J model requires about 16MB of memory per input token. -// -bool llama_eval( - const llama_model & model, - const int n_threads, - const int n_past, - const std::vector & embd_inp, - std::vector & embd_w, - size_t & mem_per_token) { - const int N = embd_inp.size(); +bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype) { + ggml_type type = GGML_TYPE_Q4_1; - const auto & hparams = model.hparams; + switch (itype) { + case 2: type = GGML_TYPE_Q4_0; break; + case 3: type = GGML_TYPE_Q4_1; break; + default: fprintf(stderr, "%s: invalid quantization type %d\n", __func__, itype); return 1; + }; - const int n_embd = hparams.n_embd; - const int n_layer = hparams.n_layer; - const int n_ctx = hparams.n_ctx; - const int n_head = hparams.n_head; - const int n_vocab = hparams.n_vocab; - const int n_rot = hparams.n_embd/hparams.n_head; + if (type != GGML_TYPE_Q4_0 && type != GGML_TYPE_Q4_1) { + fprintf(stderr, "%s: invalid quantization type %d\n", __func__, type); + return false; + } - const int d_key = n_embd/n_head; + gpt_vocab vocab; - // TODO: check if this size scales with n_ctx linearly and remove constant. somehow I feel it wasn't the case - // static size_t buf_size = hparams.n_ctx*1024*1024; - static size_t buf_size = 512u*1024*1024; - static void * buf = malloc(buf_size); + printf("%s: loading model from '%s'\n", __func__, fname_inp.c_str()); - if (mem_per_token > 0 && mem_per_token*N > buf_size) { - const size_t buf_size_new = 1.1*(mem_per_token*N); // add 10% to account for ggml object overhead - //fprintf(stderr, "\n%s: reallocating buffer from %zu to %zu bytes\n", __func__, buf_size, buf_size_new); + auto finp = std::ifstream(fname_inp, std::ios::binary); + if (!finp) { + fprintf(stderr, "%s: failed to open '%s' for reading\n", __func__, fname_inp.c_str()); + return false; + } - // reallocate - buf_size = buf_size_new; - buf = realloc(buf, buf_size); - if (buf == nullptr) { - fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size); + auto fout = std::ofstream(fname_out, std::ios::binary); + if (!fout) { + fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_out.c_str()); + return false; + } + + // verify magic + { + uint32_t magic; + finp.read((char *) &magic, sizeof(magic)); + if (magic != 0x67676d6c) { + fprintf(stderr, "%s: invalid model file '%s' (bad magic)\n", __func__, fname_inp.c_str()); return false; } + + fout.write((char *) &magic, sizeof(magic)); } - struct ggml_init_params params = { - /*.mem_size =*/ buf_size, - /*.mem_buffer =*/ buf, - }; + llama_hparams hparams; - struct ggml_context * ctx0 = ggml_init(params); - ggml_cgraph gf = {}; - gf.n_threads = n_threads; + // load hparams + { + finp.read((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); + //finp.read((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); + finp.read((char *) &hparams.n_embd, sizeof(hparams.n_embd)); + finp.read((char *) &hparams.n_mult, sizeof(hparams.n_mult)); + finp.read((char *) &hparams.n_head, sizeof(hparams.n_head)); + finp.read((char *) &hparams.n_layer, sizeof(hparams.n_layer)); + finp.read((char *) &hparams.n_rot, sizeof(hparams.n_rot)); + finp.read((char *) &hparams.f16, sizeof(hparams.f16)); - struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N); - memcpy(embd->data, embd_inp.data(), N*ggml_element_size(embd)); + printf("%s: n_vocab = %d\n", __func__, hparams.n_vocab); + printf("%s: n_ctx = %d\n", __func__, hparams.n_ctx); + printf("%s: n_embd = %d\n", __func__, hparams.n_embd); + printf("%s: n_mult = %d\n", __func__, hparams.n_mult); + printf("%s: n_head = %d\n", __func__, hparams.n_head); + printf("%s: n_layer = %d\n", __func__, hparams.n_layer); + printf("%s: f16 = %d\n", __func__, hparams.f16); - struct ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embeddings, embd); + fout.write((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); + //fout.write((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); + fout.write((char *) &hparams.n_embd, sizeof(hparams.n_embd)); + fout.write((char *) &hparams.n_mult, sizeof(hparams.n_mult)); + fout.write((char *) &hparams.n_head, sizeof(hparams.n_head)); + fout.write((char *) &hparams.n_layer, sizeof(hparams.n_layer)); + fout.write((char *) &hparams.n_rot, sizeof(hparams.n_rot)); + fout.write((char *) &itype, sizeof(hparams.f16)); + } - for (int il = 0; il < n_layer; ++il) { - struct ggml_tensor * inpSA = inpL; + // load vocab + { + const int32_t n_vocab = hparams.n_vocab; - struct ggml_tensor * cur; + if (n_vocab != hparams.n_vocab) { + fprintf(stderr, "%s: invalid model file '%s' (bad vocab size %d != %d)\n", + __func__, fname_inp.c_str(), n_vocab, hparams.n_vocab); + return false; + } - // norm - { - cur = ggml_rms_norm(ctx0, inpL); + std::string word; + for (int i = 0; i < n_vocab; i++) { + uint32_t len; + finp.read ((char *) &len, sizeof(len)); + fout.write((char *) &len, sizeof(len)); - // cur = attention_norm*cur - cur = ggml_mul(ctx0, - ggml_repeat(ctx0, model.layers[il].attention_norm, cur), - cur); + word.resize(len); + finp.read ((char *) word.data(), len); + fout.write((char *) word.data(), len); + + vocab.token_to_id[word] = i; + vocab.id_to_token[i] = word; } + } - // self-attention - { - struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, model.layers[il].wq, cur); - struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur); - struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, model.layers[il].wv, cur); + // load weights + { + size_t total_size_org = 0; + size_t total_size_new = 0; - // store key and value to memory - if (N >= 1) { - struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_k, N*n_embd, (ggml_element_size(model.memory_k)*n_embd)*(il*n_ctx + n_past)); - struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_v, N*n_embd, (ggml_element_size(model.memory_v)*n_embd)*(il*n_ctx + n_past)); + std::vector work; - ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Kcur, k)); - ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcur, v)); - } + std::vector data_u8; + std::vector data_f16; + std::vector data_f32; - // Q = Qcur.contiguous().view(n_embd/n_head, n_head, N).permute(0, 2, 1, 3) - struct ggml_tensor * Q = - ggml_permute(ctx0, - ggml_rope(ctx0, - ggml_cpy(ctx0, - Qcur, - ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd/n_head, n_head, N)), - n_past, n_rot, 0), - 0, 2, 1, 3); + std::vector hist_all(1 << 4, 0); - // K = Kmem.view(n_embd/n_head, n_head, n_past + N).permute(0, 2, 1, 3) - struct ggml_tensor * K = - ggml_permute(ctx0, - ggml_rope(ctx0, - ggml_reshape_3d(ctx0, - ggml_view_1d(ctx0, model.memory_k, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_k)*n_embd), - n_embd/n_head, n_head, n_past + N), - n_past, n_rot, 1), - 0, 2, 1, 3); + while (true) { + int32_t n_dims; + int32_t length; + int32_t ftype; - // K * Q - struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + finp.read(reinterpret_cast(&n_dims), sizeof(n_dims)); + finp.read(reinterpret_cast(&length), sizeof(length)); + finp.read(reinterpret_cast(&ftype), sizeof(ftype)); - // KQ_scaled = KQ / sqrt(n_embd/n_head) - struct ggml_tensor * KQ_scaled = - ggml_scale(ctx0, - KQ, - ggml_new_f32(ctx0, 1.0f/sqrt(float(n_embd)/n_head)) - ); + if (finp.eof()) { + break; + } - // KQ_masked = mask_past(KQ_scaled) - struct ggml_tensor * KQ_masked = ggml_diag_mask_inf(ctx0, KQ_scaled, n_past); + int32_t nelements = 1; + int32_t ne[2] = { 1, 1 }; + for (int i = 0; i < n_dims; ++i) { + finp.read (reinterpret_cast(&ne[i]), sizeof(ne[i])); + nelements *= ne[i]; + } - // KQ = soft_max(KQ_masked) - struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctx0, KQ_masked); + std::string name(length, 0); + finp.read (&name[0], length); - // V_trans = Vmem.view(n_embd/n_head, n_head, n_past + N).permute(1, 2, 0, 3).contiguous() - struct ggml_tensor * V_trans = - ggml_permute(ctx0, - ggml_reshape_3d(ctx0, - ggml_view_1d(ctx0, model.memory_v, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_v)*n_embd), - n_embd/n_head, n_head, n_past + N), - 1, 2, 0, 3); + { + static const char * ftype_str[] = { "f32", "f16", "q4_0", "q4_1", }; + printf("%48s - [%5d, %5d], type = %6s ", name.data(), ne[0], ne[1], ftype_str[ftype]); + } - // KQV = transpose(V) * KQ_soft_max - struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V_trans, KQ_soft_max); + // regexes of tensor names to be quantized + const std::vector k_names = { + ".*weight", + }; - // KQV_merged = KQV.permute(0, 2, 1, 3) - struct ggml_tensor * KQV_merged = ggml_permute(ctx0, KQV, 0, 2, 1, 3); + bool quantize = false; + for (const auto & s : k_names) { + if (std::regex_match(name, std::regex(s))) { + quantize = true; + break; + } + } + + // quantize only 2D tensors + quantize &= (n_dims == 2); + + if (quantize) { + if (ftype != 0 && ftype != 1) { + fprintf(stderr, "%s: unsupported ftype %d for integer quantization\n", __func__, ftype); + return false; + } + + if (ftype == 1) { + data_f16.resize(nelements); + finp.read(reinterpret_cast(data_f16.data()), nelements * sizeof(ggml_fp16_t)); + data_f32.resize(nelements); + for (int i = 0; i < nelements; ++i) { + data_f32[i] = ggml_fp16_to_fp32(data_f16[i]); + } + } else { + data_f32.resize(nelements); + finp.read(reinterpret_cast(data_f32.data()), nelements * sizeof(float)); + } + + ftype = itype; + } else { + const int bpe = (ftype == 0) ? sizeof(float) : sizeof(uint16_t); + + data_u8.resize(nelements*bpe); + finp.read(reinterpret_cast(data_u8.data()), nelements * bpe); + } + + fout.write(reinterpret_cast(&n_dims), sizeof(n_dims)); + fout.write(reinterpret_cast(&length), sizeof(length)); + fout.write(reinterpret_cast(&ftype), sizeof(ftype)); + for (int i = 0; i < n_dims; ++i) { + fout.write(reinterpret_cast(&ne[i]), sizeof(ne[i])); + } + fout.write(&name[0], length); + + if (quantize) { + printf("quantizing .. "); + work.resize(nelements); // for quantization + + size_t cur_size = 0; + std::vector hist_cur(1 << 4, 0); + + switch (type) { + case GGML_TYPE_Q4_0: + { + cur_size = ggml_quantize_q4_0(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); + } break; + case GGML_TYPE_Q4_1: + { + cur_size = ggml_quantize_q4_1(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); + } break; + default: + { + fprintf(stderr, "%s: unsupported quantization type %d\n", __func__, type); + return false; + } + } + + fout.write(reinterpret_cast(work.data()), cur_size); + total_size_new += cur_size; + + printf("size = %8.2f MB -> %8.2f MB | hist: ", nelements * sizeof(float)/1024.0/1024.0, cur_size/1024.0/1024.0); + for (int i = 0; i < hist_cur.size(); ++i) { + hist_all[i] += hist_cur[i]; + } + + for (int i = 0; i < hist_cur.size(); ++i) { + printf("%5.3f ", hist_cur[i] / (float)nelements); + } + printf("\n"); + } else { + printf("size = %8.3f MB\n", data_u8.size()/1024.0/1024.0); + fout.write(reinterpret_cast(data_u8.data()), data_u8.size()); + total_size_new += data_u8.size(); + } + + total_size_org += nelements * sizeof(float); + } + + printf("%s: model size = %8.2f MB\n", __func__, total_size_org/1024.0/1024.0); + printf("%s: quant size = %8.2f MB\n", __func__, total_size_new/1024.0/1024.0); + + { + int64_t sum_all = 0; + for (int i = 0; i < hist_all.size(); ++i) { + sum_all += hist_all[i]; + } + + printf("%s: hist: ", __func__); + for (int i = 0; i < hist_all.size(); ++i) { + printf("%5.3f ", hist_all[i] / (float)sum_all); + } + printf("\n"); + } + } + + finp.close(); + fout.close(); + + return true; +} + + +/* External API */ + +const std::vector& llama_context_get_embd(const llama_context& ctx) { + return ctx.state->embd; +} +gpt_vocab& llama_context_get_vocab(llama_context& ctx) { + return ctx.vocab; +} +bool llama_context_not_finished(const llama_context& ctx) +{ + return ctx.state->remaining_tokens > 0; +} +const std::vector llama_tokenize_text(const llama_context& ctx, const std::string& text) { + return llama_tokenize(ctx.vocab, text, true); +} +const std::vector& llama_context_get_last_n_tokens(const llama_context& ctx) { + return ctx.state->last_n_tokens; +} +llama_context* llama_init_from_params(const gpt_params& params) { + llama_model model{}; + gpt_vocab vocab{}; + + // Compute time taken to load model + const int64_t t_start = ggml_time_us(); + bool ret = llama_model_load(params.model, model, vocab, 1024); + const int64_t t_end = ggml_time_us(); + if(!ret) + { + return nullptr; + } + llama_context* ctx = new llama_context(std::move(model), std::move(vocab), params); + ctx->t_load_us = t_end - t_start; + return ctx; +} +void llama_free_context(llama_context* ctx) { + delete ctx; +} + +const char * llama_print_system_info(void) { + static std::string s; + + s = ""; + s += "AVX = " + std::to_string(ggml_cpu_has_avx()) + " | "; + s += "AVX2 = " + std::to_string(ggml_cpu_has_avx2()) + " | "; + s += "AVX512 = " + std::to_string(ggml_cpu_has_avx512()) + " | "; + s += "FMA = " + std::to_string(ggml_cpu_has_fma()) + " | "; + s += "NEON = " + std::to_string(ggml_cpu_has_neon()) + " | "; + s += "ARM_FMA = " + std::to_string(ggml_cpu_has_arm_fma()) + " | "; + s += "F16C = " + std::to_string(ggml_cpu_has_f16c()) + " | "; + s += "FP16_VA = " + std::to_string(ggml_cpu_has_fp16_va()) + " | "; + s += "WASM_SIMD = " + std::to_string(ggml_cpu_has_wasm_simd()) + " | "; + s += "BLAS = " + std::to_string(ggml_cpu_has_blas()) + " | "; + s += "SSE3 = " + std::to_string(ggml_cpu_has_sse3()) + " | "; + s += "VSX = " + std::to_string(ggml_cpu_has_vsx()) + " | "; + + return s.c_str(); +} + +void llama_print_context_info(const llama_context& ctx) +{ + const gpt_params& params = ctx.params; + const std::vector& embd_inp = ctx.state->embd_inp; + { + fprintf(stderr, "\n"); + fprintf(stderr, "system_info: n_threads = %d / %d | %s\n", + params.n_threads, std::thread::hardware_concurrency(), llama_print_system_info()); + } + fprintf(stderr, "\n"); + fprintf(stderr, "%s: prompt: '%s'\n", __func__, params.prompt.c_str()); + fprintf(stderr, "%s: number of tokens in prompt = %zu\n", __func__, embd_inp.size()); + for (int i = 0; i < (int) embd_inp.size(); i++) { + fprintf(stderr, "%6d -> '%s'\n", embd_inp[i], ctx.vocab.id_to_token.at(embd_inp[i]).c_str()); + } + fprintf(stderr, "\n"); +} + +void llama_print_end_stats(const llama_context& ctx) +{ + const llama_state& state = *ctx.state; + fprintf(stderr, "\n\n"); + fprintf(stderr, "%s: mem per token = %8zu bytes\n", __func__, state.mem_per_token); + fprintf(stderr, "%s: load time = %8.2f ms\n", __func__, ctx.t_load_us/1000.0f); + fprintf(stderr, "%s: sample time = %8.2f ms\n", __func__, state.t_sample_us/1000.0f); + fprintf(stderr, "%s: predict time = %8.2f ms / %.2f ms per token\n", __func__, state.t_predict_us/1000.0f, state.t_predict_us/1000.0f/state.n_past); +} +// evaluate the transformer +// +// - model: the model +// - n_threads: number of threads to use +// - n_past: the context size so far +// - embd_inp: the embeddings of the tokens in the context +// - embd_w: the predicted logits for the next token +// +// The GPT-J model requires about 16MB of memory per input token. +// +bool llama_eval( + const llama_model & model, + const int n_threads, + const int n_past, + const std::vector & embd_inp, + std::vector & embd_w, + size_t & mem_per_token) { + const int N = embd_inp.size(); + + const auto & hparams = model.hparams; + + const int n_embd = hparams.n_embd; + const int n_layer = hparams.n_layer; + const int n_ctx = hparams.n_ctx; + const int n_head = hparams.n_head; + const int n_vocab = hparams.n_vocab; + const int n_rot = hparams.n_embd/hparams.n_head; + + const int d_key = n_embd/n_head; + + // TODO: check if this size scales with n_ctx linearly and remove constant. somehow I feel it wasn't the case + // static size_t buf_size = hparams.n_ctx*1024*1024; + static size_t buf_size = 512u*1024*1024; + static void * buf = malloc(buf_size); + + if (mem_per_token > 0 && mem_per_token*N > buf_size) { + const size_t buf_size_new = 1.1*(mem_per_token*N); // add 10% to account for ggml object overhead + //fprintf(stderr, "\n%s: reallocating buffer from %zu to %zu bytes\n", __func__, buf_size, buf_size_new); + + // reallocate + buf_size = buf_size_new; + buf = realloc(buf, buf_size); + if (buf == nullptr) { + fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size); + return false; + } + } + + struct ggml_init_params params = { + /*.mem_size =*/ buf_size, + /*.mem_buffer =*/ buf, + }; + + struct ggml_context * ctx0 = ggml_init(params); + ggml_cgraph gf = {}; + gf.n_threads = n_threads; + + struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N); + memcpy(embd->data, embd_inp.data(), N*ggml_element_size(embd)); + + struct ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embeddings, embd); + + for (int il = 0; il < n_layer; ++il) { + struct ggml_tensor * inpSA = inpL; + + struct ggml_tensor * cur; + + // norm + { + cur = ggml_rms_norm(ctx0, inpL); + + // cur = attention_norm*cur + cur = ggml_mul(ctx0, + ggml_repeat(ctx0, model.layers[il].attention_norm, cur), + cur); + } + + // self-attention + { + struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, model.layers[il].wq, cur); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, model.layers[il].wv, cur); + + // store key and value to memory + if (N >= 1) { + struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_k, N*n_embd, (ggml_element_size(model.memory_k)*n_embd)*(il*n_ctx + n_past)); + struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_v, N*n_embd, (ggml_element_size(model.memory_v)*n_embd)*(il*n_ctx + n_past)); + + ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Kcur, k)); + ggml_build_forward_expand(&gf, ggml_cpy(ctx0, Vcur, v)); + } + + // Q = Qcur.contiguous().view(n_embd/n_head, n_head, N).permute(0, 2, 1, 3) + struct ggml_tensor * Q = + ggml_permute(ctx0, + ggml_rope(ctx0, + ggml_cpy(ctx0, + Qcur, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd/n_head, n_head, N)), + n_past, n_rot, 0), + 0, 2, 1, 3); + + // K = Kmem.view(n_embd/n_head, n_head, n_past + N).permute(0, 2, 1, 3) + struct ggml_tensor * K = + ggml_permute(ctx0, + ggml_rope(ctx0, + ggml_reshape_3d(ctx0, + ggml_view_1d(ctx0, model.memory_k, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_k)*n_embd), + n_embd/n_head, n_head, n_past + N), + n_past, n_rot, 1), + 0, 2, 1, 3); + + // K * Q + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + + // KQ_scaled = KQ / sqrt(n_embd/n_head) + struct ggml_tensor * KQ_scaled = + ggml_scale(ctx0, + KQ, + ggml_new_f32(ctx0, 1.0f/sqrt(float(n_embd)/n_head)) + ); + + // KQ_masked = mask_past(KQ_scaled) + struct ggml_tensor * KQ_masked = ggml_diag_mask_inf(ctx0, KQ_scaled, n_past); + + // KQ = soft_max(KQ_masked) + struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctx0, KQ_masked); + + // V_trans = Vmem.view(n_embd/n_head, n_head, n_past + N).permute(1, 2, 0, 3).contiguous() + struct ggml_tensor * V_trans = + ggml_permute(ctx0, + ggml_reshape_3d(ctx0, + ggml_view_1d(ctx0, model.memory_v, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_v)*n_embd), + n_embd/n_head, n_head, n_past + N), + 1, 2, 0, 3); + + // KQV = transpose(V) * KQ_soft_max + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V_trans, KQ_soft_max); + + // KQV_merged = KQV.permute(0, 2, 1, 3) + struct ggml_tensor * KQV_merged = ggml_permute(ctx0, KQV, 0, 2, 1, 3); // cur = KQV_merged.contiguous().view(n_embd, N) cur = ggml_cpy(ctx0, @@ -685,258 +1117,172 @@ bool llama_eval( return true; } -bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype) { - ggml_type type = GGML_TYPE_Q4_1; - - switch (itype) { - case 2: type = GGML_TYPE_Q4_0; break; - case 3: type = GGML_TYPE_Q4_1; break; - default: fprintf(stderr, "%s: invalid quantization type %d\n", __func__, itype); return 1; - }; - if (type != GGML_TYPE_Q4_0 && type != GGML_TYPE_Q4_1) { - fprintf(stderr, "%s: invalid quantization type %d\n", __func__, type); - return false; +bool llama_init_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing) { + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + const gpt_params& params = ctx.params; + + if (clear_existing) { + state.embd.clear(); + state.input_consumed = 0; + state.embd_inp.clear(); + state.last_n_tokens.clear(); + state.remaining_tokens = 0; + state.n_past = 0; } - gpt_vocab vocab; + std::vector line_inp = llama_tokenize_text(ctx, text); + state.embd_inp.insert(state.embd_inp.end(), line_inp.begin(), line_inp.end()); - printf("%s: loading model from '%s'\n", __func__, fname_inp.c_str()); + int n_predict = std::min(params.n_predict, model.hparams.n_ctx - (int) state.embd_inp.size()); + state.remaining_tokens = n_predict; - auto finp = std::ifstream(fname_inp, std::ios::binary); - if (!finp) { - fprintf(stderr, "%s: failed to open '%s' for reading\n", __func__, fname_inp.c_str()); + // determine the required inference memory per token: + state.mem_per_token = 0; + if(!llama_eval(model, params.n_threads, 0, { 0, 1, 2, 3 }, state.logits, state.mem_per_token)) + { + fprintf(stderr, "Failed to predict with initial prompt\n"); return false; } - auto fout = std::ofstream(fname_out, std::ios::binary); - if (!fout) { - fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_out.c_str()); - return false; - } + int last_n_size = params.repeat_last_n; + state.last_n_tokens = std::vector(last_n_size); + std::fill(state.last_n_tokens.begin(), state.last_n_tokens.end(), 0); - // verify magic - { - uint32_t magic; - finp.read((char *) &magic, sizeof(magic)); - if (magic != 0x67676d6c) { - fprintf(stderr, "%s: invalid model file '%s' (bad magic)\n", __func__, fname_inp.c_str()); - return false; - } + state.is_initialized = true; + return true; +} - fout.write((char *) &magic, sizeof(magic)); +/// @brief Injests a batch of input tokens into the context +/// @param ctx +void llama_injest_input_batch(llama_context& ctx) +{ + llama_state& state = *ctx.state; + const gpt_params& params = ctx.params; + + // Copy at most n_batch elements from embd_inp to embd + size_t num_copied = std::min((size_t) params.n_batch, state.embd_inp.size() - state.input_consumed); + std::copy(state.embd_inp.begin() + state.input_consumed, + state.embd_inp.begin() + state.input_consumed + num_copied, + std::back_inserter(state.embd)); + state.input_consumed += num_copied; + + // Copy the last `repeat_last_n` elements copied into embd to last_n_tokens + size_t num_copied_last_n = std::min(num_copied, (size_t) params.repeat_last_n); + state.last_n_tokens.erase(state.last_n_tokens.begin(), state.last_n_tokens.begin()+num_copied_last_n); + state.last_n_tokens.insert(state.last_n_tokens.end(), state.embd.end() - num_copied_last_n, state.embd.end()); +} + +/// @brief Run the prediction step on ctx.embd and store result in ctx.state.logits +/// @param ctx +/// @return +bool llama_predict(llama_context& ctx){ + const int64_t t_start_us = ggml_time_us(); + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + const gpt_params& params = ctx.params; + + if (!llama_eval(model, params.n_threads, state.n_past, state.embd, state.logits, state.mem_per_token)) { + fprintf(stderr, "Failed to predict\n"); + return false; } - llama_hparams hparams; + state.t_predict_us += ggml_time_us() - t_start_us; + return true; +} +/// @brief Sample a token from the logits +/// @param ctx +/// @return token id +gpt_vocab::id llama_sample_token(llama_context& ctx) +{ + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + const gpt_params& params = ctx.params; - // load hparams - { - finp.read((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); - //finp.read((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); - finp.read((char *) &hparams.n_embd, sizeof(hparams.n_embd)); - finp.read((char *) &hparams.n_mult, sizeof(hparams.n_mult)); - finp.read((char *) &hparams.n_head, sizeof(hparams.n_head)); - finp.read((char *) &hparams.n_layer, sizeof(hparams.n_layer)); - finp.read((char *) &hparams.n_rot, sizeof(hparams.n_rot)); - finp.read((char *) &hparams.f16, sizeof(hparams.f16)); + const float top_k = params.top_k; + const float top_p = params.top_p; + const float temp = params.temp; + const float repeat_penalty = params.repeat_penalty; - printf("%s: n_vocab = %d\n", __func__, hparams.n_vocab); - printf("%s: n_ctx = %d\n", __func__, hparams.n_ctx); - printf("%s: n_embd = %d\n", __func__, hparams.n_embd); - printf("%s: n_mult = %d\n", __func__, hparams.n_mult); - printf("%s: n_head = %d\n", __func__, hparams.n_head); - printf("%s: n_layer = %d\n", __func__, hparams.n_layer); - printf("%s: f16 = %d\n", __func__, hparams.f16); + const int n_vocab = model.hparams.n_vocab; - fout.write((char *) &hparams.n_vocab, sizeof(hparams.n_vocab)); - //fout.write((char *) &hparams.n_ctx, sizeof(hparams.n_ctx)); - fout.write((char *) &hparams.n_embd, sizeof(hparams.n_embd)); - fout.write((char *) &hparams.n_mult, sizeof(hparams.n_mult)); - fout.write((char *) &hparams.n_head, sizeof(hparams.n_head)); - fout.write((char *) &hparams.n_layer, sizeof(hparams.n_layer)); - fout.write((char *) &hparams.n_rot, sizeof(hparams.n_rot)); - fout.write((char *) &itype, sizeof(hparams.f16)); - } + gpt_vocab::id id = 0; - // load vocab { - const int32_t n_vocab = hparams.n_vocab; - - if (n_vocab != hparams.n_vocab) { - fprintf(stderr, "%s: invalid model file '%s' (bad vocab size %d != %d)\n", - __func__, fname_inp.c_str(), n_vocab, hparams.n_vocab); - return false; - } + const int64_t t_start_sample_us = ggml_time_us(); - std::string word; - for (int i = 0; i < n_vocab; i++) { - uint32_t len; - finp.read ((char *) &len, sizeof(len)); - fout.write((char *) &len, sizeof(len)); + id = llama_sample_top_p_top_k(ctx.vocab, state.logits.data() + (state.logits.size() - n_vocab), + state.last_n_tokens, repeat_penalty, top_k, top_p, temp, state.rng); - word.resize(len); - finp.read ((char *) word.data(), len); - fout.write((char *) word.data(), len); + state.last_n_tokens.erase(state.last_n_tokens.begin()); + state.last_n_tokens.push_back(id); - vocab.token_to_id[word] = i; - vocab.id_to_token[i] = word; - } + state.t_sample_us += ggml_time_us() - t_start_sample_us; } - - // load weights + return id; +} +/// @brief Injest all input (in multiple batches) into model and run call predict() +/// @param ctx +bool llama_injest_input(llama_context& ctx, const std::string& text, bool clear_existing) +{ + llama_state& state = *ctx.state; + + // Initialize context, tokenize text and clear existing state if necessary + if(!state.is_initialized && !llama_init_context_with_prompt(ctx, text, clear_existing)) { - size_t total_size_org = 0; - size_t total_size_new = 0; - - std::vector work; - - std::vector data_u8; - std::vector data_f16; - std::vector data_f32; - - std::vector hist_all(1 << 4, 0); - - while (true) { - int32_t n_dims; - int32_t length; - int32_t ftype; - - finp.read(reinterpret_cast(&n_dims), sizeof(n_dims)); - finp.read(reinterpret_cast(&length), sizeof(length)); - finp.read(reinterpret_cast(&ftype), sizeof(ftype)); - - if (finp.eof()) { - break; - } - - int32_t nelements = 1; - int32_t ne[2] = { 1, 1 }; - for (int i = 0; i < n_dims; ++i) { - finp.read (reinterpret_cast(&ne[i]), sizeof(ne[i])); - nelements *= ne[i]; - } - - std::string name(length, 0); - finp.read (&name[0], length); + return false; + } + // Injest the tokens into the model one batch at a time + while (state.has_more_input()) + { + llama_injest_input_batch(ctx); + if (state.embd.size() >= 0) { + if(!llama_predict(ctx)) { - static const char * ftype_str[] = { "f32", "f16", "q4_0", "q4_1", }; - printf("%48s - [%5d, %5d], type = %6s ", name.data(), ne[0], ne[1], ftype_str[ftype]); - } - - // regexes of tensor names to be quantized - const std::vector k_names = { - ".*weight", + return false; }; - - bool quantize = false; - for (const auto & s : k_names) { - if (std::regex_match(name, std::regex(s))) { - quantize = true; - break; - } - } - - // quantize only 2D tensors - quantize &= (n_dims == 2); - - if (quantize) { - if (ftype != 0 && ftype != 1) { - fprintf(stderr, "%s: unsupported ftype %d for integer quantization\n", __func__, ftype); - return false; - } - - if (ftype == 1) { - data_f16.resize(nelements); - finp.read(reinterpret_cast(data_f16.data()), nelements * sizeof(ggml_fp16_t)); - data_f32.resize(nelements); - for (int i = 0; i < nelements; ++i) { - data_f32[i] = ggml_fp16_to_fp32(data_f16[i]); - } - } else { - data_f32.resize(nelements); - finp.read(reinterpret_cast(data_f32.data()), nelements * sizeof(float)); - } - - ftype = itype; - } else { - const int bpe = (ftype == 0) ? sizeof(float) : sizeof(uint16_t); - - data_u8.resize(nelements*bpe); - finp.read(reinterpret_cast(data_u8.data()), nelements * bpe); - } - - fout.write(reinterpret_cast(&n_dims), sizeof(n_dims)); - fout.write(reinterpret_cast(&length), sizeof(length)); - fout.write(reinterpret_cast(&ftype), sizeof(ftype)); - for (int i = 0; i < n_dims; ++i) { - fout.write(reinterpret_cast(&ne[i]), sizeof(ne[i])); - } - fout.write(&name[0], length); - - if (quantize) { - printf("quantizing .. "); - work.resize(nelements); // for quantization - - size_t cur_size = 0; - std::vector hist_cur(1 << 4, 0); - - switch (type) { - case GGML_TYPE_Q4_0: - { - cur_size = ggml_quantize_q4_0(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); - } break; - case GGML_TYPE_Q4_1: - { - cur_size = ggml_quantize_q4_1(data_f32.data(), work.data(), nelements, ne[0], QK, hist_cur.data()); - } break; - default: - { - fprintf(stderr, "%s: unsupported quantization type %d\n", __func__, type); - return false; - } - } - - fout.write(reinterpret_cast(work.data()), cur_size); - total_size_new += cur_size; - - printf("size = %8.2f MB -> %8.2f MB | hist: ", nelements * sizeof(float)/1024.0/1024.0, cur_size/1024.0/1024.0); - for (int i = 0; i < hist_cur.size(); ++i) { - hist_all[i] += hist_cur[i]; - } - - for (int i = 0; i < hist_cur.size(); ++i) { - printf("%5.3f ", hist_cur[i] / (float)nelements); - } - printf("\n"); - } else { - printf("size = %8.3f MB\n", data_u8.size()/1024.0/1024.0); - fout.write(reinterpret_cast(data_u8.data()), data_u8.size()); - total_size_new += data_u8.size(); - } - - total_size_org += nelements * sizeof(float); } + state.n_past += state.embd.size(); + state.embd.clear(); + } + return true; +} +bool llama_inference(llama_context& ctx, gpt_vocab::id& id) { + llama_state& state = *ctx.state; - printf("%s: model size = %8.2f MB\n", __func__, total_size_org/1024.0/1024.0); - printf("%s: quant size = %8.2f MB\n", __func__, total_size_new/1024.0/1024.0); + // Tokenize text if we are starting out + if(!state.is_initialized) + { + fprintf(stderr, "State must be initialized before running inference"); + return false; + } - { - int64_t sum_all = 0; - for (int i = 0; i < hist_all.size(); ++i) { - sum_all += hist_all[i]; - } + // No more tokens to generate + if (state.remaining_tokens <= 0) { + return false; + } - printf("%s: hist: ", __func__); - for (int i = 0; i < hist_all.size(); ++i) { - printf("%5.3f ", hist_all[i] / (float)sum_all); - } - printf("\n"); + // Do prediction if we have enough tokens + if (state.embd.size() > 0) { + if(!llama_predict(ctx)) + { + return false; } } - - finp.close(); - fout.close(); - + // sample a token + id = llama_sample_token(ctx); + // add it to the context + state.embd.push_back(id); + + state.n_past += 1; + // decrement remaining sampling budget + --state.remaining_tokens; + + // end of text token + if (state.embd.back() == 2) { + state.remaining_tokens = 0; + } return true; } diff --git a/llama.h b/llama.h index a902273048ce6..387efa686b547 100644 --- a/llama.h +++ b/llama.h @@ -3,8 +3,27 @@ #include #include #include +#include #include "utils.h" +#include "ggml.h" + +#ifdef LLAMA_SHARED +# ifdef _WIN32 +# ifdef LLAMA_BUILD +# define LLAMA_API __declspec(dllexport) +# else +# define LLAMA_API __declspec(dllimport) +# endif +# else +# define LLAMA_API __attribute__ ((visibility ("default"))) +# endif +#else +# define LLAMA_API +#endif + + + // default hparams (LLaMA 7B) struct llama_hparams { @@ -18,50 +37,28 @@ struct llama_hparams { int32_t f16 = 1; }; -struct llama_layer { - // normalization - struct ggml_tensor * attention_norm; - - // attention - struct ggml_tensor * wq; - struct ggml_tensor * wk; - struct ggml_tensor * wv; - struct ggml_tensor * wo; +struct llama_context; - // normalization - struct ggml_tensor * ffn_norm; +void llama_free_context(llama_context* ctx); - // ff - struct ggml_tensor * w1; - struct ggml_tensor * w2; - struct ggml_tensor * w3; -}; +const std::vector& llama_context_get_embd(const llama_context& ctx); +gpt_vocab& llama_context_get_vocab(llama_context& ctx); +bool llama_context_not_finished(const llama_context& ctx); +const std::vector llama_tokenize_text(const llama_context& ctx, const std::string& text); -struct llama_model { - llama_hparams hparams; +const std::vector& llama_context_get_last_n_tokens(const llama_context& ctx); +bool llama_init_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing = true); - struct ggml_tensor * tok_embeddings; +// Various functions for loading a ggml LLaMA model. +llama_context* llama_init_from_params(const gpt_params& params); - struct ggml_tensor * norm; - struct ggml_tensor * output; +// Run inference on a LLaMA model using llama_context. +std::vector llama_eval(llama_context& ctx, const gpt_params& params, std::string& text); - std::vector layers; +bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype); - // key + value memory - struct ggml_tensor * memory_k; - struct ggml_tensor * memory_v; +bool llama_injest_input(llama_context& ctx, const std::string& text, bool clear_existing = true); - // - struct ggml_context * ctx; - std::map tensors; -}; - -bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx); -bool llama_eval( - const llama_model & model, - const int n_threads, - const int n_past, - const std::vector & embd_inp, - std::vector & embd_w, - size_t & mem_per_token); -bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype); +bool llama_inference(llama_context& ctx, gpt_vocab::id& model_output); +void llama_print_context_info(const llama_context& ctx); +void llama_print_end_stats(const llama_context& ctx); From 4b4d8a5d44bd5cb50ba77c984f2b8ae10927514d Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Tue, 14 Mar 2023 00:37:44 -0700 Subject: [PATCH 06/18] Update main.cpp to use new llama library --- main.cpp | 201 ++++++++++--------------------------------------------- 1 file changed, 37 insertions(+), 164 deletions(-) diff --git a/main.cpp b/main.cpp index 9e09a6a8579ee..c1872e1b0a363 100644 --- a/main.cpp +++ b/main.cpp @@ -55,25 +55,6 @@ void sigint_handler(int signo) { } #endif -const char * llama_print_system_info(void) { - static std::string s; - - s = ""; - s += "AVX = " + std::to_string(ggml_cpu_has_avx()) + " | "; - s += "AVX2 = " + std::to_string(ggml_cpu_has_avx2()) + " | "; - s += "AVX512 = " + std::to_string(ggml_cpu_has_avx512()) + " | "; - s += "FMA = " + std::to_string(ggml_cpu_has_fma()) + " | "; - s += "NEON = " + std::to_string(ggml_cpu_has_neon()) + " | "; - s += "ARM_FMA = " + std::to_string(ggml_cpu_has_arm_fma()) + " | "; - s += "F16C = " + std::to_string(ggml_cpu_has_f16c()) + " | "; - s += "FP16_VA = " + std::to_string(ggml_cpu_has_fp16_va()) + " | "; - s += "WASM_SIMD = " + std::to_string(ggml_cpu_has_wasm_simd()) + " | "; - s += "BLAS = " + std::to_string(ggml_cpu_has_blas()) + " | "; - s += "SSE3 = " + std::to_string(ggml_cpu_has_sse3()) + " | "; - s += "VSX = " + std::to_string(ggml_cpu_has_vsx()) + " | "; - - return s.c_str(); -} int main(int argc, char ** argv) { ggml_time_init(); @@ -107,41 +88,18 @@ int main(int argc, char ** argv) { int64_t t_load_us = 0; - gpt_vocab vocab; - llama_model model; - // load the model - { - const ggml_type memory_type = params.memory_f16 ? GGML_TYPE_F16 : GGML_TYPE_F32; - const int64_t t_start_us = ggml_time_us(); - if (!llama_model_load(params.model, model, vocab, params.n_ctx, memory_type)) { - fprintf(stderr, "%s: failed to load model from '%s'\n", __func__, params.model.c_str()); - return 1; - } - - t_load_us = ggml_time_us() - t_start_us; - } + llama_context* ctx_ptr = llama_init_from_params(params); + llama_context & ctx = *ctx_ptr; + gpt_vocab & vocab = llama_context_get_vocab(ctx); // print system information - { - fprintf(stderr, "\n"); - fprintf(stderr, "system_info: n_threads = %d / %d | %s\n", - params.n_threads, std::thread::hardware_concurrency(), llama_print_system_info()); - } - - int n_past = 0; - - int64_t t_sample_us = 0; - int64_t t_predict_us = 0; - - std::vector logits; + llama_print_context_info(ctx); // Add a space in front of the first character to match OG llama tokenizer behavior params.prompt.insert(0, 1, ' '); // tokenize the prompt - std::vector embd_inp = ::llama_tokenize(vocab, params.prompt, true); - - params.n_predict = std::min(params.n_predict, model.hparams.n_ctx - (int) embd_inp.size()); + std::vector embd_inp = llama_tokenize_text(ctx, params.prompt); // prefix & suffix for instruct mode const std::vector inp_pfx = ::llama_tokenize(vocab, "\n\n### Instruction:\n\n", true); @@ -154,24 +112,8 @@ int main(int argc, char ** argv) { } // tokenize the reverse prompt - std::vector> antipromptv_inp; - - for (auto antiprompt : params.antiprompt) { - antipromptv_inp.push_back(::llama_tokenize(vocab, antiprompt, false)); - } - - // enable interactive mode if reverse prompt is specified - if (!antipromptv_inp.size()) { - params.interactive = true; - } + std::vector antiprompt_inp = llama_tokenize_text(ctx, params.prompt); - fprintf(stderr, "\n"); - fprintf(stderr, "%s: prompt: '%s'\n", __func__, params.prompt.c_str()); - fprintf(stderr, "%s: number of tokens in prompt = %zu\n", __func__, embd_inp.size()); - for (int i = 0; i < (int) embd_inp.size(); i++) { - fprintf(stderr, "%6d -> '%s'\n", embd_inp[i], vocab.id_to_token.at(embd_inp[i]).c_str()); - } - fprintf(stderr, "\n"); if (params.interactive) { #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) struct sigaction sigint_action; @@ -200,16 +142,6 @@ int main(int argc, char ** argv) { fprintf(stderr, "sampling parameters: temp = %f, top_k = %d, top_p = %f, repeat_last_n = %i, repeat_penalty = %f\n", params.temp, params.top_k, params.top_p, params.repeat_last_n, params.repeat_penalty); fprintf(stderr, "\n\n"); - std::vector embd; - - // determine the required inference memory per token: - size_t mem_per_token = 0; - llama_eval(model, params.n_threads, 0, { 0, 1, 2, 3 }, logits, mem_per_token); - - int last_n_size = params.repeat_last_n; - std::vector last_n_tokens(last_n_size); - std::fill(last_n_tokens.begin(), last_n_tokens.end(), 0); - if (params.interactive) { fprintf(stderr, "== Running in interactive mode. ==\n" #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32) @@ -220,7 +152,6 @@ int main(int argc, char ** argv) { is_interacting = true; } - int input_consumed = 0; bool input_noecho = false; int remaining_tokens = params.n_predict; @@ -230,75 +161,33 @@ int main(int argc, char ** argv) { printf(ANSI_COLOR_YELLOW); } - while (remaining_tokens > 0 || params.interactive) { - // predict - if (embd.size() > 0) { - const int64_t t_start_us = ggml_time_us(); - - if (!llama_eval(model, params.n_threads, n_past, embd, logits, mem_per_token)) { - fprintf(stderr, "Failed to predict\n"); - return 1; - } - - t_predict_us += ggml_time_us() - t_start_us; - } - - n_past += embd.size(); - embd.clear(); - - if (embd_inp.size() <= input_consumed) { - // out of user input, sample next token - const float top_k = params.top_k; - const float top_p = params.top_p; - const float temp = params.temp; - const float repeat_penalty = params.repeat_penalty; - - const int n_vocab = model.hparams.n_vocab; - - gpt_vocab::id id = 0; - - { - const int64_t t_start_sample_us = ggml_time_us(); - - if (params.ignore_eos) { - // set the logit of the eos token to zero to avoid sampling it - logits[logits.size() - n_vocab + EOS_TOKEN_ID] = 0; - } - - id = llama_sample_top_p_top_k(vocab, logits.data() + (logits.size() - n_vocab), last_n_tokens, repeat_penalty, top_k, top_p, temp, rng); - - last_n_tokens.erase(last_n_tokens.begin()); - last_n_tokens.push_back(id); + if(!llama_injest_input(ctx, params.prompt)) + { + fprintf(stderr, "Failed to injest prompt\n"); + return 1; + }; - t_sample_us += ggml_time_us() - t_start_sample_us; - } + // display text + input_noecho = false; + const std::vector& embd = llama_context_get_embd(ctx); + if (!input_noecho) { + for (auto id : embd) { + printf("%s", vocab.id_to_token[id].c_str()); + } + fflush(stdout); + } - // add it to the context - embd.push_back(id); + if (!input_noecho && params.use_color) { + printf(ANSI_COLOR_RESET); + } - // echo this to console - input_noecho = false; + const std::vector& last_n_tokens = llama_context_get_last_n_tokens(ctx); - // decrement remaining sampling budget - --remaining_tokens; - } else { - // some user input remains from prompt or interaction, forward it to processing - while (embd_inp.size() > input_consumed) { - embd.push_back(embd_inp[input_consumed]); - last_n_tokens.erase(last_n_tokens.begin()); - last_n_tokens.push_back(embd_inp[input_consumed]); - ++input_consumed; - if ((int) embd.size() >= params.n_batch) { - break; - } - } - } - - // display text - if (!input_noecho) { - for (auto id : embd) { - printf("%s", vocab.id_to_token[id].c_str()); - } + while (llama_context_not_finished(ctx) > 0) { + gpt_vocab::id model_output = 0; + bool response = llama_inference(ctx, model_output); + if (response) { + printf("%s", vocab.id_to_token[model_output].c_str()); fflush(stdout); } // reset color to default if we there is no pending user input @@ -306,9 +195,10 @@ int main(int argc, char ** argv) { printf(ANSI_COLOR_RESET); } + // in interactive mode, and not currently processing queued inputs; // check if we should prompt the user for more - if (params.interactive && embd_inp.size() <= input_consumed) { + if (params.interactive) { // check for reverse prompt for (auto antiprompt_inp : antipromptv_inp) { if (antiprompt_inp.size() && std::equal(antiprompt_inp.rbegin(), antiprompt_inp.rend(), last_n_tokens.rbegin())) { @@ -337,15 +227,8 @@ int main(int argc, char ** argv) { } else { line.pop_back(); // Remove the continue character } - buffer += line + '\n'; // Append the line to the result - } while (another_line); - if (params.use_color) printf(ANSI_COLOR_RESET); - - std::vector line_inp = ::llama_tokenize(vocab, buffer, false); - embd_inp.insert(embd_inp.end(), line_inp.begin(), line_inp.end()); - - if (params.instruct) { - embd_inp.insert(embd_inp.end(), inp_sfx.begin(), inp_sfx.end()); + // Do not clear existing context in interactive mode + llama_init_context_with_prompt(ctx, buf, false); } remaining_tokens -= line_inp.size(); @@ -371,24 +254,14 @@ int main(int argc, char ** argv) { is_interacting = true; } } - -#if defined (_WIN32) - signal(SIGINT, SIG_DFL); -#endif - - // report timing + + // report timing from context { const int64_t t_main_end_us = ggml_time_us(); - - fprintf(stderr, "\n\n"); - fprintf(stderr, "%s: mem per token = %8zu bytes\n", __func__, mem_per_token); - fprintf(stderr, "%s: load time = %8.2f ms\n", __func__, t_load_us/1000.0f); - fprintf(stderr, "%s: sample time = %8.2f ms\n", __func__, t_sample_us/1000.0f); - fprintf(stderr, "%s: predict time = %8.2f ms / %.2f ms per token\n", __func__, t_predict_us/1000.0f, t_predict_us/1000.0f/n_past); + llama_print_end_stats(ctx); fprintf(stderr, "%s: total time = %8.2f ms\n", __func__, (t_main_end_us - t_main_start_us)/1000.0f); } - - ggml_free(model.ctx); + llama_free_context(ctx_ptr); if (params.use_color) { printf(ANSI_COLOR_RESET); From a81b6700a2e49be07285928aa14ae17c684542d1 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Wed, 15 Mar 2023 20:49:08 -0700 Subject: [PATCH 07/18] Add make_unique for C++11 --- llama.cpp | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/llama.cpp b/llama.cpp index 70f763fd58196..dbda0f3904f2b 100644 --- a/llama.cpp +++ b/llama.cpp @@ -15,6 +15,17 @@ // TODO: move somewhere else #define QK 32 +// Define std::make_unique if it's not available +// (e.g. on C++11) +#ifndef __cpp_lib_make_unique +template +std::unique_ptr make_unique(Args&&... args) { + return std::unique_ptr(new T(std::forward(args)...)); +} +#else +using std::make_unique; +#endif + // determine number of model parts based on the dimension static const std::map LLAMA_N_PARTS = { { 4096, 1 }, @@ -106,7 +117,7 @@ struct llama_context model(std::move(model)), vocab(std::move(vocab)), params(params), - state(std::make_unique()) + state(make_unique()) { } ~llama_context(){ From 7fb4c512d1fe6d836fa005142ef1a6175775a36c Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Thu, 16 Mar 2023 20:42:04 -0700 Subject: [PATCH 08/18] Apply suggestions to CMakeLists.txt --- CMakeLists.txt | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index ded54b4b6c351..773c713c3064c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -112,17 +112,17 @@ endif() # endif() add_executable(llama main.cpp - llama.cpp - utils.cpp - utils.h) +) add_executable(quantize quantize.cpp - utils.cpp - utils.h) +) add_library(llamalib - llama.cpp) + llama.cpp + llama.h + utils.cpp + utils.h) add_library(ggml ggml.c From b0ed03b9ab90a5ae2099a9ebc93d1232a8ff61f3 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Thu, 16 Mar 2023 20:42:12 -0700 Subject: [PATCH 09/18] Apply suggestions to llama.cpp and llama.h --- llama.cpp | 24 ++++++++++++------------ llama.h | 29 +++++++++++++++-------------- 2 files changed, 27 insertions(+), 26 deletions(-) diff --git a/llama.cpp b/llama.cpp index dbda0f3904f2b..65973eb460ec0 100644 --- a/llama.cpp +++ b/llama.cpp @@ -810,15 +810,15 @@ bool llama_model_quantize(const std::string & fname_inp, const std::string & fna /* External API */ -const std::vector& llama_context_get_embd(const llama_context& ctx) { +const std::vector& llama_context_get_embedding(const llama_context& ctx) { return ctx.state->embd; } gpt_vocab& llama_context_get_vocab(llama_context& ctx) { return ctx.vocab; } -bool llama_context_not_finished(const llama_context& ctx) +bool llama_context_is_finished(const llama_context& ctx) { - return ctx.state->remaining_tokens > 0; + return ctx.state->remaining_tokens <= 0; } const std::vector llama_tokenize_text(const llama_context& ctx, const std::string& text) { return llama_tokenize(ctx.vocab, text, true); @@ -1129,7 +1129,7 @@ bool llama_eval( return true; } -bool llama_init_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing) { +bool llama_update_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing) { llama_state& state = *ctx.state; llama_model& model = ctx.model; const gpt_params& params = ctx.params; @@ -1165,9 +1165,9 @@ bool llama_init_context_with_prompt(llama_context& ctx, const std::string& text, return true; } -/// @brief Injests a batch of input tokens into the context +/// @brief Ingests a batch of input tokens into the context /// @param ctx -void llama_injest_input_batch(llama_context& ctx) +void llama_ingest_input_batch(llama_context& ctx) { llama_state& state = *ctx.state; const gpt_params& params = ctx.params; @@ -1233,22 +1233,22 @@ gpt_vocab::id llama_sample_token(llama_context& ctx) } return id; } -/// @brief Injest all input (in multiple batches) into model and run call predict() +/// @brief Ingest all input (in multiple batches) into model and run call predict() /// @param ctx -bool llama_injest_input(llama_context& ctx, const std::string& text, bool clear_existing) +bool llama_ingest_input(llama_context& ctx, const std::string& text, bool clear_existing) { llama_state& state = *ctx.state; // Initialize context, tokenize text and clear existing state if necessary - if(!state.is_initialized && !llama_init_context_with_prompt(ctx, text, clear_existing)) + if(!state.is_initialized && !llama_update_context_with_prompt(ctx, text, clear_existing)) { return false; } - // Injest the tokens into the model one batch at a time + // ingest the tokens into the model one batch at a time while (state.has_more_input()) { - llama_injest_input_batch(ctx); + llama_ingest_input_batch(ctx); if (state.embd.size() >= 0) { if(!llama_predict(ctx)) { @@ -1260,7 +1260,7 @@ bool llama_injest_input(llama_context& ctx, const std::string& text, bool clear_ } return true; } -bool llama_inference(llama_context& ctx, gpt_vocab::id& id) { +bool llama_infer(llama_context& ctx, gpt_vocab::id& id) { llama_state& state = *ctx.state; // Tokenize text if we are starting out diff --git a/llama.h b/llama.h index 387efa686b547..b9cdeeecd0acf 100644 --- a/llama.h +++ b/llama.h @@ -39,26 +39,27 @@ struct llama_hparams { struct llama_context; -void llama_free_context(llama_context* ctx); +// Startup +llama_context* llama_init_from_params(const gpt_params& params); -const std::vector& llama_context_get_embd(const llama_context& ctx); -gpt_vocab& llama_context_get_vocab(llama_context& ctx); -bool llama_context_not_finished(const llama_context& ctx); +// Input processing and inference +bool llama_ingest_input(llama_context& ctx, const std::string& text, bool clear_existing = true); +bool llama_context_is_finished(const llama_context& ctx); +bool llama_update_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing = true); const std::vector llama_tokenize_text(const llama_context& ctx, const std::string& text); +bool llama_infer(llama_context& ctx, gpt_vocab::id& model_output); -const std::vector& llama_context_get_last_n_tokens(const llama_context& ctx); -bool llama_init_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing = true); - -// Various functions for loading a ggml LLaMA model. -llama_context* llama_init_from_params(const gpt_params& params); +// Teardown +void llama_free_context(llama_context* ctx); -// Run inference on a LLaMA model using llama_context. -std::vector llama_eval(llama_context& ctx, const gpt_params& params, std::string& text); +// Getters and setters +gpt_vocab& llama_context_get_vocab(llama_context& ctx); +const std::vector& llama_context_get_embedding(const llama_context& ctx); +const std::vector& llama_context_get_last_n_tokens(const llama_context& ctx); +// Other bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype); -bool llama_injest_input(llama_context& ctx, const std::string& text, bool clear_existing = true); - -bool llama_inference(llama_context& ctx, gpt_vocab::id& model_output); +// Stats void llama_print_context_info(const llama_context& ctx); void llama_print_end_stats(const llama_context& ctx); From 912e6246d6e4b097d3b1cad741256d47ce4a0bd5 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Thu, 16 Mar 2023 20:42:46 -0700 Subject: [PATCH 10/18] Apply suggestions to main.cpp --- main.cpp | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/main.cpp b/main.cpp index c1872e1b0a363..0165ce3481d91 100644 --- a/main.cpp +++ b/main.cpp @@ -98,8 +98,6 @@ int main(int argc, char ** argv) { // Add a space in front of the first character to match OG llama tokenizer behavior params.prompt.insert(0, 1, ' '); - // tokenize the prompt - std::vector embd_inp = llama_tokenize_text(ctx, params.prompt); // prefix & suffix for instruct mode const std::vector inp_pfx = ::llama_tokenize(vocab, "\n\n### Instruction:\n\n", true); @@ -161,15 +159,15 @@ int main(int argc, char ** argv) { printf(ANSI_COLOR_YELLOW); } - if(!llama_injest_input(ctx, params.prompt)) + if(!llama_ingest_input(ctx, params.prompt)) { - fprintf(stderr, "Failed to injest prompt\n"); + fprintf(stderr, "Failed to ingest prompt\n"); return 1; }; // display text input_noecho = false; - const std::vector& embd = llama_context_get_embd(ctx); + const std::vector& embd = llama_context_get_embedding(ctx); if (!input_noecho) { for (auto id : embd) { printf("%s", vocab.id_to_token[id].c_str()); @@ -183,9 +181,9 @@ int main(int argc, char ** argv) { const std::vector& last_n_tokens = llama_context_get_last_n_tokens(ctx); - while (llama_context_not_finished(ctx) > 0) { + while (llama_context_is_finished(ctx) != true) { gpt_vocab::id model_output = 0; - bool response = llama_inference(ctx, model_output); + bool response = llama_infer(ctx, model_output); if (response) { printf("%s", vocab.id_to_token[model_output].c_str()); fflush(stdout); @@ -195,7 +193,6 @@ int main(int argc, char ** argv) { printf(ANSI_COLOR_RESET); } - // in interactive mode, and not currently processing queued inputs; // check if we should prompt the user for more if (params.interactive) { @@ -228,7 +225,7 @@ int main(int argc, char ** argv) { line.pop_back(); // Remove the continue character } // Do not clear existing context in interactive mode - llama_init_context_with_prompt(ctx, buf, false); + llama_update_context_with_prompt(ctx, buf, false); } remaining_tokens -= line_inp.size(); From 05224ed4725314715700be37087b2a6be0fffe96 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Fri, 17 Mar 2023 18:46:27 -0700 Subject: [PATCH 11/18] Add llama_tokens_to_string() to utils.cpp - Also single token converter --- utils.cpp | 10 ++++++++++ utils.h | 5 +++++ 2 files changed, 15 insertions(+) diff --git a/utils.cpp b/utils.cpp index 08d5c6ba625f2..f810ad01fd645 100644 --- a/utils.cpp +++ b/utils.cpp @@ -347,6 +347,16 @@ std::vector llama_tokenize(const gpt_vocab & vocab, const std::st return res; } +std::string llama_tokens_to_string(const gpt_vocab & vocab, const std::vector & tokens) { + std::string res; + for (auto t : tokens) { + res += vocab.id_to_token.at(t); + } + return res; +} +std::string single_llama_token_to_string(const gpt_vocab & vocab, const gpt_vocab::id & tokens) { + return vocab.id_to_token.at(tokens); +} bool gpt_vocab_init(const std::string & fname, gpt_vocab & vocab) { printf("%s: loading vocab from '%s'\n", __func__, fname.c_str()); diff --git a/utils.h b/utils.h index 49658f7d9441e..3a1cc58a01e15 100644 --- a/utils.h +++ b/utils.h @@ -81,6 +81,11 @@ std::vector gpt_tokenize(const gpt_vocab & vocab, const std::stri // ref: https://github.com/google/sentencepiece std::vector llama_tokenize(const gpt_vocab & vocab, const std::string & text, bool bos); +// convert tokens to string +// opposite llama_tokenize +std::string llama_tokens_to_string(const gpt_vocab & vocab, const std::vector & tokens); +inline std::string single_llama_token_to_string(const gpt_vocab & vocab, const gpt_vocab::id & tokens); + // load the tokens from encoder.json bool gpt_vocab_init(const std::string & fname, gpt_vocab & vocab); From 3839a08ceed7f36093e7fdf50d4dd780c0c4c667 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Thu, 16 Mar 2023 21:03:23 -0700 Subject: [PATCH 12/18] Refactor llama.cpp and llama.h --- llama.cpp | 408 ++++++++++++++++++++++++++++++++---------------------- llama.h | 32 ++++- 2 files changed, 266 insertions(+), 174 deletions(-) diff --git a/llama.cpp b/llama.cpp index 65973eb460ec0..bab9da93c878f 100644 --- a/llama.cpp +++ b/llama.cpp @@ -72,33 +72,36 @@ struct llama_model { }; struct llama_state { - int64_t t_sample_us = 0; - int64_t t_predict_us = 0; + // Timers + struct timing { + int64_t t_load_us = 0; - std::vector logits; + int64_t t_sample_us = 0; + int64_t t_predict_us = 0; + } timing; - mutable std::mt19937 rng; + // Random number generator + std::mt19937 rng{}; + // Tokens std::vector embd{}; + std::vector embd_inp{}; + std::vector last_n_tokens{}; + // Logits from inference + std::vector logits{}; + + // Counters int input_consumed = 0; - std::vector embd_inp; - std::vector last_n_tokens; int remaining_tokens = 0; int n_past = 0; size_t mem_per_token = 0; - bool is_initialized = false; - llama_state() {} - bool has_more_input() const { - return input_consumed < embd_inp.size(); - } + // Flag set after initialization + bool is_initialized = false; }; struct llama_context { - int64_t t_load_us = 0; - int64_t t_start_us = 0; - ggml_type wtype = ggml_type::GGML_TYPE_F16; // weight type (FP32 or FP16) llama_model model{}; @@ -111,8 +114,6 @@ struct llama_context llama_context() = default; // constructor llama_context(llama_model&& model, gpt_vocab&& vocab, const gpt_params& params): - t_load_us(0), - t_start_us(0), wtype(ggml_type::GGML_TYPE_F16), model(std::move(model)), vocab(std::move(vocab)), @@ -125,6 +126,7 @@ struct llama_context } }; +/* Original code by @ggerganov */ // load the model's weights from a file bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx) { fprintf(stderr, "%s: loading model from '%s' - please wait ...\n", __func__, fname.c_str()); @@ -806,93 +808,6 @@ bool llama_model_quantize(const std::string & fname_inp, const std::string & fna return true; } - - -/* External API */ - -const std::vector& llama_context_get_embedding(const llama_context& ctx) { - return ctx.state->embd; -} -gpt_vocab& llama_context_get_vocab(llama_context& ctx) { - return ctx.vocab; -} -bool llama_context_is_finished(const llama_context& ctx) -{ - return ctx.state->remaining_tokens <= 0; -} -const std::vector llama_tokenize_text(const llama_context& ctx, const std::string& text) { - return llama_tokenize(ctx.vocab, text, true); -} -const std::vector& llama_context_get_last_n_tokens(const llama_context& ctx) { - return ctx.state->last_n_tokens; -} -llama_context* llama_init_from_params(const gpt_params& params) { - llama_model model{}; - gpt_vocab vocab{}; - - // Compute time taken to load model - const int64_t t_start = ggml_time_us(); - bool ret = llama_model_load(params.model, model, vocab, 1024); - const int64_t t_end = ggml_time_us(); - if(!ret) - { - return nullptr; - } - llama_context* ctx = new llama_context(std::move(model), std::move(vocab), params); - ctx->t_load_us = t_end - t_start; - return ctx; -} -void llama_free_context(llama_context* ctx) { - delete ctx; -} - -const char * llama_print_system_info(void) { - static std::string s; - - s = ""; - s += "AVX = " + std::to_string(ggml_cpu_has_avx()) + " | "; - s += "AVX2 = " + std::to_string(ggml_cpu_has_avx2()) + " | "; - s += "AVX512 = " + std::to_string(ggml_cpu_has_avx512()) + " | "; - s += "FMA = " + std::to_string(ggml_cpu_has_fma()) + " | "; - s += "NEON = " + std::to_string(ggml_cpu_has_neon()) + " | "; - s += "ARM_FMA = " + std::to_string(ggml_cpu_has_arm_fma()) + " | "; - s += "F16C = " + std::to_string(ggml_cpu_has_f16c()) + " | "; - s += "FP16_VA = " + std::to_string(ggml_cpu_has_fp16_va()) + " | "; - s += "WASM_SIMD = " + std::to_string(ggml_cpu_has_wasm_simd()) + " | "; - s += "BLAS = " + std::to_string(ggml_cpu_has_blas()) + " | "; - s += "SSE3 = " + std::to_string(ggml_cpu_has_sse3()) + " | "; - s += "VSX = " + std::to_string(ggml_cpu_has_vsx()) + " | "; - - return s.c_str(); -} - -void llama_print_context_info(const llama_context& ctx) -{ - const gpt_params& params = ctx.params; - const std::vector& embd_inp = ctx.state->embd_inp; - { - fprintf(stderr, "\n"); - fprintf(stderr, "system_info: n_threads = %d / %d | %s\n", - params.n_threads, std::thread::hardware_concurrency(), llama_print_system_info()); - } - fprintf(stderr, "\n"); - fprintf(stderr, "%s: prompt: '%s'\n", __func__, params.prompt.c_str()); - fprintf(stderr, "%s: number of tokens in prompt = %zu\n", __func__, embd_inp.size()); - for (int i = 0; i < (int) embd_inp.size(); i++) { - fprintf(stderr, "%6d -> '%s'\n", embd_inp[i], ctx.vocab.id_to_token.at(embd_inp[i]).c_str()); - } - fprintf(stderr, "\n"); -} - -void llama_print_end_stats(const llama_context& ctx) -{ - const llama_state& state = *ctx.state; - fprintf(stderr, "\n\n"); - fprintf(stderr, "%s: mem per token = %8zu bytes\n", __func__, state.mem_per_token); - fprintf(stderr, "%s: load time = %8.2f ms\n", __func__, ctx.t_load_us/1000.0f); - fprintf(stderr, "%s: sample time = %8.2f ms\n", __func__, state.t_sample_us/1000.0f); - fprintf(stderr, "%s: predict time = %8.2f ms / %.2f ms per token\n", __func__, state.t_predict_us/1000.0f, state.t_predict_us/1000.0f/state.n_past); -} // evaluate the transformer // // - model: the model @@ -1129,25 +1044,56 @@ bool llama_eval( return true; } -bool llama_update_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing) { - llama_state& state = *ctx.state; - llama_model& model = ctx.model; - const gpt_params& params = ctx.params; +const char * llama_print_system_info(void) { + static std::string s; - if (clear_existing) { - state.embd.clear(); - state.input_consumed = 0; - state.embd_inp.clear(); - state.last_n_tokens.clear(); - state.remaining_tokens = 0; - state.n_past = 0; - } + s = ""; + s += "AVX = " + std::to_string(ggml_cpu_has_avx()) + " | "; + s += "AVX2 = " + std::to_string(ggml_cpu_has_avx2()) + " | "; + s += "AVX512 = " + std::to_string(ggml_cpu_has_avx512()) + " | "; + s += "FMA = " + std::to_string(ggml_cpu_has_fma()) + " | "; + s += "NEON = " + std::to_string(ggml_cpu_has_neon()) + " | "; + s += "ARM_FMA = " + std::to_string(ggml_cpu_has_arm_fma()) + " | "; + s += "F16C = " + std::to_string(ggml_cpu_has_f16c()) + " | "; + s += "FP16_VA = " + std::to_string(ggml_cpu_has_fp16_va()) + " | "; + s += "WASM_SIMD = " + std::to_string(ggml_cpu_has_wasm_simd()) + " | "; + s += "BLAS = " + std::to_string(ggml_cpu_has_blas()) + " | "; + s += "SSE3 = " + std::to_string(ggml_cpu_has_sse3()) + " | "; + s += "VSX = " + std::to_string(ggml_cpu_has_vsx()) + " | "; - std::vector line_inp = llama_tokenize_text(ctx, text); - state.embd_inp.insert(state.embd_inp.end(), line_inp.begin(), line_inp.end()); + return s.c_str(); +} + +/* External API */ +/// @brief Initialize the context from a set of parameters +/// @param params +llama_context* llama_init_from_params(const gpt_params& params) { + llama_model model{}; + gpt_vocab vocab{}; + + // Compute time taken to load model + const int64_t t_start = ggml_time_us(); + bool ret = llama_model_load(params.model, model, vocab, 1024); + const int64_t t_end = ggml_time_us(); + if(!ret) + { + return nullptr; + } + llama_context* ctx = new llama_context(std::move(model), std::move(vocab), params); + ctx->state->timing.t_load_us = t_end - t_start; + ctx->state->rng = std::mt19937(params.seed); + return ctx; +} +/// @brief Prepare the context for inference +/// @param ctx +bool llama_prepare_context(llama_context& ctx) +{ + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + gpt_params& params = ctx.params; int n_predict = std::min(params.n_predict, model.hparams.n_ctx - (int) state.embd_inp.size()); - state.remaining_tokens = n_predict; + params.n_predict = n_predict; // determine the required inference memory per token: state.mem_per_token = 0; @@ -1160,10 +1106,76 @@ bool llama_update_context_with_prompt(llama_context& ctx, const std::string& tex int last_n_size = params.repeat_last_n; state.last_n_tokens = std::vector(last_n_size); std::fill(state.last_n_tokens.begin(), state.last_n_tokens.end(), 0); - state.is_initialized = true; + state.remaining_tokens = params.n_predict; + state.input_consumed = 0; return true; } +/// @brief Free the context +void llama_free_context(llama_context* ctx) { + delete ctx; +} + +/* Getters and setters */ +/// @brief Get the embedding vector for the last token +const std::vector& llama_context_get_embedding(const llama_context& ctx) { + return ctx.state->embd; +} +/// @brief Get the vector for the last token +/// @param ctx +gpt_vocab& llama_context_get_vocab(llama_context& ctx) { + return ctx.vocab; +} +/// @brief Is the context finished? +/// @param ctx +bool llama_context_is_finished(const llama_context& ctx) +{ + return ctx.state->remaining_tokens <= 0; +} +/// @brief Is the context finished? +/// @param ctx +void llama_reset_remaining_tokens(const llama_context& ctx) +{ + ctx.state->remaining_tokens = ctx.params.n_predict; +} + +/// @brief Tokenize a text into a vector of ids +/// @param ctx +/// @param text +const std::vector llama_tokenize_text(const llama_context& ctx, const std::string& text) { + // Make sure that the "beginning of string" token is not prefixed to the text + return llama_tokenize(ctx.vocab, text, false); +} +const std::vector& llama_context_get_last_n_tokens(const llama_context& ctx) { + return ctx.state->last_n_tokens; +} + +/// @brief Adds the "beginning of string" token to the model input +/// @param ctx +void llama_add_bos(llama_context& ctx){ + // Add the "bos" token into the model input + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + const gpt_params& params = ctx.params; + + const gpt_vocab::id bos_token = 1; + state.embd_inp.push_back(bos_token); +} + +/// @brief Updates the context and appends new input text +/// @param ctx +/// @param text +void llama_update_input(llama_context& ctx, const std::string& text) +{ + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + const gpt_params& params = ctx.params; + + std::vector line_inp = llama_tokenize_text(ctx, text); + + state.embd_inp.insert(state.embd_inp.end(), line_inp.begin(), line_inp.end()); + state.remaining_tokens -= line_inp.size(); +} /// @brief Ingests a batch of input tokens into the context /// @param ctx @@ -1173,7 +1185,7 @@ void llama_ingest_input_batch(llama_context& ctx) const gpt_params& params = ctx.params; // Copy at most n_batch elements from embd_inp to embd - size_t num_copied = std::min((size_t) params.n_batch, state.embd_inp.size() - state.input_consumed); + size_t num_copied = std::min((size_t) params.n_batch+1, state.embd_inp.size() - state.input_consumed); std::copy(state.embd_inp.begin() + state.input_consumed, state.embd_inp.begin() + state.input_consumed + num_copied, std::back_inserter(state.embd)); @@ -1185,21 +1197,60 @@ void llama_ingest_input_batch(llama_context& ctx) state.last_n_tokens.insert(state.last_n_tokens.end(), state.embd.end() - num_copied_last_n, state.embd.end()); } -/// @brief Run the prediction step on ctx.embd and store result in ctx.state.logits +/// @brief Returns true if there is unconsumed input in the context /// @param ctx -/// @return -bool llama_predict(llama_context& ctx){ - const int64_t t_start_us = ggml_time_us(); +bool llama_has_unconsumed_input(llama_context& ctx) +{ llama_state& state = *ctx.state; - llama_model& model = ctx.model; - const gpt_params& params = ctx.params; + return state.input_consumed < state.embd_inp.size(); +} - if (!llama_eval(model, params.n_threads, state.n_past, state.embd, state.logits, state.mem_per_token)) { - fprintf(stderr, "Failed to predict\n"); +/// @brief Ingest all input (in multiple batches) into model and run call predict() +/// @param ctx +bool llama_ingest_all_pending_input(llama_context& ctx, bool print_tokens) +{ + llama_state& state = *ctx.state; + const std::vector& embd = state.embd; + gpt_vocab& vocab = ctx.vocab; + + if(!state.is_initialized) + { + fprintf(stderr, "Context must be initialized before ingesting input"); return false; } - state.t_predict_us += ggml_time_us() - t_start_us; + // ingest the tokens into the model one batch at a time + while (llama_has_unconsumed_input(ctx)) + { + llama_ingest_input_batch(ctx); + if (print_tokens) { + std::string s = llama_tokens_to_string(vocab, embd); + printf("%s", s.c_str()); + fflush(stdout); + } + llama_eval_model(ctx); + } + return true; +} +/// @brief Evaluate the model with the current input batch +/// @param ctx +bool llama_eval_model(llama_context& ctx) +{ + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + const gpt_params& params = ctx.params; + + if (state.embd.size() > 0) { + const int64_t t_start_us = ggml_time_us(); + + if (!llama_eval(model, params.n_threads, state.n_past, state.embd, state.logits, state.mem_per_token)) { + fprintf(stderr, "Failed to predict\n"); + return false; + } + state.timing.t_predict_us += ggml_time_us() - t_start_us; + } + state.n_past += state.embd.size(); + state.embd.clear(); return true; } /// @brief Sample a token from the logits @@ -1229,37 +1280,13 @@ gpt_vocab::id llama_sample_token(llama_context& ctx) state.last_n_tokens.erase(state.last_n_tokens.begin()); state.last_n_tokens.push_back(id); - state.t_sample_us += ggml_time_us() - t_start_sample_us; + state.timing.t_sample_us += ggml_time_us() - t_start_sample_us; } return id; } -/// @brief Ingest all input (in multiple batches) into model and run call predict() -/// @param ctx -bool llama_ingest_input(llama_context& ctx, const std::string& text, bool clear_existing) -{ - llama_state& state = *ctx.state; - - // Initialize context, tokenize text and clear existing state if necessary - if(!state.is_initialized && !llama_update_context_with_prompt(ctx, text, clear_existing)) - { - return false; - } - - // ingest the tokens into the model one batch at a time - while (state.has_more_input()) - { - llama_ingest_input_batch(ctx); - if (state.embd.size() >= 0) { - if(!llama_predict(ctx)) - { - return false; - }; - } - state.n_past += state.embd.size(); - state.embd.clear(); - } - return true; -} +/// @brief Run inference for one token and return the token id +/// @param ctx +/// @param id bool llama_infer(llama_context& ctx, gpt_vocab::id& id) { llama_state& state = *ctx.state; @@ -1275,25 +1302,68 @@ bool llama_infer(llama_context& ctx, gpt_vocab::id& id) { return false; } - // Do prediction if we have enough tokens - if (state.embd.size() > 0) { - if(!llama_predict(ctx)) - { - return false; - } - } - // sample a token + // Already predicted, so we just need to sample + // sample a token id = llama_sample_token(ctx); + // add it to the context state.embd.push_back(id); - state.n_past += 1; // decrement remaining sampling budget --state.remaining_tokens; - // end of text token - if (state.embd.back() == 2) { - state.remaining_tokens = 0; + return true; +} +/// @brief Run inference for one token and return the token as a string +/// @param ctx +/// @param output +/// @param is_end_of_text +bool llama_infer(llama_context& ctx, std::string& output, bool& is_end_of_text) { + // Call overloaded llama_infer and convert to string before returning + gpt_vocab::id id_int; + is_end_of_text = false; + if(!llama_infer(ctx, id_int)){ + return false; } + + // Pass through the "end of text" token to the user + is_end_of_text = (id_int == EOS_TOKEN_ID); + + // Make sure to pass in the newly generated token to the model as well + llama_eval_model(ctx); + output = ctx.vocab.id_to_token.at(id_int); return true; } +bool llama_is_anti_prompt_present(llama_context& ctx, const std::vector& antiprompt_inp) +{ + llama_state& state = *ctx.state; + return std::equal(antiprompt_inp.rbegin(), antiprompt_inp.rend(), state.last_n_tokens.rbegin()); +} + +void llama_print_startup_stats(const llama_context& ctx) +{ + const gpt_params& params = ctx.params; + const std::vector& embd_inp = ctx.state->embd_inp; + { + fprintf(stderr, "\n"); + fprintf(stderr, "system_info: n_threads = %d / %d | %s\n", + params.n_threads, std::thread::hardware_concurrency(), llama_print_system_info()); + } + fprintf(stderr, "\n"); + fprintf(stderr, "%s: prompt: '%s'\n", __func__, params.prompt.c_str()); + fprintf(stderr, "%s: number of tokens in prompt = %zu\n", __func__, embd_inp.size()); + for (int i = 0; i < (int) embd_inp.size(); i++) { + fprintf(stderr, "%6d -> '%s'\n", embd_inp[i], ctx.vocab.id_to_token.at(embd_inp[i]).c_str()); + } + fprintf(stderr, "\n"); +} + +void llama_print_end_stats(const llama_context& ctx) +{ + const llama_state& state = *ctx.state; + fprintf(stderr, "\n\n"); + fprintf(stderr, "%s: mem per token = %8zu bytes\n", __func__, state.mem_per_token); + fprintf(stderr, "%s: load time = %8.2f ms\n", __func__, state.timing.t_load_us/1000.0f); + fprintf(stderr, "%s: sample time = %8.2f ms\n", __func__, state.timing.t_sample_us/1000.0f); + fprintf(stderr, "%s: predict time = %8.2f ms / %.2f ms per token\n", __func__, state.timing.t_predict_us/1000.0f, state.timing.t_predict_us/1000.0f/state.n_past); +} diff --git a/llama.h b/llama.h index b9cdeeecd0acf..adbfed173c029 100644 --- a/llama.h +++ b/llama.h @@ -22,7 +22,7 @@ # define LLAMA_API #endif - +static const int EOS_TOKEN_ID = 2; // default hparams (LLaMA 7B) @@ -41,13 +41,35 @@ struct llama_context; // Startup llama_context* llama_init_from_params(const gpt_params& params); +bool llama_prepare_context(llama_context& ctx); // Input processing and inference -bool llama_ingest_input(llama_context& ctx, const std::string& text, bool clear_existing = true); -bool llama_context_is_finished(const llama_context& ctx); -bool llama_update_context_with_prompt(llama_context& ctx, const std::string& text, bool clear_existing = true); +// Tokenize text (never adds BOS) const std::vector llama_tokenize_text(const llama_context& ctx, const std::string& text); +// Queues up a BOS token to the model input +void llama_add_bos(llama_context& ctx); +// Queues up input text to the model input +void llama_update_input(llama_context& ctx, const std::string& text); +// Ingests input previously added using llama_update_input() +void llama_ingest_input_batch(llama_context& ctx); +// Ingests all input previously added using llama_update_input() in multiple batches +// Batch size is determined by gpt_params::n_predict +bool llama_ingest_all_pending_input(llama_context& ctx, bool print_tokens = false); +// Checks if the model has unconsumed input to be ingested using llama_ingest_input_batch() +bool llama_has_unconsumed_input(llama_context& ctx); +// Checks if the model has an anti-prompt present its most recent output +bool llama_is_anti_prompt_present(llama_context& ctx, const std::vector& antiprompt_inp); + +// Evaluate the model on a batch of input. Must call llama_ingest_input_batch() first. +bool llama_eval_model(llama_context& ctx); +// Checks if the model has finished generating output (i.e. has generated an EOS token or remaining_tokens == 0) +bool llama_context_is_finished(const llama_context& ctx); +// Resets the remaining_tokens counter to the value specified in the gpt_params +void llama_reset_remaining_tokens(const llama_context& ctx); + +// Overloaded functions to run inference and return either the model output or the decoded text bool llama_infer(llama_context& ctx, gpt_vocab::id& model_output); +bool llama_infer(llama_context& ctx, std::string& output, bool& is_end_of_text); // Teardown void llama_free_context(llama_context* ctx); @@ -61,5 +83,5 @@ const std::vector& llama_context_get_last_n_tokens(const llama_co bool llama_model_quantize(const std::string & fname_inp, const std::string & fname_out, int itype); // Stats -void llama_print_context_info(const llama_context& ctx); +void llama_print_startup_stats(const llama_context& ctx); void llama_print_end_stats(const llama_context& ctx); From 660a4d5365f1b89ee82827bfe6e4201996506d54 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Fri, 17 Mar 2023 19:03:20 -0700 Subject: [PATCH 13/18] Refactor interactive mode in main.cpp --- main.cpp | 157 +++++++++++++++++++++++++++++++------------------------ 1 file changed, 90 insertions(+), 67 deletions(-) diff --git a/main.cpp b/main.cpp index 0165ce3481d91..f9d6f894ac35d 100644 --- a/main.cpp +++ b/main.cpp @@ -28,7 +28,6 @@ #define ANSI_COLOR_RESET "\x1b[0m" #define ANSI_BOLD "\x1b[1m" -static const int EOS_TOKEN_ID = 2; // determine number of model parts based on the dimension static const std::map LLAMA_N_PARTS = { @@ -56,6 +55,8 @@ void sigint_handler(int signo) { #endif +void process_interactive_input(llama_context& ctx, const gpt_params& params); + int main(int argc, char ** argv) { ggml_time_init(); const int64_t t_main_start_us = ggml_time_us(); @@ -86,15 +87,18 @@ int main(int argc, char ** argv) { // params.prompt = R"(// this function checks if the number n is prime //bool is_prime(int n) {)"; - int64_t t_load_us = 0; - // load the model - llama_context* ctx_ptr = llama_init_from_params(params); + llama_context* ctx_ptr = nullptr; + { + ctx_ptr = llama_init_from_params(params); + if (!ctx_ptr) { + fprintf(stderr, "%s: failed to load model from '%s'\n", __func__, params.model.c_str()); + return 1; + } + } + llama_context & ctx = *ctx_ptr; - gpt_vocab & vocab = llama_context_get_vocab(ctx); - - // print system information - llama_print_context_info(ctx); + const gpt_vocab & vocab = llama_context_get_vocab(ctx); // Add a space in front of the first character to match OG llama tokenizer behavior params.prompt.insert(0, 1, ' '); @@ -110,8 +114,9 @@ int main(int argc, char ** argv) { } // tokenize the reverse prompt - std::vector antiprompt_inp = llama_tokenize_text(ctx, params.prompt); + std::vector antiprompt_inp = llama_tokenize_text(ctx, params.antiprompt); + // Setup interactive mode if (params.interactive) { #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) struct sigaction sigint_action; @@ -150,43 +155,43 @@ int main(int argc, char ** argv) { is_interacting = true; } - bool input_noecho = false; - - int remaining_tokens = params.n_predict; - // set the color for the prompt which will be output initially if (params.use_color) { printf(ANSI_COLOR_YELLOW); } - if(!llama_ingest_input(ctx, params.prompt)) + // Prepare the context with input + // Send "beginning of string" + llama_add_bos(ctx); + + // load the input + llama_update_input(ctx, params.prompt); + + llama_print_startup_stats(ctx); + + if(!llama_prepare_context(ctx)) { - fprintf(stderr, "Failed to ingest prompt\n"); + fprintf(stderr, "%s: failed to prepare context\n", __func__); return 1; - }; - - // display text - input_noecho = false; - const std::vector& embd = llama_context_get_embedding(ctx); - if (!input_noecho) { - for (auto id : embd) { - printf("%s", vocab.id_to_token[id].c_str()); - } - fflush(stdout); } - if (!input_noecho && params.use_color) { - printf(ANSI_COLOR_RESET); - } - - const std::vector& last_n_tokens = llama_context_get_last_n_tokens(ctx); - - while (llama_context_is_finished(ctx) != true) { - gpt_vocab::id model_output = 0; - bool response = llama_infer(ctx, model_output); - if (response) { - printf("%s", vocab.id_to_token[model_output].c_str()); - fflush(stdout); + bool input_noecho = false; + bool is_end_of_text = false; + while (llama_context_is_finished(ctx) == false) { + std::string model_output{}; + + if (llama_has_unconsumed_input(ctx)) { + llama_ingest_all_pending_input(ctx, !input_noecho); + // reset color to default if we there is no pending user input + if (!input_noecho && params.use_color) { + printf(ANSI_COLOR_RESET); + } + }else{ + // Run inference if we don't have any pending input + llama_infer(ctx, model_output, is_end_of_text); + // print the single token output + printf("%s", model_output.c_str()); + input_noecho = false; } // reset color to default if we there is no pending user input if (!input_noecho && params.use_color && (int)embd_inp.size() == input_consumed) { @@ -195,10 +200,9 @@ int main(int argc, char ** argv) { // in interactive mode, and not currently processing queued inputs; // check if we should prompt the user for more - if (params.interactive) { - // check for reverse prompt - for (auto antiprompt_inp : antipromptv_inp) { - if (antiprompt_inp.size() && std::equal(antiprompt_inp.rbegin(), antiprompt_inp.rend(), last_n_tokens.rbegin())) { + if (params.interactive && !llama_has_unconsumed_input(ctx)) { + // check for reverse prompt + if (antiprompt_inp.size() && llama_is_anti_prompt_present(ctx, antiprompt_inp)) { // reverse prompt found is_interacting = true; break; @@ -206,37 +210,21 @@ int main(int argc, char ** argv) { } if (is_interacting) { if (params.instruct) { - input_consumed = embd_inp.size(); - embd_inp.insert(embd_inp.end(), inp_pfx.begin(), inp_pfx.end()); + llama_update_input(ctx, "\n\n### Instruction:\n\n"); printf("\n> "); } // currently being interactive - if (params.use_color) printf(ANSI_BOLD ANSI_COLOR_GREEN); - std::string buffer; - std::string line; - bool another_line = true; - do { - std::getline(std::cin, line); - if (line.empty() || line.back() != '\\') { - another_line = false; - } else { - line.pop_back(); // Remove the continue character - } - // Do not clear existing context in interactive mode - llama_update_context_with_prompt(ctx, buf, false); - } - - remaining_tokens -= line_inp.size(); - - input_noecho = true; // do not echo this again + process_interactive_input(ctx, params); + input_noecho = true; // do not echo this input again + is_interacting = false; } is_interacting = false; } // end of text token - if (embd.back() == EOS_TOKEN_ID) { + if (is_end_of_text) { if (params.interactive) { is_interacting = true; } else { @@ -246,23 +234,58 @@ int main(int argc, char ** argv) { } // In interactive mode, respect the maximum number of tokens and drop back to user input when reached. - if (params.interactive && remaining_tokens <= 0) { - remaining_tokens = params.n_predict; + if (params.interactive && llama_context_is_finished(ctx)) { + llama_reset_remaining_tokens(ctx); is_interacting = true; } } - // report timing from context + +#if defined (_WIN32) + signal(SIGINT, SIG_DFL); +#endif + + // report timing { const int64_t t_main_end_us = ggml_time_us(); llama_print_end_stats(ctx); fprintf(stderr, "%s: total time = %8.2f ms\n", __func__, (t_main_end_us - t_main_start_us)/1000.0f); } - llama_free_context(ctx_ptr); + + llama_free_context(ctx_ptr); if (params.use_color) { printf(ANSI_COLOR_RESET); } - return 0; } + +void process_interactive_input(llama_context& ctx, const gpt_params& params) +{ + bool another_line = true; + while (another_line) { + fflush(stdout); + char buf[256] = {0}; + int n_read; + if (params.use_color) printf(ANSI_BOLD ANSI_COLOR_GREEN); + if (scanf("%255[^\n]%n%*c", buf, &n_read) <= 0) { + // presumable empty line, consume the newline + std::ignore = scanf("%*c"); + n_read=0; + } + if (params.use_color) printf(ANSI_COLOR_RESET); + + if (n_read > 0 && buf[n_read-1]=='\\') { + another_line = true; + buf[n_read-1] = '\n'; + buf[n_read] = 0; + } else { + another_line = false; + buf[n_read] = '\n'; + buf[n_read+1] = 0; + } + + // Do not clear existing context in interactive mode + llama_update_input(ctx, buf); + } +} From edb52ab13cb672d3fb543fef5c262638903a4b98 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Sun, 19 Mar 2023 13:00:02 -0700 Subject: [PATCH 14/18] Update llama.cpp to use instruct mode --- llama.cpp | 12 ++++++++++++ llama.h | 2 ++ 2 files changed, 14 insertions(+) diff --git a/llama.cpp b/llama.cpp index bab9da93c878f..4cd69b7dd3915 100644 --- a/llama.cpp +++ b/llama.cpp @@ -1176,6 +1176,18 @@ void llama_update_input(llama_context& ctx, const std::string& text) state.embd_inp.insert(state.embd_inp.end(), line_inp.begin(), line_inp.end()); state.remaining_tokens -= line_inp.size(); } +/// @brief Updates the context and appends new input tokens (overloaded version) +/// @param ctx +/// @param tokens +void llama_update_input(llama_context& ctx, const std::vector& tokens) +{ + llama_state& state = *ctx.state; + llama_model& model = ctx.model; + const gpt_params& params = ctx.params; + + state.embd_inp.insert(state.embd_inp.end(), tokens.begin(), tokens.end()); + state.remaining_tokens -= tokens.size(); +} /// @brief Ingests a batch of input tokens into the context /// @param ctx diff --git a/llama.h b/llama.h index adbfed173c029..6a513873b7f91 100644 --- a/llama.h +++ b/llama.h @@ -50,6 +50,8 @@ const std::vector llama_tokenize_text(const llama_context& ctx, c void llama_add_bos(llama_context& ctx); // Queues up input text to the model input void llama_update_input(llama_context& ctx, const std::string& text); +// Queues up input tokens to the model input +void llama_update_input(llama_context& ctx, const std::vector& tokens); // Ingests input previously added using llama_update_input() void llama_ingest_input_batch(llama_context& ctx); // Ingests all input previously added using llama_update_input() in multiple batches From 1c1cf35694c5963d50b975cc62f52eaf75fd87d7 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Sun, 19 Mar 2023 13:00:08 -0700 Subject: [PATCH 15/18] Update main.cpp to use instruct mode --- main.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/main.cpp b/main.cpp index f9d6f894ac35d..49c43bd49aed6 100644 --- a/main.cpp +++ b/main.cpp @@ -210,13 +210,16 @@ int main(int argc, char ** argv) { } if (is_interacting) { if (params.instruct) { - llama_update_input(ctx, "\n\n### Instruction:\n\n"); - + llama_update_input(ctx, inp_pfx); printf("\n> "); } // currently being interactive process_interactive_input(ctx, params); + + if (params.instruct) { + llama_update_input(ctx, inp_sfx); + } input_noecho = true; // do not echo this input again is_interacting = false; } From b3541ceb65bbc85c90780b27f833544e385d27e8 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Sun, 19 Mar 2023 13:20:12 -0700 Subject: [PATCH 16/18] Update CMakeLists to rename targets - executable is now "main" and library is "llama" --- CMakeLists.txt | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 773c713c3064c..d13adc63f3e7c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -110,7 +110,7 @@ endif() # if (LLAMA_PERF) # set(LLAMA_EXTRA_FLAGS ${LLAMA_EXTRA_FLAGS} -DGGML_PERF) # endif() -add_executable(llama +add_executable(main main.cpp ) @@ -118,7 +118,7 @@ add_executable(quantize quantize.cpp ) -add_library(llamalib +add_library(llama llama.cpp llama.h utils.cpp @@ -133,10 +133,10 @@ target_compile_definitions(llama PUBLIC ${LLAMA_EXTRA_FLAGS}) target_compile_definitions(quantize PUBLIC ${LLAMA_EXTRA_FLAGS}) target_link_libraries(ggml PRIVATE ${LLAMA_EXTRA_LIBS}) -target_link_libraries(llamalib PRIVATE ggml ${LLAMA_EXTRA_LIBS}) +target_link_libraries(llama PRIVATE ggml ${LLAMA_EXTRA_LIBS}) target_include_directories(ggml PUBLIC .) -target_link_libraries(quantize PRIVATE ggml llamalib) -target_link_libraries(llama PRIVATE ggml llamalib) +target_link_libraries(quantize PRIVATE ggml llama) +target_link_libraries(main PRIVATE ggml llama) target_link_libraries(ggml PRIVATE Threads::Threads) From 5195fed01343107fd707077e88ff05265391cd44 Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Sun, 19 Mar 2023 13:39:06 -0700 Subject: [PATCH 17/18] Add support for multiple antiprompts --- main.cpp | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/main.cpp b/main.cpp index 49c43bd49aed6..5cfba70fb893f 100644 --- a/main.cpp +++ b/main.cpp @@ -114,7 +114,16 @@ int main(int argc, char ** argv) { } // tokenize the reverse prompt - std::vector antiprompt_inp = llama_tokenize_text(ctx, params.antiprompt); + std::vector> antipromptv_inp; + + for (auto antiprompt : params.antiprompt) { + antipromptv_inp.push_back(::llama_tokenize(vocab, antiprompt, false)); + } + + // enable interactive mode if reverse prompt is specified + if (!antipromptv_inp.size()) { + params.interactive = true; + } // Setup interactive mode if (params.interactive) { @@ -182,10 +191,6 @@ int main(int argc, char ** argv) { if (llama_has_unconsumed_input(ctx)) { llama_ingest_all_pending_input(ctx, !input_noecho); - // reset color to default if we there is no pending user input - if (!input_noecho && params.use_color) { - printf(ANSI_COLOR_RESET); - } }else{ // Run inference if we don't have any pending input llama_infer(ctx, model_output, is_end_of_text); @@ -193,15 +198,16 @@ int main(int argc, char ** argv) { printf("%s", model_output.c_str()); input_noecho = false; } - // reset color to default if we there is no pending user input - if (!input_noecho && params.use_color && (int)embd_inp.size() == input_consumed) { + // reset color to default (all input will be ingested already at this point) + if (!input_noecho && params.use_color) { printf(ANSI_COLOR_RESET); } // in interactive mode, and not currently processing queued inputs; // check if we should prompt the user for more if (params.interactive && !llama_has_unconsumed_input(ctx)) { - // check for reverse prompt + // check for reverse prompt + for (auto antiprompt_inp : antipromptv_inp) { if (antiprompt_inp.size() && llama_is_anti_prompt_present(ctx, antiprompt_inp)) { // reverse prompt found is_interacting = true; From 1c545e51ed9c8f7ebef225ee5c35a68518f6ab5c Mon Sep 17 00:00:00 2001 From: Thomas Antony Date: Sun, 19 Mar 2023 16:59:17 -0700 Subject: [PATCH 18/18] Update llama_model_load() from master branch --- llama.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/llama.cpp b/llama.cpp index 4cd69b7dd3915..bed40f13d3041 100644 --- a/llama.cpp +++ b/llama.cpp @@ -128,7 +128,8 @@ struct llama_context /* Original code by @ggerganov */ // load the model's weights from a file -bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx) { +bool llama_model_load(const std::string & fname, llama_model & model, gpt_vocab & vocab, int n_ctx, ggml_type memory_type = GGML_TYPE_F32) { + fprintf(stderr, "%s: loading model from '%s' - please wait ...\n", __func__, fname.c_str()); std::vector f_buf(1024*1024); @@ -1071,9 +1072,12 @@ llama_context* llama_init_from_params(const gpt_params& params) { llama_model model{}; gpt_vocab vocab{}; + const ggml_type memory_type = params.memory_f16 ? GGML_TYPE_F16 : GGML_TYPE_F32; + // Compute time taken to load model const int64_t t_start = ggml_time_us(); - bool ret = llama_model_load(params.model, model, vocab, 1024); + + bool ret = llama_model_load(params.model, model, vocab, params.n_ctx, memory_type); const int64_t t_end = ggml_time_us(); if(!ret) {