mcpField
R
rwkv-architecture

rwkv-architecture

github.com/davila7开发33.1k277 次安装今天更新
代码仓库

关于 rwkv-architecture

rwkv-architecture 是介绍 RWKV 架构及其使用流程的开发技能，面向模型开发者。内容涵盖流式生成、长上下文处理和微调，并说明状态管理与硬件配置。

主要功能

  1. 01流式逐词元生成
  2. 02处理超长上下文
  3. 03RWKV 模型微调
  4. 04管理循环推理状态

适用场景

  1. 01生成文本流
  2. 02分块处理长文档
  3. 03微调 RWKV 模型

安装

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

$npx skills add https://github.com/davila7/claude-code-templates --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