JAX-based library for simple language model training and RLHF?

[cool-RR]

Hi everyone!

I'm working on a research project around LLM RLHF training. I want to find a JAX-based library that allows me to do that. I found around 10 of them, with the most promising ones being MaxText and AXLearn; however, for both of them the setup is quite complicated. They are both meant for large-scale training of big LLMs on cloud hardware, while I prefer a dead-simple library that does everything locally on my NVidia GPU, and I'm okay with the language model being a lot smaller and modest.

Can anyone recommend such a library?

Thanks,
Ram.

[jakevdp]

I wonder if the examples here might set you in the right direction? https://github.com/jax-ml/jax-llm-examples

These are minimal examples of various well-known LLM architectures implemented in pure JAX, rather than using any framework.

[cool-RR]

Thank you, I'll check these out. So far I tried MaxText and Tunix, and I'm liking Tunix better.