ParaRNN：基于Triton Newton扫描实现非线性RNN（sLSTM、RWKV-7、CfC、Titans）的并行训练 [P]

我开发了ParaRNN这一开源库，它属于PyTorch/Triton系列，能够利用映射到关联前缀扫描的牛顿-拉夫森迭代机制，实现非线性循环架构的并行训练。虽然非线性RNN在推理时无需额外的KV缓存即可实现O(1)时间复杂度，但其传统训练方式却因缓慢的顺序展开过程而成为瓶颈。ParaRNN通过定制化的Triton求解器解决了这一问题：相比传统顺序BPTT算法，其在处理长序列时的速度提升了约200倍（例如，在RTX 3060上，CfC模型在序列长度为2048时的训练时间从643毫秒缩短至2.79毫秒）。该算法在131,072个token的上下文长度内仍能保持稳定收敛，实验证明其牛顿迭代过程中的$K\$值始终控制在3以内，而顺序展开方式的$\tau \le 10{-4}$的约束也得以满足。目前该库支持的RNN单元包括：sLSTM（采用4x4块雅可比矩阵及Beck式头部混合机制）、RWKV-7（Goose矩阵状态单体结构）、CfC（液态连续时间模型）、Titans（关联式惊喜梯度下降算法）、Modern Hopfield模型以及M²RNN模型。在生产环境中，用户可通过torch.compile(fullgraph=True)实现无图结构断开的优化，同时获得确定性的封闭形式VJPs计算结果（无需使用原子操作），此外还支持树外vLLM插件。更多信息请访问GitHub地址：https://github.com/bugkira/pararnn-torch，快速安装命令为pip install pararnn-torch。非常期待收到您的反馈、针对自定义数据结构的错误报告，以及在不同硬件环境下的性能测试数据！