Enable Checkpoint Conversion from Huggingface to Maxtext

[YixuanWang-99]

Description

Enable checkpoint conversion from Huggingface to Maxtext.

• Add to_maxtext.py to perform the checkpoint conversion from HF to MaxText.
• Add convert_gemma2_to_mt.sh to automate the conversion and verification.
• Add mt_hf_mutual_conversion_check.py to compare the Huggingface and MaxText checkpoints.
• Official Gemma2 models are supported.

Tests

The converted checkpoint is tested with mt_hf_mutual_conversion_check.py. It compared:

1. For given prompts, top-k predicted tokens and scores for the next token;
2. KL divergence of the full logit distributions

Tested on Gemma2-2b Model. A successful conversion example.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[hengtaoguo]

Excellent work!

[hengtaoguo]

This is nice, thanks for adding such a feature!

[hengtaoguo]

Hi @gagika ! I've heard this might be interesting to you for loading/saving HF checkpoints. Would you like to take a look when you got a chance? Thanks a lot for your time!

[shralex]

Thanks Yixuan! Added a few comments

[shralex]

Thanks for addressing the comments! I have 1 small comment and also a question -- did you test both directions -- to and from HF ? if so can you add both to the PR description testing section, currently it includes 1 example. Thanks!

[shralex]

Could you please clarify what's the difference between this file and forward_pass_logits_checker ? this file is used in all our end-to-end tests to verify logits. If there is something missing in that file, I wonder if we could change it rather than creating a new file ?

[YixuanWang-99]

Thanks for addressing the comments! I have 1 small comment and also a question -- did you test both directions -- to and from HF ? if so can you add both to the PR description testing section, currently it includes 1 example. Thanks!

from HF conversion with examples is pushed in previous PR: #1785 and #1821. And I have revised the run name.