@r1z4x/tezgah 1.0.0

package/skills/ai-research/01-model-architecture/rwkv/SKILL.md

@@ -0,0 +1,260 @@
+---
+name: rwkv-architecture
+description: RNN+Transformer hybrid with O(n) inference. Linear time, infinite context, no KV cache. Train like GPT (parallel), infer like RNN (sequential). Linux Foundation AI project. Production at Windows, Office, NeMo. RWKV-7 (March 2025). Models 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!)
+```
+
+## When to use vs alternatives
+
+**Use RWKV when**:
+- Need very long context (100K+ tokens)
+- Want constant memory usage
+- Building streaming applications
+- Need RNN efficiency with Transformer performance
+- Memory-constrained deployment
+
+**Key advantages**:
+- **Linear time**: O(n) vs O(n²) for Transformers
+- **No KV cache**: Constant memory per token
+- **Infinite context**: No fixed window limit
+- **Parallelizable training**: Like GPT
+- **Sequential inference**: Like RNN
+
+**Use alternatives instead**:
+- **Transformers**: Need absolute best performance, have compute
+- **Mamba**: Want state-space models
+- **RetNet**: Need retention mechanism
+- **Hyena**: Want convolution-based approach
+
+## Common issues
+
+**Issue: Out of memory during training**
+
+Use gradient checkpointing and DeepSpeed:
+```python
+trainer = pl.Trainer(
+    strategy='deepspeed_stage_3',  # Full ZeRO-3
+    precision='bf16'
+)
+```
+
+**Issue: Slow inference**
+
+Enable CUDA kernel:
+```python
+os.environ["RWKV_CUDA_ON"] = '1'
+```
+
+**Issue: Model not loading**
+
+Check model path and strategy:
+```python
+model = RWKV(
+    model='/absolute/path/to/model.pth',
+    strategy='cuda fp16'  # Or 'cpu fp32' for CPU
+)
+```
+
+**Issue: State management in RNN mode**
+
+Always pass state between forward calls:
+```python
+# WRONG: State lost
+out1, _ = model.forward(tokens1, None)
+out2, _ = model.forward(tokens2, None)  # No context from tokens1!
+
+# CORRECT: State preserved
+out1, state = model.forward(tokens1, None)
+out2, state = model.forward(tokens2, state)  # Has context from tokens1
+```
+
+## Advanced topics
+
+**Time-mixing and channel-mixing**: See [references/architecture-details.md](references/architecture-details.md) for WKV operation, time-decay mechanism, and receptance gates.
+
+**State management**: See [references/state-management.md](references/state-management.md) for att_x_prev, att_kv, ffn_x_prev states, and numerical stability considerations.
+
+**RWKV-7 improvements**: See [references/rwkv7.md](references/rwkv7.md) for latest architectural improvements (March 2025) and multimodal capabilities.
+
+## Hardware requirements
+
+- **GPU**: NVIDIA (CUDA 11.6+) or CPU
+- **VRAM** (FP16):
+  - 169M model: 1GB
+  - 430M model: 2GB
+  - 1.5B model: 4GB
+  - 3B model: 8GB
+  - 7B model: 16GB
+  - 14B model: 32GB
+- **Inference**: O(1) memory per token
+- **Training**: Parallelizable like GPT
+
+**Performance** (vs Transformers):
+- **Speed**: Similar training, faster inference
+- **Memory**: 1000× less for long sequences
+- **Scaling**: Linear vs quadratic
+
+## Resources
+
+- Paper (RWKV): https://arxiv.org/abs/2305.13048 (May 2023)
+- Paper (RWKV-7): https://arxiv.org/abs/2503.14456 (March 2025)
+- GitHub: https://github.com/BlinkDL/RWKV-LM ⭐ 12,000+
+- Docs: https://wiki.rwkv.com/
+- Models: https://huggingface.co/BlinkDL
+- Linux Foundation AI: Official project
+- Production: Microsoft Windows, Office integration, NeMo support
+
+

package/skills/ai-research/01-model-architecture/rwkv/references/architecture-details.md

@@ -0,0 +1,345 @@
+<!-- vendored from orchestra-research/AI-research-SKILLs@773a52944ba4747a18bd4ae9ade53fff041adcbc 01-model-architecture/rwkv/references/architecture-details.md; MIT (c) Orchestra Research -->
+# RWKV Architecture Details
+
+## Time-Mixing and Channel-Mixing Blocks
+
+RWKV alternates between **Time-Mixing** (sequence processing) and **Channel-Mixing** (feature processing) blocks.
+
+### Time-Mixing Block (WKV Operation)
+
+The core innovation is the **WKV (Weighted Key-Value)** mechanism:
+
+```python
+# Traditional Attention (O(n²))
+scores = Q @ K.T / sqrt(d)  # n×n matrix
+attention = softmax(scores)
+output = attention @ V
+
+# RWKV Time-Mixing (O(n))
+# Compute WKV in linear time using recurrence
+for t in range(T):
+    wkv[t] = (exp(w) * k[t] @ v[t] + a[t] * aa[t]) / (exp(w) * k[t] + a[t] * ab[t])
+    aa[t+1] = exp(w) * k[t] @ v[t] + exp(-u) * aa[t]
+    ab[t+1] = exp(w) * k[t] + exp(-u) * ab[t]
+```
+
+**Full Time-Mixing implementation**:
+
+```python
+class RWKV_TimeMix(nn.Module):
+    def __init__(self, d_model, n_layer):
+        super().__init__()
+        self.d_model = d_model
+
+        # Linear projections
+        self.key = nn.Linear(d_model, d_model, bias=False)
+        self.value = nn.Linear(d_model, d_model, bias=False)
+        self.receptance = nn.Linear(d_model, d_model, bias=False)
+        self.output = nn.Linear(d_model, d_model, bias=False)
+
+        # Time-mixing parameters
+        self.time_mix_k = nn.Parameter(torch.ones(1, 1, d_model))
+        self.time_mix_v = nn.Parameter(torch.ones(1, 1, d_model))
+        self.time_mix_r = nn.Parameter(torch.ones(1, 1, d_model))
+
+        # Time-decay and bonus
+        self.time_decay = nn.Parameter(torch.ones(d_model))  # w
+        self.time_first = nn.Parameter(torch.ones(d_model))  # u
+
+    def forward(self, x, state=None):
+        B, T, C = x.shape
+
+        # Time-shift mixing (interpolate with previous token)
+        if state is None:
+            state = torch.zeros(B, C, 3, device=x.device)  # [aa, ab, x_prev]
+
+        x_prev = state[:, :, 2].unsqueeze(1)  # Previous x
+        xk = x * self.time_mix_k + x_prev * (1 - self.time_mix_k)
+        xv = x * self.time_mix_v + x_prev * (1 - self.time_mix_v)
+        xr = x * self.time_mix_r + x_prev * (1 - self.time_mix_r)
+
+        # Compute k, v, r
+        k = self.key(xk)
+        v = self.value(xv)
+        r = self.receptance(xr)
+
+        # WKV computation (parallelizable or sequential)
+        wkv = self.wkv(k, v, state[:, :, :2])
+
+        # Apply receptance gate and output projection
+        out = self.output(torch.sigmoid(r) * wkv)
+
+        # Update state
+        new_state = torch.stack([state_aa, state_ab, x[:, -1]], dim=2)
+
+        return out, new_state
+
+    def wkv(self, k, v, state):
+        # Parallel implementation (training)
+        # Sequential implementation (inference) - see below
+        ...
+```
+
+### WKV Parallel Algorithm (Training)
+
+```python
+def wkv_forward(w, u, k, v):
+    """
+    Parallel WKV computation for training.
+    w: time_decay (d_model,)
+    u: time_first (d_model,)
+    k: keys (batch, seq_len, d_model)
+    v: values (batch, seq_len, d_model)
+    """
+    B, T, C = k.shape
+
+    # Compute cumulative sums with exponential decay
+    # This is the key to O(n) parallel computation
+    w = -torch.exp(w)  # Negative for decay
+
+    # Associative scan operation
+    wkv = torch.zeros(B, T, C, device=k.device)
+    state = torch.zeros(B, C, device=k.device)
+
+    for t in range(T):
+        kv = k[:, t] * v[:, t]
+        wkv[:, t] = (u * kv + state) / (u * k[:, t] + torch.exp(state_count))
+        state = w * state + kv
+
+    return wkv
+```
+
+### WKV Sequential Algorithm (Inference)
+
+```python
+def wkv_inference(w, u, k, v, state):
+    """
+    Sequential WKV for O(1) per-token inference.
+    state: (aa, ab) from previous step
+    """
+    w = -torch.exp(w)  # time_decay
+    u = torch.exp(u)   # time_first
+
+    # Unpack state
+    aa, ab = state  # aa = numerator, ab = denominator
+
+    # Compute WKV for current token
+    kv = k * v
+    wkv = (u * kv + aa) / (u * k + ab)
+
+    # Update state for next token
+    new_aa = w * aa + kv
+    new_ab = w * ab + k
+
+    return wkv, (new_aa, new_ab)
+```
+
+### Channel-Mixing Block
+
+Replaces Transformer FFN with time-shifted variant:
+
+```python
+class RWKV_ChannelMix(nn.Module):
+    def __init__(self, d_model, hidden_ratio=4):
+        super().__init__()
+        self.d_model = d_model
+        self.hidden = d_model * hidden_ratio
+
+        # Time-mixing for channel
+        self.time_mix_k = nn.Parameter(torch.ones(1, 1, d_model))
+        self.time_mix_r = nn.Parameter(torch.ones(1, 1, d_model))
+
+        # FFN layers
+        self.key = nn.Linear(d_model, self.hidden, bias=False)
+        self.receptance = nn.Linear(d_model, d_model, bias=False)
+        self.value = nn.Linear(self.hidden, d_model, bias=False)
+
+    def forward(self, x, x_prev):
+        # Time-shift mixing
+        xk = x * self.time_mix_k + x_prev * (1 - self.time_mix_k)
+        xr = x * self.time_mix_r + x_prev * (1 - self.time_mix_r)
+
+        # Channel mixing
+        k = self.key(xk)
+        k = torch.square(torch.relu(k))  # Squared ReLU activation
+        kv = self.value(k)
+
+        # Receptance gate
+        r = torch.sigmoid(self.receptance(xr))
+
+        return r * kv
+```
+
+## RWKV Block Structure
+
+```python
+class RWKV_Block(nn.Module):
+    def __init__(self, d_model, n_layer):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d_model)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.att = RWKV_TimeMix(d_model, n_layer)
+        self.ffn = RWKV_ChannelMix(d_model)
+
+    def forward(self, x, state):
+        # Time-mixing with residual
+        att_out, new_state = self.att(self.ln1(x), state)
+        x = x + att_out
+
+        # Channel-mixing with residual
+        ffn_out = self.ffn(self.ln2(x), state[:, :, 2])  # Use x_prev from state
+        x = x + ffn_out
+
+        return x, new_state
+
+# Full RWKV model
+model = nn.Sequential(
+    Embedding(...),
+    *[RWKV_Block(d_model, i) for i in range(n_layers)],
+    LayerNorm(d_model),
+    LMHead(...)
+)
+```
+
+## Time-Decay Mechanism
+
+The **time_decay** parameter `w` controls how fast information decays:
+
+```python
+# Initialization (RWKV-4)
+time_decay = torch.ones(n_layers, d_model)
+for i in range(n_layers):
+    for j in range(d_model):
+        # Logarithmic spacing
+        ratio = (i + 1) / n_layers
+        time_decay[i, j] = -5.0 + 8.0 * ratio + 0.3 * (j / d_model)
+
+# Effect on memory
+w = -exp(time_decay)  # Range: [-exp(-5), -exp(3)] ≈ [-0.007, -20]
+# Smaller w = slower decay = longer memory
+# Larger w = faster decay = shorter memory
+```
+
+**Layer-wise decay pattern**:
+- Early layers (shallow): Fast decay, capture local patterns
+- Later layers (deep): Slow decay, capture long-range dependencies
+
+## Receptance Gate
+
+The **receptance** mechanism controls information flow:
+
+```python
+r = sigmoid(receptance(x))  # Range [0, 1]
+output = r * wkv  # Gate the WKV output
+
+# High receptance (r ≈ 1): Pass information through
+# Low receptance (r ≈ 0): Block information
+```
+
+**Purpose**: Similar to LSTM forget gate, but learned per-token
+
+## RWKV-4 vs RWKV-5 vs RWKV-6 vs RWKV-7
+
+### RWKV-4 (Original)
+```python
+# Time-shift with previous token
+xx = x * time_mix + x_prev * (1 - time_mix)
+k, v, r = key(xx), value(xx), receptance(xx)
+```
+
+### RWKV-5 (2023)
+```python
+# Separate time-mix for k, v, r
+xk = x * time_mix_k + x_prev * (1 - time_mix_k)
+xv = x * time_mix_v + x_prev * (1 - time_mix_v)
+xr = x * time_mix_r + x_prev * (1 - time_mix_r)
+k, v, r = key(xk), value(xk), receptance(xr)
+```
+
+### RWKV-6 (2024)
+- Added **multi-head time-mixing** (like multi-head attention)
+- Separate time-decay per head
+- Improved stability for large models
+
+```python
+# Per-head processing
+for h in range(n_heads):
+    k_h = key[h](x)  # Separate projection per head
+    w_h = time_decay[h]  # Separate decay per head
+    wkv_h = wkv(k_h, v_h, w_h)
+output = concat(wkv_0, wkv_1, ..., wkv_H)
+```
+
+### RWKV-7 (March 2025)
+- **Multimodal support** (vision + language)
+- Improved numerical stability
+- Better scaling to 14B+ parameters
+
+## Numerical Stability
+
+### Issue: Exponential Overflow
+
+```python
+# Problem: exp(wkv) can overflow
+wkv = exp(u * kv) / exp(u * k)  # Can overflow!
+```
+
+### Solution: Log-space Computation
+
+```python
+# Stable implementation
+log_wkv_num = u + log(kv) + log(aa)
+log_wkv_den = u + log(k) + log(ab)
+wkv = exp(log_wkv_num - log_wkv_den)  # Numerically stable
+```
+
+### Gradient Clipping
+
+```python
+# Recommended for training stability
+torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+```
+
+## State Management
+
+### State Shape
+
+```python
+# For batch inference
+state = torch.zeros(
+    batch_size,
+    n_layers,
+    4,  # (att_aa, att_ab, att_x_prev, ffn_x_prev)
+    d_model
+)
+```
+
+### State Initialization
+
+```python
+# Zero initialization (standard)
+state = None  # Model creates zero state
+
+# Warm state (from previous conversation)
+_, state = model.forward(previous_context, None)
+# Use `state` for next turn
+```
+
+### State Serialization
+
+```python
+# Save conversation state
+torch.save(state, 'conversation_state.pt')
+
+# Resume conversation
+state = torch.load('conversation_state.pt')
+out, state = model.forward(new_tokens, state)
+```
+
+## Resources
+
+- Paper (RWKV): https://arxiv.org/abs/2305.13048 (May 2023)
+- Paper (RWKV-7): https://arxiv.org/abs/2503.14456 (March 2025)
+- GitHub: https://github.com/BlinkDL/RWKV-LM
+- Math derivation: https://wiki.rwkv.com/
+- CUDA kernels: https://github.com/BlinkDL/RWKV-CUDA