Add Flax implementation for BLOOM

[haileyschoelkopf]

Model description

I'm interested in adding an implementation of BLOOM in Flax.

The implementation shouldn't be too bad since the pytorch implementation can serve as a guide and a way to check correctness.

Open source status

• The model implementation is available
• The model weights are available

Provide useful links for the implementation

@younesbelkada @stas00 @patrickvonplaten

If someone is already planning to work on this then no worries, but if not I will start on this as soon as I have time!

[younesbelkada]

Hi!
Thank you very much for the contribution!
On my side it's a green light since I am not working on it, and it is not on my plans for now. Therefore, I'll be happy to review it! Let us know if you want to work on that :)

[haileyschoelkopf]

Thanks! I will open a WIP PR soon and tag you there once I do.

[patrickvonplaten]

Very cool idea - think this can also be a flagship project where we can showcase how to fine-tune BLOOM with Flax cc @patil-suraj @sanchit-gandhi

[patil-suraj]

Awesome! Would be very happy to help with it :)

[sanchit-gandhi]

Great idea! Would also be interested in getting involved, this would be a super cool model addition!

[haileyschoelkopf]

Thanks everyone for the interest! I'd love to collaborate with you all.

I'm hoping to push a rough draft of modeling code by the end of the weekend (earlier if I have time), and will tag you all when I open the PR with that. Does that sound alright?

[haileyschoelkopf]

I've opened a PR (and documented the state of the in-progress code I'm still working on) at #17761 ! We can discuss further in that PR how to collaborate / proceed.