Felipe97/llama-cpp-compiled

1#include "ggml.h"2#include "gguf.h"3 4#include "arg.h"5#include "build-info.h"6#include "common.h"7#include "llama.h"8#include "pca.hpp"9#include "mean.hpp"10 11#include <clocale>12 13#ifdef GGML_USE_CUDA14#include "ggml-cuda.h"15#endif16 17#ifdef GGML_USE_METAL18#include "ggml-metal.h"19#endif20 21#include <algorithm>22#include <climits>23#include <cstdio>24#include <cstring>25#include <fstream>26#include <iostream>27#include <string>28#include <tuple>29#include <vector>30 31 32//////////////////////////////////////////////////33// utils34 35template <class Iter>36static std::string tokens_to_str(llama_context * ctx, Iter begin, Iter end) {37    std::string ret;38    for (; begin != end; ++begin) {39        ret += common_token_to_piece(ctx, *begin);40    }41 42    return ret;43}44 45static void print_usage(int, char ** argv) {46    printf("\nexample usage:\n");47    printf("\n    CPU only:   %s -m ./llama-3.Q4_K_M.gguf\n", argv[0]);48    printf("\n    with GPU:   %s -m ./llama-3.Q4_K_M.gguf -ngl 99\n", argv[0]);49    printf("\n    advanced:   %s -m ./llama-3.Q4_K_M.gguf -ngl 99 --pca-iter 2000 --pca-batch 100\n", argv[0]);50    printf("\n    using mean: %s -m ./llama-3.Q4_K_M.gguf --method mean\n", argv[0]);51    printf("\n");52}53 54//////////////////////////////////////////////////55 56 57// cb_eval is reused for each pair of positive - negative prompt58struct callback_data {59    ggml_context * ctx_ggml = nullptr;   // holds v_pos, v_neg, v_diff_filtered60 61    int n_layers = 0;62    int n_tokens = 0;63    bool is_eval_pos = true;64 65    // each element of the vector correspond to one layer66    std::vector<struct ggml_tensor *> v_pos; // vector of matrices of size [n_embd, n_tokens]67    std::vector<struct ggml_tensor *> v_neg; // vector of matrices of size [n_embd, n_tokens]68    std::vector<struct ggml_tensor *> v_diff_filtered;   // vector of matrices of size [n_embd, n_nonzero_rows]. NOTE: n_nonzero_rows maybe different for each layer69 70    // save a tensor into either v_pos or v_neg (decided by is_eval_pos)71    void save_tensor_for_layer(struct ggml_tensor * t) {72        GGML_ASSERT(t->type == GGML_TYPE_F32);73 74        if (ctx_ggml == nullptr) {75            // alloc a new ctx_ggml if needed76            struct ggml_init_params params_ggml = {77                /*.mem_size   =*/ ggml_tensor_overhead() * n_layers * 3u,78                /*.mem_buffer =*/ NULL,79                /*.no_alloc   =*/ true,80            };81            ctx_ggml = ggml_init(params_ggml);82        }83 84        // copy tensor data85        auto n_bytes = ggml_nbytes(t);86        struct ggml_tensor * t_layer = ggml_new_tensor_2d(ctx_ggml, t->type, t->ne[0], t->ne[1]);87        t_layer->data = malloc(n_bytes); // TODO @ngxson : get rid of this malloc somehow88        ggml_backend_tensor_get(t, t_layer->data, 0, n_bytes);89        ggml_set_name(t_layer, ggml_get_name(t));90        //print_debug_tensor(t_layer);91 92        if (is_eval_pos) {93            v_pos.push_back(t_layer);94        } else {95            v_neg.push_back(t_layer);96        }97    }98 99    // calculate diff (v_pos - v_neg) and place the result back to v_pos100    // all zero rows in the diff tensor will also be removed101    // NOTE: final layer is ignored. we only have (n_layers - 1) to process102    std::vector<struct ggml_tensor *> calc_diff() {103        for (float il = 0; il < v_pos.size(); il++) {104            float * a = (float *) v_pos[il]->data;105            float * b = (float *) v_neg[il]->data;106            size_t n_elem = ggml_nelements(v_pos[il]);107            for (size_t j = 0; j < n_elem; j++) {108                a[j] -= b[j];109            }110            //print_debug_tensor(v_pos[i]);111            auto diff_filtered = filter_nonzero_rows(v_pos[il]);112            v_diff_filtered.push_back(diff_filtered);113        }114        return v_diff_filtered; // for convenient, we return the result std::vector115    }116 117    // delete zero rows from a given 2D tensor118    struct ggml_tensor * filter_nonzero_rows(struct ggml_tensor * a) {119        //printf("filter_nonzero_rows\n");120        auto is_row_all_zeros = [](struct ggml_tensor * t, int row, float eps) -> bool {121            // check if given row containing all zero elements122            int n_cols = t->ne[0]; // hint: should be equal to n_embd123            for (int col = 0; col < n_cols; ++col) {124                if (ggml_get_f32_nd(t, col, row, 0, 0) > eps) {125                    return false;126                }127            }128            return true;129        };130        std::vector<int> rows_to_copy; // the idx of non-zero cols (to be copied to row of diff_filtered)131        for (int i_row = 0; i_row < a->ne[1]; i_row++) {132            if (!is_row_all_zeros(a, i_row, 1e-6)) {133                rows_to_copy.push_back(i_row);134            }135        }136 137        // get "n_nonzero_rows" for the output "diff_filtered"138        int n_nonzero_rows = rows_to_copy.size();139        //printf("n_nonzero_rows: %d\n", n_nonzero_rows);140        int n_embd = a->ne[0];141        GGML_ASSERT(n_nonzero_rows > 0);142 143        // diff_filtered: [n_embd, n_nonzero_rows]144        struct ggml_tensor * diff_filtered = ggml_new_tensor_2d(145            ctx_ggml, GGML_TYPE_F32, n_embd, n_nonzero_rows);146        ggml_format_name(diff_filtered, "diff_filtered_%s", a->name);147        diff_filtered->data = malloc(ggml_nbytes(diff_filtered));148 149        // copy non-zero rows150        for (int dest_row = 0; dest_row < n_nonzero_rows; dest_row++) {151            int src_row = rows_to_copy[dest_row];152            for (int i = 0; i < n_embd; i++) {153                float src_elem = ggml_get_f32_nd(a, i, src_row, 0, 0);154                ggml_set_f32_nd(diff_filtered, i, dest_row, 0, 0, src_elem);155            }156        }157 158        //print_debug_tensor(diff_filtered);159 160        return diff_filtered;161    }162 163    // we don't implement destructor, because we want to reuse callback_data. we just want to free the tensors164    void reset() {165        for (auto ptr : v_pos) free(ptr->data);166        for (auto ptr : v_neg) free(ptr->data);167        for (auto ptr : v_diff_filtered) free(ptr->data);168        v_pos.clear();169        v_neg.clear();170        v_diff_filtered.clear();171        if (ctx_ggml) {172            ggml_free(ctx_ggml);173        }174        ctx_ggml = nullptr;175    }176};177 178/**179 * process_ctx is used to store the ggml context for pre-post processing the diff vectors180 * in short, input => v_diff and output => v_final181 */182struct train_context {183    ggml_context * ctx_ggml;184    int n_embd;185    int n_layers;186 187    /* pair of prompts to be used for generating final vector */188    std::vector<std::string> positive_entries;189    std::vector<std::string> negative_entries;190 191    // each element of the vector correspond to one layer192    // NOTE: the last layer is discard. therefore, we will have (n_layers - 1) elements here193    // NOTE (2): v_diff is transposed from v_diff_tmp194    std::vector<struct ggml_tensor *> v_diff;  // vector of matrices of size [m, n_embd] where m ~ n_tokens * n_completions (v_diff contains no zero-rows)195    std::vector<struct ggml_tensor *> v_final; // vector of vectors of size [n_embd] to be written to file196 197    // to easily re-alloc when concat v_diff, we temporary store v_diff in a vector instead of a tensor198    // v_diff_tmp will get converted unto v_diff later on199    std::vector<std::vector<uint8_t>> v_diff_tmp;200 201    train_context(int n_embd_, int n_layers_) {202        n_embd = n_embd_;203        n_layers = n_layers_;204        struct ggml_init_params params_ggml = {205            /*.mem_size   =*/ ggml_tensor_overhead() * (n_layers - 1) * 2u,206            /*.mem_buffer =*/ NULL,207            /*.no_alloc   =*/ true,208        };209        ctx_ggml = ggml_init(params_ggml);210        for (int il = 0; il < n_layers - 1; il++) {211            std::vector<uint8_t> empty;212            v_diff_tmp.push_back(empty);213            auto t = ggml_new_tensor_1d(ctx_ggml, GGML_TYPE_F32, n_embd);214            t->data = malloc(ggml_nbytes(t)); // TODO: get rid of malloc if possible215            v_final.push_back(t);216        }217    }218 219    // add new rows into existing tensor in v_diff_tmp220    void concat_diff_tmp(const std::vector<struct ggml_tensor *> & diff_filtered) {221        GGML_ASSERT((int) diff_filtered.size() == n_layers - 1);222        for (int il = 0; il < n_layers - 1; il++) {223            auto t = diff_filtered[il];224            auto & diff_tmp = v_diff_tmp[il];225            size_t curr_size = diff_tmp.size();226            diff_tmp.resize(curr_size + ggml_nbytes(t));227            memcpy(diff_tmp.data() + curr_size, t->data, ggml_nbytes(t));228        }229    }230 231    // build the v_diff tensors from v_diff_tmp (v_diff need to be transposed)232    // TODO @ngxson : maybe add option NOT to transpose v_diff; will be useful for "mean" method233    void build_v_diff(bool transpose) {234        printf("build_v_diff\n");235        for (int il = 0; il < n_layers - 1; il++) {236            auto & diff_tmp = v_diff_tmp[il];237            int n_elem = diff_tmp.size() / sizeof(float);238            GGML_ASSERT(n_elem % n_embd == 0);239            int n_rows = n_elem / n_embd;240            struct ggml_tensor * diff = transpose241                ? ggml_new_tensor_2d(ctx_ggml, GGML_TYPE_F32, n_rows, n_embd)242                : ggml_new_tensor_2d(ctx_ggml, GGML_TYPE_F32, n_embd, n_rows);243            ggml_set_name(diff, (std::string("diff_") + std::to_string(il)).c_str());244            diff->data = malloc(ggml_nbytes(diff)); // TODO: get rid of this malloc if possible245            if (transpose) {246                // copy data & transpose247                float * arr = (float *) diff_tmp.data();248                for (int ir = 0; ir < n_rows; ++ir) {249                    for (int ic = 0; ic < n_embd; ++ic) {250                        float f = arr[ir*n_embd + ic];251                        ggml_set_f32_nd(diff, ir, ic, 0, 0, f);252                    }253                }254            } else {255                // only copy256                memcpy(diff->data, diff_tmp.data(), ggml_nbytes(diff));257            }258            v_diff.push_back(diff);259            print_debug_tensor(diff);260            // free memory of diff_tmp261            diff_tmp.resize(0);262        }263    }264 265    ~train_context() {266        for (auto ptr : v_final) free(ptr->data);267        for (auto ptr : v_diff) free(ptr->data);268        // no need to free v_diff_tmp, since we didn't use malloc269        ggml_free(ctx_ggml);270    }271};272 273struct tokenized_prompt {274    std::vector<llama_token> tokens_pos;275    std::vector<llama_token> tokens_neg;276    size_t max_seq_len;277 278    tokenized_prompt(llama_context * ctx, std::string pos, std::string neg) {279        const llama_model * model = llama_get_model(ctx);280        const llama_vocab * vocab = llama_model_get_vocab(model);281        const bool add_bos = llama_vocab_get_add_bos(vocab);282        tokens_pos = common_tokenize(ctx, pos, add_bos, true);283        tokens_neg = common_tokenize(ctx, neg, add_bos, true);284        max_seq_len = std::max(tokens_pos.size(), tokens_neg.size());285        padding_seq(ctx, tokens_pos, max_seq_len);286        padding_seq(ctx, tokens_neg, max_seq_len);287    }288 289    void padding_seq(llama_context * ctx, std::vector<llama_token> & tokens, size_t len) {290        // TODO: customize padding token291        std::vector<llama_token> pad_tokens = common_tokenize(ctx, " ", false);292        llama_token pad_tok = pad_tokens.back();293        while (tokens.size() < len) {294            tokens.push_back(pad_tok);295        }296    }297};298 299//////////////////////////////////////////////////300 301template <typename T>302static std::string to_string(const T & val) {303    std::stringstream ss;304    ss << val;305    return ss.str();306}307 308static std::vector<std::string> ctrlvec_load_prompt_file(std::string path, bool skip_empty_lines) {309    std::vector<std::string> output;310    std::ifstream file(path);311    if (!file.is_open()) {312        fprintf(stderr, "error: unable to open file: %s\n", path.c_str());313        exit(1);314    }315    std::string line;316    while (std::getline(file, line)) {317        bool is_skip = skip_empty_lines && line.empty();318        if (!is_skip) {319            string_process_escapes(line);320            output.push_back(line);321        }322    }323    file.close();324    return output;325}326 327//////////////////////////////////////////////////328 329static bool cb_eval(struct ggml_tensor * t, bool ask, void * user_data) {330    auto * cb_data = (callback_data *) user_data;331    static const char * l_out_name = "l_out";332    const bool is_l_out = strncmp(t->name, l_out_name, strlen(l_out_name)) == 0;333 334    if (ask) {335        return is_l_out;336    }337 338    if (!is_l_out || t->ne[1] != cb_data->n_tokens) {339        return true;340    }341 342    // save the tensor to current context343    cb_data->save_tensor_for_layer(t);344    return true;345}346 347static bool get_hidden_layers(llama_context * ctx, std::vector<llama_token> & tokens) {348    llama_memory_clear(llama_get_memory(ctx), true);349    if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size()))) {350        fprintf(stderr, "%s : failed to eval\n", __func__);351        return false;352    }353    return true;354}355 356static void export_gguf(const std::vector<struct ggml_tensor *> & v_ctrl, const std::string fname, const std::string model_hint) {357    struct gguf_context * ctx = gguf_init_empty();358 359    const std::string arch = "controlvector";360    gguf_set_val_str(ctx, "general.architecture", arch.c_str());361    gguf_set_val_str(ctx, (arch + ".model_hint").c_str(), model_hint.c_str());362    gguf_set_val_i32(ctx, (arch + ".layer_count").c_str(), v_ctrl.size());363 364    for (size_t i = 0; i < v_ctrl.size(); ++i) {365        gguf_add_tensor(ctx, v_ctrl[i]);366        print_debug_tensor(v_ctrl[i]);367        printf("Added tensor: %s\n", v_ctrl[i]->name);368    }369 370    printf("%s: writing file...\n", __func__);371    gguf_write_to_file(ctx, fname.c_str(), false);372    printf("%s: wrote file '%s'\n", __func__, fname.c_str());373    gguf_free(ctx);374}375 376/**377 * Load prompt files and completion file.378 * Then format each pair of prompt + completion to make an entry.379 */380static int prepare_entries(common_params & params, train_context & ctx_train) {381    // load prompts382    std::vector<std::string> positive_prompts = ctrlvec_load_prompt_file(params.cvector_positive_file, true);383    std::vector<std::string> negative_prompts = ctrlvec_load_prompt_file(params.cvector_negative_file, true);384    if (positive_prompts.size() != negative_prompts.size()) {385        fprintf(stderr, "number of positive and negative prompts must be equal\n");386        return 1;387    }388    if (positive_prompts.empty()) {389        fprintf(stderr, "must provide at least one prompt pair\n");390        return 1;391    }392    ctx_train.positive_entries = positive_prompts;393    ctx_train.negative_entries = negative_prompts;394    return 0;395}396 397int main(int argc, char ** argv) {398    std::setlocale(LC_NUMERIC, "C");399 400    common_params params;401 402    params.out_file = "control_vector.gguf";403 404    common_init();405 406    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_CVECTOR_GENERATOR, print_usage)) {407        return 1;408    }409 410    if (params.n_pca_iterations % params.n_pca_batch != 0) {411        fprintf(stderr, "PCA iterations must by multiply of PCA batch size\n");412        return 1;413    }414 415 416    callback_data cb_data;417 418    // pass the callback to the backend scheduler419    // it will be executed for each node during the graph computation420    params.cb_eval = cb_eval;421    params.cb_eval_user_data = &cb_data;422    params.warmup = false;423 424    llama_print_build_info(llama_version());425    llama_backend_init();426    llama_numa_init(params.numa);427 428    // load the model to get hparams429    auto llama_init = common_init_from_params(params);430 431    auto * model = llama_init->model();432    auto * ctx   = llama_init->context();433 434    // int n_ctx = llama_n_ctx(ctx);435    int n_layers = llama_model_n_layer(model);436    int n_embd = llama_model_n_embd(model);437 438    // get model hint param (a.k.a model arch name)439    char model_hint[128];440    llama_model_meta_val_str(model, "general.architecture", model_hint, 128);441 442    // init train_context443    train_context ctx_train(n_embd, n_layers);444 445    // load and prepare entries for training446    prepare_entries(params, ctx_train);447 448    // we have to pretokenize everything because otherwise we don't know how much overhead to allocate ctx_diffs_wrapped449    std::vector<tokenized_prompt> tokenized_prompts;450    size_t n_total_tokens = 0;451    for (size_t i = 0; i < ctx_train.positive_entries.size(); ++i) {452        tokenized_prompt t(ctx, ctx_train.positive_entries[i], ctx_train.negative_entries[i]);453        n_total_tokens += 2 * t.max_seq_len;454        tokenized_prompts.push_back(std::move(t));455    }456 457    std::cout << "n_total_tokens: " << n_total_tokens << std::endl;458 459    for(size_t i = 0; i < ctx_train.positive_entries.size(); ++i) {460        bool success = false;461        tokenized_prompt t = tokenized_prompts[i];462        cb_data.n_layers = n_layers;463        cb_data.n_tokens = t.max_seq_len;464 465        printf("Evaluating prompt[%d/%d]: \"%s\" - \"%s\" (%d tokens)\n",466            (int) i+1, (int) ctx_train.positive_entries.size(),467            tokens_to_str(ctx, t.tokens_pos.cbegin(), t.tokens_pos.cend()).c_str(),468            tokens_to_str(ctx, t.tokens_neg.cbegin(), t.tokens_neg.cend()).c_str(),469            (int) t.max_seq_len);470 471        cb_data.is_eval_pos = true;472        success = get_hidden_layers(ctx, t.tokens_pos);473        if (!success) break;474 475        cb_data.is_eval_pos = false;476        success = get_hidden_layers(ctx, t.tokens_neg);477        if (!success) break;478 479        // calculate diff and remove all zero rows480        auto v_diff_filtered = cb_data.calc_diff();481 482        // save & concat the filtered v_diff to ctx_train483        ctx_train.concat_diff_tmp(v_diff_filtered);484 485        // reset for next iteration486        cb_data.reset();487    }488 489    // done with the model, we can now free it to make gain some memory490    printf("Done evaluate prompts, unload model...\n");491 492    bool use_pca = params.cvector_dimre_method == DIMRE_METHOD_PCA;493 494    // prepare ctx_train for PCA495    ctx_train.build_v_diff(use_pca);496 497    if (use_pca) {498        // run PCA499        PCA::pca_params pca_params;500        pca_params.n_threads    = params.cpuparams.n_threads;501        pca_params.n_batch      = params.n_pca_batch;502        pca_params.n_iterations = params.n_pca_iterations;503        PCA::run_pca(pca_params, ctx_train.v_diff, ctx_train.v_final);504    } else {505        // run mean506        mean::run(ctx_train.v_diff, ctx_train.v_final);507    }508 509    // write output vectors to gguf510    export_gguf(ctx_train.v_final, params.out_file, model_hint);511 512    llama_backend_free();513 514    return 0;515}516