RWKV 模型全参微调教程（入门版）

配置虚拟环境（conda）

RWKV 默认训练虚拟环境为 conda ，如果不知道如何配置 Linux 系统的 conda ，请参考此文章 (opens in a new tab)。

安装微调训练所需的软件

要获得 RWKV 模型的最佳训练性能，请使用 python 3.10、torch 2.1.2+cu121（或最新版本）、cuda 12.3+、最新的 deepspeed，但保持 pytorch-lightning ==1.9.5 。

在 Linux 环境中依次运行以下命令：

pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121

pip install pytorch-lightning==1.9.5 transformers ninja einops tqdm gpustat nvitop bitsandbytes wandb datasets triton==2.2.0 --upgrade

第一条命令通过指定 url 安装 cu121 版本的 torch ,第二条命令安装 RWKV State tuning 所需的其他软件包，并指定安装 1.9.5 版本的 PyTorch Lightning、2.2.0 版本的 triton。

-- upgrade 参数意味着如果环境中已经安装了对应的软件包，则将其升级到最新版本。

检查 CUDA 环境

依次运行以下命令，检查 CUDA 库是否可用：

python3 # 启动 python
 
import torch # 导入 torch 包
 
torch.cuda.is_available() # 检查 CUDA 是否可用

如果返回值为 True ，证明 CUDA 环境可用。如果返回 False ，请重新安装 CUDA 环境（cuda 12.3+）。

整理 jsonl 语料

首先确认“你希望 RWKV 模型学习哪些知识”。这里的知识可以是某一领域的资料，比如法律问答、金融知识等；也可以是某类任务的文本，如材料总结、角色扮演等。总而言之，你需要收集此类数据，并将其转化为特定的 jsonl 格式。

下面是各类内容的 jsonl 格式参考：

单轮问答： {"text": "User: 问题\n\nAssistant: 答案"}
 
多轮对话例子：{"text": "User: 问题一\n\nAssistant: 答案一\n\nUser: 问题二\n\nAssistant: 答案二"}
 
材料问答：{"text": "Instruction: 问题\n\nInput: 材料\n\nResponse: 答案"}
 
内容生成/续写：{"text": "将每篇文章的内容变成 JSONL 的一行，即使是一百万字的小说也变成一行。"} 或者 {"text": "《标题》\n正文内容"}
 
小说扩写：{"text": "User: 段落的开头\n\nAssistant: 完整的段落"}

整理完成的 jsonl 文件应该是这样的：

将 jsonl 文件转化为 binidx 文件

在 RWKV-LM/RWKV-v5 目录中运行以下命令，可以将 data 文件夹中的 jsonl 文件转成 binidx 文件 ：

python3 tools/preprocess_data.py --input ./data/sample.jsonl --output-prefix ./data/sample --vocab ./rwkv_vocab_v20230424.txt --dataset-impl mmap --tokenizer-type RWKVTokenizer --append-eod

如果你在使用自己的 jsonl 数据集，这条命令有两个需要修改的参数：

--input ./data/sample.jsonl # 作为输入的 jsonl 文件路径
--output-prefix ./data/sample # 输出的 bin 和 idx 文件路径

将 sample.jsonl 改成你自己的 jsonl 文件名称，其他参数保持默认即可。

调整路径参数

demo-training-run.sh 文件前三行是文件路径参数：

• load_model： 基底 RWKV 模型的路径
• proj_dir：训练日志和模型文件的输出路径
• data_file：训练数据集的路径，注意路径中不需要带 bin 和 idx 后缀，仅需文件名称。

调整 n_layer 和 n_embd 参数

基底 RWKV 模型的参数大小不同，训练时使用的 n_layer 和 n_embd 数值也不一样。以下是不同参数的 RWKV 模型和 n_layer 和 n_embd 数值的对照列表：

调整重要训练参数

调整其他训练参数

下面列出了脚本中其他可修改的训练参数，及其修改的效果。

参数调整完成后，请记得保存 demo-training-run.sh 文件。