mcpField
R
rwkv-architecture

rwkv-architecture

github.com/ovachiever开发5648 次安装49 周前更新
代码仓库

关于 rwkv-architecture

RWKV 是一种结合 Transformer 并行训练与 RNN 推理效率的模型架构，面向模型开发和部署人员。它支持逐 Token 流式生成，并以恒定的每 Token 推理内存处理长上下文。

主要功能

  1. 01并行化模型训练
  2. 02逐 Token 流式生成
  3. 03恒定每 Token 推理内存
  4. 04跨分块传递状态
  5. 05支持微调流程

适用场景

  1. 01生成流式文本
  2. 02处理超长文档序列
  3. 03在内存受限环境进行推理

安装

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

$npx skills add https://github.com/ovachiever/droid-tings --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