Support for RWKV

[philpax]

So this is a pretty immense task and I'd start with #45, but...

RWKV is an RNN with Transformer-level LLM performance, which can also be directly trained like a GPT transformer (parallelizable). And it's 100% attention-free. You only need the hidden state at position t to compute the state at position t+1. You can use the "GPT" mode to quickly compute the hidden state for the "RNN" mode.

So it's combining the best of RNN and transformer - great performance, fast inference, saves VRAM, fast training, "infinite" ctx_len, and free sentence embedding (using the final hidden state).

It's entirely open-source, so not legally burdened like LLaMA, and (from what I've seen) is more powerful than BLOOM at the same parameter count.

I asked the RWKV Discord which implementation would be worth looking at, and this is what I was told:

RWKV-LM/RWKV-v4neo/src/model.py is the implementation that's actually used to train the large models, it's cuda only and has tons of features you probably don't need.
rwkv_pip_package only implements inference, but is a good implementation and worth a look, recently got a lot more complex due to supporting more and more strategies and including various optimizations.
ChatRWKV/src/model_run is an older version, but haven't played with it so not sure how good it is. Might be worth a look since it's basically an older version of the one in rwkv_pip_package.
RWKV_in_150_lines.py I still haven't fully checked out, but I know it doesn't support GPT mode, so that may or may not be less useful
Also worth a look is RWKV-v4neo/src/model_run.py, which is a small inference-only impl capable of loading the large RWKV checkpoints
I'm not sure if it has GPT-mode, though

So it sounds like rwkv_pip_package is the way to go as source material:

https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py

The following articles are very useful for understanding how RWKV works:

• https://johanwind.github.io/2023/03/23/rwkv_overview.html
• https://johanwind.github.io/2023/03/23/rwkv_details.html

An interesting detail from the latter is the following:

The largest number a 16-bit floating point number (float16) can represent is 65 504, anything above that overflows, which is bad. Most of the code has no problems with this, partially because the Layer Normalizations keep values in a reasonable range. However, the RWKV attention contains exponentially large numbers (exp(bonus + k)). In practice, the RWKV attention is implemented in a way where we factor out an exponential factor from num and den to keep everything within float16 range. See for example the time_mixing function in RWKV in 150 lines.

This may pose issues for the GGML 4-bit quantisation format, which is non-optimal. We would likely want GPTQ quantisation.

[KerfuffleV2]

RWKV is really interesting, but the Python code is absolutely horrendous. I made an attempt to clean it up a little but the author wasn't interested.

I'd look at this one if you were interested in trying to make the Rust code just a frontend to a Python backend: https://github.com/harrisonvanderbyl/rwkvstic

But I sort of don't know what the benefit would be of doing that for a CLI app, because there's already stuff like Oobabooga's text thingy that handles a lot of different model formats with a front end.

I believe the rwkvstic repo I linked is on a path to get added to the actual Transformers repo as one of the supported models. If you wanted to do a Python frontend to stuff, maybe it would make more sense to look at making the frontend for Transformers rather than RWKV specifically. Then you'll just pretty much get that for free one support is merged.

huggingface/transformers#20737

huggingface/transformers#20809

---

Sort of related, but the TextSynth server (closed source) actually supports/publishes some 4bit RWKV models: https://bellard.org/ts_server/

It can do pretty fast inference on CPU. I contact the author and asked if he'd be willing to disclose the model format for those files and he said he probably would post a Python converter soon which would document the format as well.

[philpax]

Nah, we'd want to port the actual inference code to Rust, similarly to what we did for LLaMA itself. The less Python we have in the codebase, the better 🦀

[KerfuffleV2]

After a whole lot of struggling because I have essentially no idea of what I'm doing with NN or math stuff: https://github.com/KerfuffleV2/smolrsrwkv

Obviously that's not useful as anything more than an example of the basics, but it does work and generates tokens not all that much slower than the simple Torch version.

[philpax]

Oh wow, nice one! That's an excellent start - do you want to try porting it to ggml-rs? (#81 exposes it as a separate library, so you should be able to use a Git dependency)

[philpax]

(By the way, are you in the Discord? It'd be great to chat synchronously if need be)

[KerfuffleV2]

Thanks, I didn't really do much more than port the Python example here though: https://johanwind.github.io/2023/03/23/rwkv_details.html

do you want to try porting it to ggml-rs?

I'd have to say probably no. I don't think I'm qualified to write a version that's actually worth using in real software. So if you or anyone else wants to use my code as an example or starting point, you're certainly very welcome.

I may mess around with trying to get it to work on backends other than ndarray but I still don't think it'll be suitable for production. Also, sad to say, I have a very short attention span. Just generally speaking, you're probably better off not depending on me for anything long term.

[philpax]

I don't think I'm qualified to write a version that's actually worth using in real software.

Eh, none of us are really LLM experts - it's more a matter of engineering and refinement. I think you're just as qualified as the rest of us. Happy to help any efforts get across the line.

Also, sad to say, I have a very short attention span. Just generally speaking, you're probably better off not depending on me for anything long term.

Don't worry, I know the feeling very well. No worries at all; a driveby contribution is fine as long as it's maintainable, and from what I've seen of your implementation, you're doing fine in that regard.

I'd encourage you to give it a try, but no pressure at all - even if you don't do it, your port will still be a great resource for those coming after (especially with those sweet static types!)

[KerfuffleV2]

@philpax

I think you're just as qualified as the rest of us.

I don't know know about that! I also really don't know any math stuff either. So I think I'd probably more focus on trying to make it a good example for understanding how the model works in its basic form rather than something that could be used directly.

I'd encourage you to give it a try, but no pressure at all

I did look, but one thing I noticed is it doesn't seem like GGML supports the operations that are required (or I don't know how to work around/refactor based on it). For example, how to do stuff like subtract tensors from a value, or map something like exponent to the elements? It doesn't seem obvious how to do that with the API it provides.

