maxtext.checkpoint_conversion.utils package

Submodules

• maxtext.checkpoint_conversion.utils.hf_model_configs module
  • DeepseekV32Config
    • DeepseekV32Config.model_type
• 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
    • MemoryMonitorTqdm.format_meter()
  • load_orbax_checkpoint()
  • extract_nnx_weights()
  • extract_linen_weights()
  • detect_and_extract_checkpoint()
  • load_hf_dict_from_transformers()
  • load_hf_dict_from_safetensors()