RWKV

RWKV-LM 预训练教程

本教程基于社区维护的 RWKV-LM-V7，适合快速复现约 191M～3B 的 RWKV-7 预训练。该仓库仍处于 WIP，官方参考实现请查看 RWKV-LM 的 RWKV-v7/train_temp。

预训练案例

  • RWKV-7-Arithmetic-0.1B：基于 RWKV-7 架构预训练的 0.1B 模型，实现基础加减法运算和加减法方程求解功能。

视频教程

高画质视频请跳转到 B 站观看。

硬件需求

在相同模型规模和上下文长度下，预训练通常比 LoRA、MiSS 等参数高效微调需要更多显存。

以下是基于 HEAD_CHUNK=4096，GRAD_CP=1，KERNEL="@rwkv3" 参数，在不同 CTX_LEN 下预训练 RWKV 模型所需的显存：

模型参数层数维度M_BSZ=1 显存其他 M_BSZ 显存
0.1B127686786MB（6.63GB）BSZ8：11412MB（11.14GB）
BSZ48：40556MB（39.61GB） × 4
0.4B24102412938MB（12.63GB）BSZ6：18828MB（18.39GB）
BSZ8：20366MB（19.89GB）
1.5B24204837300MB（36.43GB）BSZ4：21340MB（20.84GB） × 4
3B32256067766MB（66.18GB）-
7B32409655424MB（54.13GB） × 4-
14B6140964 × Pro 6000 OOM-
模型参数层数维度M_BSZ=1 显存其他 M_BSZ 显存
0.1B127686446MB（6.29GB）BSZ8：9018MB（8.81GB）
0.4B24102412374MB（12.08GB）BSZ6：14894MB（14.54GB）
1.5B24204836114MB（35.27GB）-
3B32256066132MB（64.58GB）-
7B32409654864MB × 4（53.58GB × 4）-
14B6140964 × Pro 6000 OOM-

表格中标注 × 4 指通过 GPU_PER_NODE=4 指定四张 GPU 并行训练时，每张 GPU 所需的显存。

克隆仓库

git clone https://github.com/RWKV-Vibe/RWKV-LM-V7.git
# 或者国内镜像仓库
git clone https://gitee.com/rwkv-vibe/RWKV-LM-V7.git

无法克隆？复制并依次运行以下命令，下载 RWKV-LM-V7 仓库压缩包并解压：

wget -O rwkv.zip https://github.com/RWKV-Vibe/RWKV-LM-V7/archive/refs/heads/main.zip && \
unzip -q rwkv.zip && \
rm rwkv.zip

使用 zip 下载的代码仓库不含 git 信息，无法使用 git pull 更新仓库

准备训练环境

  1. 安装 CUDA Toolkit

请参考 CUDA Toolkit 安装文档，根据您的操作系统安装 CUDA Toolkit。

以下是 Ubuntu 24.04 x86_64 的安装示例：

# 下载 CUDA 仓库的 GPG 密钥包 
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
# 安装 GPG 密钥包，使系统信任 NVIDIA 的软件仓库
sudo dpkg -i cuda-keyring_1.1-1_all.deb
# 更新 apt 软件包列表，更新 NVIDIA CUDA 软件源
sudo apt-get update
# 安装 CUDA Toolkit 13.0（包含 nvcc 编译器、开发库、工具链等）
sudo apt-get -y install cuda-toolkit-13-0
  1. 使用 miniforge 等 conda 兼容包管理器，创建一个全新的虚拟环境

以下是 MiniConda-Linux x86_64 的安装和环境创建示例：

# 下载最新的 MiniConda 安装包
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
# 运行 MiniConda 安装包的安装脚本，安装期间一定要注意进展，及时输入 yes  ！
sh Miniconda3-latest-Linux-x86_64.sh
# 重启环境变量，激活 Conda 环境
source ~/.bashrc
# 创建名为 rwkv-lm-v7 的 conda 环境, 并指定 python 版本为 3.12
conda create -n rwkv-lm-v7 python=3.12
# 激活 rwkv 环境
conda activate rwkv-lm-v7
  1. 安装下列依赖，注意 pytorch-lightning 固定使用了 1.9.5 版本，此为本仓库特性，请不要升级此依赖包。
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
# 安装 CUDA 版本的 PyTorch，这里的 cu128 指的是 CUDA 12.8 版本，可以根据你自己的 CUDA 版本进行调整
pip3 install -r requirements.txt

