中文

RWKV-LM: an RNN

RWKV is an RNN that can also be directly trained like a GPT transformer (parallelizable).

Verified

Core facts

Read by a script from the official API, not rewritten by a model.

License
Apache-2.0
Languages
Python
Maintainer
BlinkDL
Verified

What it is

RWKV is an RNN that can also be directly trained like a GPT transformer (parallelizable). RWKV is 100% attention-free. RWKV is inspired by Apple's AFT. Only the hidden state at position t is needed to compute the state at position t+1.

Main features

  1. RWKV is parallelizable because the time-decay of each channel is data-independent and trainable.

  2. The project says inference is very fast, using only matrix-vector multiplications, even on CPUs.

  3. The project says RWKV-7 is very stable and spike-free, verified for 0.1/0.4/1.5/2.9b.

  4. RWKV-7 can do math.

  5. The Head-QK trick adds extra Q and K to the final output so the model can directly copy or avoid tokens in the context.

Getting started

  1. To train RWKV-4 from scratch, run train.py, which by default uses the enwik8 dataset.

  2. In train.py, an epoch is a mini-epoch, not a real epoch, and 1 mini-epoch equals 40320 times ctx_len tokens.

  3. The trainer auto-exits after the number of tokens set by the exit-tokens option.

  4. For state-tuning, rwkv 0.8.26+ is used to auto-load the trained time_state.

From the 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

Who it is for

In RWKV, R, K and V are trainable matrices, and W is a trainable vector, the time-decay factor for each channel. RWKV-7 has no good simple layer, because a pytorch layer cannot make sure it uses correct init and hyperparameters. The project says that to use RWKV-7 for another task, one should study the train_temp code (only several hundred lines) and change it as needed.

For state-tuning, the implementation is currently unoptimized and takes the same VRAM as full SFT. For RWKV 14B, the state has 200 vectors, 5 vectors for each block. For finetuning RWKV-5 models, the data uses .jsonl format.

The project says RWKV-4 can extrapolate, so training with ctxLen 1024 can work for ctxLen of 2500+. The project proposes a new sampling method called top-p-x. The project also proposes a new sampling method called top-a.

The 1.5B RWKV-3 is trained with the Adam optimizer on 8 * A100 40G, with no weight decay and no dropout.

Sources

This page is written from the official sources below. Open them to check the original text.

Verified ·

This page is written by AI from the official sources above and checked automatically. This site has not tested it. Found a mistake? Tell us.