mcpField
R
rwkv-architecture

rwkv-architecture

github.com/orchestra-research开发14.1k952 次安装19 周前更新
代码仓库

关于 rwkv-architecture

rwkv-architecture 介绍 RWKV 架构及其使用流程，面向需要运行、训练或微调 RWKV 模型的开发者。内容涵盖流式生成和长序列处理，并说明其逐 Token 推理的常量内存特性。

主要功能

  1. 01流式逐 Token 文本生成
  2. 02分块处理超长序列
  3. 03RWKV 模型微调流程
  4. 04状态传递与管理
  5. 05CUDA 推理配置

适用场景

  1. 01生成文本并逐步输出
  2. 02处理超长文档序列
  3. 03微调 RWKV 模型
  4. 04排查推理状态或内存问题

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/orchestra-research/ai-research-skills --skill rwkv-architecture

安装后在 Claude Code 中输入/rwkv-architecture确认已加载。

包含的文件

除 SKILL.md 外，安装时一起下载的参考资料和脚本。

references/
architecture-details.md
rwkv7.md
state-management.md

详细介绍

RWKV - 接收度加权键值

快速入门

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 的情况：

  • 需要超长上下文（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

高级主题

时间混合与通道混合：请参阅 references/architecture-details.md，了解 WKV 运算、时间衰减机制和接收度门控。

状态管理：请参阅 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 一样可并行

性能（与 Transformers 相比）：

  • 速度：训练速度相近，推理更快
  • 内存：处理长序列时减少 1000×
  • 扩展性：线性而非二次

资源

评论

支持 Markdown