Exclude optimizer state when running to_huggingface

[hengtaoguo]

Description

• Restore only model parameters during load_orbax_checkpoint (only used by to_huggingface conversion).
• Exclude unnecessary optimizer state to reduce memory and I/O usage.
• Use Orbax partial restore with the filtered checkpoint tree.
• Preserve base and LoRA checkpoint merging behavior.
• (opt_state for Linen, optimizer for NNX).

Result: Peak memory cut down in half for llama3.1-8b: 123.64 GB -> 62.60 GB

Fixes: b/538004926

Tests

roundtrip_conversion test still passes.

I generated train checkpoint using:

python -m maxtext.trainers.pre_train.train base_output_directory=gs://hengtaoguo-maxtext-logs/train/llama3.1-8b/2026-07-28-13-52 dataset_path=gs://maxtext-dataset tokenizer_type=huggingface load_parameters_path=gs://maxtext-model-checkpoints/llama3.1-8b/2025-01-23-19-04/unscanned/checkpoints/0/items per_device_batch_size=1 run_name=llama31_8b_train max_target_length=2048 steps=5 async_checkpointing=false checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False model_name=llama3.1-8b scan_layers=false use_multimodal=false ici_tensor_parallelism=4 hf_access_token=xxx

Checkpoint conversion for to_huggingface (downloaded gcs ckpt to local dir for faster debugging):

python -m maxtext.checkpoint_conversion.to_huggingface maxtext/configs/base.yml model_name=llama3.1-8b load_parameters_path=/dev/shm/hengtaoguo/llama3.1-8b/items base_output_directory=/dev/shm/hengtaoguo/llama3.1-8b/hf hf_access_token=xxx scan_layers=false override_model_config=true 

Before the fix:

# logs
I0728 21:24:03.478436 139669224541312 utils.py:671] ✅ Model and tokenizer (if provided) successfully processed for /dev/shm/hengtaoguo/llama3.1-8b/hf
I0728 21:24:03.478489 139669224541312 to_huggingface.py:435] ✅ MaxText model successfully saved in HuggingFace format at /dev/shm/hengtaoguo/llama3.1-8b/hf
I0728 21:24:03.478526 139669224541312 to_huggingface.py:437] Elapse for transform and save: 0.33 min
I0728 21:24:03.480079 139669224541312 to_huggingface.py:552] Overall Elapse: 1.04 min
I0728 21:24:03.480337 139669224541312 utils.py:786] Peak Memory: 123.64 GB

After the fix:

# logs
I0728 21:47:19.765015 140550674111616 utils.py:671] ✅ Model and tokenizer (if provided) successfully processed for /dev/shm/hengtaoguo/llama3.1-8b/hf
I0728 21:47:19.765071 140550674111616 to_huggingface.py:435] ✅ MaxText model successfully saved in HuggingFace format at /dev/shm/hengtaoguo/llama3.1-8b/hf
I0728 21:47:19.765115 140550674111616 to_huggingface.py:437] Elapse for transform and save: 0.23 min
I0728 21:47:19.766717 140550674111616 to_huggingface.py:552] Overall Elapse: 0.56 min
I0728 21:47:19.766989 140550674111616 utils.py:786] Peak Memory: 62.60 GB

After conversion, validated the correctness

python -m tests.utils.forward_pass_logit_checker maxtext/configs/base.yml model_name=llama3.1-8b tokenizer_path=/dev/shm/hengtaoguo/llama3.1-8b/hf load_parameters_path=/dev/shm/hengtaoguo/llama3.1-8b/items use_multimodal=false per_device_batch_size=1 scan_layers=false --max_kl_div=0.1 --run_hf_model=true --hf_model_path=/dev/shm/hengtaoguo/llama3.1-8b/hf

Forward pass logits checker against itself

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 8.16e-04
INFO:absl:Per-token KL Divergences: 
['-2.29e-04', '4.90e-03', '-7.08e-04', '-7.00e-04']
INFO:absl:
Max KL divergence for a single token in the set: 4.90e-03
INFO:absl:
--- Prompt: What is the ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 6811       | difference           | 9.8750     |
| 1888       | best                 | 9.3125     |
| 19463      | probability          | 8.6875     |
| 1455       | most                 | 8.4375     |
| 836        | name                 | 8.4375     |
| 6864       | capital              | 8.3750     |
| 7580       | purpose              | 8.1250     |
| 3560       | role                 | 8.0000     |
| 5133       | relationship         | 7.9688     |
| 7438       | meaning              | 7.8125     |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 6811       | difference           | 9.8750     |
| 1888       | best                 | 9.3750     |
| 19463      | probability          | 8.6875     |
| 6864       | capital              | 8.4375     |
| 836        | name                 | 8.4375     |
| 1455       | most                 | 8.3750     |
| 7580       | purpose              | 8.0625     |
| 3560       | role                 | 8.0000     |
| 5133       | relationship         | 7.9062     |
| 7438       | meaning              | 7.8125     |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 80.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 2.10e-03
INFO:absl:Per-token KL Divergences: 
['-2.29e-04', '1.20e-03', '5.12e-03', '2.30e-03']
INFO:absl:
Max KL divergence for a single token in the set: 5.12e-03

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	10 Missing ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.