vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

integrated new version of rwkv from upstream

Browse source
This commit is contained in:
Concedo 2023-05-03 23:26:39 +08:00
parent 4857739ab5
commit 105f818d45
6 changed files with 315 additions and 12479 deletions
Download patch file Download diff file Expand all files Collapse all files

13
Makefile
View file

@ -235,9 +235,6 @@ 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 $@
@ -263,19 +260,19 @@ main: examples/main/main.cpp build-info.h ggml.o llama.o common.o $(OBJS)
@echo '==== Run ./main -h for help. ===='
@echo
koboldcpp: ggml.o ggml_rwkv.o ggml_v1.o expose.o common.o gpttype_adapter.o
koboldcpp: ggml.o ggml_v1.o expose.o common.o gpttype_adapter.o
$(DEFAULT_BUILD)
koboldcpp_openblas: ggml_openblas.o ggml_rwkv.o ggml_v1.o expose.o common.o gpttype_adapter.o
koboldcpp_openblas: ggml_openblas.o ggml_v1.o expose.o common.o gpttype_adapter.o
$(OPENBLAS_BUILD)
koboldcpp_noavx2: ggml_noavx2.o ggml_rwkv.o ggml_v1_noavx2.o expose.o common.o gpttype_adapter.o
koboldcpp_noavx2: ggml_noavx2.o ggml_v1_noavx2.o expose.o common.o gpttype_adapter.o
$(NOAVX2_BUILD)
koboldcpp_openblas_noavx2: ggml_openblas_noavx2.o ggml_rwkv.o ggml_v1_noavx2.o expose.o common.o gpttype_adapter.o
koboldcpp_openblas_noavx2: ggml_openblas_noavx2.o ggml_v1_noavx2.o expose.o common.o gpttype_adapter.o
$(OPENBLAS_NOAVX2_BUILD)
koboldcpp_clblast: ggml_clblast.o ggml_rwkv.o ggml_v1.o expose.o common.o gpttype_adapter.o ggml-opencl.o
koboldcpp_clblast: ggml_clblast.o ggml_v1.o expose.o common.o gpttype_adapter.o ggml-opencl.o
$(CLBLAST_BUILD)
quantize_llama: examples/quantize/quantize.cpp ggml.o llama.o

11588
otherarch/ggml_rwkv.c
View file

File diff suppressed because it is too large Load diff

645
otherarch/ggml_rwkv.h
View file

@ -1,645 +0,0 @@
#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

57
otherarch/otherarch.h
View file

@ -271,61 +271,4 @@ struct stablelm_model {
std::map<std::string, struct ggml_tensor *> tensors;
};
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
};

468
otherarch/rwkv.cpp
View file

