社区项目ROSA-Tuning：验证RWKV-8 ROSA效果
本项目来自 RWKV 社区开发者 zyaaa-ux ，项目链接：https://github.com/zyaaa-ux/ROSA-Tuning。 本项目为社区提出的一种 ROSA 实现，不代表 RWKV-8 ROSA 的实际性能，效果供参考。 本项目提出 ROSA-Tuning，一种通过检索回忆机制增强预训练模型长上下文建模能力的方法。该方法在传统注意力机制之外并行引入基于 CPU 的 ROSA（RWKV Online Suffix Automaton）检索模块，从长上下文中高效定位与当前查询相关的历史位置，并以可训练的方式将检索到的信息注入模型状态，故随后的加权融合可由状态受限的高效注意力完成。 为实现端到端训练，本文设计了二值离散化策略与反事实梯度算法，并通过 CPU–GPU 异步流水线进一步优化整体执行效率。 在 Qwen3-Base-1.7B 上的系统性评估结果表明，ROSA-Tuning 能够显著恢复窗口注意力模型的长上下文建模能力，在 LongBench 等基准上取得接近甚至匹配全局注意力的性能，同时保持与窗口注意力方法几乎相当的计算效率与显存占用，为高效长上下文处理提供了一条新的技术路径。 性能测试 困惑度 (PPL) 对比 在 PG-19 数据集上的测试显示，ROSA 适配器成功修复了窗口注意力的 PPL 劣化，甚至优于全局注意力。 Model PPL (越低越好) Global Attention 18.96 Windowed Attention 74.50 Windowed + ROSA 17.63 实验设置: 基础模型：具有全局注意力或窗口注意力的 Qwen3-Base-0.6B 训练：28,000 个样本，在 PG-19 训练集上训练 1 个轮次，原始模型冻结，仅训练 ROSA 适配器 评估：PG-19 测试集，序列长度 16k，窗口大小 1024 长文本能力 (LongBench) 在大海捞针 (NIAH) 测试，ROSA 实现了 100% 的召回率。综合评分恢复至全局注意力的 96.5%。 Task / Metric Global Attention Windowed (2048) Windowed + ROSA NIAH (大海捞针) 100.00 6.20 100.00 TriviaQA 86.20 61.56 84.34 Multi_news 23.23 10.43 23.76 Samsum 42.04 32.51 40.53 TREC 72.67 52.67 68.00 Gov_report 31.11 13.08 26.19 LongBench 平均分 59.21 29.41 57.14 实验设置: 基础模型：Qwen3-1.7B-Base，具有全局注意力或窗口化注意力（窗口大小2048） 训练数据：约 37B tokens，其中约 30B 来自 prolong，约 7B 来自其他上下文推理数据集，且不与测试集重叠 使用方法 项目作者进行了非常多的实验，本次介绍 2025 年 12 月 29 日更新的 2025.12.29 qkv_update.py 的使用方法。 注意在本地准备 Hugging Face datasets 库保存到磁盘的格式（Arrow 格式）的数据。 环境准备和代码获取 首先运行下列代码安装需要的库： pip install torch transformers datasets deepspeed numba numpy 可选安装 flash-attn 库，该库能够提升代码运行速度，但首次安装时需要编译。 然后运行下列命令，获取项目代码： git clone https://github.com/zyaaa-ux/ROSA-Tuning 准备 DeepSpeed 配置文件 项目使用了 DeepSpeed 进行加速，因此需要在本地创建一个 deepspeed_config.json 文件，示例如下： { "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 } 如果显存足够，可以删除 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 命令启动。 deepspeed --num_gpus=1 2025.12.29 qkv_update.py 启动成功后，会输出以下内容： 该图为使用 200 条长度为 128 的数据在单卡 4090 上进行流程测试的示例，实际训练 16k 长度数据时需要很大的显存。 🧠 原理概述 加入 RWKV 社区 欢迎大家加入 RWKV 社区，可以从 RWKV 中文官网了解 RWKV 模型，也可以加入 RWKV 论坛、QQ 频道和 QQ 群聊，一起探讨 RWKV 模型。 📖 RWKV 中文文档：https://www.rwkv.cn 💬 RWKV 论坛：https://community.rwkv.cn/ 🐧 QQ 频道：https://pd.qq.com/s/9n21eravc | QQ 交流群：224287095 📺 BiliBili 视频教程：https://space.bilibili.com/3546689096910933 欢迎大家基于 RWKV-7 进行创业、科研，我们也会为基于 RWKV 的项目提供技术支持。 如果您的团队正在基于 RWKV 创业或开展研究，请联系我们！（在“RWKV元始智能”微信公众号留言您的联系方式，或发送邮件到“contact@rwkvos.com”。）
