Skip to content

Update Jupyter Notebooks to convert MaxText checkpoint to Hugging Face#3959

Merged
copybara-service[bot] merged 1 commit into
mainfrom
yixuannwang-dev-demo-nb
May 21, 2026
Merged

Update Jupyter Notebooks to convert MaxText checkpoint to Hugging Face#3959
copybara-service[bot] merged 1 commit into
mainfrom
yixuannwang-dev-demo-nb

Conversation

@YixuanWang-99
Copy link
Copy Markdown
Collaborator

@YixuanWang-99 YixuanWang-99 commented May 20, 2026

Description

Add a new cell of checkpoint conversion from maxtext to huggingface in SFT and RL Jupyter Notebooks.

The main changes:

  • A new cell to convert matext checkpoint to hugging face format, in both Jupyter Notebooks
  • An updated config initialization in to_huggingface.py to be compatible with the new pydantic configs.

Tests

  • The updated SFT and RL Jupyter Notebooks are tested on TPU v6e-8, outputs are stored in GCS
Screenshot 2026-05-20 at 2 10 50 PM Screenshot 2026-05-20 at 11 24 19 AM
  • The updated to_huggingface.py is tested using test_qwen3_to_hf.sh.

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.

@codecov
Copy link
Copy Markdown

codecov Bot commented May 20, 2026

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...rc/maxtext/checkpoint_conversion/to_huggingface.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

Comment thread src/maxtext/examples/rl_llama3_demo.ipynb Outdated
Comment thread src/maxtext/examples/rl_llama3_demo.ipynb Outdated
Comment thread src/maxtext/examples/rl_llama3_demo.ipynb
Comment thread src/maxtext/examples/sft_llama3_demo_tpu.ipynb Outdated
Comment thread src/maxtext/examples/sft_llama3_demo_tpu.ipynb
Comment thread src/maxtext/examples/sft_llama3_demo_tpu.ipynb
Copy link
Copy Markdown
Collaborator

@hengtaoguo hengtaoguo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@YixuanWang-99 YixuanWang-99 force-pushed the yixuannwang-dev-demo-nb branch from d5b3d0d to 171f932 Compare May 21, 2026 01:04
@copybara-service copybara-service Bot merged commit fb79a9e into main May 21, 2026
34 of 36 checks passed
@copybara-service copybara-service Bot deleted the yixuannwang-dev-demo-nb branch May 21, 2026 19:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants