Skip to content

Conversation

@RissyRan
Copy link
Collaborator

@RissyRan RissyRan commented Nov 5, 2025

Description

Move all legacy checkpoint scripts from folder src/MaxText into folder src/MaxText/utils/ckpt_scripts to align with new checkpoint utility src/MaxText/utils/ckpt_conversion. We have plan to migrate all legacy scripts into conversion utility in Q1 2026 timeline.

This PR helps us to consolidate all scripts, and resolves b/452391826.

Tests

  • Expect all tests passing
  • Successfully verified a small model with new path gemma3-4b - link
export CHKPT_BUCKET=gs://maxtext-gemma/gemma3/flax
python3 -m MaxText.utils.ckpt_scripts.convert_gemma3_chkpt --base_model_path ${CHKPT_BUCKET}/4b --maxtext_model_path gs://runner-maxtext-logs/ranran/gemma3 --model_size 4b

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.

@github-actions
Copy link

github-actions bot commented Nov 5, 2025

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

Copy link

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

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

📋 Review Summary

This Pull Request successfully consolidates legacy checkpoint scripts into a new src/MaxText/utils/ckpt_scripts folder. The changes are well-executed, primarily involving path updates across documentation, shell scripts, and Python import statements.

🔍 General Feedback

  • All file renames and path updates are consistent and correctly reflect the new directory structure.
  • The changes align with the stated goal of organizing checkpoint utilities.

@shuningjin
Copy link
Collaborator

shuningjin commented Nov 5, 2025

Do we also move dequantize_mxfp4.py here, which we use for gpt-oss dequantization?

  • It won't be a legacy script though. I see you moved deepseek_fp8_to_bf16.py.

python3 dequantize_mxfp4.py --input-path=<input_path> --output-path=<output_path>

@RissyRan
Copy link
Collaborator Author

RissyRan commented Nov 6, 2025

dequantize_mxfp4

Good call! I will move that together.

Copy link
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!

Copy link
Collaborator

@parambole parambole left a comment

Choose a reason for hiding this comment

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

LGTM!

@copybara-service copybara-service bot merged commit 2ec8b35 into main Nov 6, 2025
39 checks passed
@copybara-service copybara-service bot deleted the migrate_script branch November 6, 2025 20:21
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.

6 participants