Felipe97/llama-cpp-compiled

1#include "ggml.h"2#include "ggml-alloc.h"3#include "gguf.h"4 5#include "arg.h"6#include "common.h"7 8#include <clocale>9#include <map>10#include <vector>11#include <string>12#include <fstream>13 14static bool g_verbose = false;15 16struct tensor_transformation {17    struct ggml_tensor * in;18    struct ggml_tensor * out;19    bool is_copy;20};21 22static std::string get_kv_str(struct gguf_context * ctx_gguf, const std::string & key){23    int id = gguf_find_key(ctx_gguf, key.c_str());24    return id < 0 ? "" : std::string(gguf_get_val_str(ctx_gguf, id));25}26 27static float get_kv_f32(struct gguf_context * ctx_gguf, const std::string & key) {28    int id = gguf_find_key(ctx_gguf, key.c_str());29    return id < 0 ? 0.0f : gguf_get_val_f32(ctx_gguf, id);30}31 32static void zeros(std::ofstream & file, size_t n) {33    char zero = 0;34    for (size_t i = 0; i < n; ++i) {35        file.write(&zero, 1);36    }37}38 39static std::string ggml_ne_string(const ggml_tensor * t) {40    std::string str;41    for (int i = 0; i < GGML_MAX_DIMS; ++i) {42        str += std::to_string(t->ne[i]);43        if (i + 1 < GGML_MAX_DIMS) {44            str += ", ";45        }46    }47    return str;48}49 50static struct gguf_context * load_gguf(std::string & fname, struct ggml_context ** ctx_ggml) {51    struct gguf_init_params params = {52        /*.no_alloc = */ true,53        /*.ctx      = */ ctx_ggml,54    };55    struct gguf_context * ctx_gguf = gguf_init_from_file(fname.c_str(), params);56    if (!ctx_gguf) {57        throw std::runtime_error("failed to load input GGUF from " + fname);58    }59    return ctx_gguf;60}61 62struct file_input {63    struct ggml_context * ctx_meta = nullptr;64    struct gguf_context * ctx_gguf = nullptr;65    std::ifstream f_in;66    std::map<std::string, ggml_tensor *> tensors;67    float alpha;68    float scale;69 70    file_input(std::string & fname, float scale): f_in(fname, std::ios::binary), scale(scale) {71        if (!f_in.is_open()) {72            throw std::runtime_error("failed to open input gguf from " + fname);73        }74 75        ctx_gguf = load_gguf(fname, &ctx_meta);76        alpha = get_kv_f32(ctx_gguf, "adapter.lora.alpha");77        printf("%s: loaded gguf from %s\n", __func__, fname.c_str());78 79        for (ggml_tensor * cur = ggml_get_first_tensor(ctx_meta); cur; cur = ggml_get_next_tensor(ctx_meta, cur)) {80            std::string name(cur->name);81            tensors[name] = cur;82            if (g_verbose) {83                printf("%s: %s\n", __func__, cur->name);84            }85        }86    }87 88    ggml_tensor * get_tensor(std::string name) {89        if (tensors.find(name) == tensors.end()) {90            return nullptr;91        }92        return tensors[name];93    }94 95    void read_tensor_data(std::string name, std::vector<uint8_t> & buf) {96        if (tensors.find(name) == tensors.end()) {97            throw std::runtime_error("cannot find tensor with name: " + name);98        }99        auto len = ggml_nbytes(tensors[name]);100        if (buf.size() < len) {101            buf.resize(len);102        }103        auto i_tensor_in = gguf_find_tensor(ctx_gguf, name.c_str()); // idx of tensor in the input file104        auto offset = gguf_get_data_offset(ctx_gguf) + gguf_get_tensor_offset(ctx_gguf, i_tensor_in);105        f_in.seekg(offset);106        f_in.read((char* )buf.data(), len);107    }108 109    ~file_input() {110        gguf_free(ctx_gguf);111        ggml_free(ctx_meta);112    }113};114 115struct lora_merge_ctx {116    // input base model + adapters117    file_input base_model;118    std::vector<std::unique_ptr<file_input>> adapters;119 120    // for computing merged tensor121    int n_threads;122    ggml_backend_t backend = nullptr;123    ggml_gallocr_t allocr = nullptr;124    std::vector<uint8_t> read_buf;125 126    // output file127    struct gguf_context * ctx_out;128    struct ggml_context * ctx_out_ggml;129    std::ofstream fout;130 131    lora_merge_ctx(132            std::string & base_fname,133            std::vector<common_adapter_lora_info> & lora_files,134            std::string & outfile,135            int n_threads) : base_model(base_fname, 0), n_threads(n_threads), fout(outfile, std::ios::binary) {136        fout.exceptions(std::ofstream::failbit); // fail fast on write errors137 138        if (gguf_find_key(base_model.ctx_gguf, LLM_KV_SPLIT_COUNT) >= 0) {139            throw std::runtime_error("split model is not yet supported");140        }141 142        for (auto & lora_inp : lora_files) {143            auto fname = lora_inp.path;144            auto scale = lora_inp.scale;145            std::unique_ptr<file_input> adapter(new file_input(fname, scale));146            check_metadata_lora(adapter.get());147            adapters.push_back(std::move(adapter));148        }149 150        ctx_out = gguf_init_empty();151        struct ggml_init_params params = {152            /*.mem_size   =*/ static_cast<size_t>(gguf_get_n_tensors(base_model.ctx_gguf)*ggml_tensor_overhead()),153            /*.mem_buffer =*/ NULL,154            /*.no_alloc   =*/ true,155        };156        ctx_out_ggml = ggml_init(params);157        backend = ggml_backend_cpu_init();158        allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));159    }160 161    void check_metadata_lora(file_input * adapter) {162        auto general_type = get_kv_str(adapter->ctx_gguf, "general.type");163        if (general_type != "adapter") {164            throw std::runtime_error("expect general.type to be 'adapter', but got: " + general_type);165        }166 167        auto adapter_type = get_kv_str(adapter->ctx_gguf, "adapter.type");168        if (adapter_type != "lora") {169            throw std::runtime_error("expect adapter.type to be 'lora', but got: " + adapter_type);170        }171 172        auto general_arch_base = get_kv_str(base_model.ctx_gguf, "general.architecture");173        auto general_arch_lora = get_kv_str(adapter->ctx_gguf,   "general.architecture");174        if (general_arch_base != general_arch_lora) {175            throw std::runtime_error("model arch and LoRA arch mismatch");176        }177    }178 179    ggml_type get_out_tensor_type(struct ggml_tensor * t) {180        if (t->type == GGML_TYPE_F32) {181            return GGML_TYPE_F32;182        } else {183            return GGML_TYPE_F16;184        }185    }186 187    void run_merge() {188        // prepare metadata189        gguf_set_kv(ctx_out, base_model.ctx_gguf);190        // output is forced to f16 for now191        gguf_set_val_u32(ctx_out, "general.file_type", LLAMA_FTYPE_MOSTLY_F16);192 193        // check if all lora adapters have the same tensors194        // TODO: remove this when we can support merging subset of adapters. Ref: https://github.com/ggml-org/llama.cpp/pull/8607#discussion_r1686027777195        static const char * err_no_subset_adapter = "Input adapters do not have the same list of tensors. This is not yet supported. Please merge the adapter one-by-one instead of merging all at once.";196        if (adapters.size() > 1) {197            for (size_t i = 1; i < adapters.size(); ++i) {198                if (adapters[0]->tensors.size() != adapters[i]->tensors.size()) {199                    throw std::runtime_error(err_no_subset_adapter);200                }201                for (auto & it : adapters[i]->tensors) {202                    if (adapters[0]->get_tensor(it.first) == nullptr) {203                        throw std::runtime_error(err_no_subset_adapter);204                    }205                }206            }207        }208 209        // mapping base tensor to out tensor (same shape with base, but different type)210        std::vector<tensor_transformation> trans;211        for (auto & it : base_model.tensors) {212            bool t_a = true;213            bool t_b = true;214            for (auto & adapter : adapters) {215                t_a &= nullptr != adapter->get_tensor(it.first + ".lora_a");216                t_b &= nullptr != adapter->get_tensor(it.first + ".lora_b");217            }218            auto base_tensor = it.second;219            if (!t_a && !t_b) {220                // only copy221                struct ggml_tensor * cpy_tensor = ggml_dup_tensor(ctx_out_ggml, base_tensor);222                ggml_set_name(cpy_tensor, base_tensor->name);223                trans.push_back({224                    cpy_tensor,225                    cpy_tensor,226                    true,227                });228                gguf_add_tensor(ctx_out, cpy_tensor);229            } else if (t_a && t_b) {230                // need merging231                struct ggml_tensor * out_tensor = ggml_new_tensor(232                    ctx_out_ggml, get_out_tensor_type(base_tensor), GGML_MAX_DIMS, base_tensor->ne);233                ggml_set_name(out_tensor, base_tensor->name);234                trans.push_back({235                    base_tensor,236                    out_tensor,237                    false,238                });239                gguf_add_tensor(ctx_out, out_tensor);240            } else {241                throw std::runtime_error("tensor " + it.first + " missing either lora_a or lora_b");242            }243        }244 245        // placeholder for the meta data246        {247            size_t meta_size = gguf_get_meta_size(ctx_out);248            zeros(fout, meta_size);249        }250 251        // process base model tensors252        size_t n_merged = 0;253        for (auto & it : trans) {254            if (!it.is_copy) {255                merge_tensor(it.in, it.out);256                n_merged++;257            } else {258                copy_tensor(it.in);259            }260        }261 262        // write output metadata263        {264            std::vector<uint8_t> data(gguf_get_meta_size(ctx_out));265            gguf_get_meta_data(ctx_out, data.data());266            fout.seekp(0);267            fout.write((const char *)data.data(), data.size());268        }269 270        printf("%s : merged %zu tensors with lora adapters\n", __func__, n_merged);271        printf("%s : wrote %zu tensors to output file\n", __func__, trans.size());272    }273 274    void copy_tensor(struct ggml_tensor * base) {275        printf("%s :  %s [%s]\n", __func__, base->name, ggml_ne_string(base).c_str());276        size_t len = ggml_nbytes(base);277        base_model.read_tensor_data(base->name, read_buf);278        fout.write((char* )read_buf.data(), len);279        zeros(fout, GGML_PAD(len, GGUF_DEFAULT_ALIGNMENT) - len);280    }281 282    void merge_tensor(struct ggml_tensor * base, struct ggml_tensor * out) {283        std::string name_base(base->name);284        std::string name_lora_a = name_base + ".lora_a";285        std::string name_lora_b = name_base + ".lora_b";286 287        printf("%s : %s [%s]\n", __func__, base->name, ggml_ne_string(base).c_str());288 289        // context for input tensor290        std::vector<struct ggml_tensor *> inp_a(adapters.size());291        std::vector<struct ggml_tensor *> inp_b(adapters.size());292        struct ggml_init_params params {293            /*.mem_size   =*/ ggml_tensor_overhead()*(2+adapters.size()*2),294            /*.mem_buffer =*/ NULL,295            /*.no_alloc   =*/ true,296        };297        struct ggml_context * ctx = ggml_init(params);298 299        // alloc tensors300        struct ggml_tensor * inp_base = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, base->ne);301        for (size_t i = 0; i < adapters.size(); ++i) {302            auto t_a = adapters[i]->get_tensor(name_lora_a);303            auto t_b = adapters[i]->get_tensor(name_lora_b);304            // TODO: add support for quantized lora305            if (ggml_is_quantized(t_a->type) || ggml_is_quantized(t_b->type)) {306                throw std::runtime_error("quantized LoRA adapters is not supported, please retry with f16 or f32");307            }308            inp_a[i] = ggml_dup_tensor(ctx, t_a);309            inp_b[i] = ggml_dup_tensor(ctx, t_b);310        }311        ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);312 313        // load base tensor to backend buffer314        base_model.read_tensor_data(name_base, read_buf);315        if (base->type != GGML_TYPE_F32) {316            // optionally dequantize it317            printf("%s :   + dequantize base tensor from %s to F32\n", __func__, ggml_type_name(base->type));318            auto nels = ggml_nelements(inp_base);319            const auto * qtype = ggml_get_type_traits(base->type);320            std::vector<uint8_t> dequant_buf(nels * sizeof(float));321            qtype->to_float(read_buf.data(), (float *)dequant_buf.data(), nels);322            ggml_backend_tensor_set(inp_base, dequant_buf.data(), 0, dequant_buf.size());323        } else {324            ggml_backend_tensor_set(inp_base, read_buf.data(), 0, ggml_nbytes(inp_base));325        }326 327        // load lora tensors to backend buffer328        for (size_t i = 0; i < adapters.size(); ++i) {329            adapters[i]->read_tensor_data(name_lora_a, read_buf);330            ggml_backend_tensor_set(inp_a[i], read_buf.data(), 0, ggml_nbytes(inp_a[i]));331            adapters[i]->read_tensor_data(name_lora_b, read_buf);332            ggml_backend_tensor_set(inp_b[i], read_buf.data(), 0, ggml_nbytes(inp_b[i]));333        }334 335        // build graph336        struct ggml_cgraph * gf;337        {338            static size_t buf_size = ggml_tensor_overhead()*GGML_DEFAULT_GRAPH_SIZE + ggml_graph_overhead();339            static std::vector<uint8_t> buf(buf_size);340            struct ggml_init_params params0 = {341                /*.mem_size   =*/ buf_size,342                /*.mem_buffer =*/ buf.data(),343                /*.no_alloc   =*/ true,344            };345            struct ggml_context * ctx0 = ggml_init(params0);346            gf = ggml_new_graph(ctx0);347            struct ggml_tensor * cur = inp_base;348            for (size_t i = 0; i < adapters.size(); ++i) {349                struct ggml_tensor * delta;350                bool is_tok_embd = string_starts_with(name_base, "token_embd");351                if (is_tok_embd) {352                    printf("%s :     detected token embeddings tensor\n", __func__);353                    delta = ggml_mul_mat(ctx0,354                        ggml_cast(ctx0, inp_b[i], GGML_TYPE_F32),355                        ggml_cast(ctx0, inp_a[i], GGML_TYPE_F32));356                } else {357                    delta = ggml_mul_mat(ctx0,358                        ggml_cont(ctx0, ggml_transpose(ctx0, ggml_cast(ctx0, inp_a[i], GGML_TYPE_F32))),359                        ggml_cast(ctx0, inp_b[i], GGML_TYPE_F32));360                }361                // scale362                const float alpha = adapters[i]->alpha;363                const float rank  = (float) inp_b[i]->ne[0];364                const float scale = alpha ? adapters[i]->scale * alpha / rank : adapters[i]->scale;365                delta = ggml_scale(ctx0, delta, scale);366                cur = ggml_add(ctx0, delta, cur);367                printf("%s :   + merging from adapter[%zu] type=%s\n", __func__, i, ggml_type_name(inp_a[i]->type));368                printf("%s :     input_scale=%f calculated_scale=%f rank=%d\n", __func__, adapters[i]->scale, scale, (int) inp_b[i]->ne[0]);369            }370            cur = ggml_cast(ctx0, cur, out->type);371            printf("%s :   + output type is %s\n", __func__, ggml_type_name(out->type));372            ggml_build_forward_expand(gf, cur);373            ggml_free(ctx0);374        }375 376        // compute377        {378            ggml_gallocr_alloc_graph(allocr, gf);379            ggml_backend_cpu_set_n_threads(backend, n_threads);380            ggml_backend_graph_compute(backend, gf);381        }382 383        // write data to output file384        {385            auto * result = ggml_graph_node(gf, -1);386            size_t len = ggml_nbytes(result);387            if (read_buf.size() < len) {388                read_buf.resize(len);389            }390            ggml_backend_tensor_get(result, read_buf.data(), 0, len);391            fout.write((char* )read_buf.data(), len);392            zeros(fout, GGML_PAD(len, GGUF_DEFAULT_ALIGNMENT) - len);393        }394 395        ggml_free(ctx);396        ggml_backend_buffer_free(buffer);397    }398 399    ~lora_merge_ctx() {400        ggml_gallocr_free(allocr);401        ggml_backend_free(backend);402        gguf_free(ctx_out);403        ggml_free(ctx_out_ggml);404    }405};406 407static void print_usage(int, char ** argv) {408    printf("\nexample usage:\n");409    printf("\n  %s -m base-model.gguf --lora lora-file.gguf -o merged-model-f16.gguf\n", argv[0]);410    printf("\nNOTE: output model is F16\n");411    printf("\n");412}413 414int main(int argc, char ** argv) {415    std::setlocale(LC_NUMERIC, "C");416 417    common_params params;418 419    params.out_file = "ggml-lora-merged-f16.gguf";420 421    common_init();422 423    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_EXPORT_LORA, print_usage)) {424        return 1;425    }426 427    g_verbose = (params.verbosity > 1);428    try {429        lora_merge_ctx ctx(params.model.path, params.lora_adapters, params.out_file, params.cpuparams.n_threads);430        ctx.run_merge();431    } catch (const std::exception & err) {432        fprintf(stderr, "%s\n", err.what());433        exit(EXIT_FAILURE);434    }435 436    printf("done, output file is %s\n", params.out_file.c_str());437 438    return 0;439}440