echodict/llama.cpp

1#include "arg.h"2#include "common.h"3#include "sampling.h"4#include "speculative.h"5#include "log.h"6#include "llama.h"7 8#include <clocale>9#include <cstdio>10#include <cstring>11#include <string>12#include <vector>13 14int main(int argc, char ** argv) {15    std::setlocale(LC_NUMERIC, "C");16 17    common_params params;18 19    common_init();20 21    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) {22        return 1;23    }24 25    if (params.n_predict < -1) {26        LOG_ERR("%s: --n-predict must be >= -1\n", __func__);27        return 1;28    }29 30    if (params.speculative.mparams_dft.path.empty()) {31        LOG_ERR("%s: --model-draft is required\n", __func__);32        return 1;33    }34 35    // init llama.cpp36    llama_backend_init();37    llama_numa_init(params.numa);38 39    llama_model * model_tgt = NULL;40 41    llama_context * ctx_tgt = NULL;42 43    // load the target model44    auto llama_init_tgt = common_init_from_params(params);45 46    model_tgt = llama_init_tgt->model();47    ctx_tgt   = llama_init_tgt->context();48 49    const llama_vocab * vocab = llama_model_get_vocab(model_tgt);50 51    // load the draft model52    llama_model_ptr model_dft;53 54    // TODO: simplify this logic55    {56        const auto & params_spec = params.speculative;57 58        auto params_dft = params;59 60        params_dft.n_parallel   = 1;61        params_dft.n_ctx        = params_spec.n_ctx;62        params_dft.n_batch      = llama_n_ctx_seq(ctx_tgt);63        params_dft.devices      = params_spec.devices;64        params_dft.model        = params_spec.mparams_dft;65        params_dft.n_gpu_layers = params_spec.n_gpu_layers;66 67        if (params_spec.cpuparams.n_threads > 0) {68            params_dft.cpuparams.n_threads       = params.speculative.cpuparams.n_threads;69            params_dft.cpuparams_batch.n_threads = params.speculative.cpuparams_batch.n_threads;70        }71 72        params_dft.tensor_buft_overrides = params.speculative.tensor_buft_overrides;73 74        auto mparams_dft = common_model_params_to_llama(params_dft);75 76        model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft));77        if (model_dft == nullptr) {78            LOG_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str());79            return 1;80        }81 82        params.speculative.model_dft = model_dft.get();83        params.speculative.cparams_dft = common_context_params_to_llama(params_dft);84    }85 86    // Tokenize the prompt87    std::vector<llama_token> inp;88    inp = common_tokenize(ctx_tgt, params.prompt, true, true);89 90    if (llama_n_ctx(ctx_tgt) < (uint32_t) inp.size()) {91        LOG_ERR("%s: the prompt exceeds the context size (%d tokens, ctx %d)\n", __func__, (int) inp.size(), llama_n_ctx(ctx_tgt));92 93        return 1;94    }95 96    if (llama_n_batch(ctx_tgt) < (uint32_t) inp.size()) {97        LOG_ERR("%s: the prompt exceeds the batch size (%d tokens, batch %d)\n", __func__, (int) inp.size(), llama_n_batch(ctx_tgt));98 99        return 1;100    }101 102    LOG("\n\n");103 104    for (auto id : inp) {105        LOG("%s", common_token_to_piece(ctx_tgt, id).c_str());106    }107 108    int n_predict = 0;109    int n_drafted = 0;110    int n_accept  = 0;111 112    // used to determine end of generation113    bool has_eos = false;114 115    // ================================================116    // everything until here is standard initialization117    // the relevant stuff for speculative decoding starts here118 119    const auto t_enc_start = ggml_time_us();120 121    // target model sampling context122    struct common_sampler * smpl = common_sampler_init(model_tgt, params.sampling);123 124    // eval the prompt125    llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1));126 127    // note: keep the last token separate!128    llama_token id_last = inp.back();129 130    // all tokens currently in the target context131    llama_tokens prompt_tgt(inp.begin(), inp.end() - 1);132    prompt_tgt.reserve(llama_n_ctx(ctx_tgt));133 134    int n_past = inp.size() - 1;135 136    // init the speculator137    const auto & params_spec = params.speculative;138 139    struct common_speculative * spec = common_speculative_init(params.speculative, ctx_tgt);140 141    common_speculative_begin(spec, prompt_tgt);142 143    llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1);144 145    const auto t_enc_end = ggml_time_us();146 147    const auto t_dec_start = ggml_time_us();148 149    while (true) {150        // optionally, generate draft tokens that can be appended to the target batch151        //152        // this is the most important part of the speculation. the more probable tokens that are provided here153        // the better the performance will be. in theory, this computation can be performed asynchronously and even154        // offloaded to a remote device. it doesn't even have to be based on an LLM. instead, it can provide tokens155        // from a cache or lookup tables.156        //157        llama_tokens draft = common_speculative_draft(spec, params_spec, prompt_tgt, id_last);158 159        //LOG_DBG("draft: %s\n", string_from(ctx_dft, draft).c_str());160 161        // always have a token to evaluate from before - id_last162        common_batch_clear(batch_tgt);163        common_batch_add  (batch_tgt, id_last, n_past++, { 0 }, true);164 165        // evaluate the target model on [id_last, draft0, draft1, ..., draftN-1]166        {167            // do not waste time on small drafts168            if (draft.size() < (size_t) params_spec.n_min) {169                draft.clear();170            }171 172            for (size_t i = 0; i < draft.size(); ++i) {173                common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true);174            }175 176            //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str());177 178            llama_decode(ctx_tgt, batch_tgt);179        }180 181        // sample from the full target batch and return the accepted tokens based on the target sampler182        //183        // for each token to be accepted, the sampler would have to sample that same token184        // in such cases, instead of decoding the sampled token as we normally do, we simply continue with the185        // available logits from the batch and sample the next token until we run out of logits or the sampler186        // disagrees with the draft187        //188        const auto ids = common_sampler_sample_and_accept_n(smpl, ctx_tgt, draft);189 190        //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str());191 192        GGML_ASSERT(ids.size() > 0); // there will always be at least one accepted token193 194        n_past    += ids.size() - 1;195        n_drafted += draft.size(); // note: we ignore the discarded small drafts196        n_accept  += ids.size() - 1;197        n_predict += ids.size();198 199        // process the accepted tokens and update contexts200        //201        // this is the standard token post-processing that we normally do202        // in this case, we do it for a group of accepted tokens at once203        //204        for (size_t i = 0; i < ids.size(); ++i) {205            prompt_tgt.push_back(id_last);206 207            id_last = ids[i];208 209            if (llama_vocab_is_eog(vocab, id_last)) {210                has_eos = true;211                break;212            }213 214            const std::string token_str = common_token_to_piece(ctx_tgt, id_last);215 216            if (params.use_color && i + 1 < ids.size()) {217                LOG("\u001b[%dm%s\u001b[37m", (36 - 0 % 6), token_str.c_str());218            } else {219                LOG("%s", token_str.c_str());220            }221        }222 223        LOG_DBG("accepted %d/%d draft tokens, the last target token is: (%d)\n", (int) ids.size() - 1, (int) draft.size(), id_last);224 225        {226            LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past);227 228            llama_memory_seq_rm(llama_get_memory(ctx_tgt), 0, n_past, -1);229        }230 231        if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) {232            break;233        }234    }235 236    auto t_dec_end = ggml_time_us();237 238    const int n_input = inp.size();239 240    LOG("\n\n");241 242    LOG_INF("encoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_input,   (t_enc_end - t_enc_start) / 1e6f, inp.size() / ((t_enc_end - t_enc_start) / 1e6f));243    LOG_INF("decoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_predict, (t_dec_end - t_dec_start) / 1e6f, n_predict  / ((t_dec_end - t_dec_start) / 1e6f));244 245    LOG_INF("\n");246    LOG_INF("n_draft   = %d\n", params_spec.n_max);247    LOG_INF("n_predict = %d\n", n_predict);248    LOG_INF("n_drafted = %d\n", n_drafted);249    LOG_INF("n_accept  = %d\n", n_accept);250    LOG_INF("accept    = %.3f%%\n", 100.0f * n_accept / n_drafted);251 252    LOG_INF("\n");253    LOG_INF("draft:\n\n");254 255    LOG_INF("\n");256    LOG_INF("target:\n\n");257    common_perf_print(ctx_tgt, smpl);258 259    llama_batch_free(batch_tgt);260 261    common_sampler_free(smpl);262    common_speculative_free(spec);263 264    llama_backend_free();265 266    LOG("\n\n");267 268    return 0;269}270