Gemma2 - 27B model with JAX/Flax implementation

[mini-goel]

Hello, is there a Gemma 2 - 27B model with JAX/Flax implementation available which is fine-tunable. The kaggle link has one but it mentions that it's 'not' fine-tunable. https://www.kaggle.com/models/google/gemma-2/flax

[Gopi-Uppari]

Hi @mini-goel,

There is no official JAX/Flax implementation for fine-tuning Gemma 2 - 27B yet. However, there are existing resources on fine-tuning smaller Gemma models using JAX and Flax. These can be really useful for understanding the fine-tuning process and could potentially be adapted for larger models in the future.

1. Fine-tuning Gemma using JAX and Flax: To fine-tune the Gemma 2B Instruct model for tasks like English-French translation using the gemma library alongside JAX and Flax.

2. Fine-tuning RecurrentGemma using JAX and Flax: his guide focuses on the RecurrentGemma 2B Instruct model, providing insights into fine-tuning with JAX and Flax for similar translation tasks.

Thank you.

[Conchylicultor]

Hey @mini-goel , yes, this repository is the official implementation for fine-tuning Gemma 2. We have a few examples:

• I would recommend you to start by this tutorial, that you can adapt for Gemma 27B: https://gemma-llm.readthedocs.io/en/latest/colab_finetuning.html
• If using LoRA, you can see our examples in the documentation: https://gemma-llm.readthedocs.io/en/latest/colab_lora_finetuning.html
• Finally, for a end-to-end example to run locally, you can look at our examples folder, like: https://github.com/google-deepmind/gemma/blob/main/examples/next_token_prediction.py

Let us know if something isn't working for you!