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 を使う場合：

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

主な利点：

• 線形時間：Transformer の O(n²) に対して O(n)
• KV キャッシュなし：トークンごとのメモリ使用量が一定
• 無限コンテキスト：固定ウィンドウの制限なし
• 並列トレーニング可能：GPT のように並列化できます
• 逐次推論：RNN のように逐次処理します

代わりに以下を使う場合：

• Transformers：最高の性能が必要で、計算リソースがある場合
• Mamba：状態空間モデルを使いたい場合
• RetNet：保持機構が必要な場合
• 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

発展的なトピック

タイムミキシングとチャネルミキシング：WKV オペレーション、時間減衰メカニズム、レセプタンスゲートについては、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 のように並列化可能

性能（Transformers との比較）：

• 速度：トレーニングは同程度、推論はより高速
• メモリ：長いシーケンスでは 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 のサポート

エージェントの作業手順

1. 環境を確認します。 GPU の種類、メモリ、ドライバー／CUDA のバージョン（nvidia-smi）、フレームワークのバージョン、データセットの場所とサイズを確認します。
2. 入力を特定します。 データまたはユーザーに確認して、形式、識別子、パラメーターを確定します。結果を左右する値は推測せず、質問してください。
3. まず小規模に実行します。 スモークテストを行います。小さなモデルまたはサブセットを使い、少ないステップ数で実行して、損失が減少し、チェックポイントが保存されることを確認します。
4. 上記の手順と参考情報を使って、タスク全体を実行します。
5. 結果を検証します。 保留データで指標を追跡し、ベースラインと比較して、シード、設定、ハードウェアを記録します。
6. 報告します。 実行内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の事項を明記します。

発生した問題	対処方法
CUDA のメモリ不足	バッチサイズを小さくする、勾配累積／チェックポイントまたは混合精度を有効にする、あるいはモデルをシャーディングします。
損失が NaN になる、または発散する	学習率を下げ、データに無効な値がないか確認し、勾配クリッピングを有効にします。
これらの手順に記載された関数、フラグ、またはエンドポイントがインストール済みのバージョンにない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認し、適宜変更して、ユーザーに伝えます。API をでっち上げてはいけません。
必須の入力、識別子、またはパラメーターが曖昧	ユーザーに質問するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計を決して捏造しないでください。実行または検証できない場合は、その旨を明確に伝えます。
• ログなしに学習結果を主張しないでください。大規模なジョブを開始する前に、計算コストを見積もります。
• ここに記載されたバージョン固有の情報は古くなっている可能性があるため、インストール済みバージョンの公式ドキュメントで確認します。
• 費用が発生する、共有 GPU やクラウドのクォータを消費する、個人データや患者データに触れる、または取り消せない操作を行う前に、許可を求めてください。

関連スキル

• mamba-architecture：Mamba の選択的状態空間モデル（state-spaces/mamba package、d_state=16 の Mamba-1、およびマルチヘッド str…を使用する Mamba-2）の使い方を説明します。
• optimizing-attention-flash：PyTorch のネイティブな scaled_dot_product_attention（PyTorch 2.2+）または flash-attn ライブラリを使って、Transformer モデルで Flash Attention を有効にします。
• speculative-decoding：投機的デコーディング、Medusa の複数ヘッド、先読みデコーディングの手法を使って LLM の推論を高速化します。