Could be it provides a higher level way to accomplish the same stuff, but unfortunately I don't have a high level understanding of how it works so that's kind of beyond my ability.

Anyway, maybe if you or someone showed me how to do that, porting it to GGML would be more practical. Here's an example of the kind of thing I mean: https://github.com/KerfuffleV2/smolrsrwkv/blob/c21cb8008b51aa10fb2c0eaa2a81714e8c27f76f/src/model.rs#LL125C70-L125C70 (that's elementwise map).

even if you don't do it, your port will still be a great resource for those coming after (especially with those sweet static types!)

Haha, yeah, I honestly found the Python example very hard to follow since there were no types and it was using weird NumPy shortcuts like sometensors[a > b] = 0 — indexing with a boolean expression, what? Figuring out silly stuff like that actually took as much time as writing the actual logic.

I cleaned up/reorganized the Rust version a bit, so hopefully it's easier to understand now. The state part was particularly bad in the initial version.

[philpax]

I also really don't know any math stuff either.

You know your way around a type system, that's good enough for me :P

I did look, but one thing I noticed is it doesn't seem like GGML supports the operations that are required (or I don't know how to work around/refactor based on it). For example, how to do stuff like subtract tensors from a value, or map something like exponent to the elements? It doesn't seem obvious how to do that with the API it provides.

Oh... yeah. Just had some free time to look at this, and yeah, ggml doesn't support some critical operators (it doesn't seem to support any form of exponentiation...).

For subtracting tensors with a value, you'd use op_sub, potentially with new_f32 (if I'm interpreting what you're asking correctly). That being said, there's definitely non-zero challenges here - I was hoping that it was just a case of operations that we left out of the bindings, but it seems like there's actual holes here.

ggml's API is based around building up a computation graph and then executing all the computation at once, but it means that if you need an operation that it doesn't support you'll need to add it yourself.

weird NumPy shortcuts like sometensors[a > b] = 0 — indexing with a boolean expression, what?

Unfortunately, this is something I recognize from my MATLAB days - it's convenient, but it is absolutely bewildering the first time you see it 😅

I cleaned up/reorganized the Rust version a bit, so hopefully it's easier to understand now. The state part was particularly bad in the initial version.

Yeah, nice, appreciate it! Given that you can't actually implement a ggml version of this, I don't have any immediate suggestions for you, but you can explore the other Rust ML libraries if you're curious.

Still, tremendously cool to have a RWKV implementation in easy-to-understand Rust - you should be proud!

[KerfuffleV2]

After profiling it, 80-90% or more of the time is just spent doing the matrix multiplication. The rest isn't really significant. How crazy would it be to just use ggml for only that one operation. Even having to copy the data every time it probably still would be worth it.

It may be a lot faster, but that still wouldn't really make that code suitable for general use though since it does everything in f32 and there isn't any easy way around that with the current approach.

[philpax]

I'm not convinced that ggml would be faster for the f32 case; I think most of its benefits come through with quantized values, which would be hard to test with RWKV.

[KerfuffleV2]

Well, the big deal is currently it only uses a single thread. Supposedly ndarray supports threading but I've failed to get it to actually use more than one.

I also messed around with Rayon just running the 3 and 2 groups of .dot() operations in parallel, but it didn't really make a difference because there's one place it's forced to a single thread and one place where it can only run two in parallel. I only got like a 20-30% speed increase at best, when I have 6 cores and 12 threads available.

So just running those matmul ops in parallel would make a huge, huge difference I think.

[philpax]

Hmm, yeah, makes sense. Do the matmuls happen at the end or do they need to be copied back and forth?

[KerfuffleV2]

Probably have to be copied back and forth. You can look here: https://github.com/KerfuffleV2/smolrsrwkv/blob/main/src/model.rs

1. channel_mixing (3) — k and r are independent, but vk depends on k.
2. time_mixing (3) — All independent.
3. evaluate_layer (1) — Only happens once per step rather than per layer.

Each group is basically all together, but they need data from other calculations and further calculations depend on those results as well.

(By the way, I'm getting reasonably happy with how readable it is now that it's all split into logical modules. The main function is looking pretty clean. If you have time to skim it, what do you think of the current state as an example for other projects?)