Add unified tool to inspect checkpoint structures

[shuningjin]

Description

Unified script to inspect checkpoint structures for HF/MaxText/Orbax, to help with bringup and debugging.

Fix: b/484416862

Feature: Inspection Tool

A unified command-line tool to inspect checkpoint structures and tensor shapes. It is designed to be lightweight and minimize memory/compute costs across three supported modes:

1. HuggingFace Checkpoints (Locally downloaded file)

• Input: local path (safetensors or pth).
• Cost - safetensors: Lightweight with near-zero cost. Parses JSON file headers to read metadata instantly without allocating RAM.
• Cost - pth: Per-tensor load with small cost. Deserializes the file to load weights per-tensor (via mmap). Incurs a time cost for deserialization, but the host memory footprint is strictly constrained.
• Prints: Param key / shape.

2. MaxText Model Architecture (On-the-fly)

• Input: model name, scan mode. Run on CPU or TPU.
• Cost: Lightweight with near-zero cost. Traces JAX shapes abstractly to evaluate theoretical parameters dynamically without executing compute or allocating VRAM/HBM. Mechanism: jax.eval_shape
• Prints: Param key / shape + total parameter count.

3. Orbax Checkpoints (Pre-saved disk/GCS file)

• Input: GCS path or local path. Run on CPU or TPU.
• Cost: Lightweight with near-zero cost. Reads the structural metadata tree instantly without pulling the underlying heavy TensorStore data chunks. Mechanism: Read structural metadata.
• Prints: Param key / shape.

Usage Examples

[Mode 1: HuggingFace]   
  python src/maxtext/checkpoint_conversion/inspect_checkpoint.py hf \
      --path <local_hf_path> --format <safetensors | pth>

[Mode 2: MaxText Architecture] 
  python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext \
      model_name=<maxtext_model_name> scan_layers=<True | False>
      (Optional: other maxtext config)

[Mode 3: Orbax]        
  python src/maxtext/checkpoint_conversion/inspect_checkpoint.py orbax \
      --path <local_orbax_path | gcs_orbax_path>

Additional Features:

• --check_dtype: bool, default to False. Appends the data type of each tensor to the output.
• --output_file: str, default to empty. Path to save the output structure.

Add Documentation: When to use Tool

1. docs/guides/checkpointing_solutions/convert_checkpoint.md

• In Section Development and Troubleshooting, introduced a new ### Inspection Tools section documenting the inspect_checkpoint.py utility. It provides clear command-line examples for checking the three supported checkpoint formats.
• In Section Adding New Models, introduced using Inspection Tools to create mappings.
• In Section Common Errors, introduced using Inspection Tools to debug. The most common error (Type ShapeDtypeStruct is not a valid JAX type) is caused by structural mismatches.

2. docs/guides/model_bringup.md

• In Section 3.1 Create Mapping, introduced using Inspection Tools to create mappings.

Tests

1 HF checkpoint, locally downloaded

safetensors: no cost

python -m maxtext.checkpoint_conversion.inspect_checkpoint hf \
--path ~/deepseek2-16b/hf-bf16 \
--format safetensors \
--check_dtype True --output_file ~/checkpoint/test.txt  # Optional

• log: https://paste.googleplex.com/6521569454194688
• output file: https://paste.googleplex.com/5392498586419200
• Peak Memory: 1.09 GB

python -m maxtext.checkpoint_conversion.inspect_checkpoint hf \
--path /mnt/disks/ds32/deepseek3.2/hf-bf16 \
--format safetensors \
--check_dtype True --output_file ~/checkpoint/test.txt  # Optional

• log: https://paste.googleplex.com/4574738088329216
• Peak Memory: 1.11 GB

pth: low cost with per-tensor load

python -m maxtext.checkpoint_conversion.inspect_checkpoint hf \
--path ~/llama2-7b/hf-bf16 \
--format pth \
--check_dtype True --output_file ~/checkpoint/test.txt  # Optional

