feat(backend): add some initial decoding steps

This commit is contained in:
Morgan Funtowicz 2024-10-23 00:09:10 +02:00
parent 098c66920d
commit 45d5a6a8c5
3 changed files with 147 additions and 25 deletions

View File

@ -2,20 +2,23 @@
// Created by Morgan Funtowicz on 9/28/2024.
//
#include <expected>
#include <filesystem>
#include <span>
#include <ggml.h>
#include <llama.h>
#include <fmt/chrono.h>
#include <fmt/format.h>
#include <fmt/std.h>
#include <spdlog/spdlog.h>
#include "backend.hpp"
namespace huggingface::tgi::backends::llama {
std::expected<std::unique_ptr<TgiLlamaCppBackend>, TgiLlamaCppBackendError>
CreateLlamaCppBackend(const std::filesystem::path& modelPath) {
SPDLOG_INFO(FMT_STRING("Loading model from {}"), modelPath);
SPDLOG_DEBUG(FMT_STRING("Loading model from {}"), modelPath);
llama_backend_init();
llama_numa_init(ggml_numa_strategy::GGML_NUMA_STRATEGY_NUMACTL);
@ -28,39 +31,109 @@ namespace huggingface::tgi::backends::llama {
auto* model = llama_load_model_from_file(modelPath.c_str(), params);
auto* context = llama_new_context_with_model(model, {
.n_batch = 1,
.n_threads = 16,
.attention_type = llama_attention_type::LLAMA_ATTENTION_TYPE_CAUSAL,
.flash_attn = true,
.flash_attn = false,
});
return std::make_unique<huggingface::tgi::backends::llama::TgiLlamaCppBackend>(model, context);
}
huggingface::tgi::backends::llama::TgiLlamaCppBackend::TgiLlamaCppBackend(llama_model *const model, llama_context *const ctx)
: model(model), ctx(ctx), batch() {
char modelName[128];
llama_model_meta_val_str(model, "general.name", modelName, sizeof(modelName));
: model(model), ctx(ctx) {
#ifndef NDEBUG
char modelName[256];
llama_model_meta_val_str(llama_get_model(ctx), "general.name", modelName, sizeof(modelName));
SPDLOG_DEBUG(FMT_STRING("Created llama.cpp backend for model: '{}'"), std::string_view(modelName));
#endif
}
huggingface::tgi::backends::llama::TgiLlamaCppBackend::~TgiLlamaCppBackend() {
if (model) {
SPDLOG_DEBUG("Freeing llama.cpp model");
llama_free_model(model);
}
if (ctx) {
SPDLOG_DEBUG("Freeing llama.cpp context");
llama_free(ctx);
}
if(model) {
SPDLOG_DEBUG("Freeing llama.cpp model");
llama_free_model(model);
}
}
void huggingface::tgi::backends::llama::TgiLlamaCppBackend::schedule() {
std::vector<llama_token> tokens;
std::vector<TgiLlamaCppBackend::TokenId> TgiLlamaCppBackend::Tokenize(const std::string &text) const {
std::vector<TgiLlamaCppBackend::TokenId> tokens(llama_n_seq_max(ctx));
if(auto nTokens = llama_tokenize(model, text.c_str(), text.length(), tokens.data(), tokens.capacity(), true, true); nTokens < 0){
tokens.resize(-nTokens);
llama_tokenize(model, text.c_str(), text.length(), tokens.data(), tokens.capacity(), true, true);
} else {
tokens.resize(nTokens);
}
SPDLOG_DEBUG(FMT_STRING("Tokenized input with {:d} tokens"), tokens.size());
return tokens;
}
namespace impl {
class LlamaCppBackendImpl {
std::unique_ptr<llama_sampler *> TgiLlamaCppBackend::GetSamplerFromArgs(
const uint32_t topK, const float_t topP, const float_t frequencyPenalty, const float_t repetitionPenalty, const uint64_t seed) {
auto *sampler = llama_sampler_chain_init({.no_perf = false});
};
// Penalties
llama_sampler_chain_add(sampler, llama_sampler_init_penalties(
llama_n_vocab(model),
llama_token_eos(model),
llama_token_nl (model),
0.0f,
repetitionPenalty,
frequencyPenalty,
0.0f,
false,
false
));
llama_sampler_chain_add(sampler, llama_sampler_init_top_k(static_cast<int32_t>(topK)));
if(0 < topP && topP < 1) {
llama_sampler_chain_add(sampler, llama_sampler_init_top_p(topP, 1));
}
llama_sampler_chain_add(sampler, llama_sampler_init_dist(seed));
return std::make_unique<llama_sampler*>(sampler);
}
std::vector<TgiLlamaCppBackend::TokenId> huggingface::tgi::backends::llama::TgiLlamaCppBackend::Generate(
std::span<const TokenId> tokens, const uint32_t topK, const float_t topP, const uint32_t maxNewTokens) {
SPDLOG_DEBUG(FMT_STRING("Received {:d} tokens to schedule"), tokens.size());
// Allocate generation result
std::vector<TgiLlamaCppBackend::TokenId> generated;
generated.reserve(llama_n_seq_max(ctx) - tokens.size());
// Retrieve decoding context
auto batch = llama_batch_get_one(const_cast<int32_t *>(tokens.data()), static_cast<int32_t>(tokens.size()));
auto sampler = GetSamplerFromArgs(topK, topP, 1.0, 1.0, 2014);
// Decode
for(auto [generating, nDecoded] = std::pair{true, 0uz}; generating && nDecoded < maxNewTokens; ++nDecoded) {
#ifndef NDEBUG
const auto start = std::chrono::steady_clock::now();
const auto status = llama_decode(ctx, batch);
const auto end = std::chrono::steady_clock::now();
const auto latency = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
SPDLOG_DEBUG(FMT_STRING("Successfully decoded {:d} token(s) in {}"), batch.n_tokens, latency);
#else
const auto status = llama_decode(ctx, batch);
#endif
if (status == LLAMA_SUCCESS) {
// Sample the new token
auto new_token_id = llama_sampler_sample(*sampler, ctx, -1);
generated.emplace_back(new_token_id);
generating = !llama_token_is_eog(model, new_token_id);
// Next iteration
batch = llama_batch_get_one(&new_token_id, 1);
}
}
generated.shrink_to_fit();
return generated;
}
}

View File

@ -4,28 +4,61 @@
#ifndef TGI_LLAMA_CPP_BACKEND_BACKEND_HPP
#define TGI_LLAMA_CPP_BACKEND_BACKEND_HPP
#include <cmath>
#include <expected>
#include <filesystem>
#include <memory>
#include <llama.h>
namespace huggingface::tgi::backends::llama {
// const char* TGI_BACKEND_LLAMA_CPP_NAME = "llama.cpp";
#define LLAMA_SUCCESS 0
namespace huggingface::tgi::backends::llama {
enum TgiLlamaCppBackendError {
MODEL_FILE_DOESNT_EXIST = 1
};
class TgiLlamaCppBackend {
using TokenId = int32_t;
private:
llama_model* model;
llama_context* ctx;
llama_batch batch;
/**
*
* @param topK
* @param topP
* @return
*/
std::unique_ptr<llama_sampler *> GetSamplerFromArgs(
uint32_t topK, float_t topP, float_t frequencyPenalty, float_t repetitionPenalty, uint64_t seed);
public:
TgiLlamaCppBackend(llama_model *model, llama_context *ctx);
~TgiLlamaCppBackend();
void schedule();
/**
*
* @param text
* @return
*/
[[nodiscard]] std::vector<TgiLlamaCppBackend::TokenId> Tokenize(const std::string& text) const;
/**
*
* @param tokens
* @param topK
* @param topP
* @param maxNewTokens
* @return
*/
[[nodiscard]] std::vector<TgiLlamaCppBackend::TokenId> Generate(
std::span<const TokenId> tokens,
uint32_t topK,
float_t topP = 1.0f,
uint32_t maxNewTokens = std::numeric_limits<uint32_t>::max()
);
};
std::expected<std::unique_ptr<TgiLlamaCppBackend>, TgiLlamaCppBackendError>

View File

@ -3,21 +3,37 @@
//
#include <string_view>
#include <fmt/color.h>
#include <fmt/format.h>
#include <fmt/std.h>
#include <fmt/color.h>
#include <fmt/ranges.h>
#include <spdlog/spdlog.h>
#include "../csrc/backend.hpp"
int main(int argc, char** argv) {
if(argc < 2) {
if (argc < 2) {
fmt::print("No model folder provider");
return 1;
}
spdlog::set_level(spdlog::level::debug);
const auto prompt = "My name is Morgan";
const auto modelPath = absolute(std::filesystem::path(argv[1]));
if(auto backend = huggingface::tgi::backends::llama::CreateLlamaCppBackend(modelPath); backend.has_value())
fmt::print(fmt::emphasis::bold | fg(fmt::color::yellow), "Successfully initialized llama.cpp model from {}\n", modelPath);
}
if (auto maybeBackend = huggingface::tgi::backends::llama::CreateLlamaCppBackend(modelPath); maybeBackend.has_value()) {
// Retrieve the backend
const auto& backend = *maybeBackend;
// Generate
const auto promptTokens = backend->Tokenize(prompt);
const auto out = backend->Generate(promptTokens, 30, 1.0, 32);
fmt::print(FMT_STRING("Generated: {}"), out);
} else {
switch (maybeBackend.error()) {
case huggingface::tgi::backends::llama::TgiLlamaCppBackendError::MODEL_FILE_DOESNT_EXIST:
fmt::print(fmt::emphasis::bold | fg(fmt::color::red), "Specified file {} doesnt exist", modelPath);
return maybeBackend.error();
}
}
}