Package: rwkvstic
RWKVSTICRwkvstic, pronounced however you want to, is a library for interfacing and using the RWKV-V4 based models.Rwkvstic does not autoinstall its dependencies, as its main purpose is to be dependency agnostic, able to be used by whatever library you would prefer.When using BlinkDLs pretrained models, it would advised to have thetorchpackage installed.Some options, when left blank, will elicit a prompt asking you to choose a value.
for this purpose, please ensure you have theinquirerpackage installed.Noteas of RWKVSTIC 2.0, the default mode is GPU with FASTQUANT, my own custom implementation of strategy="cuda fp32i8".Please checkout the strategy section onRWKVfor other strategies, or look at the advanced modes below.Tables and graphsRwkv-4 models -> recomended vramrwkvstic vram
Model | 8bit | bf16/fp16 | fp32
14B   | 16GB | 28GB      | >50GB
7B    | 8GB  | 14GB      | 28GB
3B    | 2.8GB| 6GB       | 12GB
1b5   | 1.3GB| 3GB       | 6GBInstallationpipinstallrwkvsticBasic Usagefromrwkvstic.loadimportRWKV# Load the model (supports full path, relative path, and remote paths)model=RWKV("https://huggingface.co/BlinkDL/rwkv-4-pile-3b/resolve/main/RWKV-4-Pile-3B-Instruct-test1-20230124.pth")model.loadContext(newctx=f"Q: who is Jim Butcher?\n\nA:")output=model.forward(number=100)["output"]print(output)# Q: who is Jim Butcher?# A: Jim Butcher is a very popular American author of fantasy novels. He’s known for the Dresden Files series of novels.
RWKV wrapperYou can use any compatible rwkv strategy string to overwrite the default behavior with the original BlinkDL packagemodel = RWKV(
    "https://huggingface.co/BlinkDL/rwkv-4-pile-3b/resolve/main/RWKV-4-Pile-3B-Instruct-test1-20230124.pth",
    strategy="cuda fp32"
)ExportingYou can export the default FASTQUANT mode for quick downloading and loading, as it has a smaller file size and uses less Ram and Disk Spacemodel=RWKV("https://huggingface.co/BlinkDL/rwkv-4-pile-3b/resolve/main/RWKV-4-Pile-3B-Instruct-test1-20230124.pth",export="myfilename")# exported model as myfilename.rwkvmodel=RWKV("myfile.rwkv",)Advanced UsageStep 1: load the model with your choice of poisonPytorchfromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportTORCH# this is the dtype used for trivial operations, such as vector->vector operations and is the dtype that will determine the accuracy of the modelruntimedtype=torch.float32# torch.float64, torch.bfloat16# this is the dtype used for matrix-vector operations, and is the dtype that will determine the performance and memory usage of the modeldtype=torch.bfloat16# torch.float32, torch.float64, torch.bfloat16useGPU=True# Falsemodel=RWKV("path/to/model.pth",mode=TORCH,useGPU=useGPU,runtimedtype=runtimedtype,dtype=dtype)JAXfromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportJAX# Jax will automatically use the GPU if available, and will use the CPU if not availablemodel=RWKV("path/to/model.pth",mode=JAX)TensorFlowfromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportTFuseGPU=True# Falsemodel=RWKV("path/to/model.pth",mode=TF,useGPU=useGPU)Numpyfromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportNUMPY# you masochistic bastardmodel=RWKV("path/to/model.pth",mode=NUMPY)StreamingTrade vram usage for performancefromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportTORCH_STREAM# this is the dtype used for trivial operations, such as vector->vector operations and is the dtype that will determine the accuracy of the modelruntime_dtype=torch.float32# torch.float64, torch.bfloat16# this is the dtype used for matrix-vector operations, and is the dtype that will determine the performance and memory usage of the modeldtype=torch.bfloat16# torch.float32, torch.float64, torch.bfloat16# this is the amount of GB you want to use for matrix storage, if the model is too large, matrixes will be stored in ram and moved to the GPU as neededtarget=4# Pin Memory is used to speed up the transfer of data to the GPU, but will use more memory, both on the GPU and on the CPUpin_memory=Truemodel=RWKV("path/to/model.pth",mode=TORCH_STREAM,runtimedtype=runtime_dtype,dtype=dtype,target=target,pinMem=pin_memory)Multi-GPUModel weights are split(sharded) across multiple GPUsfromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportTORCH_SPLIT# this is the dtype used for trivial operations, such as vector->vector operations and is the dtype that will determine the accuracy of the modelruntime_dtype=torch.float32# torch.float64, torch.bfloat16# this is the dtype used for matrix-vector operations, and is the dtype that will determine the performance and memory usage of the modeldtype=torch.bfloat16# torch.float32, torch.float64, torch.bfloat16model=RWKV("path/to/model.pth",mode=TORCH_SPLIT,runtimedtype=runtime_dtype,dtype=dtype)QuantizationUses close to half the memory of float16, but is slightly less accurate, and is about 4x slowerfromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportTORCH_QUANT# this is the dtype used for trivial operations, such as vector->vector operations and is the dtype that will determine the accuracy of the modelruntime_dtype=torch.float32# torch.float64, torch.bfloat16# this is the amount of chunks to split the matrix rows into pre-row-quantization, the more chunks, the more accurate the model will be, but with some minor trade offschunksize=4useGPU=True# False# this is the amount of GB you want to use for matrix storage, if the model is too large, matrixes will be stored in ram and moved to the GPU as needed, same as streamtarget=4model=RWKV("path/to/model.pth",mode=TORCH_QUANT,runtimedtype=runtime_dtype,chunksize=chunksize,useGPU=useGPU,target=target)Step 2: State managementThe stateThe state is a vectorized value that is a representation of all the previous inputs and outputs of the model. It is used basically the memory of the model, and is used to generate the next output.The model has an internal state, so the following is useful in that regards.model=RWKV("path/to/model.pth")emptyState=model.emptyStatemodel.setState(emptyState)currentMem=model.getState()Step 3: Injecting contextInjecting contextWhen you want to influence the output of the model, you can inject context into the model. This is done by using theloadContextfunction.model=RWKV("path/to/model.pth")model.loadContext(newctx="Q: who is Jim Butcher?\n\nA:")print(model.forward(number=100)["output"])model.loadContext(newctx="Can you tell me more?\n\nA:")Step 4: Generating outputGenerating outputWhen you want to generate output, you can use theforwardfunction.model=RWKV("path/to/model.pth")number=100# the number of tokens to generatestopStrings=["\n\n"]# When read, the model will stop generating outputstopTokens=[0]# advanced, when the model has generated any of these tokens, it will stop generating outputtemp=1# the temperature of the model, higher values will result in more random output, lower values will result in more predictable outputtop_p=0.9# the top_p of the model, higher values will result in more random output, lower values will result in more predictable outputdefprogressLambda(properties):# "logits", "state", "output", "progress", "tokens", "total", "current"print("progress:",properties["progress"]/properties["total"])output=model.forward(number=number,stopStrings=stopStrings,stopTokens=stopTokens,temp=temp,top_p=top_p,progressLambda=progressLambda)print(output["output"])# the generated outputprint(output["state"])# the state of the model after generationprint(output["logits"])# the logits of the model after generation, before samplingImplementation DetailsThe RWKVOP objectHere is a base class, when overwritten, will allow the swapout of operations with their equivilents in different frameworks. Ill show you the JAX one, as its relatively simpleclassRWKVJaxOps(RWKVOp.module):def__init__(self,layers,embed,preJax=False):fromjaximportnumpyasnpjaxsuper().__init__(layers,embed)# convert from torch to jaxself.initTensor=lambdax:npjax.array(x.float().cpu().numpy())# jax math functionsself.sqrt=lambdax:npjax.sqrt(x)self.mean=lambdax:npjax.mean(x)self.relu=lambdax:npjax.maximum(x,0)self.exp=lambdax:npjax.exp(x)self.matvec=npjax.matmulself.lerp=lambdax,y,z:x*(1-z)+y*(z)self.minimum=lambdax,y:npjax.minimum(x,y)self.log=npjax.logdefln(x,w,b):xee2=x-self.mean(x)x2=self.sqrt(self.mean(xee2*xee2)+0.000009999999747378752)returnw*(xee2/x2)+bself.layernorm=ln# constants and stuffself.klimit=npjax.array([18]*embed)self.stack=lambdax:x# module defself.module=object# function overwrites (used for advanced stuff)self.initfunc=lambdax:xself.layerdef=lambdax:xself.mainfunc=lambdax:x# The empty stateself.emptyState=npjax.array([[0.01]*embed]*4*layers)This can then be used to construct and infer the model.Stream, Split And QuantThe stream, split and quant backends are all pytorch varients that use some tricks to use less, or distribute memory usage across multiple GPUs.Ill show you the important stuff, usually consisting of how the matrixes are constructed, and how they are used to create a matvec.(Disclaimer, just similar to the actual code, not the actual code, actual code is messy and gross)Stream# Pinning memory allows for faster transfer between CPU and GPU, but uses more memorydefpinmem(x):returnx.pin_memory()ifpinMemandx.device=="cpu"elsexdefinitMatrix(x):# if more memory is used then the target specified, then it is sent to the cpuiftorch.cuda.max_memory_reserved(0)/1024/1024/1024>target:x=x.cpu()else:x=x.cuda(non_blocking=True)returnpinmem(x)# for the matvec, it just brings it to the correct device as neededdefmatvec(z,y):returnz.to(y.device,non_blocking=True)@ySplitdefinitMatrix(x):devices=[torch.device("cuda",i)foriinrange(torch.cuda.device_count())]# split the matrix into the number of devicesx=torch.split(x,x.shape[0]//len(devices),dim=0)# send each part to a different devicex=[i.to(devices[i],non_blocking=True)foriinrange(len(x))]returnx# for the matvec, split the vector into the number of devices, and then send each part to the correct devicedefmatvec(z,y):devices=[torch.device("cuda",i)foriinrange(torch.cuda.device_count())]y=torch.split(y,y.shape[0]//len(devices),dim=0)y=[i.to(devices[i],non_blocking=True)foriinrange(len(y))]# do the matvec on each partz=[z[i].mv(y[i])foriinrange(len(z))]# put them all on one devicez=[i.to(devices[0],non_blocking=True)foriinz]# add them all togetherz=torch.sum(torch.stack(z),dim=0)returnzQuantdefQuantizeMatrix(x,runtimeDtype,device):rang=255ran,mini=(x.max(0)[0]-x.min(0)[0])/rang,x.min(0)[0]x=x.double()x=((x-mini)/ran)x=x.to(dtype=torch.uint8,non_blocking=True,device=device)returnx,ran.to(runtimeDtype).to(device=device),mini.to(runtimeDtype).to(device=device)defMatVec(x,y,runtimedtype):# resize y into a 2d arrayy=y.reshape(chunksize,-1)# retrieve the  quantized matrix, the spread, and the offsetrx,spread,zpoint=x# spread the y vector across the spread matrixyy=y*spread# convert the quantized matrix back to the runtime dtyperx=rx.to(dtype=runtimedtype)# we can use matmul to do a batched matvec for each split matrixxmain=rx.matmul(yy.reshape(yy.shape[0],-1,1)).sum(0).squeeze()# the offset is added to the resultreturnxmain+torch.tensordot(zpoint,y)definitMatrix(x):# by splitting the matrix before quantizing, it allows for much better resultssplitmatrices=torch.chunk(x,chunksize,1)xx=[QuantizeMatrix(x,runtimedtype,dev)forxinsplitmatrices]xxo=torch.stack([x[0]forxinxx])xx1=torch.stack([x[1]forxinxx])xx2=torch.stack([x[2]forxinxx])returnxxo,xx1,xx2PreQuantizationYou can prequantize the matrixes to save loading time, and bandwidth when downloading model.cd/path/to/folder/with/model
python3-mrwkvstic--pq# what model to prequantize?# -> model.pthls# model.pth# model.pqthYou can load these pre-quantized models as you would a normal file.fromrwkvstic.loadimportRWKVmodel=RWKV("model.pqth")Onnx exportYou can export the model to onnx, and then use onnxruntime/rwkvstic to infer the model.fromrwkvstic.loadimportRWKVfromrwkvstic.agnostic.backendsimportONNX_EXPORTimporttorchmodel=RWKV("model.pth",backend=ONNX_EXPORT,dtype=torch.float16)# or torch.float32# the model is exported to model_{layers}_{embed}.onnx# the external data is stored in model_{layers}_{embed}.binrwkvstic onnx runningfromrwkvstic.loadimportRWKVmodel=RWKV("model_12_768.onnx")