💡 Feature Request: QLoRA Training

Hi!
I just discovered this library and I have to say this is the best thing I have seen in LLMs so far! Not only is JAX so much more elegant and well written then PyTorch, this whole library is feature packed with so much stuff I would need several, badly integrating ones in the HuggingFace / PyTorch eco-system. Although I'm missing one thing: bitsandbytes. I know that you have 4 bit inference built into this thing but I really like the ability in Transformers to train in 4 bits as well (really neat for LoRA fine tuning).
Do you think it would be possible to add this here as well? The groundwork seems to be there, but I'm running into incompatible types errors.
I might be able to contribute this, but I'll need some pointers.