rwkv-architecture

RWKV - Receptance Weighted Key Value

クイックスタート

RWKV（RwaKuv）は、Transformer の並列化（学習）と RNN の効率性（推論）を組み合わせたものです。

インストール:

# 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

基本的な使い方（GPT モード + RNN モード）:

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!

よく使うワークフロー

ワークフロー 1: テキスト生成（ストリーミング）

トークンごとの効率的な生成:

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)

主な利点: トークンごとのメモリ使用量が一定です（KV キャッシュが増え続けません）

ワークフロー 2: 長いコンテキストの処理（無限コンテキスト）

100万トークンのシーケンスを処理:

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)!)

ワークフロー 3: RWKV のファインチューニング

標準的なファインチューニングのワークフロー:

# 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)

ワークフロー 4: RWKV と Transformer の比較

メモリの比較（100万トークンのシーケンス）:

# 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!

速度の比較（推論）:

# 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!)

使用すべき場面と代替手段

RWKV を使うべき場合:

• 非常に長いコンテキスト（10万トークン以上）が必要な場合
• メモリ使用量を一定に保ちたい場合
• ストリーミングアプリケーションを構築する場合
• Transformer 並みの性能と RNN の効率性の両方が必要な場合
• メモリに制約のあるデプロイメントの場合

主な利点:

• 線形時間: Transformer の O(n²) に対して O(n)
• KV キャッシュなし: トークンごとのメモリ使用量が一定
• 無限コンテキスト: 固定のウィンドウ長による制限がない
• 並列化可能な学習: GPT と同様
• 逐次的な推論: RNN と同様

代わりに使うべき手段:

• Transformer: 絶対的に最高の性能が必要で、計算資源がある場合
• Mamba: 状態空間モデルを使いたい場合
• RetNet: Retention メカニズムが必要な場合
• Hyena: 畳み込みベースのアプローチを使いたい場合

よくある問題

問題: 学習中にメモリ不足になる

勾配チェックポイントと DeepSpeed を使用します:

trainer = pl.Trainer(
    strategy='deepspeed_stage_3',  # Full ZeRO-3
    precision='bf16'
)

問題: 推論が遅い

CUDA カーネルを有効にします:

os.environ["RWKV_CUDA_ON"] = '1'

問題: モデルが読み込めない

モデルのパスとストラテジーを確認します:

model = RWKV(
    model='/absolute/path/to/model.pth',
    strategy='cuda fp16'  # Or 'cpu fp32' for CPU
)

問題: RNN モードでの状態管理

forward 呼び出しの間では、必ず状態を渡してください:

# 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

発展的なトピック

Time-mixing と channel-mixing: WKV 演算、時間減衰メカニズム、receptance ゲートについては references/architecture-details.md を参照してください。

状態管理: att_x_prev、att_kv、ffn_x_prev の各状態と、数値安定性に関する考慮事項については references/state-management.md を参照してください。

RWKV-7 の改良点: 最新のアーキテクチャ改良（2025 年 3 月）とマルチモーダル機能については references/rwkv7.md を参照してください。

ハードウェア要件

• GPU: NVIDIA（CUDA 11.6 以降）または CPU
• VRAM（FP16）:
  • 169M モデル: 1GB
  • 430M モデル: 2GB
  • 1.5B モデル: 4GB
  • 3B モデル: 8GB
  • 7B モデル: 16GB
  • 14B モデル: 32GB
• 推論: トークンごとに O(1) のメモリ
• 学習: GPT と同様に並列化可能

性能（Transformer との比較）:

• 速度: 学習は同程度、推論は高速
• メモリ: 長いシーケンスで 1000 倍少ない
• スケーリング: 線形 対 二次

リソース

• 論文（RWKV）: https://arxiv.org/abs/2305.13048（2023 年 5 月）
• 論文（RWKV-7）: https://arxiv.org/abs/2503.14456（2025 年 3 月）
• GitHub: https://github.com/BlinkDL/RWKV-LM ⭐ 12,000+
• ドキュメント: https://wiki.rwkv.com/
• モデル: https://huggingface.co/BlinkDL
• Linux Foundation AI: 公式プロジェクト
• 本番利用: Microsoft Windows、Office 統合、NeMo サポート