Felipe97/llama-cpp-compiled

1#include "common.h"2#include "llama.h"3#include "ggml.h"4 5#ifdef GGML_USE_CUDA6#include "ggml-cuda.h"7#endif8 9#ifdef GGML_USE_METAL10#include "ggml-metal.h"11#endif12 13#include <cstdio>14#include <ctime>15#include <random>16#include <string>17#include <vector>18 19#define DEBUG_POS 520 21static void print_debug_tensor(struct ggml_tensor * t, bool with_data = true) {22    printf("%s: %s (%s): [%d, %d]\n", __func__, t->name, ggml_type_name(t->type), (int) t->ne[0], (int) t->ne[1]);23    if (!with_data) return;24    printf("%s: %s[0] = [", __func__, t->name);25    for (size_t i = 0; i <= DEBUG_POS; i++) {26        printf(" %f,", ggml_get_f32_nd(t, i, 0, 0, 0));27    }28    printf(" ... ]\n");29}30 31namespace PCA {32 33// input params for PCA computations34struct pca_params {35    int n_threads = 1;36    int n_batch = 20; // number of iterations do to in one batch. larger the batch, more memory is used37    int n_iterations = 1000;38    float tolerance = 1e-7;39 40    // for debugging41    int i_layer = 0;42    int n_layers = 0;43};44 45// result from each iteration46struct pca_result {47    struct ggml_tensor * calculated_square = NULL;48    std::vector<struct ggml_tensor *> eigenvectors;49    std::vector<float> distances;50};51 52struct pca_model {53    ggml_backend_t backend = NULL;54    ggml_backend_buffer_t buffer;55    struct ggml_context * ctx;      // context to compute graph on target device56    struct ggml_context * ctx_host; // host context to store results57 58    // tensors on target device59    struct ggml_tensor * dev_input;60    struct ggml_tensor * dev_square;61    struct ggml_tensor * dev_eigenvector;62 63    pca_model(struct ggml_tensor * t_input) {64#ifdef GGML_USE_CUDA65        fprintf(stderr, "%s: using CUDA backend\n", __func__);66        backend = ggml_backend_cuda_init(0); // init device 067        if (!backend) {68            fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__);69        }70#endif71 72// TODO: enable Metal support when support for GGML_OP_SQRT is added73// #ifdef GGML_USE_METAL74//         fprintf(stderr, "%s: using Metal backend\n", __func__);75//         backend = ggml_backend_metal_init();76//         if (!backend) {77//             fprintf(stderr, "%s: ggml_backend_metal_init() failed\n", __func__);78//         }79// #endif80 81        // if there aren't GPU Backends fallback to CPU backend82        if (!backend) {83            backend = ggml_backend_cpu_init();84        }85 86        const int num_tensors = 4;87        struct ggml_init_params params {88            /*.mem_size   =*/ ggml_tensor_overhead() * num_tensors,89            /*.mem_buffer =*/ NULL,90            /*.no_alloc   =*/ true,91        };92        ctx = ggml_init(params);93 94        auto n_samples = t_input->ne[0];95        auto n_embd    = t_input->ne[1];96 97        dev_input       = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_samples, n_embd);98        dev_square      = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd,    n_embd);99        dev_eigenvector = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);100 101        ggml_set_name(dev_input,       "dev_input");102        ggml_set_name(dev_square,      "dev_square");103        ggml_set_name(dev_eigenvector, "dev_eigenvector");104        buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);105        ggml_backend_tensor_set(dev_input, t_input->data, 0, ggml_nbytes(t_input));106 107        // initialize eigenvector to random normalized vector108        {109            std::vector<float> random_vec(ggml_nelements(dev_eigenvector), 0.0);110            std::default_random_engine generator(static_cast<unsigned int>(std::time(0)));111            std::uniform_real_distribution<float> distribution(0.0, 1.0);112            float sum_sqr = 0.0; // for normalizing random_vec113            for (size_t i = 0; i < random_vec.size(); ++i) {114                float f = distribution(generator);115                sum_sqr += f * f;116                random_vec[i] = f;117            }118            // normalize it119            float random_vec_norm = std::sqrt(sum_sqr);120            for (size_t i = 0; i < random_vec.size(); ++i) {121                random_vec[i] /= random_vec_norm;122            }123            ggml_backend_tensor_set(dev_eigenvector, random_vec.data(), 0, ggml_nbytes(dev_eigenvector));124        }125    }126 127    ~pca_model() {128        ggml_free(ctx);129        ggml_backend_buffer_free(buffer);130        ggml_backend_free(backend);131    }132};133 134static struct ggml_cgraph * build_graph_piter(135        const struct pca_params & params,136        const pca_model & model,137        bool calc_square = false) {138    GGML_ASSERT(params.n_batch > 0);139    // TODO: buf_size must be able to scale with params.n_batch140    static size_t buf_size = ggml_tensor_overhead()*GGML_DEFAULT_GRAPH_SIZE + ggml_graph_overhead();141    static std::vector<uint8_t> buf(buf_size);142 143    struct ggml_init_params params0 = {144        /*.mem_size   =*/ buf_size,145        /*.mem_buffer =*/ buf.data(),146        /*.no_alloc   =*/ true, // the tensors will be allocated later by ggml_allocr_alloc_graph()147    };148    // create a temporally context to build the graph149    struct ggml_context * ctx0 = ggml_init(params0);150    struct ggml_cgraph * gf = ggml_new_graph(ctx0);151 152    // turn v_diff_original into square matrix if needed153    struct ggml_tensor * tmp_square;154    if (calc_square) {155        tmp_square = ggml_mul_mat(ctx0, model.dev_input, model.dev_input);156        ggml_set_name(tmp_square, "tmp_square");157    }158 159    struct ggml_tensor * b_tensor;160    struct ggml_tensor * distance;161    struct ggml_tensor * old_eigen    = model.dev_eigenvector;162    struct ggml_tensor * input_square = calc_square ? tmp_square : model.dev_square;163 164    for (int i = 0; i < params.n_batch; ++i) {165        // b_tensor = square * eigenvector^T166        b_tensor = ggml_mul_mat(ctx0, input_square, old_eigen);167        ggml_set_name(b_tensor, "b_tensor");168 169        // normalize170        b_tensor = ggml_div_inplace(ctx0,171            b_tensor,172            ggml_sqrt_inplace(ctx0, ggml_sum_rows(ctx0, ggml_sqr(ctx0, b_tensor)))173        );174        ggml_format_name(b_tensor, "b_tensor_norm_%d", i);175 176        // calculate distance(new eigenvector - old eigenvector)177        // we don't use ggml_sub because it may not be implemented on GPU backend178        struct ggml_tensor * new_sub_old = ggml_add(ctx0, old_eigen, ggml_scale(ctx0, b_tensor, -1));179        distance = ggml_sqrt_inplace(ctx0,180            ggml_sum_rows(ctx0, ggml_sqr_inplace(ctx0, new_sub_old)));181        ggml_format_name(distance, "distance_%d", i);182 183        old_eigen = b_tensor;184 185        // build operations nodes186        ggml_build_forward_expand(gf, distance);187    }188 189    // delete the temporally context used to build the graph190    ggml_free(ctx0);191    return gf;192}193 194static ggml_status compute_piter(195        const struct pca_params & params,196        const pca_model & model,197        struct ggml_cgraph * gf,198        ggml_gallocr_t allocr,199        struct pca_result & result) {200    // allocate tensors201    ggml_gallocr_alloc_graph(allocr, gf);202 203    if (ggml_backend_is_cpu(model.backend)) {204        ggml_backend_cpu_set_n_threads(model.backend, params.n_threads);205    }206 207    ggml_status res = ggml_backend_graph_compute(model.backend, gf);208    if (res == GGML_STATUS_SUCCESS) {209        auto extract_i = [](std::string prefix, std::string str) -> int {210            int i = -1;211            if (str.rfind(prefix, 0) == 0) {212                sscanf(str.c_str(), (prefix + "%d").c_str(), &i);213            }214            return i;215        };216        result.calculated_square = NULL;217        result.eigenvectors.clear();218        result.distances.clear();219        result.eigenvectors.resize(params.n_batch);220        result.distances.resize(params.n_batch);221        // get output nodes222        for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {223            auto node = ggml_graph_node(gf, i);224            int iter = -1;225            // find b_tensor (without copying data from device)226            if ((iter = extract_i("b_tensor_norm_", node->name)) > -1) {227                result.eigenvectors[iter] = node;228            }229            // find distances, then copy data from device230            if ((iter = extract_i("distance_", node->name)) > -1) {231                float d;232                ggml_backend_tensor_get(node, &d, 0, sizeof(float));233                result.distances[iter] = d;234                // std::cout << node->name << " = " << d << "\n";235            }236            // find tmp_square if it exists (without copying data from device)237            if (std::string(node->name) == "tmp_square") {238                result.calculated_square = node;239            }240        }241    }242    return res;243}244 245static void power_iteration(246        const struct pca_params & params,247        struct ggml_tensor * input, // shape of input: [n_samples, n_embd]248        struct ggml_tensor * output) {249    //printf("in power iteration\n");250    struct pca_model model(input);251 252    ggml_gallocr_t allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend));253    struct pca_result result;254    struct ggml_tensor * last_eigenvector = NULL;255 256    int n_iters = params.n_iterations / params.n_batch; // more batch, fewer iterations257    for (int iter = 0; iter < n_iters; ++iter) {258        bool calc_square = (iter == 0); // only need to calculate square for first iteration259        struct ggml_cgraph * gf = build_graph_piter(params, model, calc_square);260        // ggml_graph_dump_dot(gf, nullptr, "/tmp/_cgraph.dot");261        compute_piter(params, model, gf, allocr, result);262 263        for (size_t k = 0; k < result.distances.size(); ++k) {264            last_eigenvector = result.eigenvectors[k];265            if (result.distances[k] < params.tolerance) {266                break; // done267            }268        }269 270        if (calc_square) {271            // copy and store the square matrix if needed272            GGML_ASSERT(result.calculated_square != NULL);273            ggml_backend_tensor_copy(result.calculated_square, model.dev_square);274        }275 276        {277            // copy last eigen vector and store as input for next iteration278            GGML_ASSERT(last_eigenvector != NULL);279            ggml_backend_tensor_copy(last_eigenvector, model.dev_eigenvector);280        }281 282        printf("%s: layer %d/%d, iteration: %d / total: %d (batch = %d) ...\n",283            __func__, params.i_layer+1, params.n_layers, iter+1, n_iters, params.n_batch);284    }285 286    // get output tensor287    GGML_ASSERT(last_eigenvector);288    ggml_backend_tensor_get(last_eigenvector, output->data, 0, ggml_nbytes(last_eigenvector));289    //print_debug_tensor(output);290    ggml_gallocr_free(allocr);291 292    // TODO @ngxson : The output vector is randomly inverted293    // Solution: https://github.com/ggml-org/llama.cpp/pull/8069#issuecomment-2185328171294}295 296static void run_pca(297        struct pca_params & params,298        const std::vector<struct ggml_tensor *> & v_input, // shape of v_input[0]: [n_samples, n_embd]299        const std::vector<struct ggml_tensor *> & v_output) {300    printf("%s: Running PCA...\n", __func__);301    for (size_t il = 0; il < v_input.size(); ++il) {302 303        // prepare output vector304        struct ggml_tensor * ctrl_out = v_output[il];305        ggml_format_name(ctrl_out, "direction.%zu", il+1);306 307        // run power_iteration308        params.i_layer = il;309        params.n_layers = v_input.size();310        power_iteration(params, v_input[il], ctrl_out);311        printf("%s: Done layer %d / %d\n", __func__, (int) il+1, (int) v_input.size());312    }313}314 315}316