下载缓慢？尝试添加 -i https://mirrors.aliyun.com/pypi/simple 参数，使用阿里源加快下载速度。

安装完成后，使用以下命令检查 PyTorch 版本以及 PyTorch 能否使用 CUDA：

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

若输出类似 2.7.1+cu128 True，表示当前 PyTorch 已经识别 CUDA GPU。这个结果不能用于判断系统 CUDA Toolkit 或 nvcc 是否可用。

准备训练数据

请参考准备训练数据 文档，准备一个 jsonl 格式的预训练数据集，并将其放在 RWKV-LM-V7/data 目录下。

在 RWKV-LM-V7/data 目录中运行以下命令，将 jsonl 数据文件转成 binidx 格式：

python make_data.py demo.jsonl 30 4096

命令中的 30 表示复制次数，4096 表示上下文长度。make_data.py 将执行以下操作：

  • 对 demo.jsonl 进行 30 次复制和打乱
  • 加载复制后的 demo.jsonl 并基于 rwkv_vocab_v20230424 词表进行分词
  • 把数据保存为 binidx 格式文件：demo.bin 和 demo.idx
  • 针对 4096 上下文长度，计算出预训练所需的 my_exit_tokens 和 magic_prime 参数

make-data-py

务必保存命令行输出的 --my_exit_tokens、--magic_prime 和 --ctx_len 参数，这些参数会在后续的预训练步骤中频繁使用。

在此示例中，我们将使用一个混合了中文/大写中文、英文、半角/全角阿拉伯数字、正负数、小数的加减法多轮对话数据集，预训练一个 RWKV-7 0.1B 加减法模型。

pretrain-showcase-dataset

初始化 RWKV 模型

在开始预训练之前，我们需要初始化一个 RWKV 模型，作为预训练的起点。

使用文本编辑器打开 RWKV-LM-V7 目录的 demo-training-prepare.sh 脚本，修改以下初始化参数：

参数参数解释
MODEL_TYPE="x070"训练的模型版本，当前 RWKV-LM-V7 使用 RWKV-7 的 x070 架构
N_LAYER="12" N_EMBD="768"模型层数和维度，维度和层数决定预训练模型的参数大小，建议参考硬件需求修改，N_EMBD 必须 64 的倍数
CTX_LEN="4096"预训练的上下文长度，必须是 16 的倍数，并与准备训练数据时使用的 --ctx_len 参数相同
data_file预训练的数据路径，使用已转换的 bin 和 idx 数据，无需文件名后缀
--my_exit_tokens训练数据集的总 token 数，会在全部 token 训练完毕后退出，在 make_data.py 中计算得到
--magic_prime训练数据集的magic_prime 值，在 make_data.py 中计算得到

其余参数请保持默认值，修改完毕保存 demo-training-prepare.sh 文件，然后在 RWKV-LM-V7 目录打开终端，并在终端运行 sh demo-training-prepare.sh 命令初始化 RWKV 模型。

demo-training-prepare.sh

预训练阶段

调整训练参数

初始化完成后，使用文本编辑器打开 RWKV-LM-V7 目录的 demo-training-run.sh 脚本，修改训练参数。

这些训练参数必须和初始化阶段保持一致：

参数参数解释
MODEL_TYPE="x070"训练的模型版本，当前 RWKV-LM-V7 使用 RWKV-7 的 x070 架构
N_LAYER="12" N_EMBD="768"模型层数和维度，维度和层数决定预训练模型的参数大小，建议参考硬件需求修改，N_EMBD 必须 64 的倍数
CTX_LEN="4096"预训练的上下文长度，必须是 16 的倍数，并与初始化模型、准备训练数据时使用的值一致
data_file预训练的数据路径，使用已转换的 bin 和 idx 数据，无需文件名后缀
--my_exit_tokens训练数据集的总 token 数，会在全部 token 训练完毕后退出，在 make_data.py 中计算得到
--magic_prime训练数据集的magic_prime 值，在 make_data.py 中计算得到

