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：文本生成（流式）

高效的逐 token 生成：

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)

主要优势：每个 token 的内存占用恒定（KV 缓存不会增长）

工作流 2：长上下文处理（无限上下文）

处理百万 token 序列：

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 对比

内存对比（1M token 序列）：

# 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 而非替代方案

适合使用 RWKV 的场景：

• 需要非常长的上下文（100K+ token）
• 希望内存用量保持恒定
• 构建流式应用
• 需要兼具 RNN 的效率和 Transformer 的性能
• 部署环境受内存限制

主要优势：

• 线性时间：O(n)，而 Transformer 为 O(n²)
• 无需 KV 缓存：每个 token 的内存占用恒定
• 无限上下文：没有固定窗口限制
• 可并行训练：类似 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 操作、时间衰减机制和 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
• 推理：每个 token 的内存占用为 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 支持

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 确认输入。 根据数据或用户提供的信息确认格式、标识符和参数。任何会改变结果的值都应先询问，不要猜测。
3. 先运行小规模版本。 进行冒烟测试：使用小型模型或子集，运行少量步骤，并确认损失下降且检查点已保存。
4. 使用上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 汇报。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果发生这种情况	请这样做
CUDA 显存不足	减小批次大小，启用梯度累积/检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否存在无效值，并启用梯度裁剪。
这些说明中提到的函数、标志或端点在已安装的版本中不存在	查阅已安装版本自身的文档（help()、--help、官方文档），进行适配，并告知用户。切勿捏造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明假设。

诚信准则

• 切勿捏造结果、参数、标识符、引文或统计数据。如果无法运行或验证某项内容，请直言说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前应估算计算成本。
• 请将此处与版本相关的细节视为可能已过时，并根据已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• mamba-architecture：介绍如何使用 Mamba 选择性状态空间模型（state-spaces/mamba 包、具有 d_state=16 的 Mamba-1，以及具有多头结构的 Mamba-2……
• optimizing-attention-flash：使用 PyTorch 原生 scaled_dot_product_attention（PyTorch 2.2+）或 flash-attn 库，为 Transformer 模型启用 Flash Attention……
• speculative-decoding：使用推测解码、Medusa 多头和前瞻解码技术加速 LLM 推理。