Implement Gemma 2

[monatis]

It might be a good idea to split work for different Gemma versions. There's an example in the Flax repo that we can make use of a basis. We can also get inspired from that implementation in other ways, i.e., about some utilities such as saw_lib etc. for the class in #10.

[keshan]

It will be bit challenging but I will give this a try.

[mdda]

The flax.nxx examples version for gemma2-2b was proven out (i.e. layer outputs match the GDM gemma2 version) in my fork. My PR hasn't been incorporated into the main flax repo, because (apparently) the Google team's update deserved a 2 week head-start, and then I got pulled into other things.

The key updates in my fork are listed here:

• New TransformerConfig for gemma2-2b model
  • Renames existing configs to make them uniform
• Added additional params reading key adjustment (key in gemma2-2b needs remapping from kaggle download)
• Adds GQA to Attention module
• Reorders the operations in Block module so that logits output from overall Transformer are not gibberish logits confirmed to (approximately) match those from GDE gemma (flax linen) model (this was also fixed in the Google team's update, AFAICT)

If you'd like to collaborate, please let me know (though I'm going to be a bit tied up with ICLR next week...)

[monatis]

Hi @mdda, thanks for joining here and congrats on your PR. It seems like a great update. We would definitely like to collaborate on it. Would you like to integrate it in this repo yourself, or should we move on and acknowledge your work / request your review?

[monatis]

Architecture implemented in #29 but needs a little work to fix weight conversion.