Merge pull request #3139 from AI-Hypercomputer:shuningjin-ckpt-structure

PiperOrigin-RevId: 934514537

Author: Google-ML-Automation

docs/guides/checkpointing_solutions/convert_checkpoint.md

@@ -240,14 +240,42 @@ Max KL divergence for a single token in the set: 0.003497
 
 ______________________________________________________________________
 
-## Troubleshooting and Development
+## Development and Troubleshooting
+
+### Inspection Tools
+
+We provide a unified inspection tool In [`inspect_checkpoint.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/checkpoint_conversion/inspect_checkpoint.py) to help you view model structures. This is highly useful for both development and troubleshooting.
+
+- **Lightweight & low-overhead:** Operates on either CPU or TPU, with a near-zero memory footprint (RAM/HBM) and negligible compute costs.
+- **Detailed outputs:** Prints tensor keys and shapes. Optionally appends data types with `--check_dtype=True` flag.
+- **Save to file:** Optionally saves the printed layout structure to a text file with `--output_file=<path>` flag.
+
+**Tool 1 (HF Inspector)**. To inspect the HuggingFace checkpoint structure, from either `.safetensors` or `.pth` files:
+
+```
+python -m maxtext.checkpoint_conversion.inspect_checkpoint hf --path <local_hf_path> --format <safetensors | pth>
+```
+
+**Tool 2 (MaxText Inspector)**. To inspect the MaxText model structure:
+
+```
+python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext model_name=<maxtext_model_name> scan_layers=<True | False>
+```
+
+**Tool 3 (Orbax Inspector)**. To inspect the Orbax checkpoint:
+
+```
+python -m maxtext.checkpoint_conversion.inspect_checkpoint orbax --path <local_orbax_path | gcs_orbax_path>
+```
 
 ### Adding New Models
 
 To extend conversion support to a new model architecture, you must define its specific parameter and configuration mappings. The conversion logic is decoupled, so you only need to modify the mapping files.
 
 1. **Add parameter mappings**:
 
+- As the first step, we inspect the checkpoint structures. To see the HuggingFace checkpoint structure, use **Tool 1 (HF Inspector)**. To see the MaxText model structure, use **Tool 2 (MaxText Inspector)**.
+
 - In [`utils/param_mapping.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/checkpoint_conversion/utils/param_mapping.py), add the parameter name mappings(`def {MODEL}_MAXTEXT_TO_HF_PARAM_MAPPING`). This is the 1-to-1 mappings of parameters names per layer.
 
 - In [`utils/param_mapping.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/checkpoint_conversion/utils/param_mapping.py), add the `hook_fn` logic (`def {MODEL}_MAXTEXT_TO_HF_PARAM_HOOK_FN`). This is the transformation needed per layer.
@@ -262,11 +290,22 @@ Here is an example [PR to add support for gemma3 multi-modal model](https://gith
 
 ### Common Errors
 
-- "Type ShapeDtypeStruct is not a valid JAX type" or generic **PyTree structure/shape mismatches** (e.g., Orbax reporting `"X/Y paths matched"`, such as `143/145 paths`):
-  This is almost always caused by a mismatch in the `scan_layers` configuration between the checkpoint conversion script (e.g., `to_maxtext.py` or `to_huggingface.py`) and the trainer/inference runner (e.g., `train.py`).
+- **Error:** When loading a converted checkpoint, `Type ShapeDtypeStruct is not a valid JAX type`.
+
+  - **Cause (most common): Structure mismatch** between the converted checkpoint and the MaxText model.
+
+  - **Solution:** To see the MaxText model structure, use **Tool 2 (MaxText Inspector)**. To inspect the checkpoint, use **Tool 3 (Orbax Inspector)**.
+
+- **Error:** `Type ShapeDtypeStruct is not a valid JAX type` or generic **PyTree structure/shape mismatches** (e.g., Orbax reporting `"X/Y paths matched"`, such as `143/145 paths`).
+
+  - **Cause: Configuration mismatch** (e.g., `scan_layers`) between the checkpoint conversion script (e.g., `to_maxtext.py` or `to_huggingface.py`) and the trainer/inference runner (e.g., `train.py`).
 
   - **Solution:** Ensure the `scan_layers` flag is set to the exact same value (`True` or `False`) in both the conversion command and your training/execution command.
 
-- If the converted checkpoint loads without errors but produces nonsensical output, likely an error in the Q/K/V weight reshaping logic during conversion.
+- **Error:** The converted checkpoint loads without errors but produces nonsensical output.
+
+  - **Cause:** There is likely an error in the Q/K/V weight reshaping logic during conversion.
+
+- **Error:** The model generates repetitive text sequences.
 
-- If the model generates repetitive text sequences, check if layer normalization parameters were mapped correctly.
+  - **Cause:** Layer normalization parameters were likely not mapped correctly.

docs/guides/model_bringup.md

@@ -42,7 +42,7 @@ This step can be bypassed if the current MaxText codebase already supports all c
 
 ## 3. Checkpoint Conversion
 
-While most open-source models are distributed in Safetensors or PyTorch formats, MaxText requires conversion to the [Orbax](https://orbax.readthedocs.io/en/latest/) format.
+While most open-source models are distributed in Safetensors or PyTorch formats, MaxText requires conversion to the [Orbax](https://orbax.readthedocs.io/en/latest) format.
 
 There are [two primary formats](checkpoints) for Orbax checkpoints within MaxText, and while both are technically compatible with training and inference, we recommend following these performance-optimized guidelines:
 
@@ -51,10 +51,27 @@ There are [two primary formats](checkpoints) for Orbax checkpoints within MaxTex
 
 ### 3.1 Create Mapping
 
-Success starts with a clear map. You must align the parameter names from your source checkpoints (Safetensors/PyTorch) with the corresponding MaxText internal names.
+To successfully convert a model, you must define the exact mapping between the parameter names in your source checkpoints (Safetensors/PyTorch) and the corresponding MaxText internal names.
 
-- You can print out the keys and shapes of your original `.safetensors` or `.pth` files.
-- To see the target structure, you can initiate a pre-training run to save a randomly initialized checkpoint for inspection.
+We provide a unified utility [`inspect_checkpoint.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/checkpoint_conversion/inspect_checkpoint.py) to help you view and compare these structures.
+
+(1) **HF Inspector**. To see the HuggingFace checkpoint structure, you can print out the keys and shapes of your original `.safetensors` or `.pth` files.
+
+```
+python -m maxtext.checkpoint_conversion.inspect_checkpoint hf --path <local_hf_path> --format <safetensors | pth>
+```
+
+(2) **MaxText Inspector**. View the expected parameter structure of the target MaxText model:
+
+```
+python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext model_name=<maxtext_model_name> scan_layers=<True | False>
+```
+
+(3) **Orbax Inspector** (Optional). If you have already saved an Orbax checkpoint during pretraining, you can inspect its structure directly:
+
+```
+python -m maxtext.checkpoint_conversion.inspect_checkpoint orbax --path <local_orbax_path | gcs_orbax_path>
+```
 
 ### 3.2 Write Script