• log: https://paste.googleplex.com/4895937938980864
• output file: https://paste.googleplex.com/5891138853666816
• Peak Memory: 1.09 GB

2 MaxText model, on-fly, no cost

python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \
model_name=deepseek2-16b scan_layers=True \
--check_dtype True --output_file ~/checkpoint/test.txt  # Optional

• log: https://paste.googleplex.com/5636796997304320
• output file: https://paste.googleplex.com/5253647192686592

python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \
model_name=deepseek2-16b scan_layers=False \
--check_dtype True --output_file ~/checkpoint/test.txt  # Optional

• log: https://paste.googleplex.com/5392016862216192
• output file: https://paste.googleplex.com/5745144996429824

3 Orbax checkpoint, pre-saved, no cost

python -m maxtext.checkpoint_conversion.inspect_checkpoint orbax  \
--path gs://maxtext-deepseek/deepseek3-671b/2025-03-31/scanned/0/items \
--check_dtype True --output_file ~/checkpoint/test.txt  # Optional

• log: https://paste.googleplex.com/5267549800497152
• output file: https://paste.googleplex.com/5626295399612416

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]

Codecov Report

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

Files with missing lines	Patch %	Lines
...axtext/checkpoint_conversion/inspect_checkpoint.py	0.00%	144 Missing ⚠️

📢 Thoughts on this report? Let us know!

[RissyRan]

Thanks for making our process better!

[github-actions]

🤖 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.

[github-actions]

## 📋 Review Summary

This Pull Request introduces a unified checkpoint inspector tool for HuggingFace, MaxText architecture, and Orbax. The tool is a great addition for debugging model bring-ups. However, there is a significant bug in the MaxText mode where layer indices are ignored during path flattening, leading to incomplete output. Additionally, there are opportunities to improve memory efficiency when inspecting large checkpoints.

🔍 General Feedback

• Bug Fix Required: The MaxText architecture inspection currently collapses all layers into a single set of keys because it ignores SequenceKey indices. This needs to be addressed to provide a complete view of the model structure.
• Memory Efficiency: For safetensors, using get_slice instead of get_tensor avoids unnecessary data loading. For PyTorch checkpoints, better handling of large files and common state-dict wrappers would make the tool more robust.
• Consistency: Standardizing separators across different modes (e.g., using . everywhere) would improve the user experience.

[parambole]

1. JAX path flattening currently ignores SequenceKey (layer indices in lists) and GetAttrKey because it filters strictly with hasattr(k, "key"). This collapses all layers (e.g., layer_0 and layer_1 become just layer), causing incorrect parameter counts and output. We should use the existing helper param_key_parts_from_path(path_tuple) to format keys correctly. When scan_layers=False
2. "attention=dot_product" after remaining_args overrides any user-specified CLI flags (e.g., attention=flash). We should prepend defaults before remaining_args.
3. Remove a leftover debug prints.

[parambole]

LGTM. Thank you adding this awesome tool.

[hengtaoguo]

This is awesome tool, thanks Shuning!

As a nice-to-have feature, I wonder if we could add a --save_to_txt=<path> flag. In addition to printing logs to the terminal, it would be great to dump them into a separate text file. This would allow both humans and agents to review the files and analyze the differences more easily later on. WDYT?

Does the tool currently support this? If not, we can definitely treat this as a follow-up item.

[shuningjin]

As a nice-to-have feature, I wonder if we could add a --save_to_txt=<path> flag. In addition to printing logs to the terminal, it would be great to dump them into a separate text file. This would allow both humans and agents to review the files and analyze the differences more easily later on. WDYT?

Thanks @hengtaoguo for the suggestion! I have enabled saving to output, with flag --output_file.

Tests: I have rerun the tests and included the saved file in the PR description.

• HF - output file: https://paste.googleplex.com/5392498586419200
• MaxText - output file: https://paste.googleplex.com/5253647192686592
• Orbax - output file: https://paste.googleplex.com/5626295399612416