SuperAPI AI 工具集
合作ENAI 行业资讯收藏
Skill 广场rwkv-architecture

rwkv-architecture

智能体与模型

查看项目

更新于 10-15 20:21

技能介绍

用途

RWKV（RwaKuv，Receptance Weighted Key Value）将 Transformer 的并行化能力（用于训练）与 RNN 的推理效率结合在一起。它适用于需要超长上下文、恒定内存占用、流式生成，或在显存受限环境下部署的场景。

主要能力

  • 线性时间推理：O(n)，相比 Transformer 的 O(n²) 更高效。
  • 无 KV cache：每个 token 的内存占用为常数（O(1)）。
  • 无限上下文：不受固定窗口长度限制。
  • 可并行训练：训练方式类似 GPT。
  • 顺序推理：推理方式类似 RNN。
  • 支持 GPT 模式（并行处理）与 RNN 模式（顺序处理，结果一致）。

使用方式

安装

先安装 PyTorch：

pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121

再安装依赖：

pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade

最后安装 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 cache）。

工作流 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 的对比

内存对比（100 万 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 的情况：

  • 需要超长上下文（10 万 token 以上）。
  • 希望内存占用恒定。
  • 构建流式应用。
  • 需要兼具 RNN 效率与 Transformer 性能。
  • 显存受限的部署环境。

关键优势：线性时间（O(n)，Transformer 为 O(n²)）；无 KV cache（每 token 内存恒定）；无限上下文（无固定窗口限制）；可并行训练（类似 GPT）；顺序推理（类似 RNN）。

可考虑替代方案的情况：

  • Transformers：需要绝对最佳性能且算力充足。
  • Mamba：需要状态空间模型。
  • RetNet：需要 retention 机制。
  • Hyena：需要基于卷积的方法。

常见问题与限制

  • 训练时内存不足：使用梯度检查点与 DeepSpeed，例如 strategy='deepspeed_stage_3'、precision='bf16'。
  • 推理缓慢：启用 CUDA kernel，设置 os.environ["RWKV_CUDA_ON"] = '1'。
  • 模型无法加载：检查模型路径与 strategy，例如 cuda fp16，或 CPU 下用 cpu fp32。
  • RNN 模式下的状态管理：必须在 forward 调用之间传递 state。错误做法是每次都传入 None，会丢失上下文；正确做法是 out1, state = model.forward(tokens1, None) 后 out2, state = model.forward(tokens2, state)。

进阶主题

  • 时间混合与通道混合：参见 references/architecture-details.md，了解 WKV 操作、时间衰减机制与 receptance gate。
  • 状态管理：参见 references/state-management.md，了解 att_x_prev、att_kv、ffn_x_prev 状态及数值稳定性考量。
  • RWKV-7 改进：参见 references/rwkv7.md，了解最新架构改进（2025 年 3 月）与多模态能力。

硬件需求

GPU：NVIDIA（CUDA 11.6+）或 CPU。

显存（FP16）：

  • 169M 模型：1GB
  • 430M 模型：2GB
  • 1.5B 模型：4GB
  • 3B 模型：8GB
  • 7B 模型：16GB
  • 14B 模型：32GB

推理：每 token 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+ stars）
  • 文档：https://wiki.rwkv.com/
  • 模型：https://huggingface.co/BlinkDL
  • Linux Foundation AI：官方项目
  • 生产应用：Microsoft Windows、Office 集成、NeMo 支持