Checkpoint Conversion Utilities

This guide provides instructions to use checkpoint conversion scripts to convert model checkpoints bidirectionally between Hugging Face and MaxText formats.

Supported models

The following models are supported:

Model Family	Sizes	HF \(\to\) Orbax (scan)	HF \(\to\) Orbax (unscan)	Orbax (scan) \(\to\) HF	Orbax (unscan) \(\to\) HF
Gemma2	2B, 9B, 27B	√	√	√	√
Gemma3 (Multimodal)	4B, 12B, 27B	√	√	√	√
Gemma4	26B (MoE), 31B	√	√	√	√
Gemma4 Small	E2B, E4B	-	√	-	√
Llama3.1	8B, 70B, 450B	√	√	√	√
Qwen2.5	1.5B, 7B, 14B	√	√	√	√
Qwen3	0.6B, 4B, 8B, 14B, 32B	√	√	√	√
Qwen3 MoE	30B, 235B, 480B	√	√	√	√
Mixtral	8x7B, 8x22B	√	√	√	√
GPT-OSS	20B, 120B	√	√	√	√
DeepSeek2	16B	√	√	√	√
DeepSeek3	671B	√	√	√	√
DeepSeek3.2	671B	√	√	-	-
Qwen3 Next	80B	√	√	√	√

Note: Gemma4 Small (E2B / E4B) uses the gemma4_small decoder block, which has per-layer KV sharing and is incompatible with nn.scan. Only unscanned (scan_layers=False) conversion is supported for these variants.

Prerequisites

• MaxText must be installed in a Python virtual environment using the maxtext[tpu] option. For instructions on installing MaxText on your VM, please refer to the official installation documentation.

• Hugging Face model checkpoints are cached locally at $HOME/.cache/huggingface/hub before conversion. Ensure you have sufficient disk space.

• Authenticate via the Hugging Face CLI if using private or gated models.

Hugging Face to MaxText

Use the to_maxtext.py script to convert a Hugging Face model checkpoint into a MaxText checkpoint. The script will automatically download the specified model from the Hugging Face Hub, perform conversion, and save converted checkpoints to the given output directory.

Note: For more information, checkout qwen3-4b example script and gemma3-4b example script.

Note (Multimodal Models / Gemma 3): When converting multimodal models to maxtext, you must set use_multimodal=true and disable lazy loading with --lazy_load_tensors=False (lazy loading of HF tensors is not supported for multimodal models yet).

Furthermore, for Gemma 3, you must explicitly specify --eager_load_method='transformers' (the default is 'safetensors') to ensure key structure compatibility with Hugging Face Transformers.

Note (Gemma 4 E2B / E4B): The gemma4-e2b and gemma4-e4b variants use the gemma4_small decoder block, which has per-layer KV sharing and is incompatible with nn.scan. Conversion must therefore be run with scan_layers=False (unscanned), and any downstream training/inference run must also use scan_layers=False. Additionally, if you want to convert the base (pretrained) model rather than the instruction-tuned default, pass --hf_model_path=google/gemma-4-E4B (or google/gemma-4-E2B) explicitly — HF_IDS defaults to the -it repository. See convert_gemma4_base.sh for a full example.

Setup Environment

# Setup environment variables
# -- Model configuration --
# The MaxText model name. See `src/maxtext/configs/types.py` for `ModelName` for a
# full list of supported models.
export MODEL=<HF_MODEL> # e.g. 'llama3.1-8b-Instruct'

# -- MaxText configuration --
# The directory where the converted Orbax checkpoint will be stored (GCS bucket or local directory).
export BASE_OUTPUT_DIRECTORY=<CKPT_PATH> # e.g., gs://my-bucket/my-checkpoint-directory

# Set to 1 if you intend to use Pathways for training, 0 for McJAX.
export USE_PATHWAYS=0

Run Conversion

# Optional: If you run out of disk space when downloading Hugging Face safetensors,
# customize your "HF_HOME" to redirect the cache to a larger or mounted disk (e.g., on a TPU VM).
# export HF_HOME="/dev/shm/huggingface_tmp"

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name=${MODEL?} \
    base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
    scan_layers=True \
    use_multimodal=false \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --save_dtype=bfloat16

You can find your converted checkpoint files under ${BASE_OUTPUT_DIRECTORY}/0/items.

Key Parameters

• model_name: The specific model identifier. It must match a supported entry in the MaxText globals.py.

• scan_layers: Controls whether the output uses a scanned (scan_layers=true) or unscanned (scan_layers=false) checkpoint format. Refer to the Checkpoints guide for more information. Note: When resuming or loading a checkpoint in MaxText, this setting will automatically be loaded from the checkpoint metadata, meaning you do not have to manually specify it during model execution. If you do explicitly specify a value for scan_layers, it must match the checkpoint’s saved configuration, or a ValueError mismatch error will be raised.

• use_multimodal: Indicates if multimodality is used. Note: When use_multimodal=true, you must pass --lazy_load_tensors=False because lazy loading of tensors is not yet supported for multimodal models.

• base_output_directory: The path where the converted Orbax checkpoint will be stored; it can be Google Cloud Storage (GCS) or local.

• hardware=cpu: The conversion script runs on a CPU machine.

• checkpoint_storage_use_zarr3 and checkpoint_storage_use_ocdbt: These storage flags enable McJAX compatibility when set to True (the default). For Pathways, these should be False.

• --lazy_load_tensors (Optional / Required for multimodal): Controls on-demand loading of HF tensors to minimize RAM usage. Defaults to True for text models. Must be set to False (--lazy_load_tensors=False) for multimodal models.

• --eager_load_method (Optional): Backend used when eager loading (--lazy_load_tensors=False). Choices are safetensors (default) or transformers. Must be set to transformers (--eager_load_method='transformers') for Gemma 3 models to ensure key mapping compatibility.

• --hf_model_path (Optional): Specifies a customized remote directory or local directory containing the model weights. If unspecified, we use the default Hugging Face repository ID (e.g., openai/gpt-oss-20b). This is necessary for locally dequantized models like GPT-OSS or DeepSeek.

• --save_dtype (Optional): Specifies the data type of saved model weights. Default to bfloat16 to save memory.

MaxText to Hugging Face

Use the to_huggingface.py script to convert a MaxText checkpoint into the Hugging Face format. This is useful for sharing your models or integrating them with the Hugging Face ecosystem.

Note: For more information, checkout qwen3-4b example script.

Note (Gemma 4 E2B / E4B): The gemma4-e2b and gemma4-e4b variants use the gemma4_small decoder block and must be converted with scan_layers=False. If your MaxText checkpoint was fine-tuned from the base model, also pass --hf_model_path=google/gemma-4-E4B (or google/gemma-4-E2B) so the exported checkpoint bundles the base tokenizer — otherwise to_huggingface sources the tokenizer from HF_IDS[gemma4-e4b] = google/gemma-4-E4B-it.

Setup Environment

# Setup environment variables
# -- Model configuration --
# The MaxText model name. See `src/maxtext/configs/types.py` for `ModelName` for a
# full list of supported models.
export MODEL=<MODEL_NAME> # e.g. 'qwen3-4b'

# The path to the MaxText compatible model checkpoint (Orbax format).
export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items

# -- MaxText configuration --
# The directory where the converted Hugging Face checkpoint will be stored (GCS bucket,
# local directory, or Hugging Face Hub path).
export BASE_OUTPUT_DIRECTORY=<HF_CKPT_PATH> # e.g., gs://my-bucket/my-checkpoint-directory

Run Conversion