@ -1,10 +1,11 @@
//adapted from RWKV.cpp repo under MIT license
// https://github.com/saharNooby/rwkv.cpp
#include "rwkv.h"
#include "ggml_rwkv.h"
#include "otherarch.h"
#include "rwkv.h"
#include "ggml.h"
#include <string>
#include <vector>
#include <thread>
@ -17,14 +18,10 @@
#include <iostream>
#include <unordered_map>
#include "model_adapter.h"
#include "rwkv_vocab.cpp"
// --- 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 { \
@ -51,22 +48,86 @@ bool read_int32(FILE * file, int32_t * dest) {
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
#define GGML_TYPE_UNKNOWN GGML_TYPE_COUNT
#define FORMAT_TYPE_COUNT 10
static const ggml_type FORMAT_TYPE_TO_GGML_TYPE[FORMAT_TYPE_COUNT] = {
GGML_TYPE_F32,
GGML_TYPE_F16,
GGML_TYPE_Q4_0,
GGML_TYPE_Q4_1,
GGML_TYPE_UNKNOWN, // Unused
GGML_TYPE_Q4_2,
GGML_TYPE_UNKNOWN, // Unused
GGML_TYPE_Q5_0,
GGML_TYPE_Q5_1,
GGML_TYPE_Q8_0
};
static int32_t format_name_to_format_type(const char * format_name) {
if (strcmp(format_name, "Q4_0") == 0) return 2;
if (strcmp(format_name, "Q4_1") == 0) return 3;
if (strcmp(format_name, "Q4_2") == 0) return 5;
if (strcmp(format_name, "Q5_0") == 0) return 7;
if (strcmp(format_name, "Q5_1") == 0) return 8;
if (strcmp(format_name, "Q8_0") == 0) return 9;
return -1;
}
// --- Model definition and loading utilities ---
struct rwkv_layer {
struct ggml_tensor * ln1_weight;
struct ggml_tensor * ln1_bias;
// RWKV, also called "attention" by the author.
struct ggml_tensor * att_time_mix_k;
struct ggml_tensor * att_time_mix_v;
struct ggml_tensor * att_time_mix_r;
struct ggml_tensor * att_time_first;
struct ggml_tensor * att_time_decay;
struct ggml_tensor * att_key;
struct ggml_tensor * att_value;
struct ggml_tensor * att_receptance;
struct ggml_tensor * att_output;
struct ggml_tensor * ln2_weight;
struct ggml_tensor * ln2_bias;
// FFN.
struct ggml_tensor * ffn_time_mix_k;
struct ggml_tensor * ffn_time_mix_r;
struct ggml_tensor * ffn_key;
struct ggml_tensor * ffn_value;
struct ggml_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_tensor * emb;
struct ggml_tensor * ln0_weight;
struct ggml_tensor * ln0_bias;
std::vector<rwkv_layer> layers;
struct ggml_tensor * ln_out_weight;
struct ggml_tensor * ln_out_bias;
struct ggml_tensor * head;
};
// 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];
bool set_parameter(std::unordered_map<std::string, struct ggml_tensor *> * parameters, char * key, struct ggml_tensor ** dest) {
struct ggml_tensor * parameter = (*parameters)[key];
RWKV_ASSERT_FALSE(parameter != NULL, "Parameter %s not found in model file", key);
*dest = parameter;
return true;
@ -74,7 +135,7 @@ bool set_parameter(std::unordered_map<std::string, struct ggml_rwkv_tensor *> *
// 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) {
bool set_block_parameter(std::unordered_map<std::string, struct ggml_tensor *> * parameters, int32_t block_index, char * key, struct ggml_tensor ** dest) {
char full_key[128];
sprintf(full_key, "blocks.%d.%s", block_index, key);
return set_parameter(parameters, full_key, dest);
@ -82,18 +143,70 @@ bool set_block_parameter(std::unordered_map<std::string, struct ggml_rwkv_tensor
// --- 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) {
void rwkv_exp_impl(const int n_cols, float * dest, const float * src) {
for (int i = 0; i < n_cols; i++) {
dest[i] = expf(src[i]);
}
}
void rwkv_1_minus_x_impl(const int n_cols, float * dest, const float * src) {
for (int i = 0; i < n_cols; i++) {
dest[i] = 1.0F - src[i];
}
}
void rwkv_sigmoid_impl(const int n_cols, float * dest, const float * src) {
for (int i = 0; i < n_cols; i++) {
dest[i] = 1.0F / (1.0F + expf(-src[i]));
}
}
void rwkv_max_impl(const int n_cols, float * dest, const float * src0, const float * src1) {
for (int i = 0; i < n_cols; i++) {
dest[i] = fmaxf(src0[i], src1[i]);
}
}
struct ggml_tensor * rwkv_exp(ggml_context * ctx, struct ggml_tensor * x) {
return ggml_map_unary_f32(ctx, x, rwkv_exp_impl);
}
struct ggml_tensor * rwkv_1_minus_x(ggml_context * ctx, struct ggml_tensor * x) {
return ggml_map_unary_f32(ctx, x, rwkv_1_minus_x_impl);
}
struct ggml_tensor * rwkv_sigmoid(ggml_context * ctx, struct ggml_tensor * x) {
return ggml_map_unary_f32(ctx, x, rwkv_sigmoid_impl);
}
struct ggml_tensor * rwkv_max(ggml_context * ctx, struct ggml_tensor * x, struct ggml_tensor * y) {
return ggml_map_binary_f32(ctx, x, y, rwkv_max_impl);
}
struct ggml_tensor * rwkv_layer_norm(ggml_context * ctx, struct ggml_tensor * x, struct ggml_tensor * weight, struct ggml_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);
// Looks like ggml_norm does the first part, we only need to apply weight & bias.
x = ggml_norm(ctx, x);
x = ggml_mul(ctx, x, weight);
x = ggml_add(ctx, x, bias);
return x;
}
// --- Implementation ---
struct rwkv_context {
struct rwkv_model * model;
struct ggml_tensor * token_index;
struct ggml_tensor * state;
struct ggml_tensor ** state_parts;
struct ggml_tensor * logits;
struct ggml_context * ctx;
struct ggml_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
};
struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_threads) {
FILE * file = fopen(file_path, "rb");
@ -119,14 +232,16 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
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 < FORMAT_TYPE_COUNT, "Unsupported model data type %d", 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
model->data_type != 4,
"Models in Q4_1_O format cannot be loaded anymore because the format was removed. You need to quantize the model into another format"
);
RWKV_ASSERT_NULL(
model->data_type != 6,
"Models in Q4_3 format cannot be loaded anymore because the format was removed. You need to quantize the model into another format"
);
// Parameter tensors would take at least this amount in memory.
@ -152,12 +267,13 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
size_t(256) * 1024 * 1024;
// Initialize ggml
struct ggml_rwkv_init_params params;
struct ggml_init_params params;
params.mem_size = memory_required;
params.mem_buffer = NULL;
struct ggml_rwkv_context * ctx = ggml_rwkv_init(params);
params.no_alloc = false;
struct ggml_context * ctx = ggml_init(params);
std::unordered_map<std::string, struct ggml_rwkv_tensor *> parameters;
std::unordered_map<std::string, struct ggml_tensor *> parameters;
while (true) {
int32_t dim_count;
@ -176,30 +292,24 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
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
);
RWKV_ASSERT_NULL(data_type >= 0 && data_type < FORMAT_TYPE_COUNT, "Unsupported parameter data type %d", data_type);
ggml_rwkv_type ggml_rwkv_data_type = FORMAT_TYPE_TO_GGML_RWKV_TYPE[data_type];
ggml_type ggml_data_type = FORMAT_TYPE_TO_GGML_TYPE[data_type];
struct ggml_rwkv_tensor * tensor;
RWKV_ASSERT_NULL(ggml_data_type != GGML_TYPE_UNKNOWN, "Unsupported parameter data type %d", data_type);
struct ggml_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);
tensor = ggml_new_tensor_1d(ctx, ggml_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);
tensor = ggml_new_tensor_2d(ctx, ggml_data_type, x, y);
} else {
abort();
}
@ -207,7 +317,7 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
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");
RWKV_ASSERT_NULL(fread(tensor->data, 1, ggml_nbytes(tensor), file) == ggml_nbytes(tensor), "Failed to read parameter data");
parameters[key] = tensor;
}
@ -255,26 +365,26 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
set_parameter(&parameters, "head.weight", &(model->head));
// Verify order of dimensions
struct ggml_rwkv_tensor * emb = model->emb;
struct ggml_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]);
RWKV_ASSERT_NULL(emb->ne[0] == model->n_embed, "Unexpected dimension of embedding matrix %lld", emb->ne[0]);
RWKV_ASSERT_NULL(emb->ne[1] == model->n_vocab, "Unexpected dimension of embedding matrix %lld", 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);
struct ggml_tensor * state = ggml_new_tensor_1d(ctx, GGML_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);
struct ggml_tensor * token_index = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
struct ggml_tensor * x = ggml_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];
struct ggml_tensor ** state_parts = new ggml_tensor * [n_layer * 5];
for (int i = 0; i < n_layer; i++) {
auto layer = model->layers[i];
@ -282,99 +392,99 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
// 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);
struct ggml_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);
struct ggml_tensor * x_prev = ggml_view_1d(ctx, state, n_embed, (5 * i + 1) * n_embed * sizeof(float));
// 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(
struct ggml_tensor * xk = ggml_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))
ggml_mul(ctx, x0, layer.att_time_mix_k),
ggml_mul(ctx, x_prev, rwkv_1_minus_x(ctx, layer.att_time_mix_k))
);
struct ggml_rwkv_tensor * xv = ggml_rwkv_add(
struct ggml_tensor * xv = ggml_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))
ggml_mul(ctx, x0, layer.att_time_mix_v),
ggml_mul(ctx, x_prev, rwkv_1_minus_x(ctx, layer.att_time_mix_v))
);
struct ggml_rwkv_tensor * xr = ggml_rwkv_add(
struct ggml_tensor * xr = ggml_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))
ggml_mul(ctx, x0, layer.att_time_mix_r),
ggml_mul(ctx, x_prev, 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(
struct ggml_tensor * r = rwkv_sigmoid(
ctx,
ggml_rwkv_mul_mat(ctx, layer.att_receptance, xr)
ggml_mul_mat(ctx, layer.att_receptance, xr)
);
// k = kw @ xk
struct ggml_rwkv_tensor * k = ggml_rwkv_mul_mat(ctx, layer.att_key, xk);
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.att_key, xk);
// v = vw @ xv
struct ggml_rwkv_tensor * v = ggml_rwkv_mul_mat(ctx, layer.att_value, xv);
struct ggml_tensor * v = ggml_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);
struct ggml_tensor * aa = ggml_view_1d(ctx, state, n_embed, (5 * i + 2) * n_embed * sizeof(float));
struct ggml_tensor * bb = ggml_view_1d(ctx, state, n_embed, (5 * i + 3) * n_embed * sizeof(float));
struct ggml_tensor * pp = ggml_view_1d(ctx, state, n_embed, (5 * i + 4) * n_embed * sizeof(float));
// ww = time_first + k
struct ggml_rwkv_tensor * ww = ggml_rwkv_add(ctx, layer.att_time_first, k);
struct ggml_tensor * ww = ggml_add(ctx, layer.att_time_first, k);
// qq = torch.maximum(pp, ww)
struct ggml_rwkv_tensor * qq = ggml_rwkv_max(ctx, pp, ww);
struct ggml_tensor * qq = 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));
struct ggml_tensor * e1 = rwkv_exp(ctx, ggml_sub(ctx, pp, qq));
// e2 = torch.exp(ww - qq)
struct ggml_rwkv_tensor * e2 = ggml_rwkv_exp(ctx, ggml_rwkv_sub(ctx, ww, qq));
struct ggml_tensor * e2 = rwkv_exp(ctx, ggml_sub(ctx, ww, qq));
// a = e1 * aa + e2 * v
struct ggml_rwkv_tensor * a = ggml_rwkv_add(
struct ggml_tensor * a = ggml_add(
ctx,
ggml_rwkv_mul(ctx, e1, aa),
ggml_rwkv_mul(ctx, e2, v)
ggml_mul(ctx, e1, aa),
ggml_mul(ctx, e2, v)
);
// b = e1 * bb + e2
struct ggml_rwkv_tensor * b = ggml_rwkv_add(
struct ggml_tensor * b = ggml_add(
ctx,
ggml_rwkv_mul(ctx, e1, bb),
ggml_mul(ctx, e1, bb),
e2
);
// wkv = a / b
struct ggml_rwkv_tensor * wkv = ggml_rwkv_div(ctx, a, b);
struct ggml_tensor * wkv = ggml_div(ctx, a, b);
// ww = pp + time_decay
ww = ggml_rwkv_add(ctx, pp, layer.att_time_decay);
ww = ggml_add(ctx, pp, layer.att_time_decay);
// qq = torch.maximum(ww, k)
qq = ggml_rwkv_max(ctx, ww, k);
qq = rwkv_max(ctx, ww, k);
// e1 = torch.exp(ww - qq)
e1 = ggml_rwkv_exp(ctx, ggml_rwkv_sub(ctx, ww, qq));
e1 = rwkv_exp(ctx, ggml_sub(ctx, ww, qq));
// e2 = torch.exp(k - qq)
e2 = ggml_rwkv_exp(ctx, ggml_rwkv_sub(ctx, k, qq));
e2 = rwkv_exp(ctx, ggml_sub(ctx, k, qq));
// state[5 * i + 2] = e1 * aa + e2 * v
state_parts[5 * i + 2] = ggml_rwkv_add(
state_parts[5 * i + 2] = ggml_add(
ctx,
ggml_rwkv_mul(ctx, e1, aa),
ggml_rwkv_mul(ctx, e2, v)
ggml_mul(ctx, e1, aa),
ggml_mul(ctx, e2, v)
);
// state[5 * i + 3] = e1 * bb + e2
state_parts[5 * i + 3] = ggml_rwkv_add(
state_parts[5 * i + 3] = ggml_add(
ctx,
ggml_rwkv_mul(ctx, e1, bb),
ggml_mul(ctx, e1, bb),
e2
);
// state[5 * i + 4] = qq
state_parts[5 * i + 4] = qq;
// ow @ (r * wkv)
x = ggml_rwkv_add(
x = ggml_add(
ctx,
x,
ggml_rwkv_mul_mat(
ggml_mul_mat(
ctx,
layer.att_output,
ggml_rwkv_mul(ctx, r, wkv)
ggml_mul(ctx, r, wkv)
)
);
}
@ -382,42 +492,42 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
// 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);
struct ggml_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);
struct ggml_tensor * x_prev = ggml_view_1d(ctx, state, n_embed, (5 * i + 0) * n_embed * sizeof(float));
// 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(
struct ggml_tensor * xk = ggml_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))
ggml_mul(ctx, x0, layer.ffn_time_mix_k),
ggml_mul(ctx, x_prev, rwkv_1_minus_x(ctx, layer.ffn_time_mix_k))
);
struct ggml_rwkv_tensor * xr = ggml_rwkv_add(
struct ggml_tensor * xr = ggml_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))
ggml_mul(ctx, x0, layer.ffn_time_mix_r),
ggml_mul(ctx, x_prev, 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(
struct ggml_tensor * r = rwkv_sigmoid(
ctx,
ggml_rwkv_mul_mat(ctx, layer.ffn_receptance, xr)
ggml_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(
struct ggml_tensor * k = ggml_sqr(ctx, ggml_relu(
ctx,
ggml_rwkv_mul_mat(ctx, layer.ffn_key, xk)
ggml_mul_mat(ctx, layer.ffn_key, xk)
));
// r * (vw @ k)
x = ggml_rwkv_add(
x = ggml_add(
ctx,
x,
ggml_rwkv_mul(
ggml_mul(
ctx,
r,
ggml_rwkv_mul_mat(ctx, layer.ffn_value, k)
ggml_mul_mat(ctx, layer.ffn_value, k)
)
);
}
@ -427,14 +537,14 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr
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_tensor * logits = ggml_mul_mat(ctx, model->head, x);
struct ggml_rwkv_cgraph * graph = (struct ggml_rwkv_cgraph *) calloc(1, sizeof(struct ggml_rwkv_cgraph));
struct ggml_cgraph * graph = (struct ggml_cgraph *) calloc(1, sizeof(struct ggml_cgraph));
*graph = ggml_rwkv_build_forward(logits);
*graph = ggml_build_forward(logits);
for (int i = 0; i < n_layer * 5; i++) {
ggml_rwkv_build_forward_expand(graph, state_parts[i]);
ggml_build_forward_expand(graph, state_parts[i]);
}
graph->n_threads = n_threads;
@ -468,51 +578,59 @@ bool rwkv_eval(struct rwkv_context * ctx, int32_t token, float * state_in, float
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);
ggml_set_i32_1d(ctx->token_index, 0, token);
if (state_in == NULL) {
ggml_rwkv_set_f32(ctx->state, 0.0F);
ggml_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),
ggml_set_f32(
ggml_view_1d(ctx->ctx, ctx->state, n_embed, (5 * i + 4) * n_embed * sizeof(float)),
-1e30F
);
}
} else {
memcpy(ctx->state->data, state_in, ctx->state->ne[0] * FP32_SIZE);
memcpy(ctx->state->data, state_in, ctx->state->ne[0] * sizeof(float));
}
ggml_rwkv_graph_compute(ctx->ctx, ctx->graph);
ggml_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];
struct ggml_tensor * part = ctx->state_parts[i];
memcpy(state_out + i * n_embed, part->data, part->ne[0] * FP32_SIZE);
memcpy(state_out + i * n_embed, part->data, part->ne[0] * sizeof(float));
}
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);
memcpy(logits_out, ctx->logits->data, ctx->logits->ne[0] * sizeof(float));
return true;
}
void rwkv_free(struct rwkv_context * ctx) {
ggml_rwkv_free(ctx->ctx);
delete ctx->model;
delete ctx->state_parts;
delete ctx;
ctx->model->layers.~vector();
free(ctx->model);
delete[] ctx->state_parts;
ggml_free(ctx->ctx);
free(ctx->graph);
free(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);
bool rwkv_quantize_model_file(const char * model_file_path_in, const char * model_file_path_out, const char * format_name) {
int32_t format_type = format_name_to_format_type(format_name);
ggml_rwkv_type type = FORMAT_TYPE_TO_GGML_RWKV_TYPE[q_type];
RWKV_ASSERT_FALSE(format_type != -1, "Unsupported format \"%s\"", format_name);
ggml_type type = FORMAT_TYPE_TO_GGML_TYPE[format_type];
RWKV_ASSERT_FALSE(type != GGML_TYPE_UNKNOWN, "Unsupported format \"%s\"", format_name);
// Needed to initialize FP16 lookup table
{
struct ggml_init_params params = { 0, NULL, false };
struct ggml_context * ctx = ggml_init(params);
ggml_free(ctx);
}
printf("Loading model from '%s'\n", model_file_path_in);
@ -546,7 +664,7 @@ bool rwkv_quantize_model_file(const char * model_file_path_in, const char * mode
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;
data_type = format_type;
fout.write((char *) &n_vocab, sizeof(n_vocab));
fout.write((char *) &n_embed, sizeof(n_embed));
@ -562,7 +680,7 @@ bool rwkv_quantize_model_file(const char * model_file_path_in, const char * mode
std::vector<float> work;
std::vector<uint8_t> data_u8;
std::vector<ggml_rwkv_fp16_t> data_f16;
std::vector<ggml_fp16_t> data_f16;
std::vector<float> data_f32;
std::vector<int64_t> hist_all(1 << 4, 0);
@ -580,6 +698,12 @@ bool rwkv_quantize_model_file(const char * model_file_path_in, const char * mode
break;
}
RWKV_ASSERT_FALSE(parameter_data_type >= 0 && parameter_data_type < FORMAT_TYPE_COUNT, "Invalid parameter data type %d", parameter_data_type);
ggml_type parameter_ggml_type = FORMAT_TYPE_TO_GGML_TYPE[parameter_data_type];
RWKV_ASSERT_FALSE(parameter_ggml_type != GGML_TYPE_UNKNOWN, "Invalid parameter data type %d", parameter_data_type);
int32_t nelements = 1;
int32_t ne[2] = { 1, 1 };
for (int i = 0; i < n_dims; ++i) {
@ -591,16 +715,9 @@ bool rwkv_quantize_model_file(const char * model_file_path_in, const char * mode
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]);
printf("%48s - [%5d, %5d], type = %6s ", name.data(), ne[0], ne[1], ggml_type_name(parameter_ggml_type));
total_size_orig += (size_t) (nelements * ggml_rwkv_type_sizef(FORMAT_TYPE_TO_GGML_RWKV_TYPE[parameter_data_type]));
total_size_orig += (size_t) (nelements * ggml_type_sizef(parameter_ggml_type));
}
// Quantize only 2D tensors, except embedding and head matrices.
@ -619,17 +736,17 @@ bool rwkv_quantize_model_file(const char * model_file_path_in, const char * mode
if (parameter_data_type == 1) {
data_f16.resize(nelements);
finp.read(reinterpret_cast<char *>(data_f16.data()), nelements * sizeof(ggml_rwkv_fp16_t));
finp.read(reinterpret_cast<char *>(data_f16.data()), nelements * sizeof(ggml_fp16_t));
data_f32.resize(nelements);
for (int i = 0; i < nelements; ++i) {
data_f32[i] = ggml_rwkv_fp16_to_fp32(data_f16[i]);
data_f32[i] = ggml_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;
parameter_data_type = format_type;
} else {
const int bytes_per_element = (parameter_data_type == 0) ? sizeof(float) : sizeof(uint16_t);
data_u8.resize(nelements * bytes_per_element);
@ -655,20 +772,25 @@ bool rwkv_quantize_model_file(const char * model_file_path_in, const char * mode
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:
{
case GGML_TYPE_Q4_0:
cur_size = ggml_quantize_q4_0(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
break;
case GGML_TYPE_Q4_1:
cur_size = ggml_quantize_q4_1(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
break;
case GGML_TYPE_Q4_2:
cur_size = ggml_quantize_q4_2(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
break;
case GGML_TYPE_Q5_0:
cur_size = ggml_quantize_q5_0(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
break;
case GGML_TYPE_Q5_1:
cur_size = ggml_quantize_q5_1(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
break;
case GGML_TYPE_Q8_0:
cur_size = ggml_quantize_q8_0(data_f32.data(), work.data(), nelements, ne[0], hist_cur.data());
break;
default: {
fprintf(stderr, "unsupported quantization type %d\n", type);
return false;
}
@ -697,7 +819,7 @@ bool rwkv_quantize_model_file(const char * model_file_path_in, const char * mode
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);
printf("compression ratio = %8.2f\n", 1.0 * total_size_orig / total_size_new);
{
int64_t sum_all = 0;
@ -726,18 +848,18 @@ 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()) + " | ";
s += "AVX = " + std::to_string(ggml_cpu_has_avx()) + " | ";
s += "AVX2 = " + std::to_string(ggml_cpu_has_avx2()) + " | ";
s += "AVX512 = " + std::to_string(ggml_cpu_has_avx512()) + " | ";
s += "FMA = " + std::to_string(ggml_cpu_has_fma()) + " | ";
s += "NEON = " + std::to_string(ggml_cpu_has_neon()) + " | ";
s += "ARM_FMA = " + std::to_string(ggml_cpu_has_arm_fma()) + " | ";
s += "F16C = " + std::to_string(ggml_cpu_has_f16c()) + " | ";
s += "FP16_VA = " + std::to_string(ggml_cpu_has_fp16_va()) + " | ";
s += "WASM_SIMD = " + std::to_string(ggml_cpu_has_wasm_simd()) + " | ";
s += "BLAS = " + std::to_string(ggml_cpu_has_blas()) + " | ";
s += "SSE3 = " + std::to_string(ggml_cpu_has_sse3()) + " | ";
s += "VSX = " + std::to_string(ggml_cpu_has_vsx()) + " | ";
return s.c_str();
}

13
otherarch/rwkv.h
View file

@ -52,12 +52,19 @@ extern "C" {
// 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.
// Quantizes FP32 or FP16 model to one of quantized 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);
// - format_name: must be one of available format names below.
// Available format names:
// - Q4_0
// - Q4_1
// - Q4_2
// - Q5_0
// - Q5_1
// - Q8_0
RWKV_API bool rwkv_quantize_model_file(const char * model_file_path_in, const char * model_file_path_out, const char * format_name);
// Returns system information string.
RWKV_API const char * rwkv_get_system_info_string(void);