⚡

RWKV架构Skill rwkv-architecture

RWKV是一种结合Transformer和RNN的混合神经网络架构，具有线性推理复杂度、无限上下文处理能力和高效内存使用。适用于长序列处理、流式应用和大模型训练与推理，特别适合AI领域的模型开发和优化。关键词：RWKV, Transformer, RNN, 线性复杂度, 无限上下文, 高效推理, 机器学习, AI模型, 大模型微调

大模型微调 0 次安装 0 次浏览 更新于 3/21/2026

name: rwkv-architecture description: RNN+Transformer混合架构，具有O(n)推理复杂度。线性时间，无限上下文，无KV缓存。训练类似GPT（并行），推理类似RNN（顺序）。Linux Foundation AI项目。生产应用于Windows、Office、NeMo。RWKV-7（2025年3月）。模型参数高达140亿。 version: 1.0.0 author: Orchestra Research license: MIT tags: [RWKV, 模型架构, RNN, Transformer混合, 线性复杂度, 无限上下文, 高效推理, Linux Foundation, 替代架构] dependencies: [rwkv, torch, transformers]

RWKV - 接收加权键值

快速入门

RWKV（RwaKuv）结合了Transformer的并行化（训练）与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'  # 使用CUDA内核加速

# 加载模型
model = RWKV(
    model='/path/to/RWKV-4-Pile-1B5-20220903-8040',
    strategy='cuda fp16'
)

# GPT模式（并行处理）
out, state = model.forward([187, 510, 1563, 310, 247], None)
print(out.detach().cpu().numpy())  # 逻辑值

# RNN模式（顺序处理，相同结果）
out, state = model.forward([187, 510], None)  # 前2个令牌
out, state = model.forward([1563], state)      # 下一个令牌
out, state = model.forward([310, 247], state)  # 最后令牌
print(out.detach().cpu().numpy())  # 与上方相同的逻辑值！

常见工作流

工作流1：文本生成（流式）

高效令牌逐令牌生成:

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")

# 初始提示
prompt = "人工智能的未来是"
state = None

# 逐令牌生成
for token in prompt:
    out, state = pipeline.model.forward(pipeline.encode(token), state)

# 继续生成
for _ in range(100):
    out, state = pipeline.model.forward(None, state)
    token = pipeline.sample_logits(out)
    print(pipeline.decode(token), end='', flush=True)

关键优势: 每个令牌恒定内存（无增长的KV缓存）

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

处理百万令牌序列:

model = RWKV(model='RWKV-4-Pile-14B', strategy='cuda fp16')

# 处理非常长的文档
state = None
long_document = load_document()  # 例如，100万令牌

# 流式通过整个文档
for chunk in chunks(long_document, chunk_size=1024):
    out, state = model.forward(chunk, state)

# 状态现在包含整个100万令牌文档的信息
# 内存使用: O(1)（恒定，非O(n)!）

工作流3：微调RWKV

标准微调工作流:

# 训练脚本
import pytorch_lightning as pl
from rwkv.model import RWKV
from rwkv.trainer import RWKVTrainer

# 配置模型
config = {
    'n_layer': 24,
    'n_embd': 1024,
    'vocab_size': 50277,
    'ctx_len': 1024
}

# 设置训练器
trainer = pl.Trainer(
    accelerator='gpu',
    devices=8,
    precision='bf16',
    strategy='deepspeed_stage_2',
    max_epochs=1
)

# 训练
model = RWKV(config)
trainer.fit(model, train_dataloader)

工作流4：RWKV与Transformer比较

内存比较（100万令牌序列）:

# Transformer（GPT）
# 内存: O(n²) 用于注意力
# KV缓存: 100万 × 隐藏维度 × 层数 × 2（键 + 值）
# 示例: 100万 × 4096 × 24 × 2 = ~400GB（不切实际!）

# RWKV
# 内存: O(1) 每个令牌
# 状态: 隐藏维度 × 层数 = 4096 × 24 = ~400KB
# 效率提升1000000倍！

速度比较（推理）:

# Transformer: O(n) 每个令牌（整体二次）
# 第一个令牌: 1次计算
# 第二个令牌: 2次计算
# ...
# 第1000个令牌: 1000次计算

# RWKV: O(1) 每个令牌（整体线性）
# 每个令牌: 1次计算
# 第1000个令牌: 1次计算（与第一个相同!）

何时使用与替代方案

使用RWKV当:

  • 需要非常长上下文（10万+令牌）
  • 希望恒定内存使用
  • 构建流式应用
  • 需要RNN效率与Transformer性能
  • 内存受限部署

关键优势:

  • 线性时间: O(n) 对比 Transformers 的 O(n²)
  • 无KV缓存: 每个令牌恒定内存
  • 无限上下文: 无固定窗口限制
  • 可并行化训练: 类似GPT
  • 顺序推理: 类似RNN

使用替代方案替代:

  • Transformers: 需要绝对最佳性能，有计算资源
  • Mamba: 想要状态空间模型
  • RetNet: 需要保留机制
  • Hyena: 想要基于卷积的方法

常见问题

问题: 训练期间内存不足

使用梯度检查点和DeepSpeed:

trainer = pl.Trainer(
    strategy='deepspeed_stage_3',  # 完全ZeRO-3
    precision='bf16'
)

问题: 推理慢

启用CUDA内核:

os.environ["RWKV_CUDA_ON"] = '1'

问题: 模型未加载

检查模型路径和策略:

model = RWKV(
    model='/absolute/path/to/model.pth',
    strategy='cuda fp16'  # 或 'cpu fp32' 用于CPU
)

问题: RNN模式中的状态管理

始终在向前调用之间传递状态:

# 错误: 状态丢失
out1, _ = model.forward(tokens1, None)
out2, _ = model.forward(tokens2, None)  # 没有tokens1的上下文！

# 正确: 状态保留
out1, state = model.forward(tokens1, None)
out2, state = model.forward(tokens2, state)  # 具有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
  • VRAM（FP16）:
    • 1.69亿模型: 1GB
    • 4.30亿模型: 2GB
    • 15亿模型: 4GB
    • 30亿模型: 8GB
    • 70亿模型: 16GB
    • 140亿模型: 32GB
  • 推理: O(1) 每个令牌内存
  • 训练: 可并行化类似GPT

性能（对比Transformers）:

  • 速度: 训练相似，推理更快
  • 内存: 长序列内存减少1000倍
  • 扩展: 线性对比二次

资源