feat: Auto-detect and convert Tunix post-training checkpoints on load

[hsuan-lun-chiang]

Overview

This PR updates the checkpoint loader in MaxText to auto-detect and dynamically convert Tunix post-training checkpoints on load, allowing downstream MaxText pre-training runs to restore from Tunix DPO/SFT checkpoints seamlessly.

Key Changes

• Auto-Detection: Identifies Tunix-formatted checkpoints by inspecting the checkpoint layout (presence of model_params collection).
• On-the-fly PyTree Reshaping:
  • Dynamically unwraps Optax inject_hyperparams wrapper states (_drop_inject_hyperparams).
  • Strips LoRA adapter layers (_drop_adapter_level) when mapping from adapter states to base pre-training states, handling dicts, lists, tuples, and namedtuples.
  • Correctly repacks base adapter levels (_add_adapter_level) when querying metadata or loading checkpoints with has_base=True.
• model_creation_utils.from_pretrained: Directs parameter inspection and single-item PyTree restore to model_params when loading Tunix checkpoints.

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request implements on-the-fly detection and conversion of Tunix-formatted checkpoints during the pre-training load process, avoiding manual checkpoint conversion. The feedback highlights several critical issues: first, setting metadata to None on exception in model_creation_utils.py will cause an AttributeError when accessing metadata.item_metadata later; second, _drop_adapter_level must recurse into lists, tuples, and namedtuples to handle nested optimizer states; and third, an _add_adapter_level helper is required to correctly wrap optimizer state fields in base when restoring adapter checkpoints to prevent structure mismatches.

[codecov]

Codecov Report

❌ Patch coverage is 26.74419% with 126 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/checkpointing.py	20.14%	104 Missing and 7 partials ⚠️
src/maxtext/utils/model_creation_utils.py	51.61%	8 Missing and 7 partials ⚠️

📢 Thoughts on this report? Let us know!