maxtext.integration.tunix.weight_mapping package

Provides a centralized access point for vLLM weight mappings.

This module defines the StandaloneVllmWeightMapping class, which acts as a dispatcher to retrieve the correct weight mapping configuration for a given model name. This allows for easy extension to support new models.

class maxtext.integration.tunix.weight_mapping.StandaloneVllmWeightMapping

Bases: object

Mapping MaxText model weights to vLLM’s model weights.

Submodules

• maxtext.integration.tunix.weight_mapping.deepseek3 module
  • DEEPSEEK_VLLM_MAPPING
    • DEEPSEEK_VLLM_MAPPING.to_hf_hook_fns()
    • DEEPSEEK_VLLM_MAPPING.to_hf_transpose_keys()
    • DEEPSEEK_VLLM_MAPPING.lora_to_hf_mappings()
    • DEEPSEEK_VLLM_MAPPING.to_hf_mapping()
• maxtext.integration.tunix.weight_mapping.gpt_oss module
  • GPT_OSS_VLLM_MAPPING
    • GPT_OSS_VLLM_MAPPING.lora_to_hf_mappings()
    • GPT_OSS_VLLM_MAPPING.to_hf_hook_fns()
    • GPT_OSS_VLLM_MAPPING.to_hf_transpose_keys()
    • GPT_OSS_VLLM_MAPPING.to_hf_mapping()
• maxtext.integration.tunix.weight_mapping.llama3 module
  • LLAMA3_VLLM_MAPPING
    • LLAMA3_VLLM_MAPPING.to_hf_hook_fns()
    • LLAMA3_VLLM_MAPPING.to_hf_transpose_keys()
    • LLAMA3_VLLM_MAPPING.lora_to_hf_mappings()
    • LLAMA3_VLLM_MAPPING.to_hf_mapping()
• maxtext.integration.tunix.weight_mapping.qwen2 module
  • QWEN2_VLLM_MAPPING
    • QWEN2_VLLM_MAPPING.to_hf_hook_fns()
    • QWEN2_VLLM_MAPPING.to_hf_transpose_keys()
    • QWEN2_VLLM_MAPPING.lora_to_hf_mappings()
    • QWEN2_VLLM_MAPPING.to_hf_mapping()
• maxtext.integration.tunix.weight_mapping.qwen3 module
  • QWEN3_VLLM_MAPPING
    • QWEN3_VLLM_MAPPING.to_hf_hook_fns()
    • QWEN3_VLLM_MAPPING.to_hf_transpose_keys()
    • QWEN3_VLLM_MAPPING.lora_to_hf_mappings()
    • QWEN3_VLLM_MAPPING.to_hf_mapping()