LoRA 微调是什么？
LORA（Low-Rank Adaptation）是一种针对大型预训练模型的微调技术。它不改变原始模型大部分参数，而是调整模型的部分权重，以此实现对特定任务的优化。
本文的 LoRA 微调方法来自 RWKV 社区微调项目 RWKV-PEFT (opens in a new tab) 。
在开始 LoRA 微调之前，请确保你拥有一个 Linux 工作区，以及支持 CUDA 的 NVIDIA 显卡。
LoRA 显存参考
RWKV LoRA 微调的显存（GPU VRAM）需求可参考下表：
| 模型参数 | bf16 | int8 | nf4 |
|---|---|---|---|
| RWKV-6-1.6B | 7.3GB | 5.9GB | 5.4GB |
| RWKV-6-3B | 11.8GB | 9.4GB | 8.1GB |
| RWKV-6-7B | 23.7GB | 17.3GB | 14.9GB |
上表的数据基于以下训练参数：
随着训练参数的变更，RWKV LoRA 微调所需显存会发生改变。
收集训练数据
你需要使用收集更适合训练 RWKV 的 binidx 数据，具体方法可参考 准备微调数据集 (opens in a new tab) 。
配置训练环境
要训练 RWKV 模型，首先要配置 conda 等训练环境。具体过程请参考 RWKV 微调环境配置 (opens in a new tab) 板块。
克隆仓库并安装依赖
在 Linux 或 WSL 中，使用 git 命令克隆 RWKV-PEFT 仓库：
git clone https://github.com/JL-er/RWKV-PEFT.git
克隆完成后，使用
cd RWKV-PEFT
命令进入 RWKV-PEFT 目录。并运行以下命令，安装项目所需依赖：
pip install -r requirements.txt修改训练参数
使用任意文本编辑器（如 vscode）打开
RWKV-PEFT/scripts
目录下的
run_lora.sh
文件，可以修改训练参数，进而控制微调的训练过程和训练效果：
以下是一次 LoRA 微调的调参过程：
调整路径参数
调整 n_layer 和 n_embd 参数
以下 RWKV 模型参数对应的 n_layer/n_embd 值：
| 模型参数 | n_layer | n_embd |
|---|---|---|
| 0.1B | 12 | 768 |
| 0.4B | 24 | 1024 |
| 1.5B | 24 | 2048 |
| 3B | 32 | 2560 |
| 7B | 32 | 4096 |
| 14B | 61 | 4096 |
调整重要训练参数
| 参数 | 描述 |
|---|---|
micro_bsz=1
|
微批次大小，根据显存大小调整，微调时从 1 开始逐渐增大 |
epoch_save=5
|
每隔多少个训练轮次保存一次 LoRA 文件，注意存储空间是否充足 |
epoch_steps=1000
|
每个训练轮次的步数，增加会拉长单个 epoch 的训练时间 |
ctx_len=512
|
微调模型的上下文长度，建议根据语料长度修改 |
调整 LoRA 相关参数
| 参数 | 描述 |
|---|---|
| "lora_load":"" | LoRA 文件路径，代表从哪个 LoRA checkpoint 开始微调。若从头训练 LoRA 可不填 |
| "lora_r":32 | LoRA 微调的 rank 参数，值越大效果越好，但训练速度越慢/显存需求越高，一般训练使用 32 或者 64 即可 |
| "lora_alpha":32 | LoRA 微调的 alpha 参数（缩放因子），建议保持 lora_r 的两倍 |
| "lora_dropout":0.01 | LoRA 微调的丢弃率，建议使用 0.01 |
调整其他训练参数
| 参数 | 描述 |
|---|---|
--data_type binidx
|
训练语料的文件格式，支持："utf-8", "utf-16le", "numpy", "binidx", "dummy", "wds_img", "uint16" |
--vocab_size 65536
|
词表大小，默认为 65536，设置为 0 表示模型自动确定词汇表大小 |
--epoch_count 5
|
总训练轮次 |
--epoch_begin 0
|
初始训练轮次，即从第 N 个训练轮次开始加载 |
--pre_ffn 0
|
用 ffn 替换第一个 att 层，通常保持默认值 0 |
--head_qk 0
|
通常保持默认值 0，即关闭状态 |
--lr_init 5e-5
|
初始学习率，LoRA 建议 5e-5 ，最大不超过 1e-4 |
--lr_final 5e-5
|
最终学习率，建议和初始学习率保持一致 |
--warmup_steps 0
|
预热步骤数，默认 0，加载模型时可尝试改成 50 |
--beta1 0.9
|
Adam 优化器的 beta1 参数，保持默认值 |
--beta2 0.99
|
Adam 优化器的 beta2 参数，保持默认值 |
--adam_eps 1e-8
|
Adam 优化器的 epsilon 参数，保持默认值 |
--accelerator gpu
|
使用的加速器类型，目前主要支持 gpu， cpu 基本不支持训练 |
--devices 1
|
显卡数量，单显卡填 1，多卡按实际数量填写 |
--precision bf16
|
训练精度，建议保持默认值 bf16，支持："fp32", "tf32", "fp16", "bf16" |
--strategy deepspeed_stage_1
|
lightning 训练策略参数，微调推荐使用 deepspeed_stage_1，设备显存太小可将 1 改成 2 |
--grad_cp 1
|
梯度累积步数，0 训练更快但需更多显存，1 训练较慢但节省显存 |
--my_testing "x060"
|
训练的 RWKV 模型版本，v5 选 x052，v6 选 x060 |
--dataload pad
|
数据加载选项，"pad" 支持
bsz>1
，"only" 则限制
bsz=1
|
--loss_mask pad
|
在数据末尾进行 padding，可改成 "qa" 以对 QA 任务中的问题部分进行屏蔽，防止模型根据问题来记忆答案，从而增强模型的泛化能力。 |
--fla
|
是否启用 fla 以减小显存需求，bsz 小于 8 时建议开启 |
--peft lora
|
训练类型，LoRA 微调填 lora，支持："lora", "state", "pissa", "bone" |
--quant int8/nf4
|
RWKV 默认使用 bf16 训练精度，但支持 int8 和 nf4 两种量化训练类型，推荐使用精度损失较小的 int8 |
--train_parts
|
指定微调的模型区域，如“emb”、“head”、“time”、“ln”。训练时建议移除此参数，以使用默认值 ["time", "ln"] |
--wandb RWKV-PEFT-LORA
|
是否使用 wandb 可视化记录训练日志，需提前配置 wandb (opens in a new tab) 账号 |
附录：run_lora.sh 配置参考
load_model='/home/rwkv/models/basemodel/3b.pth'
proj_dir='/home/rwkv/RWKV-PEFT/output_lora'
data_file='/home/rwkv/RWKV-PEFT/data/lora-test'
n_layer=32
n_embd=2560
micro_bsz=1
epoch_save=1
epoch_steps=500
ctx_len=1024
lora_config='{"lora_load":"","lora_r":32,"lora_alpha":64,"lora_dropout":0.01}'
python train.py --load_model $load_model \
--proj_dir $proj_dir --data_file $data_file \
--data_type binidx --vocab_size 65536 \
--ctx_len $ctx_len --epoch_steps $epoch_steps --epoch_count 5 --epoch_begin 0 --epoch_save $epoch_save --micro_bsz $micro_bsz \
--n_layer $n_layer --n_embd $n_embd \
--pre_ffn 0 --head_qk 0 --lr_init 5e-5 --lr_final 5e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
--accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 1 \
--my_testing "x060" \
--dataload pad --loss_mask pad \
--peft lora --lora_config $lora_config \
--wandb RWKV-PEFT-LORA
# 以下是可选项
# --fla
# --quant int8/nf4 （是否量化训练）
# --train_parts ["time", "ln"]（是否指定训练范围，不建议添加此参数）
在 RWKV-PEFT 目录，运行
sh scripts/run_lora.sh
命令，开启 LoRA 微调 。
正常开始训练后，应当是如下画面：
训练完毕后，应当可以在输出文件夹中找到训练好的 LoRA 权重文件（
.pth
格式）和训练日志（
.txt
文件）:
如何使用 LoRA 权重文件
得到 LoRA 权重文件后，你需要将它合并到基底 RWKV 模型中，获得一个完整的 LoRA 微调模型。
这个过程可以使用
scripts
目录的
merge_lora.sh
脚本文件来完成。
修改合并参数
使用文本编辑器打开
scripts
目录下的
merge_lora.sh
脚本文件，并修改脚本中的合并参数：
| 参数 | 描述 |
|---|---|
| base_model | LoRA 微调的基底 RWKV 模型路径，参考 run_lora.sh |
| lora_checkpoint | 训练得到的某个 LoRA 检查点文件路径 |
| output | 合并后的 LoRA 模型输出路径（包含模型命名） |
| TYPE='lora' | 合并类型：LoRA，无需修改 |
| Lora_alpha | 跟随训练时的 Lora_alpha 参数 |
| --quant | 跟随训练时的量化参数，"nf4"或者"int8"。如果训练时未使用量化，则无需添加此参数 |
合并参数参考：
base_model='/home/rwkv/models/basemodel/3b.pth'
lora_checkpoint='/home/rwkv/RWKV-PEFT/output_lora/rwkv-2.pth'
output='/home/rwkv/RWKV-PEFT/output_lora/peft-demo.pth'
TYPE='lora'
Lora_alpha=64
python merge/merge.py --base_model $base_model \
--lora_checkpoint $lora_checkpoint \
--output $output \
--type $TYPE \
--lora_alpha $Lora_alpha
# --quant int8/nf4 （和训练时的量化参数同步）使用 LoRA 权重文件
修改完成后保存文件，在
RWKV-PEFT
主目录运行
sh scripts/merge_lora.sh
命令，进行 LoRA 模型合并：