add Gemma2 support to MaxText

[ZhaoyueCheng]

add Gemma2 support to MaxText on top of PR

• add gemma Decoder block which merges [local_sliding_attention, global_attention] combination into one Decoder layer with post_attn_norm, post_ffw_norm support
• add convert_gemma2_chkpt.py to convert checkpoint from Gemma2 architecture
  • merges one [local_sliding_attention, global_attention] combination into one Decoder layer
  • adds post_attn_norm, post_ffw_norm, transpose_gating_einsum and query_pre_attn_scalar support which is new from Gemma1
• enabling post_attn_norm, post_ffw_norm in the Gemma2 Decoder architecture
• add final_logits_soft_cap support
• add gemma2-2b gemma2-9b gemma2-27b config yaml file and 2b end_to_end test script and test golden logits dumped from flax

[gagika]

Some quick comments

[gagika]

Do we want to open source all the notebooks?

[ZhaoyueCheng]

I think we had the original gemma export notebook and other notebooks (llama, mixtral, etc) open sourced so added the gemma2 notebook to scratch_code folder as well with other open source notebooks

[salrowili]

Please note that if we merge local_sliding_attention, global_attention and assigned half the decoder layer to "base_num_decoder_layers", the TFLOPS calculation will report half the actual TFLOPS . This is due to how maxtext_utils calcuate the TFLOPS
https://github.com/google/maxtext/blob/644eb87ae90dd8b210ce17f1c16ca7a54e80fceb/MaxText/maxtext_utils.py#L139

[gobbleturk]

Please note that if we merge local_sliding_attention, global_attention and assigned half the decoder layer to "base_num_decoder_layers", the TFLOPS calculation will report half the actual TFLOPS . This is due to how maxtext_utils calcuate the TFLOPS

https://github.com/google/maxtext/blob/644eb87ae90dd8b210ce17f1c16ca7a54e80fceb/MaxText/maxtext_utils.py#L139

Good catch! Thank you for noting this

[ZhaoyueCheng]

Please note that if we merge local_sliding_attention, global_attention and assigned half the decoder layer to "base_num_decoder_layers", the TFLOPS calculation will report half the actual TFLOPS . This is due to how maxtext_utils calcuate the TFLOPS

https://github.com/google/maxtext/blob/644eb87ae90dd8b210ce17f1c16ca7a54e80fceb/MaxText/maxtext_utils.py#L139

updated, thanks for the comment!

[khatwanimohit]

Thanks for adding this! Left a few nits.

QQ: Are we going to add logit checker for Gemma9B in a separate PR? Since you have the code for golden logits here, can you also add golden logits jsonl file for Gemma2-9B in this PR

[khatwanimohit]

nit: Update the copyright to 2024

[ZhaoyueCheng]

updated, thanks for the review!

[gobbleturk]

Overall looks very good, thank you for adding this! Just want to clean up the flops calculation.

[gobbleturk]

Awesome, thank you for fixing the tflops calculation!

[ZhaoyueCheng]

Awesome, thank you for fixing the tflops calculation!

Thanks for the detailed review and suggestions!!

[salrowili]

@ZhaoyueCheng I would like to thank you for adding Gemma 2 support which really is appreciated.

I have pre-trained Gemma 2 using maxtext but now i stuck since i could not convert MaxText to hugging face format to do the SFT stage. Gemma2, unlike Llama, Gemma and Mistral, uses local and global attention . Is there any thing you can do to "MaxText/llama_or_mistral_ckpt.py" script to include Gemma 2 checkpoint conversion to HF format? The global and local attention used in the following weights:

mlp_global
mlp_local
post_ffw_norm_global
post_ffw_norm_local
post_self_attention_norm_global
post_self_attention_norm_local
pre_ffw_norm_global
pre_ffw_norm_local
pre_self_attention_norm_global
pre_self_attention_norm_local
self_attention_global

This related to #829