本项目来自 RWKV 社区开辟者 zyaaa-ux ,项目链接:https://github.com/zyaaa-ux/ROSA-Tuning。
本项目为社区提出的一种 ROSA 实现,不代表 RWKV-8 ROSA 的实际机能,后果供参考。
本项目提出 ROSA-Tuning,一种经由过程检索回想机制加强预练习模型长高低文建模才能的办法。该办法在传统留意力机制之外并行引入基于 CPU 的 ROSA(RWKV>pip install torch transformers datasets deepspeed numba numpy
![]()
![]()
参加 RWKV 社区
然后运行下列敕令,获取项目代码:
git clone https://github.com/zyaaa-ux/ROSA-Tuning预备 DeepSpeed 设备文件
项目应用了 DeepSpeed 进行加快,是以须要在本地创建一个
deepspeed_config.json文件,示例如下:假如显存足够,可以删除
offload_optimizer中的pin_memory参数,并将device的值修改为 none,来获得更快的运行速度。
修改设备
2025.12.29 qkv_update.py 的 68~73 行定义了路径参数,须要修改路径参数为你本地的路径。
MODEL_LOCAL_DIR = "/path/to/base/model/" # 本地基本模型路径MODEL_DIR = "/path/to/checkpoint/" # 模型检查点保存路径DATASET_DIR = "/path/to/processed/dataset/" # 数据集路径OUTPUT_DIR = "/path/to/output/" # 输前程径DEEPSPEED_CONFIG_PATH = "/path/to/deepspeed/config.json" # DeepSpeed 设备文件路径
假如须要加倍节俭显存,可以修改代码第 119 行动 True,打开梯度累计:
GRADIENT_CHECKPOINTING = True # 源代码是 False
假如本地无 flash-attn 库,可以修改第 78 行的代码,封闭 flash-attn 的应用:
USE_FLASH_ATTN = False # 本来是 True
运行启动敕令
因为应用了 DeepSpeed 和分布式练习逻辑(is_main_process 等检查),推荐 deepspeed 敕令启动。
可选安装 flash-attn 库,该库可以或许晋升代码运行速度,但初次安装时须要编译。
deepspeed --num_gpus=1 2025.12.29 qkv_update.py
启动成功后,会输出以下内容:

该图为应用 200 条长度为 128 的数据在单卡 4090 长进行流程测试的示例,实际练习 16k 长度数据时须要很大年夜的显存。
🧠 道理概述
迎接大年夜家参加 RWKV 社区,可以从 RWKV 中文官网懂得 RWKV 模型,也可以参加 RWKV 论坛、QQ 频道和 QQ 群聊,一路商量 RWKV 模型。
- 📖 RWKV 中文文档:https://www.rwkv.cn
- 💬 RWKV 论坛:https://community.rwkv.cn/
- 🐧 QQ 频道:https://pd.qq.com/s/9n21eravc
- 📺 BiliBili 视频教程:https://space.bilibili.com/3546689096910933
{ "fp16": { "enabled": "auto", "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, "hysteresis": 2, "min_loss_scale": 1 }, "bf16": { "enabled": "auto" }, "zero_optimization": { "stage": 2, "allgather_partitions": true, "allgather_bucket_size": 200000000, "overlap_comm": true, "reduce_scatter": true, "reduce_bucket_size": 200000000, "contiguous_gradients": true, "offload_optimizer": { "device": "cpu", "pin_memory": true }, "offload_param": { "device": "none" } }, "gradient_accumulation_steps": "auto", "train_batch_size": "auto", "train_micro_batch_size_per_gpu": "auto", "gradient_clipping": "auto", "steps_per_print": 20, "wall_clock_breakdown": false}迎接大年夜家基于 RWKV-7 进行创业、科研,我们也会为基于 RWKV 的项目供给技巧支撑。
假如您的团队正在基于 RWKV 创业或开展研究,请接洽我们!(在“RWKV元始智能”微信"大众,"号留言您的接洽方法,或发送邮件到“contact@rwkvos.com”。)


发表评论 取消回复