Skip to content

123mbcz123/RWKV-Tensorflow2.0

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

11 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

RWKV4-Tensorflow2.0

A RWKV written in TensorFlow 2.0,only tested under FP32 precision in the World model.
RWKV_MODEL.py is the TensorFlow version implementation of RWKV, while the remaining files were copied from chatRWKV for testing accuracy with this repository.

一个使用Tensorflow2.0编写的RWKV模型加载器.目前只支持FP32下加载模型,模型方面我只测试了World 0.1B/0.4B模型

RWKV_MODEL.py是RWKV的完整实现,其余的文件由ChatRWKV库拷贝编辑而得.其余文件仅仅是为了测试tf版本实现的正确性,一般情况下无需下载.

此仓库与ChatRWKV v2(经过一点修改)后误差平均值小于1.3238086e-05

这是我翻译模型的一部分,如果没有咕咕的话会继续更新.

下一步我会为RWKV添加训练部分,然后再尝试把注意力机制带回RWKV

RWKV_MODEL 是第一版实现,TimeMix基于RNN实现,但是存在注释混乱的问题,代码经过拆分但是注释没动
RWKV_MODEL_Attention 是第二版实现,现在基本是不可用的存在. 这个版本加入了MultiQueryAttention,但是未经过测试,这个版本着重引入了并行版本的RWKV,但是存在精度问题,需要下一个版本解决.

上两版都不建议使用,更加完善的第三版将在解决并行状态下精度问题后提交,到时候也会对Attention机制进行拆分.Attention将采用类似插件的方式或者集成的方式挂载到RWKV上.

暂时找不到在不写算子的情况下,可以使RWKV4高效并行训练的方法,想使用可以使用第一版,串行推理的效率与ChatRWKV接近。

About

A RWKV written in TensorFlow 2.0,only tested under FP32 precision in the World model.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages