Progress / Capabilities

[jasonmhead]

How far along is this project?
What can be done with the code so far?

[tensorpro]

The code checked in has been validated against official torch implementations, with the softmaxes aligned with the torch version.

It can be used to train on TPUs, I'm able to fine tune the 1.5B param model on a v3-8 with model parallelism.

For training on GPUs, I would recommend using the official implementation since the Cuda WKV kernel is more performant than the lax.scan version I user here.