RWKV无限上下文训练器

如果你是RWKV新手，建议先通过我们的wiki了解更多信息：https://wiki.rwkv.com/

RWKV训练器具备以下特点：

无训练上下文限制（通过BPTT实现）
支持DeepSpeed 3
集成HuggingFace数据集

使用这个实现，你可以在（近乎）恒定的显存消耗下训练任意长的上下文；以RWKV 7B为例，每增加1024/2048个标记（取决于你选择的ctx_len），显存消耗仅增加约2MB，这使得训练超过100万个标记的序列成为可能。

顺便说一下，训练代码已经大幅重构，现在使用PyTorch 2.0、Lightning 2.0和DeepSpeed 2.0，启动脚本现在依赖于LightningCLI，因此你会看到config-example.yaml文件包含所有开关，大多数是Lightning自身处理的标准开关。还有一些针对RWKV和数据集解析器的新开关。

要使用这个仓库，进入RWKV-v4neo目录并执行：

python3 lightning_trainer.py fit -c {your_config}.yaml

记得根据自己的需求修改配置。

有关各种选项的文档，请参见RWKV-v4neo/config-example.yaml

注意：由于当前实现不完整，没有状态梯度，bptt_truncate被强制设为true

环境设置

注意：在撰写本文时，CUDA 12.0与多GPU存在已知问题。请至少升级到CUDA 12.1或12.2，或者降级到11.8

以下是使用conda的虚拟环境设置，请根据你的使用情况进行相应修改：

# 新训练器需要ninja-build
sudo apt-get install ninja-build

# 更新conda及其包列表
conda update conda

# 创建虚拟环境，使用Python 3.10
# Python 3.11在torch.compile / h100s方面存在问题
# 如果你想使用3.11，需要安装nightly版本
conda create -n rwkv-infctx python=3.11 pip
conda activate rwkv-infctx

# 安装PyTorch（>=2.1.2）
conda install -y pytorch==2.1.2 torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
python -m pip install lightning==2.1.3 deepspeed==0.12.6

# 目前，为了让torch.compile + 3.11在某些平台上工作，你可能需要尝试nightly版本
# 这被认为是高度"不稳定"的
# ---
# conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch-nightly -c nvidia
# python -m pip install lightning==2.0.5 deepspeed==0.10.0

# 验证你的PyTorch版本
python -c "import torch; print(torch.__version__)"

# 安装所有其他各种依赖
# 注：我们使用python -m pip而不是直接使用pip，因为它可以解决虚拟环境不加载正确pip的问题
python -m pip install datasets transformers 
python -m pip install ninja numexpr jsonargparse 'jsonargparse[signatures]'
python -m pip install lm-dataformat ftfy sentencepiece tokenizers wandb

# 可选依赖，对运行notebooks等有用
python -m pip install papermill

或者，你可以使用requirements.txt（这可能无法正确安装pytorch-cuda，并且被发现与conda环境不兼容）

python3 -m pip install -r requirements.txt

由于Windows上的deepspeed问题，仅支持Linux环境。不推荐使用带Windows的WSL2，因为在这个过程中会有严重的性能损失（无法使用deepspeed卸载，速度慢约50%）

整体训练流程

初始化新模型或下载现有模型
- 要初始化新模型，使用python3 ./init_model.py --n_layer {层数} --n_embd {嵌入大小} --vocab_size {词汇量大小/neox/world} --skip-if-exists ../model/file/path.pth
设置config.yaml文件，根据你的基础模型/微调用例进行自定义
使用python3 preload_datapath.py {你的配置}.yaml预加载数据集
使用python3 lightning_trainer.py fit -c {你的配置}.yaml开始训练过程
训练完成后，使用python3 export_checkpoint.py ../path/to/checkpoint/last.ckpt/ ../path/to/export/model.pth导出检查点
可选：运行dragon prompt作为快速理智检查python3 dragon_test.py ../path/to/export/model.pth

总结一下代码，从训练器目录（例如RWKV-v4neo）开始：

# 初始化空白模型（或下载预训练模型）
python3 init_model.py --n_layer {层数} --n_embd {嵌入大小} --vocab_size {词汇量大小/neox/world} --skip-if-exists ../model/file/path.pth

# 预加载你的数据集
python3 preload_datapath.py {你的配置}.yaml

# 运行训练过程
python3 lightning_trainer.py fit -c {你的配置}.yaml

# 将检查点导出为模型代码
python3 export_checkpoint.py ../path/to/checkpoint/last.ckpt/ ../path/to/export/model.pth

# 使用dragon prompt快速测试模型
python3 dragon_test.py ../path/to/export/model.pth

# @TODO，将模型转换为bf16格式（而不是现在的巨大fp32格式）
#        目前你需要使用RWKV pip包通过Python代码来实现这一点：
#        https://pypi.org/project/rwkv/

配置文件示例

你可以在以下位置找到相关笔记本/示例：

各种配置的完整注释位于 ./RWKV-v4neo/config-example.py
最小配置示例位于 ./RWKV-v4neo/config-example.py
各种数据集用例的配置/笔记本在这里
@TODO: 特定训练场景的示例

对于配置问题，请先查看上述列出的示例，然后再在Discord上提问。

你可以在我们的Discord上找到训练频道：https://discord.com/channels/992359628979568762/992362252269256815

infctx lightning训练器的重要注意事项

确保你的主机不运行CUDA 12.0（使用11.8或>=12.1），因为已知这会导致冻结问题
从检查点恢复时，估计时间不准确。参见：https://github.com/Lightning-AI/lightning/issues/18220
注意一些术语可能令人困惑，以下是简要术语表
- 下方进度条中的"step"指的是每个GPU的1个数据样本
- 经典Transformer的一个批次在wandb中是"trainer/global_step"
- wandb中的"substep"指单个数据样本 -"(累积梯度批次 * GPU数量) substeps = 1 trainer/global_step"

我应该使用官方RWKV-LM训练器还是infctx训练器？

通常，如果你从头开始训练基础模型 - 具有固定上下文大小，并且需要在多个节点（例如，10个装满A100服务器的节点）上获得绝对最高的吞吐量，官方训练器会表现更好（根据设置可能快2倍）

如果你需要DeepSpeed 3支持，或者处理动态数据集，这个训练器在几乎所有其他用例中都更加灵活。

随着时间推移，我们优化infctx训练器，与官方训练器的差距应该会缩小，但这不是最高优先级（infctx能工作 > 绝对速度）

一些长期架构目标

CUDA应该是可选的
- 展望未来，这使我们可能在其他架构上进行训练（即使性能有所损失），如AMD ROCM、TPU或Apple M1架构。
不依赖官方RWKV pip包
- 这是一个有意识的选择，以帮助促进#rwkv-x开发中模型架构的轻松迭代。这样，设计更改的整个训练-测试-验证过程可以在此存储库中完成。

现有限制

以下功能尚不支持（可能存在于blink的原始仓库中）

numpy文件数据集
模型调整权重（从较小模型初始化到较大模型）
向现有模型添加新令牌的辅助脚本
不支持torch compile，因为在nightly构建中这一直不稳定
尚不支持LoRA，现在请使用https://github.com/blealtan/RWKV-LM-LoRA

指定维护者

@picocreator - 是项目的当前维护者，如果你对这个项目有任何问题，可以在RWKV discord上联系他

致谢（v4neo和v5代码）

第一个infctx训练器的大部分内容最初由@Blealtan重写，位于：https://github.com/Blealtan/RWKV-LM-LoRA/tree/dev-infctx
RWKV-LM和原始训练器代码归功于@BlinkDL，位于：https://github.com/BlinkDL/RWKV-LM
特别感谢我们RWKV discord上的@Yuzaboto和@bananaman，他们的帮助对调试和修复仓库以分别适用于RWKVv4和RWKVv5代码至关重要。
@picocreator使项目功能完整，适用于RWKV主线发布