rwkv-architecture

1# RWKV - Receptance Weighted Key Value
2 
3## Quick start
4 
5RWKV (RwaKuv) combines Transformer parallelization (training) with RNN efficiency (inference).
6 
7**Installation**:
8```bash
9# Install PyTorch
10pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121
11 
12# Install dependencies
13pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade
14 
15# Install RWKV
16pip install rwkv
17```
18 
19**Basic usage** (GPT mode + RNN mode):
20```python
21import os
22from rwkv.model import RWKV
23 
24os.environ["RWKV_JIT_ON"] = '1'
25os.environ["RWKV_CUDA_ON"] = '1'  # Use CUDA kernel for speed
26 
27# Load model
28model = RWKV(
29    model='/path/to/RWKV-4-Pile-1B5-20220903-8040',
30    strategy='cuda fp16'
31)
32 
33# GPT mode (parallel processing)
34out, state = model.forward([187, 510, 1563, 310, 247], None)
35print(out.detach().cpu().numpy())  # Logits
36 
37# RNN mode (sequential processing, same result)
38out, state = model.forward([187, 510], None)  # First 2 tokens
39out, state = model.forward([1563], state)      # Next token
40out, state = model.forward([310, 247], state)  # Last tokens
41print(out.detach().cpu().numpy())  # Same logits as above!
42```
43 
44## Common workflows
45 
46### Workflow 1: Text generation (streaming)
47 
48**Efficient token-by-token generation**:
49```python
50from rwkv.model import RWKV
51from rwkv.utils import PIPELINE
52 
53model = RWKV(model='RWKV-4-Pile-14B-20230313-ctx8192-test1050', strategy='cuda fp16')
54pipeline = PIPELINE(model, "20B_tokenizer.json")
55 
56# Initial prompt
57prompt = "The future of AI is"
58…