Rwkv architecture

--- name: rwkv-architecture description: RNN+Transformer hybrid with O(n) inference. Linear time, infinite context, no KV cache. Trains like GPT (parallel), infers like RNN (sequential). Linux Foundation AI project. In production at Windows, Office, and NeMo. RWKV-7 (March 2025). Models with up to 14B parameters. version: 1.0.0 author: Orchestra Research license: MIT tags: [RWKV, Model Architecture, RNN, Transformer Hybrid, Linear Complexity, Infinite Context, Efficient Inference, Linux Foundation, Alternative Architecture] dependencies: [rwkv, torch, transformers] --- # RWKV - Receptance Weighted Key Value ## Quick Start RWKV (RwaKuv) combines Transformer parallelization (training) with RNN efficiency (inference). Installation: ``bash # Install PyTorch pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121 # Install dependencies pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade # Install RWKV pip install rwkv ` **Basic usage** (GPT mode + RNN mode): `python import os from rwkv.model import RWKV os.environ["RWKV_JIT_ON"] = '1' os.environ["RWKV_CUDA_ON"] = '1' # Use CUDA kernel for speed # Load model model = RWKV( model='/path/to/RWKV-4-Pile-1B5-20220903-8040', strategy='cuda fp16' ) # GPT mode (parallel processing) out, state = model.forward([187, 510, 1563, 310, 247], None) print(out.detach().cpu().numpy()) # Logits # RNN mode (sequential processing, same result) out, state = model.forward([187, 510], None) # First 2 tokens out, state = model.forward([1563], state) # Next token out, state = model.forward([310, 247], state) # Last tokens print(out.detach().cpu().numpy()) # Same logits as above! ` ## Common workflows ### Workflow 1: Text generation (streaming) **Efficient token-by-token generation**: `python from rwkv.model import RWKV from rwkv.utils import PIPELINE model = RWKV(model='RWKV-4-Pile-14B-20230313-ctx8192-test1050', strategy='cuda fp16') pipeline = PIPELINE(model, "20B_tokenizer.json") # Initial prompt prompt = "The future of AI is" state = None # Generate token by token for token in prompt: out, state = pipeline.model.forward(pipeline.encode(token), state) # Continue generation for _ in range(100): out, state = pipeline.model.forward(None, state) token = pipeline.sample_logits(out) print(pipeline.decode(token), end='', flush=True) ` **Key advantage**: Constant memory per token (no growing KV cache) ### Workflow 2: Long context processing (infinite context) **Process million-token sequences**: `python model = RWKV(model='RWKV-4-Pile-14B', strategy='cuda fp16') # Process very long document state = None long_document = load_document() # e.g., 1M tokens # Stream through entire document for chunk in chunks(long_document, chunk_size=1024): out, state = model.forward(chunk, state) # State now contains information from entire 1M token document # Memory usage: O(1) (constant, not O(n)!) ` ### Workflow 3: Fine-tuning RWKV **Standard fine-tuning workflow**: `python # Training script import pytorch_lightning as pl from rwkv.model import RWKV from rwkv.trainer import RWKVTrainer # Configure model config = { 'n_layer': 24, 'n_embd': 1024, 'vocab_size': 50277, 'ctx_len': 1024 } # Setup trainer trainer = pl.Trainer( accelerator='gpu', devices=8, precision='bf16', strategy='deepspeed_stage_2', max_epochs=1 ) # Train model = RWKV(config) trainer.fit(model, train_dataloader) ` ### Workflow 4: RWKV vs Transformer comparison **Memory comparison** (1M token sequence): `python # Transformer (GPT) # Memory: O(n²) for attention # KV cache: 1M × hidden_dim × n_layers × 2 (keys + values) # Example: 1M × 4096 × 24 × 2 = ~400GB (impractical!) # RWKV # Memory: O(1) per token # State: hidden_dim × n_layers = 4096 × 24 = ~400KB # 1,000,000× more efficient! ` **Speed comparison** (inference): `python # Transformer: O(n) per token (quadratic overall) # First token: 1 computation # Second token: 2 computations # ... # 1000th token: 1000 computations # RWKV: O(1) per token (linear overall) # Every token: 1 computation # 1000th token: 1 computation (same as first!) `` ## Wh