要训练 RWKV 模型，请先遵循以下步骤，为你的 Linux 系统配置训练环境。
视频教程
高画质视频请跳转到 B 站 (opens in a new tab)观看。
配置 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 环境。
安装训练所需的软件
要获得 RWKV 模型的最佳训练性能，推荐安装如下软件：
- torch 2.1.2+cu121（或最新版本的 torch）
- 最新的 deepspeed
- pytorch-lightning ==1.9.5
在终端的 Conda 环境中依次运行以下命令，安装 RWKV 训练环境：
# 通过指定 url 安装 CUDA 12.1 版本的 torch
pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121
# 指定安装 1.9.5 版本的 PyTorch Lightning 和其他常用的机器学习工具包
# --upgrade 参数意味着如果环境中已经安装了对应的软件包，则将其升级到最新版本
pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade
# 以下是使用 RWKV-PEFT 微调时会用到的一些其他工具包，建议提前安装
pip install bitsandbytes einops triton rwkv-fla rwkv transformers GPUtil plotly gradio datasets推荐的软件版本为 RWKV 训练的最佳实践，也可以选择安装其他版本的软件，保证兼容性即可。
检查 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: True 。则可能需要重新安装 CUDA 环境（建议 cuda 12.1 或更新版本），安装步骤请参阅 CUDA Toolkit 安装文档 (opens in a new tab)。