From 9caffac6eeab72ab478136c49bec4fedbd381a40 Mon Sep 17 00:00:00 2001 From: Shuning Jin Date: Thu, 18 Jun 2026 18:41:26 +0000 Subject: [PATCH] Add unified tool to inspect checkpoint structures --- .../convert_checkpoint.md | 49 ++- docs/guides/model_bringup.md | 25 +- .../inspect_checkpoint.py | 385 ++++++++++++++++++ .../checkpoint_conversion/to_huggingface.py | 2 +- 4 files changed, 451 insertions(+), 10 deletions(-) create mode 100644 src/maxtext/checkpoint_conversion/inspect_checkpoint.py diff --git a/docs/guides/checkpointing_solutions/convert_checkpoint.md b/docs/guides/checkpointing_solutions/convert_checkpoint.md index a178b8ddbe..e65627e794 100644 --- a/docs/guides/checkpointing_solutions/convert_checkpoint.md +++ b/docs/guides/checkpointing_solutions/convert_checkpoint.md @@ -240,7 +240,33 @@ 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=` 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 --format +``` + +**Tool 2 (MaxText Inspector)**. To inspect the MaxText model structure: + +``` +python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext model_name= scan_layers= +``` + +**Tool 3 (Orbax Inspector)**. To inspect the Orbax checkpoint: + +``` +python -m maxtext.checkpoint_conversion.inspect_checkpoint orbax --path +``` ### Adding New Models @@ -248,6 +274,8 @@ To extend conversion support to a new model architecture, you must define its sp 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. diff --git a/docs/guides/model_bringup.md b/docs/guides/model_bringup.md index 2388ceaab7..34971cb776 100644 --- a/docs/guides/model_bringup.md +++ b/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 --format +``` + +(2) **MaxText Inspector**. View the expected parameter structure of the target MaxText model: + +``` +python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext model_name= scan_layers= +``` + +(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 +``` ### 3.2 Write Script diff --git a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py new file mode 100644 index 0000000000..c63f2e1161 --- /dev/null +++ b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py @@ -0,0 +1,385 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A unified command-line tool to inspect checkpoint structures and tensor shapes. + +Lightweight and low-overhead, this tool operates on either CPU or TPU, +with a near-zero memory footprint (RAM/HBM) and negligible compute costs. + +Supported formats and frameworks: + 1. HuggingFace Checkpoints (Locally downloaded directory) + - Input: Local path containing `.safetensors` or `.pth` files. + - `safetensors` Cost: *Lightweight with near-zero cost*. Parses JSON file headers + to read metadata instantly without allocating RAM. + - `pth` Cost: *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: Parameter keys and shapes. + + 2. MaxText Model Architecture (On-the-fly) + - Input: MaxText model name, scan mode, and optional config args. + - Cost: *Lightweight with near-zero cost*. Traces JAX shapes abstractly + (via `jax.eval_shape`) to evaluate theoretical parameters dynamically + without executing compute or allocating RAM/HBM. + - Prints: Parameter keys and shapes, along with total parameter count. + + 3. Orbax Checkpoints (Pre-saved disk/GCS file) + - Input: Local or GCS path to the checkpoint. + - Cost: *Lightweight with near-zero cost*. Reads the structural metadata tree + instantly without pulling the underlying heavy TensorStore data chunks + into memory. + - Prints: Parameter keys and shapes. + +Usage Examples: + [Mode 1: HuggingFace] + python -m maxtext.checkpoint_conversion.inspect_checkpoint hf \ + --path --format + + [Mode 2: MaxText Architecture] + python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \ + model_name= scan_layers= + (Optional: other maxtext config) + + [Mode 3: Orbax] + python -m maxtext.checkpoint_conversion.inspect_checkpoint orbax \ + --path + +Optional Flags: + `--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. +""" + +import argparse +import sys +import os +import re +import pathlib +import absl +from maxtext.inference.inference_utils import str2bool +from maxtext.checkpoint_conversion.utils.utils import print_peak_memory + + +def natural_sort_key(s: str): + """Sorts strings containing numbers naturally (e.g., 1, 2, 10 instead of 1, 10, 2).""" + return [int(text) if text.isdigit() else text for text in re.split(r"(\d+)", str(s))] + + +def print_structure(data_dict, output_file=""): + """Utility to format and print sorted keys and shapes from a flattened dictionary.""" + if output_file: + # Save command + save_lines = [f"# {" ".join(sys.orig_argv)}", ""] + + for key in sorted(data_dict.keys(), key=natural_sort_key): + line = f"key: {key} | {data_dict[key]}" + print(line) + + if output_file: + # Save structure + save_lines.append(line) + + if output_file: + pathlib.Path(output_file).write_text("\n".join(save_lines) + "\n", encoding="utf-8") + print(f"\nStructure saved to {output_file}") + + +# ============================================================================== +# Mode 1: HuggingFace Checkpoint (Locally downloaded, safetensors or pth) +# ============================================================================== + + +# pylint: disable=import-outside-toplevel +def _inspect_safetensors(ckpt_paths, check_dtype): + """Inspects HuggingFace checkpoints, from local .safetensors file. + + Lightweight with near-zero cost: Bypasses heavy RAM allocation and disk I/O. + It strictly parses the JSON header of the locally downloaded safetensors file + to extract shapes and dtypes instantly. + + Optimization: Read metadata rather than weights via `f.get_slice(k)`. + """ + # Defer imports to avoid overhead when running in other modes. + try: + from safetensors import safe_open + except ImportError: + sys.exit("Error: 'safetensors' is required. Run `pip install safetensors`") + + param_dict = {} + for i, ckpt_path in enumerate(ckpt_paths): + print(f"Loading {ckpt_path.name} ({i+1}/{len(ckpt_paths)})...") + with safe_open(ckpt_path, framework="pt") as f: + for k in f.keys(): + # Optimization: `f.get_slice(k)` only reads the file header. + # This provides instant access to metadata. Using `f.get_tensor(k)` + # instead would force the actual tensor values into host RAM. + slice_obj = f.get_slice(k) + shape = slice_obj.get_shape() + param_dict[k] = f"shape: {shape}" + + if check_dtype: + dtype = slice_obj.get_dtype() + param_dict[k] += f" | dtype: {dtype}" + + return param_dict + + +# pylint: disable=import-outside-toplevel +def _inspect_pth(ckpt_paths, check_dtype): + """Inspects HuggingFace checkpoints, from local .pth file. + + Per-tensor load with small cost: Memory is allocated per-tensor rather than + loading the entire massive file into RAM at once. Because it deserializes + PyTorch's pickled format, it incurs a time cost just to extract keys and shapes. + + Optimization: Per-tensor load via `mmap=True`. + """ + # Defer imports to avoid overhead when running in other modes. + try: + import torch + except ImportError: + sys.exit("Error: 'torch' is required. Run `pip install torch`") + + param_dict = {} + for i, ckpt_path in enumerate(ckpt_paths): + print(f"Loading {ckpt_path.name} ({i+1}/{len(ckpt_paths)})...") + + # `mmap=True`: ensures the file is memory-mapped, keeping the RAM footprint + # strictly bounded to the per-tensor level rather than pulling the whole file. + # `weights_only=True`: prevents arbitrary code execution. + checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True, mmap=True) + + # Flattening logic may be necessary depending on the .pth structure. + # This assumes a standard state_dict; wrapper keys may need manual handling if present. + if isinstance(checkpoint, dict): + for k, v in checkpoint.items(): + if hasattr(v, "shape"): + param_dict[k] = f"shape: {v.shape}" + + if check_dtype and hasattr(v, "dtype"): + param_dict[k] += f" | dtype: {v.dtype}" + + return param_dict + + +def inspect_hf(args): + """Inspects the structure and tensor shapes of HuggingFace checkpoints. + + Note: The checkpoint structure typically matches the standard `from_pretrained` + format, with minor exceptions for specific architectures (e.g., multimodal Gemma). + Use this to verify structure compatibility with `to_maxtext` conversion pipelines. + """ + print(f"\n--- Inspecting {args.format} files in {args.path} ---") + + ckpt_paths = sorted(pathlib.Path(args.path).glob(f"[!.]*.{args.format}")) + if not ckpt_paths: + sys.exit(f"No files with extension .{args.format} found in {args.path}") + + if args.format == "safetensors": + param_dict = _inspect_safetensors(ckpt_paths, args.check_dtype) + elif args.format == "pth": + param_dict = _inspect_pth(ckpt_paths, args.check_dtype) + else: + sys.exit(f"Unsupported format: {args.format}") + + print("\n=== Structure ===") + print_structure(param_dict, args.output_file) + + +# ============================================================================== +# Mode 2: MaxText Architecture (On-the-fly) +# ============================================================================== + + +# pylint: disable=import-outside-toplevel +def inspect_maxtext(args, remaining_args): + """Inspects MaxText model architecture, from on-the-fly execution. + + Lightweight with near-zero cost: Uses JAX's abstract tracing + to evaluate the model's parameters dynamically. This instantly computes + theoretical shapes and dtypes without executing actual FLOPs or allocating + device memory (VRAM/TPU HBM). + + Optimization: Uses `jax.eval_shape`. + """ + # Defer imports to avoid overhead when running in other modes. + import jax + from maxtext.utils import max_utils, maxtext_utils + from maxtext import pyconfig + from maxtext.utils.globals import MAXTEXT_PKG_DIR + from maxtext.layers import quantizations + from maxtext.models import models + from maxtext.checkpoint_conversion.utils.utils import param_key_parts_from_path + + Transformer = models.transformer_as_linen + + # Configure the PyConfig environment. + # The first argument in argv is typically the script name. + argv = ( + [None, os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")] + + remaining_args + # The parameter structure is identical across attention implementations. + # For CPU compatibility, force `dot_product` attention. + # On CPU, default option `flash` will attempt to use GPU-specific Pallas kernels, causing potential failure + + ["attention=dot_product"] + # For CPU compatibability + + ["skip_jax_distributed_system=true"] + ) + print(argv) + config = pyconfig.initialize(argv) + + print(f"\n--- Inspecting MaxText Architecture: {config.model_name} (Scan: {config.scan_layers}) ---") + devices_array = maxtext_utils.create_device_mesh(config) + mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) + quant = quantizations.configure_quantization(config) + model = Transformer(config, mesh=mesh, quant=quant) + + # Extract abstract parameters. This returns a PyTree of `ShapeDtypeStruct` + # objects, without materializing weights. + abstract_param = maxtext_utils.get_abstract_param(model, config) + + # Calculate and display the total parameter count based purely on abstract shapes. + num_params = max_utils.calculate_num_params_from_pytree(abstract_param) + print(f"\nTotal Parameters: {num_params} (~{num_params/1e9:.2f} B)") + + print("\n=== Structure ===") + + abstract_params_flat, _ = jax.tree_util.tree_flatten_with_path(abstract_param) + + param_dict = {} + # Example of abstract_leaf_value format: ShapeDtypeStruct(shape=(128, 58), dtype=float32) + for path_tuple, abstract_leaf_value in abstract_params_flat: + key_parts = param_key_parts_from_path(path_tuple) + + # Construct a MaxText-style parameter key (e.g., "params.params.layer.weight"). + param_key = "params." + ".".join(key_parts) + + shape = abstract_leaf_value.shape + param_dict[param_key] = f"shape: {shape}" + + if args.check_dtype: + dtype = abstract_leaf_value.dtype + param_dict[param_key] += f" | dtype: {dtype}" + + print_structure(param_dict, args.output_file) + + +# ============================================================================== +# Mode 3: Orbax Checkpoint (Pre-saved) +# ============================================================================== + + +# pylint: disable=import-outside-toplevel +def inspect_orbax(args): + """Inspects Orbax checkpoint, from pre-saved disk/GCS file. + + Lightweight with near-zero cost: Orbax separates metadata from + heavy payload data. This method strictly parses the structural metadata tree, + bypassing the need to pull the underlying TensorStore data chunks from disk or GCS. + + Optimization: Reads metadata rather than weights. + """ + print(f"\n--- Inspecting Orbax Checkpoint: {args.path} ---") + + # Defer imports to avoid overhead when running in other modes. + import orbax.checkpoint as ocp + from etils import epath + + path = epath.Path(args.path) + + # Retrieve structural metadata. Note: Depending on the Orbax version, metadata + # access might vary slightly. This logic aligns with StandardCheckpointer usage. + metadata = ocp.StandardCheckpointer().metadata(path) + if hasattr(metadata, "item_metadata"): + metadata = metadata.item_metadata + + # Flatten the nested metadata tree into a standard dictionary. + dictionary = ocp.tree.to_flat_dict(metadata) + + # Filter strictly for parameter keys and format them. + param_dict = {} + for k, v in dictionary.items(): + # `k` is a tuple representing the path hierarchy; join it into a single string. + # `v` is a metadata object containing `.shape` and `.dtype`. + param_key = ".".join(k) + if not param_key.startswith("params"): + continue + + shape = v.shape + param_dict[param_key] = f"shape: {shape}" + + if args.check_dtype: + dtype = v.dtype + param_dict[param_key] += f" | dtype: {dtype}" + + print("\n=== Structure ===") + print_structure(param_dict, args.output_file) + + +# ============================================================================== +# Main CLI Driver +# ============================================================================== +def main(): + + # Shared parser for arguments common across all modes. + shared_parser = argparse.ArgumentParser(add_help=False) + shared_parser.add_argument( + "--check_dtype", type=str2bool, required=False, default=False, help="Whether to append dtype info to the output" + ) + shared_parser.add_argument( + "--output_file", type=str, required=False, default="", help="Path to save the output structure" + ) + + # Main parser and sub-parsers for distinct inspection modes. + parser = argparse.ArgumentParser(description="Consolidated Model Checkpoint Inspector") + subparsers = parser.add_subparsers(dest="mode", required=True, help="Inspection mode: hf, maxtext, orbax") + + # Mode 1: HuggingFace + parser_hf = subparsers.add_parser("hf", parents=[shared_parser], help="Inspect .safetensors or .pth files") + parser_hf.add_argument("--path", type=str, required=True, help="Directory containing checkpoint files") + parser_hf.add_argument( + "--format", + type=str, + required=False, + choices=["safetensors", "pth"], + default="safetensors", + help="Checkpoint file format", + ) + + # Mode 2: MaxText Architecture + subparsers.add_parser("maxtext", parents=[shared_parser], help="Inspect MaxText theoretical architecture") + + # Mode 3: Orbax + parser_orbax = subparsers.add_parser("orbax", parents=[shared_parser], help="Inspect saved Orbax checkpoint metadata") + parser_orbax.add_argument("--path", type=str, required=True, help="Path to checkpoint items (local or GCS)") + + args, remaining_args = parser.parse_known_args() + + # For output file, create parent directory if not exists + if args.output_file: + pathlib.Path(args.output_file).parent.mkdir(parents=True, exist_ok=True) + + if args.mode == "hf": + inspect_hf(args) + elif args.mode == "maxtext": + # remaining_args accepts maxtext config, like `model_name= scan_layers=` + inspect_maxtext(args, remaining_args) + elif args.mode == "orbax": + inspect_orbax(args) + + absl.logging.set_verbosity(absl.logging.INFO) + print_peak_memory() + + +if __name__ == "__main__": + main() diff --git a/src/maxtext/checkpoint_conversion/to_huggingface.py b/src/maxtext/checkpoint_conversion/to_huggingface.py index 0c498d7c3a..8467b95e2c 100644 --- a/src/maxtext/checkpoint_conversion/to_huggingface.py +++ b/src/maxtext/checkpoint_conversion/to_huggingface.py @@ -57,7 +57,7 @@ To merge a LoRA adapter into a base model and save as a full HF model: export HF_AUTH_TOKEN="hf_YOUR_TOKEN" - python src/maxtext/checkpoint_conversion/to_huggingface.py \ + python -m maxtext.checkpoint_conversion.to_huggingface \ src/maxtext/configs/base.yml \ model_name="gemma3-4b" \ load_parameters_path="/path/to/base/checkpoint/" \