EN

RWKV-LM:一种 RNN

RWKV 是一种 RNN,也可以像 GPT Transformer 那样直接训练(可并行)。

核验日期

核心信息

以下信息由程序从官方接口读取，不经模型改写。

许可证
Apache-2.0
语言
Python
维护方
BlinkDL
核验日期

它是什么

RWKV 是一种 RNN,也可以像 GPT Transformer 那样直接训练(可并行)。RWKV 是 100% 无注意力的。RWKV 受到 Apple 的 AFT 启发。计算 t+1 位置的状态,只需要 t 位置的隐藏状态。

主要能力

  1. RWKV 之所以可并行,是因为每个通道的时间衰减与数据无关且可训练。

  2. 项目方称,推理非常快,只使用矩阵-向量乘法,即使在 CPU 上也是如此。

  3. 项目方称,RWKV-7 非常稳定且无尖峰,已在 0.1/0.4/1.5/2.9b 上验证。

  4. RWKV-7 可以做数学。

  5. Head-QK 技巧在最终输出中加入额外的 Q 和 K,使模型能够直接复制或避开上下文中的 token。

怎么上手

  1. 从零训练 RWKV-4 时,运行 train.py,它默认使用 enwik8 数据集。

  2. 在 train.py 中,一个 epoch 是 mini-epoch,不是真正的 epoch,1 个 mini-epoch 等于 40320 乘以 ctx_len 个 token。

  3. 训练器会在达到退出 token 选项所设定的 token 数后自动退出。

  4. 对于 state-tuning,使用 rwkv 0.8.26+ 自动加载训练好的 time_state。

摘自 README
# you can use latest torch + latest cuda (not limited to cu121)
pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121
pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade

# train RWKV-7
cd RWKV-v7/train_temp/ 

# download minipile .bin .idx to train_temp/data first (check demo-training-prepare.sh)
# this will generate the initial weight rwkv-init.pth in out/....../
sh ./demo-training-prepare.sh

# this will load rwkv-init.pth and train the model. you may want to log in to wandb first
sh ./demo-training-run.sh

your out/....../train_log.txt should have losses similar to:
0 4.875856 131.0863 0.00059975 2025-04-24 02:23:42.481256 0
1 4.028621 56.1834 0.00059899 2025-04-24 02:28:16.674463 1
2 3.801625 44.7739 0.00059773 2025-04-24 02:32:51.059568 2
3 3.663070 38.9808 0.00059597 2025-04-24 02:37:25.409892 3
4 3.578974 35.8368 0.00059371 2025-04-24 02:41:59.711315 4
5 3.510906 33.4786 0.00059096 2025-04-24 02:46:33.990839 5
6 3.462345 31.8917 0.00058771 2025-04-24 02:51:08.378331 6
7 3.412196 30.3318 0.00058399 2025-04-24 02:55:42.927474 7
8 3.376724 29.2747 0.00057978 2025-04-24 03:00:17.504665 8
9 3.336911 28.1321 0.00057511 2025-04-24 03:04:52.006063 9
10 3.313411 27.4787 0.00056999 2025-04-24 03:09:27.563336 10
11 3.295895 27.0016 0.00056441 2025-04-24 03:14:01.786079 11

适合谁

在 RWKV 中,R、K、V 是可训练的矩阵,W 是可训练的向量,即每个通道的时间衰减因子。RWKV-7 没有好用的简单层,因为 pytorch 层无法确保自身使用正确的初始化和超参数。项目方称,若要把 RWKV-7 用于其他任务,应研读 train_temp 代码(只有几百行)并按需修改。

对于 state-tuning,当前实现未经优化,占用的显存与完整 SFT 相同。对于 RWKV 14B,状态有 200 个向量,每个块 5 个向量。微调 RWKV-5 模型时,数据使用 .jsonl 格式。

项目方称,RWKV-4 可以外推,因此用 ctxLen 1024 训练可适用于 2500+ 的 ctxLen。项目提出了一种名为 top-p-x 的新采样方法。项目还提出了一种名为 top-a 的新采样方法。

对于 1.5B 的 RWKV-3,使用 Adam 优化器在 8 * A100 40G 上训练,不使用权重衰减,也不使用 dropout。

来源

本页内容整理自以下官方来源，点击可核对原文。

核验日期 ·

本页由 AI 依据上述官方来源整理，并经自动核对，未经本站实测。 发现错误？告诉我们。