The RWKV Language Model
https://github.com/BlinkDL/ChatRWKV
https://github.com/BlinkDL/RWKV-LM
# set these before import RWKV
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0' # if '1' then compile CUDA kernel for seq mode (much faster)
########################################################################################################
#
# Use '/' in model path, instead of '\'. Use ctx4096 models if you need long ctx.
#
# fp16 = good for GPU (!!! DOES NOT support CPU !!!)
# fp32 = good for CPU
# bf16 = worse accuracy, supports CPU
# xxxi8 (example: fp16i8) = xxx with int8 quantization to save 50% VRAM/RAM, slower, slightly less accuracy
#
# We consider [ln_out+head] to be an extra layer, so L12-D768 (169M) has "13" layers, L24-D2048 (1.5B) has "25" layers, etc.
# Strategy Examples: (device = cpu/cuda/cuda:0/cuda:1/...)
# 'cpu fp32' = all layers cpu fp32
# 'cuda fp16' = all layerscuda fp16
# 'cuda fp16i8' = all layers cuda fp16 with int8 quantization
# 'cuda fp16i8 *10 -> cpu fp32' = first 10 layers cuda fp16i8, then cpu fp32 (increase 10 for better speed)
# 'cuda:0 fp16 *10 -> cuda:1 fp16 *8 -> cpu fp32' = first 10 layers cuda:0 fp16, then 8 layers cuda:1 fp16, then cpu fp32
#
# Basic Strategy Guide:
# 100% VRAM = 'cuda fp16' # all layers cuda fp16
# 97% VRAM = 'cuda fp16i8 *1 -> cuda fp16' # first 1 layer cuda fp16i8, then cuda fp16
# 95% VRAM = 'cuda fp16i8 *2 -> cuda fp16' # first 2 layers cuda fp16i8, then cuda fp16
# 93% VRAM = 'cuda fp16i8 *3 -> cuda fp16' # first 3 layers cuda fp16i8, then cuda fp16
# ...
# 50% VRAM = 'cuda fp16i8' # all layers cuda fp16i8
# 48% VRAM = 'cuda fp16i8 -> cpu fp32 *1' # most layers cuda fp16i8, last 1 layer cpu fp32
# 46% VRAM = 'cuda fp16i8 -> cpu fp32 *2' # most layers cuda fp16i8, last 2 layers cpu fp32
# 44% VRAM = 'cuda fp16i8 -> cpu fp32 *3' # most layers cuda fp16i8, last 3 layers cpu fp32
# ...
# 0% VRAM = 'cpu fp32' # all layers cpu fp32
#
# Use '+' for STREAM mode, which can save VRAM too, and it is sometimes faster
# 'cuda fp16 *10+' = first 10 layers cuda fp16, then fp16 stream the rest to it (increase 10 for better speed)
# 'cuda fp16i8 *10 -> cuda fp16 *0+' = first 10 layers cuda fp16i8, then fp16 stream the rest to it
#
# Extreme STREAM: 3G VRAM is enough to run RWKV 14B (extremely slow. will be faster in future)
# 'cuda fp16 *0+ -> cpu fp32 *1' = stream all layers cuda fp16, last 1 layer [ln_out+head] cpu fp32
#
# ########################################################################################################
from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
# download models: https://huggingface.co/BlinkDL
model = RWKV(model='/fsx/BlinkDL/HF-MODEL/rwkv-4-pile-169m/RWKV-4-Pile-169M-20220807-8023', strategy='cpu fp32')
pipeline = PIPELINE(model, "20B_tokenizer.json") # 20B_tokenizer.json is in https://github.com/BlinkDL/ChatRWKV
ctx = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."
print(ctx, end='')
def my_print(s):
print(s, end='', flush=True)
# For alpha_frequency and alpha_presence, see "Frequency and presence penalties":
# https://platform.openai.com/docs/api-reference/parameter-details
args = PIPELINE_ARGS(temperature = 1.0, top_p = 0.7,
alpha_frequency = 0.25,
alpha_presence = 0.25,
token_ban = [0], # ban the generation of some tokens
token_stop = []) # stop generation whenever you see any token here
pipeline.generate(ctx, token_count=200, args=args, callback=my_print)
print('\n')
out, state = model.forward([187, 510, 1563, 310, 247], None)
print(out.detach().cpu().numpy()) # get logits
out, state = model.forward([187, 510], None)
out, state = model.forward([1563], state) # RNN has state (use deepcopy to clone states)
out, state = model.forward([310, 247], state)
print(out.detach().cpu().numpy()) # same result as above
print('\n')
1 maintainer
Metadata
1 maintainer
Release files for rwkv 0.0.8
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| rwkv-0.0.8.tar.gz | 17.3 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| rwkv-0.0.8-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 41.0 kB
Release files / rwkv-0.0.8.tar.gz
| Download URL | rwkv-0.0.8.tar.gz |
|---|---|
| Size | 17.3 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
84b5f88a6613d759cc891189b4e9fe460b72835a5a2593f43abec0a868da30f8
|
|
BLAKE2b-256 checksum How to use checksums |
1254be480502aa4f21a6851ba8bcfddbf07f3866e4e295eee3ecf2137189a7f5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.2 CPython/3.9.12
|
Release files / rwkv-0.0.8-py3-none-any.whl
| Download URL | rwkv-0.0.8-py3-none-any.whl |
|---|---|
| Size | 23.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
70a7e61542fd4b285659e56ec2327a6a2c34a9f3485afc87467e53607a148d31
|
|
BLAKE2b-256 checksum How to use checksums |
9186e0c4940db39cf8566e70368e27440e073a2a0932b0775837ebb15fadb685
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.2 CPython/3.9.12
|
