微调环境配置
要训练 RWKV 模型,请先遵循以下步骤,为你的 Linux 系统配置训练环境。
视频教程
高画质视频请跳转到 B 站观看。
安装 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 配置 Conda 虚拟环境
RWKV 模型的微调训练默认虚拟环境为 Conda。此处以 MiniConda 为例,介绍如何配置 Linux 系统的 Conda 环境。
在你的 Linux 或 WSL 工作区,依次运行以下命令:
# 下载最新的 MiniConda 安装包
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
# 运行 MiniConda 安装包的安装脚本,安装期间一定要注意进展,及时输入 yes !
sh Miniconda3-latest-Linux-x86_64.sh -u
# 重启环境变量,激活 Conda 环境
source ~/.bashrc此时我们处于 Conda 默认的 base 环境中。为了避免软件版本冲突,我们需要新建并激活一个名为 rwkv 的 conda 环境。依次运行以下命令:
# 创建名为 rwkv 的 conda 环境, 并指定 python 版本为 3.10
conda create -n rwkv python=3.10
# 激活 rwkv 环境
conda activate rwkv
至此,我们已经配置了 RWKV 微调所需的 Conda 环境。
安装训练所需的软件
在终端的 Conda 环境中依次运行以下命令,安装 RWKV 训练环境:
# 通过指定 url 安装 CUDA 12.8 版本的最新 torch
pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu128
# 安装 PyTorch Lightning 和其他常用的机器学习工具包,--upgrade 参数将软件升级到最新版本
pip install lightning deepspeed wandb ninja --upgrade
# 以下是使用 RWKV-PEFT 微调时会用到的一些其他工具包,建议提前安装
pip install bitsandbytes einops triton rwkv-fla rwkv transformers GPUtil plotly datasets --upgrade推荐的软件版本为 RWKV 微调的最佳实践,也可以选择安装其他版本的软件,保证兼容性即可。
下载缓慢?尝试在命令后添加 -i https://mirrors.aliyun.com/pypi/simple 参数,使用阿里源加快下载速度。
检查 CUDA 环境
上述依赖安装完毕后,你应该已经同时安装好最新版本的 Torch 和 CUDA 12.1 工具包。
依次运行以下命令,检查 PyTorch 的版本和 CUDA 工具是否可用:
python3 # 启动 python3
import torch # 导入 torch 包
# 打印 PyTorch 版本
print(f"PyTorch version: {torch.__version__}")
# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")
# 打印 CUDA 版本
print(f"CUDA version: {torch.version.cuda}")
# 打印可用的 GPU 数量
print(f"Number of GPUs: {torch.cuda.device_count()}")如果 print(f"CUDA available: {torch.cuda.is_available()}") 命令返回 CUDA available: True ,证明 CUDA 环境可用:

如果返回 False ,则使用 print(f"CUDA version: {torch.version.cuda}") 命令,检查是否正确安装了 CUDA 版本的 torch。
如果安装了 CUDA 版本的 torch ,仍然返回 CUDA available: False 。则可能需要重新安装 CUDA 环境(建议 cuda 12.1 或更新版本),安装步骤请参阅 CUDA Toolkit 安装文档。
这份文档对您有帮助吗?