maxtext.checkpoint_conversion package

Subpackages

• maxtext.checkpoint_conversion.utils package
  • Submodules
    • maxtext.checkpoint_conversion.utils.hf_model_configs module
      • DeepseekV32Config
    • maxtext.checkpoint_conversion.utils.hf_shape module
      • GEMMA3_HF_WEIGHTS_TO_SHAPE()
      • GEMMA4_HF_WEIGHTS_TO_SHAPE()
      • GEMMA2_HF_WEIGHTS_TO_SHAPE()
      • DEEPSEEK_HF_WEIGHTS_TO_SHAPE()
      • QWEN3_NEXT_HF_WEIGHTS_TO_SHAPE()
      • GPT_OSS_HF_WEIGHTS_TO_SHAPE()
      • QWEN_HF_WEIGHTS_TO_SHAPE()
      • LLAMA31_HF_WEIGHTS_TO_SHAPE()
      • MIXTRAL_HF_WEIGHTS_TO_SHAPE()
    • maxtext.checkpoint_conversion.utils.hf_utils module
      • convert_jax_weight_to_torch()
      • check_arrays_match()
      • check_predicted_tokens_match()
      • get_logits_comparison_metrics()
    • maxtext.checkpoint_conversion.utils.param_mapping module
      • GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING()
      • GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING()
      • GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • QWEN_MAXTEXT_TO_HF_PARAM_MAPPING()
      • QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING()
      • QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING()
      • DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • DEEPSEEK_NNX_TO_VLLM_PARAM_HOOK_FN()
      • GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING()
      • GPT_OSS_TO_HF_PARAM_HOOK_FN()
      • QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING()
      • QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • QWEN3_NNX_TO_VLLM_PARAM_HOOK_FN()
      • LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING()
      • LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • LLAMA31_NNX_TO_VLLM_PARAM_HOOK_FN()
      • MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING()
      • MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING()
      • GEMMA4_MAXTEXT_TO_HF_PARAM_HOOK_FN()
      • OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING()
      • OLMO3_MAXTEXT_TO_HF_PARAM_HOOK_FN()
    • maxtext.checkpoint_conversion.utils.utils module
      • validate_and_filter_param_map_keys()
      • apply_hook_fns()
      • convert_jax_weight_to_numpy()
      • process_maxtext_param()
      • create_huggingface_hub_repo_if_not_exist()
      • save_config_file()
      • shard_checkpoint()
      • save_safetensor_file()
      • save_index_file()
      • save_weight_files()
      • get_local_save_path_manager()
      • save_model_files()
      • upload_state_dict_to_gcs()
      • upload_file_to_gcs()
      • upload_folder_to_gcs()
      • print_ram_usage()
      • print_peak_memory()
      • MemoryMonitorTqdm
      • load_orbax_checkpoint()
      • extract_nnx_weights()
      • extract_linen_weights()
      • detect_and_extract_checkpoint()
      • load_hf_dict_from_transformers()
      • load_hf_dict_from_safetensors()

Submodules

• maxtext.checkpoint_conversion.compare_hf_ckpt module
  • load_safetensors_generic()
  • get_hf_model_state_dict()
  • verify_dictionaries()
  • main()
• maxtext.checkpoint_conversion.load_and_quantize_checkpoint module
• maxtext.checkpoint_conversion.reshard_checkpoint module
• maxtext.checkpoint_conversion.to_huggingface module
• maxtext.checkpoint_conversion.to_maxtext module
  • LazyHFLoader
    • LazyHFLoader.get_tensor()
  • LazyTensor
    • LazyTensor.size
    • LazyTensor.nbytes
    • LazyTensor.itemsize
  • LazyTensorHandler
    • LazyTensorHandler.serialize()
  • get_maxtext_model_info()
  • main()