其他训练参数的解释和参考值如下：

部分参数会影响显存占用和训练速度，修改前请确保正确理解参数含义。

参数参数解释
PROJ_DIR模型输出目录，必须和预训练阶段保持一致，建议保持默认
M_BSZ=16建议为 2 的幂；显存允许时可以适当增大，以提高训练吞吐
LR_INIT="6e-4"初始学习率，计算公式为 0.45 / N_EMBD 并适当取整。 比如 L12-D768 0.1B 模型，初始学习率为 0.45/768=0.0005859375，取整为 6e-4。继续预训练的学习率和从头预训练稍有不同。
LR_FINAL="6e-5"最终学习率，计算公式为 0.04 / N_EMBD 并适当取整。
GRAD_CP=1是否启用梯度检查点；GRAD_CP=1 节省显存但训练较慢，GRAD_CP=0 训练更快但消耗更多显存
HEAD_CHUNK=4096LM HEAD；填 0 训练快但消耗的 VRAM 更多；65536 较慢但消耗的 VRAM 较少；4096 更慢，但需要的 VRAM 更少。（注意：这是 LM HEAD，非 RWKV HEAD）
KERNEL="@rwkv3"建议保持 "@rwkv3" ，运行更快，尤其是针对 H100 系列。可通过设置为 "" 调整为旧版 kernel
EPOCH_SAVE=10每隔多少个 "miniepochs" 保存一次训练模型（1 miniepoch = 40320 * ctx_len tokens）
epoch_steps会自动根据 M_BSZ 等参数进行计算，计算公式为 epoch_steps = 40320 / M_BSZ / N_NODE / GPU_PER_NODE 。
N_NODE=1节点数，通常保持默认值 1
GPU_PER_NODE=1每个节点上的 GPU 数量，单显卡填 1，多卡则改为实际数量
DS_BUCKET_MB=2DeepSpeed bucket size（单位 MB）。当前示例脚本已禁用该参数；如需启用，请先根据固定的 DeepSpeed 版本验证兼容性
--load_model "0"保持默认值 0。训练中断后，继续训练会自动检测最新检查点
--wandb "Test"建议注册一个 WandB 账号并填写你的项目名称，以便观察和比较 loss。详情查看附录 使用 WandB 监控训练过程
--train_stage 3预训练阶段，保持默认值 3。
--epoch_count 999999总训练轮次，该参数不生效。训练进程会在达到 --my_exit_tokens 指定的 tokens 数量后自动退出
--epoch_begin 0初始训练轮次，始终写 0，自动加载最新的检查点
--warmup_steps 10预热步骤，应当根据优化器和学习率的选取进行实验
--beta1 0.9Adam 优化器 beta1 参数，保持默认值
--beta2 0.99Adam 优化器 beta2 参数，保持默认值
--adam_eps 1e-18Adam 优化器的 epsilon 参数，小的 epsilon 更稳定，保持默认值即可
--data_type binidx当前训练脚本仅支持 binidx 数据
--vocab_size 65536词表大小，必须与训练数据使用的 tokenizer 一致；本教程使用 RWKV World 词表，保持 65536
--weight_decay 0.001权重衰减，保持默认值 0.001
--head_size 64头大小，保持默认值 64
--accelerator gpu加速器类型，必须是 gpu
--precision bf16训练精度，当前 RWKV-7 CUDA 算子使用 bf16，请保持该值
--strategy deepspeed_stage_2训练策略，默认 deepspeed_stage_2，更多训练策略请参考DeepSpeed 训练策略
--enable_progress_bar True是否在终端显示进度条，通常保持默认值 True

开始预训练

demo-training-run.sh 开头包含删除部分检查点的 rm 命令。如果 PROJ_DIR 中已有需要保留的检查点，请先注释这些命令并核对目标目录。

训练参数调整完毕后，保存 demo-training-run.sh 文件，并在 RWKV-LM-V7 目录下运行 sh demo-training-run.sh 命令开启预训练。

如果你在训练参数中配置了 --wandb "Test"，开启训练后可访问终端中的 WandB 链接（https://wandb.ai/xxx），可视化查看当前训练的 loss 曲线等信息。

pre-training-run

测试预训练模型

预训练完成后，在 RWKV-LM-V7/out 目录中可以找到训练的中间检查点 rwkv-0/1/2/...pth 和最终模型文件 rwkv-final.pth。

pretrain-get-model

可以使用 RWKV Runner 或者 RWKV pip - API_DEMO_CHAT.py 脚本测试预训练模型。除了最终模型文件 rwkv-final.pth，也可以测试最后几个模型检查点。

test-pretrain-model

基于官方 RWKV 模型继续预训练

初始化 RWKV 模型时会获得一个 rwkv-init.pth 初始化检查点。运行 sh demo-training-run.sh 命令时，训练程序会从这个初始化检查点开始训练。

利用这个原理，我们可以将 RWKV 官方模型改成 rwkv-init.pth，替换初始化目录中的检查点，以实现从 RWKV 官方模型继续预训练。

从 RWKV 官方模型继续训练时，需要确保 demo-training-run.sh 脚本中的训练参数 N_LAYER / N_EMBD 和 RWKV 模型的层数、维度完全对齐。

RWKV 模型参数层数 N_LAYER维度 N_EMBD
0.1B12768
0.4B241024
1.5/1.6B242048
2.9/3B322560
7/7.2B324096
13.3/14B614096

继续预训练 RWKV7-G1 模型时，请在 src/model.py 中核对 D_DECAY_LORA、D_AAA_LORA、D_MV_LORA 和 D_GATE_LORA，确保这些维度与基底模型参数完全对齐。

params0.1B0.4B1.5B2.9B7.2B13.3B
D_DECAY_LORAw64649696128192
D_AAA_LORAa64649696128192
D_MV_LORAv3232646496128
D_GATE_LORAg128128256320480384

附录

DeepSpeed 训练策略

本教程使用的 PyTorch Lightning 提供以下 DeepSpeed 策略名称：

策略名称分片内容主要取舍
deepspeed_stage_1优化器状态通信开销相对较低，GPU 显存节省有限
deepspeed_stage_2优化器状态和梯度进一步节省 GPU 显存，同时增加通信开销
deepspeed_stage_2_offloadStage 2，并将优化器状态卸载到 CPU继续降低 GPU 显存占用，速度取决于 CPU 与传输带宽
deepspeed_stage_3优化器状态、梯度和模型参数GPU 显存节省最多，但通信与调度更复杂
deepspeed_stage_3_offloadStage 3，并将优化器状态和模型参数卸载到 CPU进一步降低 GPU 显存占用，训练速度更依赖传输带宽

本页示例使用 deepspeed_stage_2。实际训练时应根据模型规模、GPU 数量、显存和互联带宽选择策略。

参考文档：

使用 WandB 监控训练过程

WandB 是一个用于实验跟踪和模型评估的工具，可以帮助我们更好地监控训练过程和 loss 数据。

使用 WandB 监控训练过程：

  1. 访问 WandB 官网，点击 "Sign up" 按钮注册一个 WandB 账号（可以使用邮箱、GitHub、Google 等方式注册）
  2. 注册成功后，点击右上角头像进入 “Settings” 页面，在 “API Keys” 栏中点击 “+ New Key” 创建一个新的 API Key，并复制下来
  3. 在你的训练设备打开终端，运行 wandb login 命令，粘贴你复制的 API Key，按回车确认
  4. 在训练脚本中添加 --wandb "Test" 参数，训练日志上传到你的 WandB 项目页面，项目名为 Test
  5. 你可以在 WandB 主页中点击该项目，实时查看训练过程中的损失（loss）曲线、学习率变化等指标

wandb-screenshot

常见问题

这份文档对您有帮助吗？
联系我们© 2026 RWKV. All rights reserved.粤ICP备2024242518号-1粤公网安备44030002010697号粤公网安备44030002010697号

On this page

预训练案例视频教程硬件需求克隆仓库准备训练环境准备训练数据初始化 RWKV 模型预训练阶段调整训练参数开始预训练测试预训练模型基于官方 RWKV 模型继续预训练附录DeepSpeed 训练策略使用 WandB 监控训练过程常见问题