arranged files, updated kobold lite, modified makefile for extra link args on linux, started RWKV implementation
This commit is contained in:
Concedo
parent
9581171a9f
commit
763ad172c0
21 changed files with 13597 additions and 46 deletions
31
Makefile
31
Makefile
|
|
@ -12,6 +12,7 @@ endif
|
|||
|
||||
ifndef ARCH_LINUX
|
||||
ARCH_LINUX := $(shell grep "Arch Linux" /etc/os-release 2>/dev/null)
|
||||
ARCH_LIKE := $(shell grep "ID_LIKE=arch" /etc/os-release 2>/dev/null)
|
||||
endif
|
||||
|
||||
CCV := $(shell $(CC) --version | head -n 1)
|
||||
|
|
@ -52,10 +53,15 @@ CXXFLAGS += -pthread -s -Wno-multichar
|
|||
ifeq ($(UNAME_S),Linux)
|
||||
CFLAGS += -pthread
|
||||
CXXFLAGS += -pthread
|
||||
ifdef ARCH_LINUX
|
||||
LDFLAGS += -lcblas
|
||||
endif
|
||||
ifdef ARCH_LINUX
|
||||
LDFLAGS += -lcblas
|
||||
else
|
||||
ifdef ARCH_LIKE
|
||||
LDFLAGS += -lcblas
|
||||
endif
|
||||
endif
|
||||
endif
|
||||
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
CFLAGS += -pthread
|
||||
CXXFLAGS += -pthread
|
||||
|
|
@ -117,7 +123,7 @@ ifdef LLAMA_OPENBLAS
|
|||
endif
|
||||
ifdef LLAMA_CLBLAST
|
||||
CFLAGS += -DGGML_USE_CLBLAST -DGGML_USE_OPENBLAS
|
||||
LDFLAGS += -lclblast -lOpenCL
|
||||
LDFLAGS += -lclblast -lOpenCL -lopenblas
|
||||
endif
|
||||
ifdef LLAMA_GPROF
|
||||
CFLAGS += -pg
|
||||
|
|
@ -202,6 +208,9 @@ ggml_v1.o: otherarch/ggml_v1.c otherarch/ggml_v1.h
|
|||
ggml_v1_noavx2.o: otherarch/ggml_v1.c otherarch/ggml_v1.h
|
||||
$(CC) $(CFLAGS) $(BONUSCFLAGS1) -c $< -o $@
|
||||
|
||||
ggml_rwkv.o: otherarch/ggml_rwkv.c otherarch/ggml_rwkv.h
|
||||
$(CC) $(CFLAGS) $(BONUSCFLAGS1) $(BONUSCFLAGS2) -c $< -o $@
|
||||
|
||||
llama.o: llama.cpp llama.h llama_util.h
|
||||
$(CXX) $(CXXFLAGS) -c $< -o $@
|
||||
|
||||
|
|
@ -226,19 +235,19 @@ main: examples/main/main.cpp ggml.o llama.o common.o
|
|||
@echo '==== Run ./main -h for help. ===='
|
||||
@echo
|
||||
|
||||
koboldcpp.dll: ggml.o ggml_v1.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
koboldcpp.dll: ggml.o ggml_rwkv.o ggml_v1.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
$(CXX) $(CXXFLAGS) $^ -shared -o $@ $(LDFLAGS)
|
||||
|
||||
koboldcpp_openblas.dll: ggml_openblas.o ggml_v1.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
koboldcpp_openblas.dll: ggml_openblas.o ggml_rwkv.o ggml_v1.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
$(OPENBLAS_BUILD)
|
||||
|
||||
koboldcpp_noavx2.dll: ggml_noavx2.o ggml_v1_noavx2.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
koboldcpp_noavx2.dll: ggml_noavx2.o ggml_rwkv.o ggml_v1_noavx2.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
$(NOAVX2_BUILD)
|
||||
|
||||
koboldcpp_openblas_noavx2.dll: ggml_openblas_noavx2.o ggml_v1_noavx2.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
koboldcpp_openblas_noavx2.dll: ggml_openblas_noavx2.o ggml_rwkv.o ggml_v1_noavx2.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
$(OPENBLAS_NOAVX2_BUILD)
|
||||
|
||||
koboldcpp_clblast.dll: ggml_clblast.o ggml_v1.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
koboldcpp_clblast.dll: ggml_clblast.o ggml_rwkv.o ggml_v1.o expose.o common.o llama_adapter.o gpttype_adapter.o
|
||||
$(CLBLAST_BUILD)
|
||||
|
||||
quantize_llama: examples/quantize/quantize.cpp ggml.o llama.o
|
||||
|
|
@ -247,10 +256,10 @@ quantize_llama: examples/quantize/quantize.cpp ggml.o llama.o
|
|||
quantize-stats: examples/quantize-stats/quantize-stats.cpp ggml.o llama.o
|
||||
$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
|
||||
|
||||
quantize_gptj: ggml.o llama.o otherarch/gptj_quantize.cpp
|
||||
quantize_gptj: ggml.o llama.o otherarch/tools/gptj_quantize.cpp
|
||||
$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
|
||||
|
||||
quantize_gpt2: ggml.o llama.o otherarch/gpt2_quantize.cpp
|
||||
quantize_gpt2: ggml.o llama.o otherarch/tools/gpt2_quantize.cpp
|
||||
$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
|
||||
|
||||
perplexity: examples/perplexity/perplexity.cpp ggml.o llama.o common.o
|
||||
|
|
|
|||
42
expose.cpp
42
expose.cpp
|
|
@ -53,16 +53,23 @@ extern "C"
|
|||
ModelLoadResult lr = gpttype_load_model(inputs, file_format);
|
||||
if (lr == ModelLoadResult::RETRY_LOAD)
|
||||
{
|
||||
file_format = FileFormat::GPTJ_2;
|
||||
printf("\n---\nRetrying as GPT-J model: (ver %d)\nAttempting to Load...\n---\n", file_format);
|
||||
lr = gpttype_load_model(inputs, file_format);
|
||||
}
|
||||
if (lr == ModelLoadResult::RETRY_LOAD)
|
||||
{
|
||||
file_format = FileFormat::GPTJ_3;
|
||||
printf("\n---\nRetrying as GPT-J model: (ver %d)\nAttempting to Load...\n---\n", file_format);
|
||||
lr = gpttype_load_model(inputs, file_format);
|
||||
}
|
||||
if(file_format==FileFormat::GPTJ_1)
|
||||
{
|
||||
//if we tried 1 first, then try 3 and lastly 2
|
||||
//otherwise if we tried 3 first, then try 2
|
||||
file_format = FileFormat::GPTJ_3;
|
||||
printf("\n---\nRetrying as GPT-J model: (ver %d)\nAttempting to Load...\n---\n", file_format);
|
||||
lr = gpttype_load_model(inputs, file_format);
|
||||
}
|
||||
|
||||
//lastly try format 2
|
||||
if (lr == ModelLoadResult::RETRY_LOAD)
|
||||
{
|
||||
file_format = FileFormat::GPTJ_2;
|
||||
printf("\n---\nRetrying as GPT-J model: (ver %d)\nAttempting to Load...\n---\n", file_format);
|
||||
lr = gpttype_load_model(inputs, file_format);
|
||||
}
|
||||
}
|
||||
|
||||
if (lr == ModelLoadResult::FAIL || lr == ModelLoadResult::RETRY_LOAD)
|
||||
{
|
||||
|
|
@ -92,6 +99,19 @@ extern "C"
|
|||
return true;
|
||||
}
|
||||
}
|
||||
else if(file_format==FileFormat::RWKV_1)
|
||||
{
|
||||
printf("\n---\nIdentified as RWKV model: (ver %d)\nAttempting to Load...\n---\n", file_format);
|
||||
ModelLoadResult lr = gpttype_load_model(inputs, file_format);
|
||||
if (lr == ModelLoadResult::FAIL || lr == ModelLoadResult::RETRY_LOAD)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
else
|
||||
{
|
||||
return true;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
printf("\n---\nIdentified as LLAMA model: (ver %d)\nAttempting to Load...\n---\n", file_format);
|
||||
|
|
@ -102,7 +122,7 @@ extern "C"
|
|||
generation_outputs generate(const generation_inputs inputs, generation_outputs &output)
|
||||
{
|
||||
if (file_format == FileFormat::GPTJ_1 || file_format == FileFormat::GPTJ_2 || file_format==FileFormat::GPTJ_3
|
||||
|| file_format==FileFormat::GPT2_1 || file_format==FileFormat::GPT2_2 )
|
||||
|| file_format==FileFormat::GPT2_1 || file_format==FileFormat::GPT2_2 || file_format==FileFormat::RWKV_1)
|
||||
{
|
||||
return gpttype_generate(inputs, output);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@
|
|||
#include "otherarch/gptj_v2.cpp"
|
||||
#include "otherarch/gpt2_v1.cpp"
|
||||
#include "otherarch/gpt2_v2.cpp"
|
||||
#include "otherarch/rwkv.cpp"
|
||||
|
||||
//return val: 0=fail, 1=(original ggml, alpaca), 2=(ggmf), 3=(ggjt)
|
||||
static FileFormat file_format = FileFormat::BADFORMAT;
|
||||
|
|
@ -25,6 +26,7 @@ static gptj_model_v1 model_v1;
|
|||
static gptj_model model_v2;
|
||||
static gpt2_v1_model model_gpt2_v1;
|
||||
static gpt2_model model_gpt2_v2;
|
||||
static rwkv_context * rwkv_context_v1;
|
||||
static gpt_params params;
|
||||
static int n_past = 0;
|
||||
static int n_threads = 4;
|
||||
|
|
@ -59,7 +61,45 @@ ModelLoadResult gpttype_load_model(const load_model_inputs inputs, FileFormat in
|
|||
params.n_ctx = inputs.max_context_length;
|
||||
model_v1.hparams.n_ctx = model_v2.hparams.n_ctx = model_gpt2_v1.hparams.n_ctx = model_gpt2_v2.hparams.n_ctx = params.n_ctx;
|
||||
|
||||
if (file_format == FileFormat::GPT2_1)
|
||||
if (file_format == FileFormat::RWKV_1)
|
||||
{
|
||||
rwkv_context_v1 = rwkv_init_from_file(modelname.c_str(), n_threads);
|
||||
|
||||
//setup buffers for rwkv state
|
||||
auto padding = 512u;
|
||||
auto statebufsiz = rwkv_get_state_buffer_element_count(rwkv_context_v1) * sizeof(float) + padding;
|
||||
auto logitbufsiz = rwkv_get_logits_buffer_element_count(rwkv_context_v1) * sizeof(float) + padding;
|
||||
|
||||
printf("\nRWKV Init: State Buffer:%u, Logit Buffer:%u\n", statebufsiz, logitbufsiz);
|
||||
rwkv_context_v1->state_out = (float *)malloc(statebufsiz);
|
||||
rwkv_context_v1->logits_out = (float *)malloc(logitbufsiz);
|
||||
rwkv_context_v1->state_in = nullptr;
|
||||
n_batch = 1;
|
||||
|
||||
std::string word;
|
||||
for (int i = 0; i < 20; i++) {
|
||||
uint32_t len;
|
||||
word = ('a'+i);
|
||||
vocab.token_to_id[word] = i;
|
||||
vocab.id_to_token[i] = word;
|
||||
}
|
||||
|
||||
int vocabsiz = vocab.token_to_id.size();
|
||||
bool testeval = rwkv_eval(rwkv_context_v1, 0, rwkv_context_v1->state_in, rwkv_context_v1->state_out, rwkv_context_v1->logits_out);
|
||||
if(!testeval)
|
||||
{
|
||||
printf("\nError: RWKV Init Eval Failed!\n");
|
||||
}
|
||||
logits.resize(vocabsiz);
|
||||
memcpy(logits.data(), rwkv_context_v1->logits_out, sizeof(float)*vocabsiz);
|
||||
|
||||
if (rwkv_context_v1 == NULL)
|
||||
{
|
||||
return ModelLoadResult::FAIL;
|
||||
}
|
||||
return ModelLoadResult::SUCCESS;
|
||||
}
|
||||
else if (file_format == FileFormat::GPT2_1)
|
||||
{
|
||||
ModelLoadResult res = legacy_gpt2_model_load(params.model, model_gpt2_v1, vocab, file_format);
|
||||
if(res==ModelLoadResult::FAIL)
|
||||
|
|
@ -209,7 +249,10 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
|
|||
std::fill(last_n_tokens.begin(), last_n_tokens.end(), 0);
|
||||
n_past = 0;
|
||||
|
||||
ContextFastForward(current_context_tokens, embd_inp, n_past, last_n_tokens, nctx, smartcontext, useSmartContext);
|
||||
if(file_format!=FileFormat::RWKV_1)
|
||||
{
|
||||
ContextFastForward(current_context_tokens, embd_inp, n_past, last_n_tokens, nctx, smartcontext, useSmartContext);
|
||||
}
|
||||
|
||||
//if using BLAS and prompt is big enough, switch to single thread and use a huge batch
|
||||
bool approved_format = (file_format!=FileFormat::GPT2_1 && file_format!=FileFormat::GPTJ_1 && file_format!=FileFormat::GPTJ_2);
|
||||
|
|
@ -228,6 +271,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
|
|||
current_context_tokens.resize(n_past);
|
||||
|
||||
int remaining_tokens = params.n_predict;
|
||||
int stopper_unused_tokens = 0;
|
||||
int input_consumed = 0;
|
||||
std::mt19937 rng(params.seed);
|
||||
std::string concat_output = "";
|
||||
|
|
@ -254,12 +298,17 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
|
|||
{
|
||||
n_vocab = model_gpt2_v2.hparams.n_vocab;
|
||||
}
|
||||
else if(file_format == FileFormat::RWKV_1)
|
||||
{
|
||||
n_vocab = vocab.id_to_token.size(); //handled seperately
|
||||
}
|
||||
else
|
||||
{
|
||||
printf("Bad format!");
|
||||
}
|
||||
|
||||
printf("\n");
|
||||
|
||||
while (remaining_tokens > 0)
|
||||
{
|
||||
gpt_vocab::id id = 0;
|
||||
|
|
@ -278,9 +327,12 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
|
|||
}
|
||||
|
||||
bool evalres = false;
|
||||
|
||||
//print_tok_vec(logits);
|
||||
if(file_format==FileFormat::GPT2_1)
|
||||
|
||||
if(file_format==FileFormat::RWKV_1)
|
||||
{
|
||||
evalres = rwkv_eval(rwkv_context_v1, embd[0], rwkv_context_v1->state_in, rwkv_context_v1->state_out, rwkv_context_v1->logits_out);
|
||||
}
|
||||
else if(file_format==FileFormat::GPT2_1)
|
||||
{
|
||||
evalres = legacy_gpt2_eval(model_gpt2_v1, params.n_threads, n_past, embd, logits, mem_per_token, file_format);
|
||||
}
|
||||
|
|
@ -326,14 +378,14 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
|
|||
}
|
||||
|
||||
{
|
||||
// set the logit of the eos token (2) to zero to avoid sampling it
|
||||
logits[50256] = (logits[50256]<0?logits[50256]:0);
|
||||
|
||||
// set the logit of the eos token (2) to zero to avoid sampling it
|
||||
if(logits.size()>50256)
|
||||
{
|
||||
logits[50256] = (logits[50256]<0?logits[50256]:0);
|
||||
}
|
||||
//gpt2 uses negative logits, so we cant zero it
|
||||
|
||||
id = gptj_sample_top_p_top_k(vocab, logits.data() + (logits.size() - n_vocab), last_n_tokens, repeat_penalty, top_k, top_p, temp, rng);
|
||||
|
||||
|
||||
last_n_tokens.erase(last_n_tokens.begin());
|
||||
last_n_tokens.push_back(id);
|
||||
current_context_tokens.push_back(id);
|
||||
|
|
@ -352,6 +404,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
|
|||
{
|
||||
if (concat_output.find(matched) != std::string::npos)
|
||||
{
|
||||
stopper_unused_tokens = remaining_tokens;
|
||||
remaining_tokens = 0;
|
||||
printf("\n(Stop sequence triggered: <%s>)",matched.c_str());
|
||||
break;
|
||||
|
|
@ -378,7 +431,8 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
|
|||
}
|
||||
time2 = timer_check();
|
||||
float pt1 = (time1*1000.0/(embd_inp_size==0?1:embd_inp_size));
|
||||
float pt2 = (time2*1000.0/(params.n_predict==0?1:params.n_predict));
|
||||
int realnpredict = params.n_predict-stopper_unused_tokens;
|
||||
float pt2 = (time2*1000.0/(realnpredict==0?1:realnpredict));
|
||||
printf("\nTime Taken - Processing:%.1fs (%.0fms/T), Generation:%.1fs (%.0fms/T), Total:%.1fs", time1, pt1, time2, pt2, (time1 + time2));
|
||||
fflush(stdout);
|
||||
output.status = 1;
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -435,7 +435,7 @@ def main(args):
|
|||
RunServerMultiThreaded(args.host, args.port, embedded_kailite)
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("Welcome to KoboldCpp - Version 1.9") # just update version manually
|
||||
print("Welcome to KoboldCpp - Version 1.10") # just update version manually
|
||||
parser = argparse.ArgumentParser(description='Kobold llama.cpp server')
|
||||
modelgroup = parser.add_mutually_exclusive_group() #we want to be backwards compatible with the unnamed positional args
|
||||
modelgroup.add_argument("--model", help="Model file to load", nargs="?")
|
||||
|
|
|
|||
|
|
@ -57,11 +57,9 @@ bool llama_load_model(const load_model_inputs inputs, FileFormat in_file_format)
|
|||
ctx_params.use_mlock = false;
|
||||
|
||||
file_format = in_file_format;
|
||||
|
||||
|
||||
ctx = llama_init_from_file(modelname.c_str(), ctx_params);
|
||||
|
||||
|
||||
if (ctx == NULL)
|
||||
{
|
||||
fprintf(stderr, "%s: error: failed to load model '%s'\n", __func__, modelname.c_str());
|
||||
|
|
@ -162,6 +160,7 @@ generation_outputs llama_generate(const generation_inputs inputs, generation_out
|
|||
current_context_tokens.resize(n_past);
|
||||
|
||||
int remaining_tokens = params.n_predict;
|
||||
int stopper_unused_tokens = 0;
|
||||
int input_consumed = 0;
|
||||
std::mt19937 rng(params.seed);
|
||||
std::string concat_output = "";
|
||||
|
|
@ -245,6 +244,7 @@ generation_outputs llama_generate(const generation_inputs inputs, generation_out
|
|||
{
|
||||
if (concat_output.find(matched) != std::string::npos)
|
||||
{
|
||||
stopper_unused_tokens = remaining_tokens;
|
||||
remaining_tokens = 0;
|
||||
printf("\n(Stop sequence triggered: <%s>)",matched.c_str());
|
||||
break;
|
||||
|
|
@ -270,7 +270,8 @@ generation_outputs llama_generate(const generation_inputs inputs, generation_out
|
|||
}
|
||||
time2 = timer_check();
|
||||
float pt1 = (time1*1000.0/(embd_inp_size==0?1:embd_inp_size));
|
||||
float pt2 = (time2*1000.0/(params.n_predict==0?1:params.n_predict));
|
||||
int realnpredict = params.n_predict-stopper_unused_tokens;
|
||||
float pt2 = (time2*1000.0/(realnpredict==0?1:realnpredict));
|
||||
printf("\nTime Taken - Processing:%.1fs (%.0fms/T), Generation:%.1fs (%.0fms/T), Total:%.1fs", time1, pt1, time2, pt2, (time1 + time2));
|
||||
fflush(stdout);
|
||||
output.status = 1;
|
||||
|
|
|
|||
|
|
@ -132,6 +132,12 @@ void print_tok_vec(std::vector<float> &embd)
|
|||
else if(magic == 0x67676d66) //v2 format ggmf
|
||||
{
|
||||
fileformat = FileFormat::GGHF;
|
||||
uint32_t temp;
|
||||
fin.read((char *)&temp, sizeof(temp)); //file version
|
||||
if(temp==100)
|
||||
{
|
||||
fileformat = FileFormat::RWKV_1;
|
||||
}
|
||||
}
|
||||
else if(magic == 0x67676a74) //v3 format ggjt
|
||||
{
|
||||
|
|
|
|||
|
|
@ -25,7 +25,9 @@ enum FileFormat
|
|||
GPTJ_3=102, //uses new ggml lib
|
||||
|
||||
GPT2_1=200,
|
||||
GPT2_2=201
|
||||
GPT2_2=201,
|
||||
|
||||
RWKV_1=300,
|
||||
};
|
||||
|
||||
enum ModelLoadResult
|
||||
|
|
|
|||
11588
otherarch/ggml_rwkv.c
Normal file
11588
otherarch/ggml_rwkv.c
Normal file
File diff suppressed because it is too large
Load diff
645
otherarch/ggml_rwkv.h
Normal file
645
otherarch/ggml_rwkv.h
Normal file
|
|
@ -0,0 +1,645 @@
|
|||
#pragma once
|
||||
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stddef.h>
|
||||
#include <stdbool.h>
|
||||
|
||||
#define GGML_RWKV_MAX_DIMS 4
|
||||
#define GGML_RWKV_MAX_NODES 4096
|
||||
#define GGML_RWKV_MAX_PARAMS 16
|
||||
#define GGML_RWKV_MAX_CONTEXTS 64
|
||||
#define GGML_RWKV_MAX_OPT 4
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
// we use the built-in 16-bit float type
|
||||
typedef __fp16 ggml_rwkv_fp16_t;
|
||||
#else
|
||||
typedef uint16_t ggml_rwkv_fp16_t;
|
||||
#endif
|
||||
|
||||
// convert FP16 <-> FP32
|
||||
float ggml_rwkv_fp16_to_fp32(ggml_rwkv_fp16_t x);
|
||||
ggml_rwkv_fp16_t ggml_rwkv_fp32_to_fp16(float x);
|
||||
|
||||
struct ggml_rwkv_object;
|
||||
struct ggml_rwkv_context;
|
||||
|
||||
enum ggml_rwkv_type {
|
||||
GGML_RWKV_TYPE_Q4_0,
|
||||
// Stores min and delta per block, does quantized matmul.
|
||||
GGML_RWKV_TYPE_Q4_1,
|
||||
// Same as Q4_1, but stores outliers separately, and matmul is done in FP32.
|
||||
// An outlier is the single absmax element in the quantized block.
|
||||
GGML_RWKV_TYPE_Q4_1_O,
|
||||
GGML_RWKV_TYPE_I8,
|
||||
GGML_RWKV_TYPE_I16,
|
||||
GGML_RWKV_TYPE_I32,
|
||||
GGML_RWKV_TYPE_F16,
|
||||
GGML_RWKV_TYPE_F32,
|
||||
GGML_RWKV_TYPE_COUNT,
|
||||
};
|
||||
|
||||
// available tensor operations:
|
||||
enum ggml_rwkv_op {
|
||||
GGML_RWKV_OP_NONE = 0,
|
||||
|
||||
GGML_RWKV_OP_DUP,
|
||||
GGML_RWKV_OP_ADD,
|
||||
GGML_RWKV_OP_SUB,
|
||||
GGML_RWKV_OP_MUL,
|
||||
GGML_RWKV_OP_DIV,
|
||||
GGML_RWKV_OP_SQR,
|
||||
GGML_RWKV_OP_SQRT,
|
||||
GGML_RWKV_OP_SUM,
|
||||
GGML_RWKV_OP_MEAN,
|
||||
GGML_RWKV_OP_REPEAT,
|
||||
GGML_RWKV_OP_ABS,
|
||||
GGML_RWKV_OP_SGN,
|
||||
GGML_RWKV_OP_NEG,
|
||||
// Element-wise exponential function `e^x`.
|
||||
// Same as `torch.exp(x)` from PyTorch.
|
||||
GGML_RWKV_OP_EXP,
|
||||
// Element-wise `1 - x`.
|
||||
GGML_RWKV_OP_1_MINUS_X,
|
||||
|
||||
// Element-wise maximum of 2 values. Argument shapes must match.
|
||||
// Same as `torch.maximum(x)` from PyTorch.
|
||||
GGML_RWKV_OP_MAX,
|
||||
|
||||
GGML_RWKV_OP_STEP,
|
||||
GGML_RWKV_OP_RELU,
|
||||
GGML_RWKV_OP_GELU,
|
||||
// Element-wise sigmoid activation `1 / (1 + e^-x)`, also called logistic function.
|
||||
// Same as `torch.sigmoid(x)` from PyTorch.
|
||||
GGML_RWKV_OP_SIGMOID,
|
||||
GGML_RWKV_OP_SILU,
|
||||
GGML_RWKV_OP_NORM, // normalize
|
||||
GGML_RWKV_OP_RMS_NORM,
|
||||
|
||||
GGML_RWKV_OP_MUL_MAT,
|
||||
|
||||
GGML_RWKV_OP_SCALE,
|
||||
GGML_RWKV_OP_CPY,
|
||||
GGML_RWKV_OP_RESHAPE,
|
||||
GGML_RWKV_OP_VIEW,
|
||||
GGML_RWKV_OP_PERMUTE,
|
||||
GGML_RWKV_OP_TRANSPOSE,
|
||||
GGML_RWKV_OP_GET_ROWS,
|
||||
GGML_RWKV_OP_DIAG_MASK_INF,
|
||||
GGML_RWKV_OP_SOFT_MAX,
|
||||
GGML_RWKV_OP_ROPE,
|
||||
GGML_RWKV_OP_CONV_1D_1S,
|
||||
GGML_RWKV_OP_CONV_1D_2S,
|
||||
|
||||
GGML_RWKV_OP_FLASH_ATTN,
|
||||
GGML_RWKV_OP_FLASH_FF,
|
||||
|
||||
GGML_RWKV_OP_COUNT,
|
||||
};
|
||||
|
||||
// n-dimensional tensor
|
||||
struct ggml_rwkv_tensor {
|
||||
enum ggml_rwkv_type type;
|
||||
|
||||
int n_dims;
|
||||
int ne[GGML_RWKV_MAX_DIMS]; // number of elements
|
||||
size_t nb[GGML_RWKV_MAX_DIMS]; // stride in bytes:
|
||||
// nb[0] = sizeof(type)
|
||||
// nb[1] = nb[0] * ne[0] + padding
|
||||
// nb[i] = nb[i-1] * ne[i-1]
|
||||
|
||||
// compute data
|
||||
enum ggml_rwkv_op op;
|
||||
|
||||
bool is_param;
|
||||
|
||||
struct ggml_rwkv_tensor * grad;
|
||||
struct ggml_rwkv_tensor * src0;
|
||||
struct ggml_rwkv_tensor * src1;
|
||||
struct ggml_rwkv_tensor * opt[GGML_RWKV_MAX_OPT];
|
||||
|
||||
// thread scheduling
|
||||
int n_tasks;
|
||||
|
||||
// performance
|
||||
int perf_runs;
|
||||
int64_t perf_cycles;
|
||||
int64_t perf_time_us;
|
||||
|
||||
void * data;
|
||||
char padding[8];
|
||||
};
|
||||
|
||||
// computation graph
|
||||
struct ggml_rwkv_cgraph {
|
||||
int n_nodes;
|
||||
int n_leafs;
|
||||
int n_threads;
|
||||
|
||||
size_t work_size;
|
||||
struct ggml_rwkv_tensor * work;
|
||||
|
||||
struct ggml_rwkv_tensor * nodes[GGML_RWKV_MAX_NODES];
|
||||
struct ggml_rwkv_tensor * grads[GGML_RWKV_MAX_NODES];
|
||||
struct ggml_rwkv_tensor * leafs[GGML_RWKV_MAX_NODES];
|
||||
|
||||
// performance
|
||||
int perf_runs;
|
||||
int64_t perf_cycles;
|
||||
int64_t perf_time_us;
|
||||
};
|
||||
|
||||
// scratch buffer
|
||||
struct ggml_rwkv_scratch {
|
||||
size_t offs;
|
||||
size_t size;
|
||||
void * data;
|
||||
};
|
||||
|
||||
struct ggml_rwkv_init_params {
|
||||
// memory pool
|
||||
size_t mem_size; // bytes
|
||||
void * mem_buffer; // if NULL, memory will be allocated internally
|
||||
};
|
||||
|
||||
void ggml_rwkv_time_init(void); // call this once at the beginning of the program
|
||||
int64_t ggml_rwkv_time_ms(void);
|
||||
int64_t ggml_rwkv_time_us(void);
|
||||
int64_t ggml_rwkv_cycles(void);
|
||||
int64_t ggml_rwkv_cycles_per_ms(void);
|
||||
|
||||
void ggml_rwkv_print_object (const struct ggml_rwkv_object * obj);
|
||||
void ggml_rwkv_print_objects(const struct ggml_rwkv_context * ctx);
|
||||
|
||||
int ggml_rwkv_nelements(const struct ggml_rwkv_tensor * tensor);
|
||||
size_t ggml_rwkv_nbytes (const struct ggml_rwkv_tensor * tensor);
|
||||
|
||||
int ggml_rwkv_blck_size (enum ggml_rwkv_type type);
|
||||
size_t ggml_rwkv_type_size (enum ggml_rwkv_type type); // size in bytes for all elements in a block
|
||||
float ggml_rwkv_type_sizef(enum ggml_rwkv_type type); // ggml_rwkv_type_size()/ggml_rwkv_blck_size() as float
|
||||
|
||||
size_t ggml_rwkv_element_size(const struct ggml_rwkv_tensor * tensor);
|
||||
|
||||
struct ggml_rwkv_context * ggml_rwkv_init(struct ggml_rwkv_init_params params);
|
||||
void ggml_rwkv_free(struct ggml_rwkv_context * ctx);
|
||||
|
||||
size_t ggml_rwkv_used_mem(const struct ggml_rwkv_context * ctx);
|
||||
|
||||
size_t ggml_rwkv_set_scratch(struct ggml_rwkv_context * ctx, struct ggml_rwkv_scratch scratch);
|
||||
|
||||
bool ggml_rwkv_mlock_supported(void);
|
||||
bool ggml_rwkv_mlock(struct ggml_rwkv_context * ctx, char ** err_p);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_new_tensor(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
enum ggml_rwkv_type type,
|
||||
int n_dims,
|
||||
const int *ne);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_new_tensor_1d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
enum ggml_rwkv_type type,
|
||||
int ne0);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_new_tensor_2d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
enum ggml_rwkv_type type,
|
||||
int ne0,
|
||||
int ne1);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_new_tensor_3d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
enum ggml_rwkv_type type,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int ne2);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_new_tensor_4d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
enum ggml_rwkv_type type,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int ne2,
|
||||
int ne3);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_new_i32(struct ggml_rwkv_context * ctx, int32_t value);
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_new_f32(struct ggml_rwkv_context * ctx, float value);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_dup_tensor (struct ggml_rwkv_context * ctx, const struct ggml_rwkv_tensor * src);
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_view_tensor(struct ggml_rwkv_context * ctx, const struct ggml_rwkv_tensor * src);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_set_zero(struct ggml_rwkv_tensor * tensor);
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_set_i32 (struct ggml_rwkv_tensor * tensor, int32_t value);
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_set_f32 (struct ggml_rwkv_tensor * tensor, float value);
|
||||
|
||||
int32_t ggml_rwkv_get_i32_1d(const struct ggml_rwkv_tensor * tensor, int i);
|
||||
void ggml_rwkv_set_i32_1d(const struct ggml_rwkv_tensor * tensor, int i, int32_t value);
|
||||
|
||||
float ggml_rwkv_get_f32_1d(const struct ggml_rwkv_tensor * tensor, int i);
|
||||
void ggml_rwkv_set_f32_1d(const struct ggml_rwkv_tensor * tensor, int i, float value);
|
||||
|
||||
void * ggml_rwkv_get_data (const struct ggml_rwkv_tensor * tensor);
|
||||
float * ggml_rwkv_get_data_f32(const struct ggml_rwkv_tensor * tensor);
|
||||
|
||||
//
|
||||
// operations on tensors with backpropagation
|
||||
//
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_dup(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_add(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_sub(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_mul(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_div(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_sqr(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_sqrt(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
// return scalar
|
||||
// TODO: compute sum along rows
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_sum(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
// mean along rows
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_mean(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
// if a is the same shape as b, and a is not parameter, return a
|
||||
// otherwise, return a new tensor: repeat(a) to fit in b
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_repeat(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_abs(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_sgn(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_neg(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_exp(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_1_minus_x(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_max(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_step(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_relu(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
// TODO: double-check this computation is correct
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_gelu(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_sigmoid(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_silu(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
// normalize along rows
|
||||
// TODO: eps is hardcoded to 1e-5 for now
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_norm(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_rms_norm(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
// A: m rows, n columns
|
||||
// B: p rows, n columns (i.e. we transpose it internally)
|
||||
// result is m columns, p rows
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_mul_mat(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
//
|
||||
// operations on tensors without backpropagation
|
||||
//
|
||||
|
||||
// in-place, returns view(a)
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_scale(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
// a -> b, return view(b)
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_cpy(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
// return view(a), b specifies the new shape
|
||||
// TODO: when we start computing gradient, make a copy instead of view
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_reshape(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
// return view(a)
|
||||
// TODO: when we start computing gradient, make a copy instead of view
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_reshape_2d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
int ne0,
|
||||
int ne1);
|
||||
|
||||
// return view(a)
|
||||
// TODO: when we start computing gradient, make a copy instead of view
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_reshape_3d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int ne2);
|
||||
|
||||
// offset in bytes
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_view_1d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
int ne0,
|
||||
size_t offset);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_view_2d(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
int ne0,
|
||||
int ne1,
|
||||
size_t nb1, // row stride in bytes
|
||||
size_t offset);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_permute(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
int axis0,
|
||||
int axis1,
|
||||
int axis2,
|
||||
int axis3);
|
||||
|
||||
// alias for ggml_rwkv_permute(ctx, a, 1, 0, 2, 3)
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_transpose(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_get_rows(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
// set elements above the diagonal to -INF
|
||||
// in-place, returns view(a)
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_diag_mask_inf(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
int n_past);
|
||||
|
||||
// in-place, returns view(a)
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_soft_max(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a);
|
||||
|
||||
// rotary position embedding
|
||||
// in-place, returns view(a)
|
||||
// if mode == 1, skip n_past elements
|
||||
// TODO: avoid creating a new tensor every time
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_rope(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
int n_past,
|
||||
int n_dims,
|
||||
int mode);
|
||||
|
||||
// padding = 1
|
||||
// TODO: we don't support extra parameters for now
|
||||
// that's why we are hard-coding the stride, padding, and dilation
|
||||
// not great ..
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_conv_1d_1s(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_conv_1d_2s(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_flash_attn(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * q,
|
||||
struct ggml_rwkv_tensor * k,
|
||||
struct ggml_rwkv_tensor * v,
|
||||
bool masked);
|
||||
|
||||
struct ggml_rwkv_tensor * ggml_rwkv_flash_ff(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * a,
|
||||
struct ggml_rwkv_tensor * b0,
|
||||
struct ggml_rwkv_tensor * b1,
|
||||
struct ggml_rwkv_tensor * c0,
|
||||
struct ggml_rwkv_tensor * c1);
|
||||
|
||||
//
|
||||
// automatic differentiation
|
||||
//
|
||||
|
||||
void ggml_rwkv_set_param(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_tensor * tensor);
|
||||
|
||||
void ggml_rwkv_build_forward_expand(struct ggml_rwkv_cgraph * cgraph, struct ggml_rwkv_tensor * tensor);
|
||||
|
||||
struct ggml_rwkv_cgraph ggml_rwkv_build_forward (struct ggml_rwkv_tensor * tensor);
|
||||
struct ggml_rwkv_cgraph ggml_rwkv_build_backward(struct ggml_rwkv_context * ctx, struct ggml_rwkv_cgraph * gf, bool keep);
|
||||
|
||||
void ggml_rwkv_graph_compute(struct ggml_rwkv_context * ctx, struct ggml_rwkv_cgraph * cgraph);
|
||||
void ggml_rwkv_graph_reset (struct ggml_rwkv_cgraph * cgraph);
|
||||
|
||||
// print info and performance information for the graph
|
||||
void ggml_rwkv_graph_print(const struct ggml_rwkv_cgraph * cgraph);
|
||||
|
||||
// dump the graph into a file using the dot format
|
||||
void ggml_rwkv_graph_dump_dot(const struct ggml_rwkv_cgraph * gb, const struct ggml_rwkv_cgraph * gf, const char * filename);
|
||||
|
||||
//
|
||||
// optimization
|
||||
//
|
||||
|
||||
// optimization methods
|
||||
enum ggml_rwkv_opt_type {
|
||||
GGML_RWKV_OPT_ADAM,
|
||||
GGML_RWKV_OPT_LBFGS,
|
||||
};
|
||||
|
||||
// linesearch methods
|
||||
enum ggml_rwkv_linesearch {
|
||||
GGML_RWKV_LINESEARCH_DEFAULT = 1,
|
||||
|
||||
GGML_RWKV_LINESEARCH_BACKTRACKING_ARMIJO = 0,
|
||||
GGML_RWKV_LINESEARCH_BACKTRACKING_WOLFE = 1,
|
||||
GGML_RWKV_LINESEARCH_BACKTRACKING_STRONG_WOLFE = 2,
|
||||
};
|
||||
|
||||
// optimization return values
|
||||
enum ggml_rwkv_opt_result {
|
||||
GGML_RWKV_OPT_OK = 0,
|
||||
GGML_RWKV_OPT_DID_NOT_CONVERGE,
|
||||
GGML_RWKV_OPT_NO_CONTEXT,
|
||||
GGML_RWKV_OPT_INVALID_WOLFE,
|
||||
GGML_RWKV_OPT_FAIL,
|
||||
|
||||
GGML_RWKV_LINESEARCH_FAIL = -128,
|
||||
GGML_RWKV_LINESEARCH_MINIMUM_STEP,
|
||||
GGML_RWKV_LINESEARCH_MAXIMUM_STEP,
|
||||
GGML_RWKV_LINESEARCH_MAXIMUM_ITERATIONS,
|
||||
GGML_RWKV_LINESEARCH_INVALID_PARAMETERS,
|
||||
};
|
||||
|
||||
// optimization parameters
|
||||
//
|
||||
// see ggml.c (ggml_rwkv_opt_default_params) for default values
|
||||
//
|
||||
struct ggml_rwkv_opt_params {
|
||||
enum ggml_rwkv_opt_type type;
|
||||
|
||||
int n_threads;
|
||||
|
||||
// delta-based convergence test
|
||||
//
|
||||
// if past == 0 - disabled
|
||||
// if past > 0:
|
||||
// stop if |f(x) - f(x_past)| < delta * max(1, |f(x)|)
|
||||
//
|
||||
int past;
|
||||
float delta;
|
||||
|
||||
// maximum number of iterations without improvement
|
||||
//
|
||||
// if 0 - disabled
|
||||
// if > 0:
|
||||
// assume convergence if no cost improvement in this number of iterations
|
||||
//
|
||||
int max_no_improvement;
|
||||
|
||||
bool print_forward_graph;
|
||||
bool print_backward_graph;
|
||||
|
||||
// ADAM parameters
|
||||
struct {
|
||||
int n_iter;
|
||||
|
||||
float alpha; // learning rate
|
||||
float beta1;
|
||||
float beta2;
|
||||
float eps; // epsilon for numerical stability
|
||||
float eps_f; // epsilon for convergence test
|
||||
float eps_g; // epsilon for convergence test
|
||||
} adam;
|
||||
|
||||
// LBFGS parameters
|
||||
struct {
|
||||
int m; // number of corrections to approximate the inv. Hessian
|
||||
int n_iter;
|
||||
int max_linesearch;
|
||||
|
||||
float eps; // convergence tolerance
|
||||
float ftol; // line search tolerance
|
||||
float wolfe;
|
||||
float min_step;
|
||||
float max_step;
|
||||
|
||||
enum ggml_rwkv_linesearch linesearch;
|
||||
} lbfgs;
|
||||
};
|
||||
|
||||
struct ggml_rwkv_opt_params ggml_rwkv_opt_default_params(enum ggml_rwkv_opt_type type);
|
||||
|
||||
// optimize the function defined by the tensor f
|
||||
enum ggml_rwkv_opt_result ggml_rwkv_opt(
|
||||
struct ggml_rwkv_context * ctx,
|
||||
struct ggml_rwkv_opt_params params,
|
||||
struct ggml_rwkv_tensor * f);
|
||||
|
||||
//
|
||||
// quantization
|
||||
//
|
||||
|
||||
size_t ggml_rwkv_quantize_q4_0(const float * src, void * dst, int n, int k, int64_t * hist);
|
||||
size_t ggml_rwkv_quantize_q4_1(const float * src, void * dst, int n, int k, int64_t * hist);
|
||||
size_t ggml_rwkv_quantize_q4_1_o(const float * src, void * dst, int n, int k, int64_t * hist);
|
||||
|
||||
//
|
||||
// system info
|
||||
//
|
||||
|
||||
int ggml_rwkv_cpu_has_avx(void);
|
||||
int ggml_rwkv_cpu_has_avx2(void);
|
||||
int ggml_rwkv_cpu_has_avx512(void);
|
||||
int ggml_rwkv_cpu_has_fma(void);
|
||||
int ggml_rwkv_cpu_has_neon(void);
|
||||
int ggml_rwkv_cpu_has_arm_fma(void);
|
||||
int ggml_rwkv_cpu_has_f16c(void);
|
||||
int ggml_rwkv_cpu_has_fp16_va(void);
|
||||
int ggml_rwkv_cpu_has_wasm_simd(void);
|
||||
int ggml_rwkv_cpu_has_blas(void);
|
||||
int ggml_rwkv_cpu_has_sse3(void);
|
||||
int ggml_rwkv_cpu_has_vsx(void);
|
||||
|
||||
// Run test suite for ggml.
|
||||
// Exits normally, if all tests pass.
|
||||
// Aborts the execution if any test did not pass.
|
||||
void ggml_rwkv_run_test_suite();
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
|
@ -213,7 +213,61 @@ struct gpt2_model {
|
|||
std::map<std::string, struct ggml_tensor *> tensors;
|
||||
};
|
||||
|
||||
// ModelLoadResult legacy_gptj_model_load(const std::string &fname, gptj_model_v1 &model, gpt_vocab &vocab, FileFormat file_format);
|
||||
// bool legacy_gptj_eval(const gptj_model_v1 &model, const int n_threads, const int n_past, const std::vector<gpt_vocab::id> &embd_inp, std::vector<float> &embd_w, size_t &mem_per_token, FileFormat file_format);
|
||||
// ModelLoadResult gptj_model_load(const std::string &fname, gptj_model &model, gpt_vocab &vocab);
|
||||
// bool gptj_eval(const gptj_model &model, const int n_threads, const int n_past, const std::vector<gpt_vocab::id> &embd_inp, std::vector<float> &embd_w, size_t &mem_per_token);
|
||||
struct rwkv_layer {
|
||||
struct ggml_rwkv_tensor * ln1_weight;
|
||||
struct ggml_rwkv_tensor * ln1_bias;
|
||||
|
||||
// RWKV, also called "attention" by the author.
|
||||
struct ggml_rwkv_tensor * att_time_mix_k;
|
||||
struct ggml_rwkv_tensor * att_time_mix_v;
|
||||
struct ggml_rwkv_tensor * att_time_mix_r;
|
||||
struct ggml_rwkv_tensor * att_time_first;
|
||||
struct ggml_rwkv_tensor * att_time_decay;
|
||||
struct ggml_rwkv_tensor * att_key;
|
||||
struct ggml_rwkv_tensor * att_value;
|
||||
struct ggml_rwkv_tensor * att_receptance;
|
||||
struct ggml_rwkv_tensor * att_output;
|
||||
|
||||
struct ggml_rwkv_tensor * ln2_weight;
|
||||
struct ggml_rwkv_tensor * ln2_bias;
|
||||
|
||||
// FFN.
|
||||
struct ggml_rwkv_tensor * ffn_time_mix_k;
|
||||
struct ggml_rwkv_tensor * ffn_time_mix_r;
|
||||
struct ggml_rwkv_tensor * ffn_key;
|
||||
struct ggml_rwkv_tensor * ffn_value;
|
||||
struct ggml_rwkv_tensor * ffn_receptance;
|
||||
};
|
||||
|
||||
struct rwkv_model {
|
||||
int32_t n_vocab;
|
||||
int32_t n_layer;
|
||||
int32_t n_embed;
|
||||
// 0 for float32, 1 for float16.
|
||||
int32_t data_type;
|
||||
|
||||
struct ggml_rwkv_tensor * emb;
|
||||
|
||||
struct ggml_rwkv_tensor * ln0_weight;
|
||||
struct ggml_rwkv_tensor * ln0_bias;
|
||||
|
||||
std::vector<rwkv_layer> layers;
|
||||
|
||||
struct ggml_rwkv_tensor * ln_out_weight;
|
||||
struct ggml_rwkv_tensor * ln_out_bias;
|
||||
|
||||
struct ggml_rwkv_tensor * head;
|
||||
};
|
||||
struct rwkv_context {
|
||||
struct rwkv_model * model;
|
||||
struct ggml_rwkv_tensor * token_index;
|
||||
struct ggml_rwkv_tensor * state;
|
||||
struct ggml_rwkv_tensor ** state_parts;
|
||||
struct ggml_rwkv_tensor * logits;
|
||||
struct ggml_rwkv_context * ctx;
|
||||
struct ggml_rwkv_cgraph * graph;
|
||||
bool freed;
|
||||
float * state_in = 0; //stores input state, or use null for a new state
|
||||
float * state_out = 0; //stores address of output state buffer
|
||||
float * logits_out = 0; //stores address of output logit buffer
|
||||
};
|
||||
|
|
|
|||
739
otherarch/rwkv.cpp
Normal file
739
otherarch/rwkv.cpp
Normal file
|
|
@ -0,0 +1,739 @@
|
|||
#include "rwkv.h"
|
||||
#include "ggml_rwkv.h"
|
||||
#include "otherarch.h"
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <thread>
|
||||
#include <cassert>
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <unordered_map>
|
||||
|
||||
|
||||
#include "model_adapter.h"
|
||||
|
||||
// --- Utilities ---
|
||||
|
||||
#define FP32_SIZE 4
|
||||
|
||||
// Checks that x is not false. If x is false, prints fancy message to stderr and returns 0.
|
||||
#define RWKV_ASSERT_FALSE(x, ...) \
|
||||
do { \
|
||||
if (!(x)) { \
|
||||
fprintf(stderr, __VA_ARGS__); \
|
||||
fprintf(stderr, "\n%s:%d: %s\n", __FILE__, __LINE__, #x); \
|
||||
return false; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// Checks that x is not false. If x is false, prints fancy message to stderr and returns NULL.
|
||||
#define RWKV_ASSERT_NULL(x, ...) \
|
||||
do { \
|
||||
if (!(x)) { \
|
||||
fprintf(stderr, __VA_ARGS__); \
|
||||
fprintf(stderr, "\n%s:%d: %s\n", __FILE__, __LINE__, #x); \
|
||||
return NULL; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// Reads single int32 value from a file.
|
||||
bool read_int32(FILE * file, int32_t * dest) {
|
||||
RWKV_ASSERT_FALSE(fread(dest, 4, 1, file) == 1, "Failed to read an int32 value from a file");
|
||||
return true;
|
||||
}
|
||||
|
||||
static const ggml_rwkv_type FORMAT_TYPE_TO_GGML_RWKV_TYPE[5] = {
|
||||
GGML_RWKV_TYPE_F32,
|
||||
GGML_RWKV_TYPE_F16,
|
||||
GGML_RWKV_TYPE_Q4_0,
|
||||
GGML_RWKV_TYPE_Q4_1,
|
||||
GGML_RWKV_TYPE_Q4_1_O
|
||||
};
|
||||
|
||||
// --- Model definition and loading utilities ---
|
||||
|
||||
|
||||
|
||||
// Finds model parameter by key and sets it into dest.
|
||||
// If the parameter was not found, returns false.
|
||||
bool set_parameter(std::unordered_map<std::string, struct ggml_rwkv_tensor *> * parameters, char * key, struct ggml_rwkv_tensor ** dest) {
|
||||
struct ggml_rwkv_tensor * parameter = (*parameters)[key];
|
||||
RWKV_ASSERT_FALSE(parameter != NULL, "Parameter %s not found in model file", key);
|
||||
*dest = parameter;
|
||||
return true;
|
||||
}
|
||||
|
||||
// Finds block parameter by block index and key and sets it into dest.
|
||||
// If the parameter was not found, returns false.
|
||||
bool set_block_parameter(std::unordered_map<std::string, struct ggml_rwkv_tensor *> * parameters, int32_t block_index, char * key, struct ggml_rwkv_tensor ** dest) {
|
||||
char full_key[128];
|
||||
sprintf(full_key, "blocks.%d.%s", block_index, key);
|
||||
return set_parameter(parameters, full_key, dest);
|
||||
}
|
||||
|
||||
// --- Operators ---
|
||||
|
||||
struct ggml_rwkv_tensor * rwkv_layer_norm(ggml_rwkv_context * ctx, struct ggml_rwkv_tensor * x, struct ggml_rwkv_tensor * weight, struct ggml_rwkv_tensor * bias) {
|
||||
// LayerNorm in RWKV is `x = (x - mean(x)) / sqrt(variance(x) + 1e-5) * weight + bias`
|
||||
// Looks like ggml_rwkv_norm does the first part, we only need to apply weight & bias.
|
||||
x = ggml_rwkv_norm(ctx, x);
|
||||
x = ggml_rwkv_mul(ctx, x, weight);
|
||||
x = ggml_rwkv_add(ctx, x, bias);
|
||||
return x;
|
||||
}
|
||||
|
||||
// --- Implementation ---
|
||||
|
||||
|
||||
|
||||
struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_threads) {
|
||||
FILE * file = fopen(file_path, "rb");
|
||||
RWKV_ASSERT_NULL(file != NULL, "Failed to open file %s", file_path);
|
||||
|
||||
int32_t magic;
|
||||
read_int32(file, &magic);
|
||||
RWKV_ASSERT_NULL(magic == RWKV_FILE_MAGIC, "Unexpected magic value %d", magic);
|
||||
|
||||
int32_t version;
|
||||
read_int32(file, &version);
|
||||
RWKV_ASSERT_NULL(version == RWKV_FILE_VERSION, "Unsupported file version %d", version);
|
||||
|
||||
struct rwkv_model * model = (struct rwkv_model *) calloc(1, sizeof(struct rwkv_model));
|
||||
|
||||
read_int32(file, &(model->n_vocab));
|
||||
RWKV_ASSERT_NULL(model->n_vocab > 0, "Non-positive n_vocab %d", model->n_vocab);
|
||||
|
||||
read_int32(file, &(model->n_embed));
|
||||
RWKV_ASSERT_NULL(model->n_embed > 0, "Non-positive n_embed %d", model->n_embed);
|
||||
|
||||
read_int32(file, &(model->n_layer));
|
||||
RWKV_ASSERT_NULL(model->n_layer > 0, "Non-positive n_layer %d", model->n_layer);
|
||||
|
||||
read_int32(file, &(model->data_type));
|
||||
RWKV_ASSERT_NULL(
|
||||
model->data_type == 0 ||
|
||||
model->data_type == 1 ||
|
||||
model->data_type == 2 ||
|
||||
model->data_type == 3 ||
|
||||
model->data_type == 4,
|
||||
"Unsupported model data type %d",
|
||||
model->data_type
|
||||
);
|
||||
|
||||
// Parameter tensors would take at least this amount in memory.
|
||||
size_t file_size;
|
||||
|
||||
{
|
||||
auto fin = std::ifstream(file_path, std::ios::binary);
|
||||
RWKV_ASSERT_NULL(fin, "Failed to open file %s", file_path);
|
||||
fin.seekg(0, fin.end);
|
||||
file_size = fin.tellg();
|
||||
fin.close();
|
||||
}
|
||||
|
||||
size_t memory_required = file_size +
|
||||
// Intermediary vectors for calculation; there are around 100 calls to ggml
|
||||
size_t(100) * model->n_embed * sizeof(float) +
|
||||
// State, in and out
|
||||
size_t(2) * 5 * model->n_layer * model->n_embed * sizeof(float) +
|
||||
// Logits
|
||||
size_t(model->n_vocab) * sizeof(float) +
|
||||
// +256 MB just for any overhead
|
||||
// TODO This is too much for smaller models; need a more proper and robust way of measuring required memory
|
||||
size_t(256) * 1024 * 1024;
|
||||
|
||||
// Initialize ggml
|
||||
struct ggml_rwkv_init_params params;
|
||||
params.mem_size = memory_required;
|
||||
params.mem_buffer = NULL;
|
||||
struct ggml_rwkv_context * ctx = ggml_rwkv_init(params);
|
||||
|
||||
std::unordered_map<std::string, struct ggml_rwkv_tensor *> parameters;
|
||||
|
||||
while (true) {
|
||||
int32_t dim_count;
|
||||
size_t elements_read = fread(&dim_count, 4, 1, file);
|
||||
|
||||
if (feof(file)) {
|
||||
break;
|
||||
}
|
||||
|
||||
RWKV_ASSERT_NULL(elements_read == 1, "Failed to read dimension count");
|
||||
RWKV_ASSERT_NULL(dim_count == 1 || dim_count == 2, "Unsupported dimension count %d", dim_count);
|
||||
|
||||
int32_t key_length;
|
||||
read_int32(file, &key_length);
|
||||
RWKV_ASSERT_NULL(key_length > 0, "Non-positive key length %d", key_length);
|
||||
|
||||
int32_t data_type;
|
||||
read_int32(file, &data_type);
|
||||
RWKV_ASSERT_NULL(
|
||||
data_type == 0 ||
|
||||
data_type == 1 ||
|
||||
data_type == 2 ||
|
||||
data_type == 3 ||
|
||||
data_type == 4,
|
||||
"Unsupported parameter data type %d",
|
||||
data_type
|
||||
);
|
||||
|
||||
ggml_rwkv_type ggml_rwkv_data_type = FORMAT_TYPE_TO_GGML_RWKV_TYPE[data_type];
|
||||
|
||||
struct ggml_rwkv_tensor * tensor;
|
||||
|
||||
int32_t x = -1;
|
||||
int32_t y = -1;
|
||||
|
||||
if (dim_count == 1) {
|
||||
read_int32(file, &x);
|
||||
tensor = ggml_rwkv_new_tensor_1d(ctx, ggml_rwkv_data_type, x);
|
||||
} else if (dim_count == 2) {
|
||||
read_int32(file, &x);
|
||||
read_int32(file, &y);
|
||||
tensor = ggml_rwkv_new_tensor_2d(ctx, ggml_rwkv_data_type, x, y);
|
||||
} else {
|
||||
abort();
|
||||
}
|
||||
|
||||
std::string key(key_length, 0);
|
||||
RWKV_ASSERT_NULL(fread(&key[0], 1, key_length, file) == uint32_t(key_length), "Failed to read parameter key");
|
||||
|
||||
RWKV_ASSERT_NULL(fread(tensor->data, 1, ggml_rwkv_nbytes(tensor), file) == ggml_rwkv_nbytes(tensor), "Failed to read parameter data");
|
||||
|
||||
parameters[key] = tensor;
|
||||
}
|
||||
|
||||
fclose(file);
|
||||
|
||||
model->layers.resize(model->n_layer);
|
||||
|
||||
set_parameter(¶meters, "emb.weight", &(model->emb));
|
||||
|
||||
set_parameter(¶meters, "blocks.0.ln0.weight", &(model->ln0_weight));
|
||||
set_parameter(¶meters, "blocks.0.ln0.bias", &(model->ln0_bias));
|
||||
|
||||
for (int i = 0; i < model->n_layer; i++) {
|
||||
rwkv_layer layer = model->layers[i];
|
||||
|
||||
set_block_parameter(¶meters, i, "ln1.weight", &(layer.ln1_weight));
|
||||
set_block_parameter(¶meters, i, "ln1.bias", &(layer.ln1_bias));
|
||||
|
||||
set_block_parameter(¶meters, i, "att.time_mix_k", &(layer.att_time_mix_k));
|
||||
set_block_parameter(¶meters, i, "att.time_mix_v", &(layer.att_time_mix_v));
|
||||
set_block_parameter(¶meters, i, "att.time_mix_r", &(layer.att_time_mix_r));
|
||||
set_block_parameter(¶meters, i, "att.time_first", &(layer.att_time_first));
|
||||
set_block_parameter(¶meters, i, "att.time_decay", &(layer.att_time_decay));
|
||||
set_block_parameter(¶meters, i, "att.key.weight", &(layer.att_key));
|
||||
set_block_parameter(¶meters, i, "att.value.weight", &(layer.att_value));
|
||||
set_block_parameter(¶meters, i, "att.receptance.weight", &(layer.att_receptance));
|
||||
set_block_parameter(¶meters, i, "att.output.weight", &(layer.att_output));
|
||||
|
||||
set_block_parameter(¶meters, i, "ln2.weight", &(layer.ln2_weight));
|
||||
set_block_parameter(¶meters, i, "ln2.bias", &(layer.ln2_bias));
|
||||
|
||||
set_block_parameter(¶meters, i, "ffn.time_mix_k", &(layer.ffn_time_mix_k));
|
||||
set_block_parameter(¶meters, i, "ffn.time_mix_r", &(layer.ffn_time_mix_r));
|
||||
set_block_parameter(¶meters, i, "ffn.key.weight", &(layer.ffn_key));
|
||||
set_block_parameter(¶meters, i, "ffn.value.weight", &(layer.ffn_value));
|
||||
set_block_parameter(¶meters, i, "ffn.receptance.weight", &(layer.ffn_receptance));
|
||||
|
||||
model->layers[i] = layer;
|
||||
}
|
||||
|
||||
set_parameter(¶meters, "ln_out.weight", &(model->ln_out_weight));
|
||||
set_parameter(¶meters, "ln_out.bias", &(model->ln_out_bias));
|
||||
|
||||
set_parameter(¶meters, "head.weight", &(model->head));
|
||||
|
||||
// Verify order of dimensions
|
||||
struct ggml_rwkv_tensor * emb = model->emb;
|
||||
RWKV_ASSERT_NULL(emb->n_dims == 2, "Unexpected dimension count of embedding matrix %d", emb->n_dims);
|
||||
RWKV_ASSERT_NULL(emb->ne[0] == model->n_embed, "Unexpected dimension of embedding matrix %d", emb->ne[0]);
|
||||
RWKV_ASSERT_NULL(emb->ne[1] == model->n_vocab, "Unexpected dimension of embedding matrix %d", emb->ne[1]);
|
||||
|
||||
int32_t n_embed = model->n_embed;
|
||||
int32_t n_layer = model->n_layer;
|
||||
|
||||
// Build graph
|
||||
struct ggml_rwkv_tensor * state = ggml_rwkv_new_tensor_1d(ctx, GGML_RWKV_TYPE_F32, n_layer * 5 * n_embed);
|
||||
|
||||
// x = self.w.emb.weight[token]
|
||||
struct ggml_rwkv_tensor * token_index = ggml_rwkv_new_tensor_1d(ctx, GGML_RWKV_TYPE_I32, 1);
|
||||
struct ggml_rwkv_tensor * x = ggml_rwkv_get_rows(ctx, model->emb, token_index);
|
||||
|
||||
// x = self.layer_norm(x, self.w.blocks[0].ln0)
|
||||
x = rwkv_layer_norm(ctx, x, model->ln0_weight, model->ln0_bias);
|
||||
|
||||
// We collect parts of new state here. Each part is (n_embed) vector.
|
||||
struct ggml_rwkv_tensor ** state_parts = new ggml_rwkv_tensor * [n_layer * 5];
|
||||
|
||||
for (int i = 0; i < n_layer; i++) {
|
||||
auto layer = model->layers[i];
|
||||
|
||||
// RWKV/time mixing
|
||||
{
|
||||
// self.layer_norm(x, self.w.blocks[i].ln1)
|
||||
struct ggml_rwkv_tensor * x0 = rwkv_layer_norm(ctx, x, layer.ln1_weight, layer.ln1_bias);
|
||||
// state[5 * i + 1]
|
||||
struct ggml_rwkv_tensor * x_prev = ggml_rwkv_view_1d(ctx, state, n_embed, (5 * i + 1) * n_embed * FP32_SIZE);
|
||||
// xk = x * time_mix_k + state[5 * i + 1] * (1 - time_mix_k)
|
||||
// xv = x * time_mix_v + state[5 * i + 1] * (1 - time_mix_v)
|
||||
// xr = x * time_mix_r + state[5 * i + 1] * (1 - time_mix_r)
|
||||
struct ggml_rwkv_tensor * xk = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, x0, layer.att_time_mix_k),
|
||||
ggml_rwkv_mul(ctx, x_prev, ggml_rwkv_1_minus_x(ctx, layer.att_time_mix_k))
|
||||
);
|
||||
struct ggml_rwkv_tensor * xv = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, x0, layer.att_time_mix_v),
|
||||
ggml_rwkv_mul(ctx, x_prev, ggml_rwkv_1_minus_x(ctx, layer.att_time_mix_v))
|
||||
);
|
||||
struct ggml_rwkv_tensor * xr = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, x0, layer.att_time_mix_r),
|
||||
ggml_rwkv_mul(ctx, x_prev, ggml_rwkv_1_minus_x(ctx, layer.att_time_mix_r))
|
||||
);
|
||||
// state[5 * i + 1] = x
|
||||
state_parts[5 * i + 1] = x0;
|
||||
|
||||
// r = torch.sigmoid(rw @ xr)
|
||||
struct ggml_rwkv_tensor * r = ggml_rwkv_sigmoid(
|
||||
ctx,
|
||||
ggml_rwkv_mul_mat(ctx, layer.att_receptance, xr)
|
||||
);
|
||||
// k = kw @ xk
|
||||
struct ggml_rwkv_tensor * k = ggml_rwkv_mul_mat(ctx, layer.att_key, xk);
|
||||
// v = vw @ xv
|
||||
struct ggml_rwkv_tensor * v = ggml_rwkv_mul_mat(ctx, layer.att_value, xv);
|
||||
|
||||
// aa = state[5 * i + 2]
|
||||
// bb = state[5 * i + 3]
|
||||
// pp = state[5 * i + 4]
|
||||
struct ggml_rwkv_tensor * aa = ggml_rwkv_view_1d(ctx, state, n_embed, (5 * i + 2) * n_embed * FP32_SIZE);
|
||||
struct ggml_rwkv_tensor * bb = ggml_rwkv_view_1d(ctx, state, n_embed, (5 * i + 3) * n_embed * FP32_SIZE);
|
||||
struct ggml_rwkv_tensor * pp = ggml_rwkv_view_1d(ctx, state, n_embed, (5 * i + 4) * n_embed * FP32_SIZE);
|
||||
|
||||
// ww = time_first + k
|
||||
struct ggml_rwkv_tensor * ww = ggml_rwkv_add(ctx, layer.att_time_first, k);
|
||||
// qq = torch.maximum(pp, ww)
|
||||
struct ggml_rwkv_tensor * qq = ggml_rwkv_max(ctx, pp, ww);
|
||||
// e1 = torch.exp(pp - qq)
|
||||
struct ggml_rwkv_tensor * e1 = ggml_rwkv_exp(ctx, ggml_rwkv_sub(ctx, pp, qq));
|
||||
// e2 = torch.exp(ww - qq)
|
||||
struct ggml_rwkv_tensor * e2 = ggml_rwkv_exp(ctx, ggml_rwkv_sub(ctx, ww, qq));
|
||||
// a = e1 * aa + e2 * v
|
||||
struct ggml_rwkv_tensor * a = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, e1, aa),
|
||||
ggml_rwkv_mul(ctx, e2, v)
|
||||
);
|
||||
// b = e1 * bb + e2
|
||||
struct ggml_rwkv_tensor * b = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, e1, bb),
|
||||
e2
|
||||
);
|
||||
// wkv = a / b
|
||||
struct ggml_rwkv_tensor * wkv = ggml_rwkv_div(ctx, a, b);
|
||||
// ww = pp + time_decay
|
||||
ww = ggml_rwkv_add(ctx, pp, layer.att_time_decay);
|
||||
// qq = torch.maximum(ww, k)
|
||||
qq = ggml_rwkv_max(ctx, ww, k);
|
||||
// e1 = torch.exp(ww - qq)
|
||||
e1 = ggml_rwkv_exp(ctx, ggml_rwkv_sub(ctx, ww, qq));
|
||||
// e2 = torch.exp(k - qq)
|
||||
e2 = ggml_rwkv_exp(ctx, ggml_rwkv_sub(ctx, k, qq));
|
||||
// state[5 * i + 2] = e1 * aa + e2 * v
|
||||
state_parts[5 * i + 2] = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, e1, aa),
|
||||
ggml_rwkv_mul(ctx, e2, v)
|
||||
);
|
||||
// state[5 * i + 3] = e1 * bb + e2
|
||||
state_parts[5 * i + 3] = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, e1, bb),
|
||||
e2
|
||||
);
|
||||
// state[5 * i + 4] = qq
|
||||
state_parts[5 * i + 4] = qq;
|
||||
// ow @ (r * wkv)
|
||||
x = ggml_rwkv_add(
|
||||
ctx,
|
||||
x,
|
||||
ggml_rwkv_mul_mat(
|
||||
ctx,
|
||||
layer.att_output,
|
||||
ggml_rwkv_mul(ctx, r, wkv)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
// FFN/channel mixing
|
||||
{
|
||||
// self.layer_norm(x, self.w.blocks[i].ln2)
|
||||
struct ggml_rwkv_tensor * x0 = rwkv_layer_norm(ctx, x, layer.ln2_weight, layer.ln2_bias);
|
||||
// state[5 * i + 0]
|
||||
struct ggml_rwkv_tensor * x_prev = ggml_rwkv_view_1d(ctx, state, n_embed, (5 * i + 0) * n_embed * FP32_SIZE);
|
||||
// xk = x * time_mix_k + state[5 * i + 0] * (1 - time_mix_k)
|
||||
// xr = x * time_mix_r + state[5 * i + 0] * (1 - time_mix_r)
|
||||
struct ggml_rwkv_tensor * xk = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, x0, layer.ffn_time_mix_k),
|
||||
ggml_rwkv_mul(ctx, x_prev, ggml_rwkv_1_minus_x(ctx, layer.ffn_time_mix_k))
|
||||
);
|
||||
struct ggml_rwkv_tensor * xr = ggml_rwkv_add(
|
||||
ctx,
|
||||
ggml_rwkv_mul(ctx, x0, layer.ffn_time_mix_r),
|
||||
ggml_rwkv_mul(ctx, x_prev, ggml_rwkv_1_minus_x(ctx, layer.ffn_time_mix_r))
|
||||
);
|
||||
// state[5 * i + 0] = x
|
||||
state_parts[5 * i + 0] = x0;
|
||||
|
||||
// r = torch.sigmoid(rw @ xr)
|
||||
struct ggml_rwkv_tensor * r = ggml_rwkv_sigmoid(
|
||||
ctx,
|
||||
ggml_rwkv_mul_mat(ctx, layer.ffn_receptance, xr)
|
||||
);
|
||||
// k = torch.square(torch.relu(kw @ xk))
|
||||
struct ggml_rwkv_tensor * k = ggml_rwkv_sqr(ctx, ggml_rwkv_relu(
|
||||
ctx,
|
||||
ggml_rwkv_mul_mat(ctx, layer.ffn_key, xk)
|
||||
));
|
||||
// r * (vw @ k)
|
||||
x = ggml_rwkv_add(
|
||||
ctx,
|
||||
x,
|
||||
ggml_rwkv_mul(
|
||||
ctx,
|
||||
r,
|
||||
ggml_rwkv_mul_mat(ctx, layer.ffn_value, k)
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// x = self.layer_norm(x, self.w.ln_out)
|
||||
x = rwkv_layer_norm(ctx, x, model->ln_out_weight, model->ln_out_bias);
|
||||
|
||||
// x = (self.w.head.weight @ x).float()
|
||||
struct ggml_rwkv_tensor * logits = ggml_rwkv_mul_mat(ctx, model->head, x);
|
||||
|
||||
struct ggml_rwkv_cgraph * graph = (struct ggml_rwkv_cgraph *) calloc(1, sizeof(struct ggml_rwkv_cgraph));
|
||||
|
||||
*graph = ggml_rwkv_build_forward(logits);
|
||||
|
||||
for (int i = 0; i < n_layer * 5; i++) {
|
||||
ggml_rwkv_build_forward_expand(graph, state_parts[i]);
|
||||
}
|
||||
|
||||
graph->n_threads = n_threads;
|
||||
|
||||
struct rwkv_context * rwkv_ctx = (struct rwkv_context *) calloc(1, sizeof(struct rwkv_context));
|
||||
rwkv_ctx->model = model;
|
||||
rwkv_ctx->token_index = token_index;
|
||||
rwkv_ctx->state = state;
|
||||
rwkv_ctx->state_parts = state_parts;
|
||||
rwkv_ctx->logits = logits;
|
||||
rwkv_ctx->ctx = ctx;
|
||||
rwkv_ctx->graph = graph;
|
||||
return rwkv_ctx;
|
||||
}
|
||||
|
||||
uint32_t rwkv_get_state_buffer_element_count(struct rwkv_context * ctx) {
|
||||
return ctx->model->n_layer * 5 * ctx->model->n_embed;
|
||||
}
|
||||
|
||||
uint32_t rwkv_get_logits_buffer_element_count(struct rwkv_context * ctx) {
|
||||
return ctx->model->n_vocab;
|
||||
}
|
||||
|
||||
bool rwkv_eval(struct rwkv_context * ctx, int32_t token, float * state_in, float * state_out, float * logits_out) {
|
||||
RWKV_ASSERT_FALSE(state_out != NULL, "state_out is NULL");
|
||||
RWKV_ASSERT_FALSE(logits_out != NULL, "logits_out is NULL");
|
||||
|
||||
int32_t n_layer = ctx->model->n_layer;
|
||||
int32_t n_embed = ctx->model->n_embed;
|
||||
int32_t n_vocab = ctx->model->n_vocab;
|
||||
|
||||
RWKV_ASSERT_FALSE(token >= 0 && token < n_vocab, "Token is out of range 0..%d", n_vocab - 1);
|
||||
|
||||
ggml_rwkv_set_i32(ctx->token_index, 0);
|
||||
ggml_rwkv_set_i32_1d(ctx->token_index, 0, token);
|
||||
|
||||
if (state_in == NULL) {
|
||||
ggml_rwkv_set_f32(ctx->state, 0.0F);
|
||||
|
||||
for (int i = 0; i < n_layer; i++) {
|
||||
// state[5 * i + 4] = -1e30
|
||||
ggml_rwkv_set_f32(
|
||||
ggml_rwkv_view_1d(ctx->ctx, ctx->state, n_embed, (5 * i + 4) * n_embed * FP32_SIZE),
|
||||
-1e30F
|
||||
);
|
||||
}
|
||||
} else {
|
||||
memcpy(ctx->state->data, state_in, ctx->state->ne[0] * FP32_SIZE);
|
||||
}
|
||||
|
||||
ggml_rwkv_graph_compute(ctx->ctx, ctx->graph);
|
||||
|
||||
for (size_t i = 0; i < size_t(n_layer * 5); i++) {
|
||||
struct ggml_rwkv_tensor * part = ctx->state_parts[i];
|
||||
|
||||
memcpy(state_out + i * n_embed, part->data, part->ne[0] * FP32_SIZE);
|
||||
}
|
||||
|
||||
memcpy(logits_out, ctx->logits->data, ctx->logits->ne[0] * FP32_SIZE);
|
||||
|
||||
// Uncomment to measure used memory for adding the value into get_memory_required_mb.
|
||||
//fprintf(stderr, "Used mem: %d MB\n", ggml_rwkv_used_mem(ctx->ctx) / 1024 / 1024);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
void rwkv_free(struct rwkv_context * ctx) {
|
||||
ggml_rwkv_free(ctx->ctx);
|
||||
|
||||
delete ctx->model;
|
||||
delete ctx->state_parts;
|
||||
delete ctx;
|
||||
}
|
||||
|
||||
bool rwkv_quantize_model_file(const char * model_file_path_in, const char * model_file_path_out, uint32_t q_type) {
|
||||
RWKV_ASSERT_FALSE(q_type == 2 || q_type == 3 || q_type == 4, "Unsupported quantization type %d", q_type);
|
||||
|
||||
ggml_rwkv_type type = FORMAT_TYPE_TO_GGML_RWKV_TYPE[q_type];
|
||||
|
||||
printf("Loading model from '%s'\n", model_file_path_in);
|
||||
|
||||
auto finp = std::ifstream(model_file_path_in, std::ios::binary);
|
||||
RWKV_ASSERT_FALSE(finp, "Failed to open %s for reading", model_file_path_in);
|
||||
|
||||
auto fout = std::ofstream(model_file_path_out, std::ios::binary);
|
||||
RWKV_ASSERT_FALSE(fout, "Failed to open %s for writing", model_file_path_out);
|
||||
|
||||
// Process header
|
||||
{
|
||||
uint32_t magic;
|
||||
finp.read((char *) &magic, sizeof(magic));
|
||||
RWKV_ASSERT_FALSE(magic == RWKV_FILE_MAGIC, "Unexpected magic value %d", magic);
|
||||
fout.write((char *) &magic, sizeof(magic));
|
||||
|
||||
uint32_t format_version;
|
||||
finp.read((char *) &format_version, sizeof(format_version));
|
||||
RWKV_ASSERT_FALSE(format_version == RWKV_FILE_VERSION, "Unsupported file version %d", format_version);
|
||||
fout.write((char *) &format_version, sizeof(format_version));
|
||||
|
||||
int32_t n_vocab;
|
||||
int32_t n_embed;
|
||||
int32_t n_layer;
|
||||
int32_t data_type;
|
||||
|
||||
finp.read((char *) &n_vocab, sizeof(n_vocab));
|
||||
finp.read((char *) &n_embed, sizeof(n_embed));
|
||||
finp.read((char *) &n_layer, sizeof(n_layer));
|
||||
finp.read((char *) &data_type, sizeof(data_type));
|
||||
|
||||
RWKV_ASSERT_FALSE(data_type == 0 || data_type == 1, "Unsupported data type %d, only FP32 and FP16 can be quantized", data_type);
|
||||
|
||||
data_type = q_type;
|
||||
|
||||
fout.write((char *) &n_vocab, sizeof(n_vocab));
|
||||
fout.write((char *) &n_embed, sizeof(n_embed));
|
||||
fout.write((char *) &n_layer, sizeof(n_layer));
|
||||
fout.write((char *) &data_type, sizeof(data_type));
|
||||
}
|
||||
|
||||
// Process parameters
|
||||
{
|
||||
size_t total_size_orig = 0;
|
||||
size_t total_size_new = 0;
|
||||
|
||||
std::vector<float> work;
|
||||
|
||||
std::vector<uint8_t> data_u8;
|
||||
std::vector<ggml_rwkv_fp16_t> data_f16;
|
||||
std::vector<float> data_f32;
|
||||
|
||||
std::vector<int64_t> hist_all(1 << 4, 0);
|
||||
|
||||
while (true) {
|
||||
int32_t n_dims;
|
||||
int32_t key_length;
|
||||
int32_t parameter_data_type;
|
||||
|
||||
finp.read(reinterpret_cast<char *>(&n_dims), sizeof(n_dims));
|
||||
finp.read(reinterpret_cast<char *>(&key_length), sizeof(key_length));
|
||||
finp.read(reinterpret_cast<char *>(¶meter_data_type), sizeof(parameter_data_type));
|
||||
|
||||
if (finp.eof()) {
|
||||
break;
|
||||
}
|
||||
|
||||
int32_t nelements = 1;
|
||||
int32_t ne[2] = { 1, 1 };
|
||||
for (int i = 0; i < n_dims; ++i) {
|
||||
finp.read (reinterpret_cast<char *>(&ne[i]), sizeof(ne[i]));
|
||||
nelements *= ne[i];
|
||||
}
|
||||
|
||||
std::string name(key_length, 0);
|
||||
finp.read(&name[0], key_length);
|
||||
|
||||
{
|
||||
static const char * parameter_data_type_str[] = {
|
||||
"F32",
|
||||
"F16",
|
||||
"Q4_0",
|
||||
"Q4_1",
|
||||
"Q4_1_O"
|
||||
};
|
||||
printf("%48s - [%5d, %5d], type = %6s ", name.data(), ne[0], ne[1], parameter_data_type_str[parameter_data_type]);
|
||||
|
||||
total_size_orig += (size_t) (nelements * ggml_rwkv_type_sizef(FORMAT_TYPE_TO_GGML_RWKV_TYPE[parameter_data_type]));
|
||||
}
|
||||
|
||||
// Quantize only 2D tensors, except embedding and head matrices.
|
||||
// Embedding and head take not too much space, especially in bigger models;
|
||||
// but they significantly increase perplexity when quantized.
|
||||
bool quantize = n_dims == 2 &&
|
||||
name != std::string("emb.weight") &&
|
||||
name != std::string("head.weight");
|
||||
|
||||
if (quantize) {
|
||||
RWKV_ASSERT_FALSE(
|
||||
parameter_data_type == 0 || parameter_data_type == 1,
|
||||
"Unsupported parameter data type %d, only FP32 and FP16 can be quantized",
|
||||
parameter_data_type
|
||||
);
|
||||
|
||||
if (parameter_data_type == 1) {
|
||||
data_f16.resize(nelements);
|
||||
finp.read(reinterpret_cast<char *>(data_f16.data()), nelements * sizeof(ggml_rwkv_fp16_t));
|
||||
data_f32.resize(nelements);
|
||||
for (int i = 0; i < nelements; ++i) {
|
||||
data_f32[i] = ggml_rwkv_fp16_to_fp32(data_f16[i]);
|
||||
}
|
||||
} else {
|
||||
data_f32.resize(nelements);
|
||||
finp.read(reinterpret_cast<char *>(data_f32.data()), nelements * sizeof(float));
|
||||
}
|
||||
|
||||
parameter_data_type = q_type;
|
||||
} else {
|
||||
const int bytes_per_element = (parameter_data_type == 0) ? sizeof(float) : sizeof(uint16_t);
|
||||
data_u8.resize(nelements * bytes_per_element);
|
||||
finp.read(reinterpret_cast<char *>(data_u8.data()), nelements * bytes_per_element);
|
||||
}
|
||||
|
||||
fout.write(reinterpret_cast<char *>(&n_dims), sizeof(n_dims));
|
||||
fout.write(reinterpret_cast<char *>(&key_length), sizeof(key_length));
|
||||
fout.write(reinterpret_cast<char *>(¶meter_data_type), sizeof(parameter_data_type));
|
||||
|
||||
for (int i = 0; i < n_dims; ++i) {
|
||||
fout.write(reinterpret_cast<char *>(&ne[i]), sizeof(ne[i]));
|
||||
}
|
||||
|
||||
fout.write(&name[0], key_length);
|
||||
|
||||
if (quantize) {
|
||||
printf("quantizing... ");
|
||||
work.resize(nelements); // for quantization
|
||||
|
||||
size_t cur_size = 0;
|
||||
// This is a histogramm of some values. If it shows single 1.0, then all 0.0, something went very wrong!
|
||||
std::vector<int64_t> hist_cur(1 << 4, 0);
|
||||
|
||||
switch (type) {
|
||||
case GGML_RWKV_TYPE_Q4_0:
|
||||
{
|
||||
cur_size = ggml_rwkv_quantize_q4_0(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
|
||||
} break;
|
||||
case GGML_RWKV_TYPE_Q4_1:
|
||||
{
|
||||
cur_size = ggml_rwkv_quantize_q4_1(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
|
||||
} break;
|
||||
case GGML_RWKV_TYPE_Q4_1_O:
|
||||
{
|
||||
cur_size = ggml_rwkv_quantize_q4_1_o(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
|
||||
} break;
|
||||
default:
|
||||
{
|
||||
fprintf(stderr, "unsupported quantization type %d\n", type);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
fout.write(reinterpret_cast<char *>(work.data()), cur_size);
|
||||
total_size_new += cur_size;
|
||||
|
||||
printf("size = %8.2f MB -> %8.2f MB | hist: ", nelements * sizeof(float) / 1024.0 / 1024.0, cur_size / 1024.0 / 1024.0);
|
||||
|
||||
for (int i = 0; i < (int) hist_cur.size(); ++i) {
|
||||
hist_all[i] += hist_cur[i];
|
||||
}
|
||||
|
||||
for (int i = 0; i < (int) hist_cur.size(); ++i) {
|
||||
printf("%5.3f ", hist_cur[i] / float(nelements));
|
||||
}
|
||||
|
||||
printf("\n");
|
||||
} else {
|
||||
printf("size = %8.3f MB\n", data_u8.size() / 1024.0 / 1024.0);
|
||||
fout.write(reinterpret_cast<char *>(data_u8.data()), data_u8.size());
|
||||
total_size_new += data_u8.size();
|
||||
}
|
||||
}
|
||||
|
||||
printf("original size = %8.2f MB\n", total_size_orig / 1024.0 / 1024.0);
|
||||
printf("quantized size = %8.2f MB\n", total_size_new / 1024.0 / 1024.0);
|
||||
printf("compression ratio = %8.2f%\n", 1.0 * total_size_orig / total_size_new);
|
||||
|
||||
{
|
||||
int64_t sum_all = 0;
|
||||
|
||||
for (int i = 0; i < (int) hist_all.size(); ++i) {
|
||||
sum_all += hist_all[i];
|
||||
}
|
||||
|
||||
printf("hist: ");
|
||||
|
||||
for (int i = 0; i < (int) hist_all.size(); ++i) {
|
||||
printf("%5.3f ", hist_all[i] / float(sum_all));
|
||||
}
|
||||
|
||||
printf("\n");
|
||||
}
|
||||
}
|
||||
|
||||
finp.close();
|
||||
fout.close();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
const char * rwkv_get_system_info_string(void) {
|
||||
static std::string s;
|
||||
|
||||
s = "";
|
||||
s += "AVX = " + std::to_string(ggml_rwkv_cpu_has_avx()) + " | ";
|
||||
s += "AVX2 = " + std::to_string(ggml_rwkv_cpu_has_avx2()) + " | ";
|
||||
s += "AVX512 = " + std::to_string(ggml_rwkv_cpu_has_avx512()) + " | ";
|
||||
s += "FMA = " + std::to_string(ggml_rwkv_cpu_has_fma()) + " | ";
|
||||
s += "NEON = " + std::to_string(ggml_rwkv_cpu_has_neon()) + " | ";
|
||||
s += "ARM_FMA = " + std::to_string(ggml_rwkv_cpu_has_arm_fma()) + " | ";
|
||||
s += "F16C = " + std::to_string(ggml_rwkv_cpu_has_f16c()) + " | ";
|
||||
s += "FP16_VA = " + std::to_string(ggml_rwkv_cpu_has_fp16_va()) + " | ";
|
||||
s += "WASM_SIMD = " + std::to_string(ggml_rwkv_cpu_has_wasm_simd()) + " | ";
|
||||
s += "BLAS = " + std::to_string(ggml_rwkv_cpu_has_blas()) + " | ";
|
||||
s += "SSE3 = " + std::to_string(ggml_rwkv_cpu_has_sse3()) + " | ";
|
||||
s += "VSX = " + std::to_string(ggml_rwkv_cpu_has_vsx()) + " | ";
|
||||
|
||||
return s.c_str();
|
||||
}
|
||||
69
otherarch/rwkv.h
Normal file
69
otherarch/rwkv.h
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
#ifndef RWKV_H
|
||||
#define RWKV_H
|
||||
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
#include <stdbool.h>
|
||||
|
||||
#ifdef RWKV_SHARED
|
||||
# if defined(_WIN32) && !defined(__MINGW32__)
|
||||
# ifdef RWKV_BUILD
|
||||
# define RWKV_API __declspec(dllexport)
|
||||
# else
|
||||
# define RWKV_API __declspec(dllimport)
|
||||
# endif
|
||||
# else
|
||||
# define RWKV_API __attribute__ ((visibility ("default")))
|
||||
# endif
|
||||
#else
|
||||
# define RWKV_API
|
||||
#endif
|
||||
|
||||
// 'ggmf' in hex.
|
||||
#define RWKV_FILE_MAGIC 0x67676d66
|
||||
#define RWKV_FILE_VERSION 100
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
struct rwkv_context;
|
||||
|
||||
// Loads the model from a file and prepares it for inference.
|
||||
// Returns NULL on any error. Error messages would be printed to stderr.
|
||||
// - model_file_path: path to model file in ggml format.
|
||||
// - n_threads: count of threads to use, must be positive.
|
||||
RWKV_API struct rwkv_context * rwkv_init_from_file(const char * model_file_path, uint32_t n_threads);
|
||||
|
||||
// Evaluates the model for a single token.
|
||||
// Returns false on any error. Error messages would be printed to stderr.
|
||||
// - token: next token index, in range 0 <= token < n_vocab.
|
||||
// - state_in: FP32 buffer of size rwkv_get_state_buffer_element_count; or NULL, if this is a first pass.
|
||||
// - state_out: FP32 buffer of size rwkv_get_state_buffer_element_count. This buffer will be written to.
|
||||
// - logits_out: FP32 buffer of size rwkv_get_logits_buffer_element_count. This buffer will be written to.
|
||||
RWKV_API bool rwkv_eval(struct rwkv_context * ctx, int32_t token, float * state_in, float * state_out, float * logits_out);
|
||||
|
||||
// Returns count of FP32 elements in state buffer.
|
||||
RWKV_API uint32_t rwkv_get_state_buffer_element_count(struct rwkv_context * ctx);
|
||||
|
||||
// Returns count of FP32 elements in logits buffer.
|
||||
RWKV_API uint32_t rwkv_get_logits_buffer_element_count(struct rwkv_context * ctx);
|
||||
|
||||
// Frees all allocated memory and the context.
|
||||
RWKV_API void rwkv_free(struct rwkv_context * ctx);
|
||||
|
||||
// Quantizes FP32 or FP16 model to one of INT4 formats.
|
||||
// Returns false on any error. Error messages would be printed to stderr.
|
||||
// - model_file_path_in: path to model file in ggml format, must be either FP32 or FP16.
|
||||
// - model_file_path_out: quantized model will be written here.
|
||||
// - q_type: set to 2 for GGML_RWKV_TYPE_Q4_0, set to 3 for GGML_RWKV_TYPE_Q4_1.
|
||||
RWKV_API bool rwkv_quantize_model_file(const char * model_file_path_in, const char * model_file_path_out, uint32_t q_type);
|
||||
|
||||
// Returns system information string.
|
||||
RWKV_API const char * rwkv_get_system_info_string(void);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif
|
||||
183
otherarch/tools/convert_hf_gpt2.py
Normal file
183
otherarch/tools/convert_hf_gpt2.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
# Convert Cerebras models to ggml format
|
||||
#
|
||||
# ref: https://www.cerebras.net/blog/cerebras-gpt-a-family-of-open-compute-efficient-large-language-models/
|
||||
#
|
||||
|
||||
import sys
|
||||
import struct
|
||||
import json
|
||||
import torch
|
||||
import numpy as np
|
||||
import re
|
||||
|
||||
from transformers import GPTJForCausalLM, AutoModelForCausalLM
|
||||
|
||||
# ref: https://github.com/openai/gpt-2/blob/master/src/encoder.py
|
||||
def bytes_to_unicode():
|
||||
"""
|
||||
Returns list of utf-8 byte and a corresponding list of unicode strings.
|
||||
The reversible bpe codes work on unicode strings.
|
||||
This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
|
||||
When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
|
||||
This is a signficant percentage of your normal, say, 32K bpe vocab.
|
||||
To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
|
||||
And avoids mapping to whitespace/control characters the bpe code barfs on.
|
||||
"""
|
||||
bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
|
||||
cs = bs[:]
|
||||
n = 0
|
||||
for b in range(2**8):
|
||||
if b not in bs:
|
||||
bs.append(b)
|
||||
cs.append(2**8+n)
|
||||
n += 1
|
||||
cs = [chr(n) for n in cs]
|
||||
return dict(zip(bs, cs))
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: convert-h5-to-ggml.py dir-model [use-f32]\n")
|
||||
sys.exit(1)
|
||||
|
||||
# output in the same directory as the model
|
||||
dir_model = sys.argv[1]
|
||||
fname_out = sys.argv[1] + "/ggml-model-f16.bin"
|
||||
|
||||
with open(dir_model + "/vocab.json", "r", encoding="utf-8") as f:
|
||||
encoder = json.load(f)
|
||||
|
||||
with open(dir_model + "/config.json", "r", encoding="utf-8") as f:
|
||||
hparams = json.load(f)
|
||||
|
||||
# use 16-bit or 32-bit floats
|
||||
use_f16 = True
|
||||
if len(sys.argv) > 2:
|
||||
use_f16 = False
|
||||
fname_out = sys.argv[1] + "/ggml-model-f32.bin"
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(dir_model, low_cpu_mem_usage=True)
|
||||
#print (model)
|
||||
|
||||
list_vars = model.state_dict()
|
||||
#print (list_vars)
|
||||
|
||||
print(hparams)
|
||||
|
||||
fout = open(fname_out, "wb")
|
||||
|
||||
fout.write(struct.pack("i", 0x67676d6c)) # magic: ggml in hex
|
||||
fout.write(struct.pack("i", hparams["vocab_size"]))
|
||||
fout.write(struct.pack("i", hparams["n_positions"]))
|
||||
fout.write(struct.pack("i", hparams["n_embd"]))
|
||||
fout.write(struct.pack("i", hparams["n_head"]))
|
||||
fout.write(struct.pack("i", hparams["n_layer"]))
|
||||
fout.write(struct.pack("i", use_f16))
|
||||
|
||||
byte_encoder = bytes_to_unicode()
|
||||
byte_decoder = {v:k for k, v in byte_encoder.items()}
|
||||
|
||||
fout.write(struct.pack("i", len(encoder)))
|
||||
|
||||
for key in encoder:
|
||||
text = bytearray([byte_decoder[c] for c in key])
|
||||
fout.write(struct.pack("i", len(text)))
|
||||
fout.write(text)
|
||||
|
||||
for name in list_vars.keys():
|
||||
data = list_vars[name].squeeze().numpy()
|
||||
print("Processing variable: " + name + " with shape: ", data.shape)
|
||||
|
||||
# rename headers to keep compatibility
|
||||
if name == "transformer.ln_f.weight":
|
||||
name = "model/ln_f/g"
|
||||
elif name == "transformer.ln_f.bias":
|
||||
name = "model/ln_f/b"
|
||||
elif name == "transformer.wte.weight":
|
||||
name = "model/wte"
|
||||
elif name == "transformer.wpe.weight":
|
||||
name = "model/wpe"
|
||||
elif name == "lm_head.weight":
|
||||
name = "model/lm_head"
|
||||
elif re.match(r"transformer.h\.\d+\.ln_1\.weight", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/ln_1/g"
|
||||
elif re.match(r"transformer.h\.\d+\.ln_1\.bias", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/ln_1/b"
|
||||
elif re.match(r"transformer.h\.\d+\.attn\.c_attn\.weight", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/attn/c_attn/w"
|
||||
elif re.match(r"transformer.h\.\d+\.attn\.c_attn\.bias", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/attn/c_attn/b"
|
||||
elif re.match(r"transformer.h\.\d+\.attn\.c_proj\.weight", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/attn/c_proj/w"
|
||||
elif re.match(r"transformer.h.\d+.attn.c_proj.bias", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/attn/c_proj/b"
|
||||
elif re.match(r"transformer.h.\d+.ln_2.weight", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/ln_2/g"
|
||||
elif re.match(r"transformer.h.\d+.ln_2.bias", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/ln_2/b"
|
||||
elif re.match(r"transformer.h.\d+.mlp.c_fc.weight", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/mlp/c_fc/w"
|
||||
elif re.match(r"transformer.h.\d+.mlp.c_fc.bias", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/mlp/c_fc/b"
|
||||
elif re.match(r"transformer.h.\d+.mlp.c_proj.weight", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/mlp/c_proj/w"
|
||||
elif re.match(r"transformer.h.\d+.mlp.c_proj.bias", name):
|
||||
i = re.findall("\d+", name)[0]
|
||||
name = f"model/h{i}/mlp/c_proj/b"
|
||||
else:
|
||||
print("Unrecognized variable name. %s", name)
|
||||
|
||||
# we don't need these
|
||||
if name.endswith("attn.masked_bias") or name.endswith(".attn.bias"):
|
||||
print(" Skipping variable: " + name)
|
||||
continue
|
||||
|
||||
n_dims = len(data.shape);
|
||||
|
||||
# ftype == 0 -> float32, ftype == 1 -> float16
|
||||
ftype = 0;
|
||||
if use_f16:
|
||||
if (name == "model/wte" or name == "model/lm_head" or name[-2:] == "/g" or name[-2:] == "/w") and n_dims == 2:
|
||||
print(" Converting to float16")
|
||||
data = data.astype(np.float16)
|
||||
ftype = 1
|
||||
else:
|
||||
print(" Converting to float32")
|
||||
data = data.astype(np.float32)
|
||||
ftype = 0
|
||||
|
||||
# for efficiency - transpose the projection matrices
|
||||
# "model/h.*/attn/c_attn/w"
|
||||
# "model/h.*/attn/c_proj/w"
|
||||
# "model/h.*/mlp/c_fc/w"
|
||||
# "model/h.*/mlp/c_proj/w"
|
||||
if name[-14:] == "/attn/c_attn/w" or \
|
||||
name[-14:] == "/attn/c_proj/w" or \
|
||||
name[-11:] == "/mlp/c_fc/w" or \
|
||||
name[-13:] == "/mlp/c_proj/w":
|
||||
print(" Transposing")
|
||||
data = data.transpose()
|
||||
|
||||
# header
|
||||
str = name.encode('utf-8')
|
||||
fout.write(struct.pack("iii", n_dims, len(str), ftype))
|
||||
for i in range(n_dims):
|
||||
fout.write(struct.pack("i", data.shape[n_dims - 1 - i]))
|
||||
fout.write(str);
|
||||
|
||||
# data
|
||||
data.tofile(fout)
|
||||
|
||||
fout.close()
|
||||
|
||||
print("Done. Output file: " + fname_out)
|
||||
print("")
|
||||
181
otherarch/tools/convert_pt_rwkv.py
Normal file
181
otherarch/tools/convert_pt_rwkv.py
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
# Converts an RWKV model checkpoint to an rwkv.cpp compatible file.
|
||||
# Usage: python convert_pytorch_to_ggml.py C:\RWKV-4-Pile-169M-20220807-8023.pth C:\rwkv.cpp-169M.bin float32
|
||||
# Get model checkpoints from https://huggingface.co/BlinkDL
|
||||
|
||||
# File format:
|
||||
#
|
||||
# RWKVModelFile {
|
||||
# // All ints and floats are in machine byte order.
|
||||
# // Magic is "ggml" string bytes.
|
||||
# int32 magic = 0x67676d66;
|
||||
# int32 version = 100;
|
||||
# int32 n_vocab;
|
||||
# int32 n_embed;
|
||||
# int32 n_layer;
|
||||
# // 0 if float32, 1 if float16, 2 if Q4_0, 3 if Q4_1, 4 if Q4_1_O.
|
||||
# int32 data_type;
|
||||
# // Read until EOF.
|
||||
# Parameter[] parameters;
|
||||
# }
|
||||
#
|
||||
# Parameter {
|
||||
# int32 dim_count;
|
||||
# int32 key_length;
|
||||
# // 0 if float32, 1 if float16, 2 if Q4_0, 3 if Q4_1, 4 if Q4_1_O.
|
||||
# int32 data_type;
|
||||
# // Compared to PyTorch's tensor.shape, dimension order is reversed here!
|
||||
# int32[dim_count] shape;
|
||||
# // Keys are like "emb.weight", "block.0.ln1.weight".
|
||||
# uint8[key_length] key_utf8;
|
||||
# // float32: 4 * element_count bytes.
|
||||
# // float16: 2 * element_count bytes.
|
||||
# // Q4_0: element_count / 32 * 20 bytes.
|
||||
# // Q4_1: element_count / 32 * 24 bytes.
|
||||
# // Q4_1_O: element_count / 32 * 24 bytes.
|
||||
# byte[] data;
|
||||
# }
|
||||
|
||||
import os
|
||||
import argparse
|
||||
import struct
|
||||
import torch
|
||||
from typing import Dict
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Convert an RWKV model checkpoint to an rwkv.cpp compatible file')
|
||||
parser.add_argument('src_path', help='Path to PyTorch checkpoint file')
|
||||
parser.add_argument('dest_path', help='Path to rwkv.cpp checkpoint file, will be overwritten')
|
||||
parser.add_argument('data_type', help='Data type, float16 or float32', type=str, choices=['float16', 'float32'], default='float32')
|
||||
return parser.parse_args()
|
||||
|
||||
def get_layer_count(state_dict: Dict[str, torch.Tensor]) -> int:
|
||||
n_layer = 0
|
||||
|
||||
while f'blocks.{n_layer}.ln1.weight' in state_dict:
|
||||
n_layer += 1
|
||||
|
||||
assert n_layer > 0
|
||||
|
||||
return n_layer
|
||||
|
||||
def write_state_dict(state_dict: Dict[str, torch.Tensor], dest_path: str, data_type: str) -> None:
|
||||
emb_weight: torch.Tensor = state_dict['emb.weight']
|
||||
|
||||
n_layer = get_layer_count(state_dict)
|
||||
n_vocab = emb_weight.shape[0]
|
||||
n_embed = emb_weight.shape[1]
|
||||
|
||||
with open(dest_path, 'wb') as out_file:
|
||||
out_file.write(struct.pack(
|
||||
# Disable padding with '='
|
||||
'=iiiiii',
|
||||
# Magic: 'ggmf' in hex
|
||||
0x67676d66,
|
||||
# llama.cpp uses file versions 1+, let's use 100+ for rwkv.cpp
|
||||
100,
|
||||
n_vocab,
|
||||
n_embed,
|
||||
n_layer,
|
||||
1 if data_type == 'float16' else 0
|
||||
))
|
||||
|
||||
for k in state_dict.keys():
|
||||
tensor = state_dict[k].float()
|
||||
|
||||
# Same processing as in "RWKV_in_150_lines.py"
|
||||
if '.time_' in k:
|
||||
# (1, 1, n_embed) -> (n_embed)
|
||||
tensor = tensor.squeeze()
|
||||
|
||||
if '.time_decay' in k:
|
||||
tensor = -torch.exp(tensor)
|
||||
|
||||
# Keep 1-dim vectors in fp32
|
||||
if data_type == 'float16' and len(tensor.shape) > 1:
|
||||
tensor = tensor.half()
|
||||
|
||||
shape = tensor.shape
|
||||
|
||||
print(f'Writing {k}, shape {shape}, type {tensor.dtype}')
|
||||
|
||||
k_encoded: bytes = k.encode('utf-8')
|
||||
|
||||
out_file.write(struct.pack(
|
||||
'=iii',
|
||||
len(shape),
|
||||
len(k_encoded),
|
||||
1 if tensor.dtype == torch.float16 else 0
|
||||
))
|
||||
|
||||
# Dimension order is reversed here:
|
||||
# * PyTorch shape is (x rows, y columns)
|
||||
# * ggml shape is (y elements in a row, x elements in a column)
|
||||
# Both shapes represent the same tensor.
|
||||
for dim in reversed(tensor.shape):
|
||||
out_file.write(struct.pack('=i', dim))
|
||||
|
||||
out_file.write(k_encoded)
|
||||
|
||||
tensor.numpy().tofile(out_file)
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
print(f'Reading {args.src_path}')
|
||||
|
||||
state_dict: Dict[str, torch.Tensor] = torch.load(args.src_path, map_location='cpu')
|
||||
|
||||
write_state_dict(state_dict, args.dest_path, args.data_type)
|
||||
|
||||
print('Done')
|
||||
|
||||
# --- Tests ---
|
||||
|
||||
def test() -> None:
|
||||
test_file_path = 'convert_pytorch_rwkv_to_ggml_test.tmp'
|
||||
|
||||
try:
|
||||
state_dict: Dict[str, torch.Tensor] = {
|
||||
'emb.weight': torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float32),
|
||||
'blocks.0.ln1.weight': torch.tensor([1], dtype=torch.float32)
|
||||
}
|
||||
|
||||
write_state_dict(state_dict, dest_path=test_file_path, data_type='float32')
|
||||
|
||||
with open(test_file_path, 'rb') as input:
|
||||
actual_bytes: bytes = input.read()
|
||||
|
||||
expected_bytes: bytes = struct.pack(
|
||||
'=iiiiii' + 'iiiii10sffffff' + 'iiii19sf',
|
||||
0x67676d66,
|
||||
100,
|
||||
3,
|
||||
2,
|
||||
1,
|
||||
0,
|
||||
# emb.weight
|
||||
2,
|
||||
10,
|
||||
0,
|
||||
2, 3,
|
||||
'emb.weight'.encode('utf-8'),
|
||||
1.0, 2.0, 3.0,
|
||||
4.0, 5.0, 6.0,
|
||||
# blocks.0.ln1.weight
|
||||
1,
|
||||
19,
|
||||
0,
|
||||
1,
|
||||
'blocks.0.ln1.weight'.encode('utf-8'),
|
||||
1.0
|
||||
)
|
||||
|
||||
assert list(actual_bytes) == list(expected_bytes), f'\nActual: {list(actual_bytes)}\nExpected: {list(expected_bytes)}'
|
||||
|
||||
print('All tests pass')
|
||||
finally:
|
||||
if os.path.isfile(test_file_path):
|
||||
os.remove(test_file_path)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -293,7 +293,7 @@ gpt_vocab::id gptj_sample_top_p_top_k(
|
|||
}
|
||||
}
|
||||
|
||||
gptj_sample_top_k(logits_id, top_k);
|
||||
gptj_sample_top_k(logits_id, top_k > 0 ? std::min(top_k, n_logits) : n_logits);
|
||||
|
||||
double maxl = -INFINITY;
|
||||
for (const auto & kv : logits_id) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue