maxtext.integration.tunix.utils module

maxtext.integration.tunix.utils module#

Utils for Tunix integration.

class maxtext.integration.tunix.utils.VllmWeightMapping(model_name, config=None, use_standalone_mappings=False)[source]#

Bases: object

Mapping MaxText model weights to vLLM’s model weights.

to_hf_mapping()[source]#

Returns a mapping from MaxText parameter names to HuggingFace parameter names.

to_hf_transpose_keys()[source]#
to_hf_hook_fns()[source]#

Returns a mapping from MaxText parameter names to transformation functions.

lora_to_hf_mappings()[source]#
convert_hf_map_to_sharding_map(hf_mapping)[source]#

Converts a MaxText-to-HF name map into a generic MaxText-to-vLLM sharding map.

Parameters:

hf_mapping (dict) – The output from QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING. - Keys are MaxText param names (e.g., “params-decoder-layers…”). - Values are HF param names (str) or lists of names (list).

Returns:

A mapping from generalized MaxText names (e.g.,

”base.decoder.layers.mlp.wi_0.kernel”) to a tuple containing: (str: generalized HF/vLLM name, tuple: sharding specification).

Return type:

dict