rwkv-architecture

# RWKV - Receptance Weighted Key Value ## 快速开始 RWKV (RwaKuv) 结合了 Transformer 的并行化（训练）与 RNN 的高效性（推理）。 **安装**： bash # 安装 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 模式)： python 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()) # Logits # RNN 模式 (顺序处理，结果相同) out, state = model.forward([187, 510], None) # 前 2 个 token out, state = model.forward([1563], state) # 下一个 token out, state = model.forward([310, 247], state) # 最后 token print(out.detach().cpu().numpy()) # 与上述结果相同的 Logits! ## 常见工作流 ### 工作流 1：文本生成 (流式) **高效的逐 token 生成**： python 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 = "The future of AI is" state = None # 逐 token 生成 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) **核心优势**：每个 token 的内存占用恒定（无增长的 KV cache） ### 工作流 2：长上下文处理 (无限上下文) **处理百万 token 序列**： python model = RWKV(model='RWKV-4-Pile-14B', strategy='cuda fp16') # 处理超长文档 state = None long_document = load_document() # 例如 1M tokens # 流式处理整个文档 for chunk in chunks(long_document, chunk_size=1024): out, state = model.forward(chunk, state) # 此时 state 包含了整个 1M token 文档的信息 # 内存占用：O(1) (恒定，而非 O(n)!) ### 工作流 3：RWKV 微调 **标准微调工作流**： python # 训练脚本 imp