The following command converts a MaxText checkpoint and saves it locally, to GCS (gs://), or uploads it directly to the Hugging Face Hub (hf://).

python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name=${MODEL?} \
    load_parameters_path=${MAXTEXT_CKPT_PATH?} \
    base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
    hardware=cpu \
    skip_jax_distributed_system=true \
    scan_layers=false \
    use_multimodal=false \
    weight_dtype=bfloat16

Key Parameters

• model_name: The specific model identifier. It must match a supported entry in the MaxText globals.py.

• load_parameters_path: The path to the MaxText Orbax checkpoint.

• scan_layers: Controls whether the output uses a scanned (scan_layers=true) or unscanned (scan_layers=false) checkpoint format. Refer to the Checkpoints guide for more information.

• use_multimodal: Indicates if multimodality is used. Set to true when converting multimodal models back to Hugging Face format.

• hardware=cpu: The conversion script runs on a CPU machine.

• base_output_directory: The path where the converted checkpoint will be stored; it can be Google Cloud Storage (GCS), Hugging Face Hub or local.

• weight_dtype: It affects the resulting Hugging Face weight dtype. Default value is float32. We recommend using bfloat16 to save memory and speed up conversion.

Verifying conversion correctness

To ensure the conversion was successful, you can use the test script. It runs a forward pass on both the original and converted models and compares the output logits to verify conversion. It is used to verify the bidirectional conversion.

Note: This correctness test will only work when MaxText is installed from source by following the installation instructions here.

Setup Environment

# Setup environment variables
# -- Model configuration --
# The MaxText model name. See `src/maxtext/configs/types.py` for `ModelName` for a
# full list of supported models.
export MODEL=<MODEL_NAME> # e.g. 'qwen3-4b'

# The path to the MaxText compatible model checkpoint (Orbax format).
export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items

# The path to the Hugging Face model checkpoint directory.
export HF_CKPT_PATH=<HF_CKPT_PATH> # e.g., gs://my-bucket/my-checkpoint-directory

# Install PyTorch (in MaxText virtual environment)
python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu

Run Correctness Test

python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml \
    load_parameters_path=${MAXTEXT_CKPT_PATH?} \
    model_name=${MODEL?} \
    skip_jax_distributed_system=true \
    scan_layers=false \
    max_prefill_predict_length=4 \
    max_target_length=8 \
    use_multimodal=false \
    --run_hf_model=True \
    --hf_model_path=${HF_CKPT_PATH?} \
    --max_kl_div=0.015

Key Parameters

• load_parameters_path: The path to the MaxText Orbax checkpoint (e.g., gs://your-bucket/maxtext-checkpoint/0/items).

• model_name: The corresponding model name in the MaxText configuration (e.g., qwen3-4b).

• scan_layers: Controls whether the output uses a scanned (scan_layers=true) or unscanned (scan_layers=false) checkpoint format. Refer to the Checkpoints guide for more information.

• use_multimodal: Indicates if multimodality is used.

• --run_hf_model (Optional): Indicates if loading Hugging Face model from the hf_model_path. If not set, it will compare the maxtext logits with pre-saved golden logits.

• --hf_model_path (Optional): The path to the Hugging Face checkpoint (if --run_hf_model=True).

• --golden_logits_path (Optional): The pre-saved golden logits. (if --run_hf_model is not set).

• --max_kl_div: Max KL divergence tolerance during comparisons.

Example of Successful Conversion Verification

Here is part of the output of forward_pass_logit_checker for the gemma2-2b.

--- Prompt: What is the ---

--- MaxText model top 10 tokens ---
| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 5830       | difference           | 27.2500    |
| 1963       | best                 | 26.6250    |
| 5316       | average              | 26.3750    |
| 2669       | change               | 26.1250    |
| 12070      | percentage           | 26.1250    |
| 1618       | value                | 25.8750    |
| 1546       | most                 | 25.7500    |
| 66202      | molar                | 25.5000    |
| 3051       | total                | 25.5000    |
| 1503       | name                 | 25.3750    |


--- HF model top 10 tokens ---
| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 5830       | difference           | 27.2500    |
| 1963       | best                 | 26.6250    |
| 5316       | average              | 26.3750    |
| 12070      | percentage           | 26.1250    |
| 2669       | change               | 26.1250    |
| 1618       | value                | 25.8750    |
| 1546       | most                 | 25.7500    |
| 66202      | molar                | 25.5000    |
| 3051       | total                | 25.5000    |
| 6187       | purpose              | 25.3750    |


--- Similarity Metrics of Top Tokens ---
| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 9/10                 |
| jaccard_similarity             | 0.8181818181818182   |
| rank_agreement_percentage      | 70.0                 |


Average KL divergence per token (D_KL(P_golden || Q_model)): 0.000409

Max KL divergence for a single token in the set: 0.003497

---

Development and Troubleshooting

Inspection Tools

We provide a unified inspection tool In 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, 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, add the hook_fn logic (def {MODEL}_MAXTEXT_TO_HF_PARAM_HOOK_FN). This is the transformation needed per layer.

2. Add Hugging Face weights Shape: In utils/hf_shape.py, define the tensor shape of Hugging Face format (def {MODEL}_HF_WEIGHTS_TO_SHAPE). This is used to ensure the tensor shape is matched after to_huggingface conversion.

3. Register model key: In utils/globals.py, add the new model key in HF_IDS.

4. Add transformer config: In utils/hf_model_configs.py, add the transformers.Config object, describing the Hugging Face model configuration (defined in src/maxtext/configs/models). This configuration must precisely match the MaxText model’s architecture.

Here is an example PR to add support for gemma3 multi-modal model.

Common Errors

• Error: ValueError: lazy loading of HF tensors is not supported for multimodal models yet.

  • Cause: to_maxtext enables --lazy_load_tensors=True by default to conserve RAM, but lazy loading of tensors is not currently supported for multimodal models.

  • Solution: Pass --lazy_load_tensors=False to disable lazy loading when converting multimodal models.

• Error: ValueError: HuggingFace key model.language_model.embed_tokens.weight not found in state_dict. (or parameter name mismatch when converting Gemma 3)

  • Cause: When eager loading is enabled (--lazy_load_tensors=False), --eager_load_method defaults to "safetensors". Gemma 3 parameter mapping expects the module hierarchy from Hugging Face Transformers (from_pretrained).

  • Solution: Pass --eager_load_method='transformers' to the to_maxtext conversion command for Gemma 3 models.

• 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). (Note: Since MaxText automatically loads scan_layers from the checkpoint’s saved metadata, you should only encounter this error if you explicitly set a mismatching value on the command line, which raises a ValueError mismatch error).

  • Solution: Omit the scan_layers parameter from your training or execution command to allow MaxText to automatically resolve it from the checkpoint metadata, or ensure any explicitly specified scan_layers parameter matches the format of the loaded checkpoint.

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

  • Cause: Layer normalization parameters were likely not mapped correctly.