State Tuning Tutorial

The State tuning method in this article comes from the RWKV community fine-tuning project RWKV-PEFT.

Before you start, make sure you have a Linux workspace and an NVIDIA graphics card that supports CUDA.

The VRAM requirements for State tuning can be referred to in the following table:

The data in the above table is based on the following test parameters:

• ctxlen=1024
• micro_bsz=1
• strategy=deepspeed_stage_1

Organize the Training Data

Collect Training Data in jsonl Format

To fine-tune the RWKV model through state tuning, you need to collect data suitable for training RWKV (in jsonl format). For specific methods, you can refer to Preparing the Training Dataset.

Configure the Training Environment

Please refer to the RWKV Training Environment Configuration section to configure the training environment such as Conda.

Clone the Repository and Install Dependencies

In Linux or WSL, use the git command to clone the RWKV-PEFT repository:

git clone https://github.com/JL-er/RWKV-PEFT.git

After the cloning is completed, use the cd RWKV-PEFT command to enter the RWKV-PEFT directory. And run the following command to install the dependencies required by the project:

pip install -r requirements.txt

Modify the Training Parameters

Open the state tuning.sh file in the RWKV-PEFT/scripts directory using any text editor (such as vscode), and modify the training parameters to control the fine-tuning training process and training effect:

The following is a state tuning parameter adjustment process:

Adjust the Path Parameters

The first three lines of the state tuning.sh file are file path parameters:

• load_model: The path of the base RWKV model
• proj_dir: The output path of the training log and the state file obtained from training
• data_file: The path of the training dataset. Note that there is no need to include the bin and idx suffixes in the path, only the file name is required.

Adjust the n_layer and n_embd Parameters

The parameter sizes of the base RWKV models are different, and the values of n_layer and n_embd used during training are also different. The following is the comparison list of the RWKV models with different parameters and the values of n_layer and n_embd:

Adjust the Important Training Parameters

Adjust Other Training Parameters

The following lists other modifiable training parameters in the script and the effects of their modification.

Appendix: State Tuning Configuration Reference

load_model='/home/rwkv/RWKV-PEFT/models/rwkv7-g1a-0.1b-20250728-ctx4096.pth'
proj_dir='/home/rwkv/RWKV-PEFT/test-output'
data_file='/home/rwkv/RWKV-PEFT/data/test.jsonl'

n_layer=12
n_embd=768

micro_bsz=1
epoch_save=1
epoch_steps=300
ctx_len=1024

python train.py --load_model $load_model \
--proj_dir $proj_dir --data_file $data_file \
--vocab_size 65536 \
--data_type jsonl \
--n_layer $n_layer --n_embd $n_embd \
--ctx_len $ctx_len --micro_bsz $micro_bsz \
--epoch_steps $epoch_steps --epoch_count 4 --epoch_save $epoch_save \
--lr_init 1e-2 --lr_final 1e-4 \
--accelerator gpu --precision bf16 \
--devices 1 --strategy deepspeed_stage_1 --grad_cp 1 \
--my_testing "x070" \
--peft state --op fla

Start the Training

In the RWKV-PEFT directory, run the command sh scripts/state tuning.sh to start state tuning.

After the training starts normally, it should be as follows:

After the training is completed, you should be able to find the trained state file (in .pth format) and the training log (.txt file) in the output folder:

How to Use the State File

After obtaining the state file, you can use it as follows:

• Use the demo-state-merge.sh tool to merge the state file into the base RWKV model to obtain a complete state fine-tuned model.

• You can choose to mount the state file separately in tools such as RWKV Runner or Ai00. (Recommended usage)

For example: This state file is fine-tuned based on the RWKV-6-World-3B-v3 model. Then, in tools such as RWKV Runner or Ai00, you must start the RWKV-6-World-3B-v3 model for the mounted state file to take effect.

Since our sample data is based on a large number of emojis and the ctx is very short, the effect of the trained State file is as follows: