Implementation:Mlc ai Mlc llm Gemma3 Loader

Overview

The Gemma3 Loader module (python/mlc_llm/model/gemma3/gemma3_loader.py) defines the parameter mapping from HuggingFace Gemma3 model weights to MLC LLM's internal parameter format. It handles the Gemma-specific weight transformation where RMS LayerNorm weights have 1 added to them, and fuses gate/up projections in the MLP layers.

Location

• File: python/mlc_llm/model/gemma3/gemma3_loader.py
• Lines: 151
• Module: mlc_llm.model.gemma3

Function: huggingface

def huggingface(model_config: Gemma3Config, quantization: Quantization) -> ExternMapping:

Returns a parameter mapping that maps from MLC LLM parameter names to HuggingFace PyTorch parameter names for the Gemma3 architecture.

Parameters:

Parameter	Type	Description
model_config	Gemma3Config	The configuration of the Gemma3 model.
quantization	Quantization	The quantization configuration.

Returns: An ExternMapping instance containing all parameter correspondences.

Initialization

The function begins by instantiating a Gemma3ForCausalLM model, optionally casting it to the quantization dtype, and exporting its named parameters via TVM:

model = Gemma3ForCausalLM(model_config)
if quantization is not None:
    model.to(quantization.model_dtype)
_, _named_params, _ = model.export_tvm(
    spec=model.get_default_spec(),
    allow_extern=True,
)
named_parameters = dict(_named_params)

Prefix Handling

Gemma3 uses a dual-prefix scheme to handle both text-only and multimodal model configurations:

mlc_prefix = "language_model."
hf_prefix = "language_model." if not model_config.is_text_model else ""

When the model is a text-only model, HuggingFace parameters have no prefix. When it is a multimodal model (e.g., vision-language), both MLC and HuggingFace parameters carry the language_model. prefix.

Per-Layer Mappings

For each of the num_hidden_layers transformer layers, the following mappings are created:

MLP Gate-Up Fusion

The separate gate_proj.weight and up_proj.weight from HuggingFace are concatenated along axis 0 into a single gate_up_proj.weight:

mlc_name = f"{mlc_prefix + mlp}.gate_up_proj.weight"
mapping.add_mapping(
    mlc_name,
    [
        f"{hf_prefix + mlp}.gate_proj.weight",
        f"{hf_prefix + mlp}.up_proj.weight",
    ],
    functools.partial(
        lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype),
        dtype=mlc_param.dtype,
    ),
)

RMS LayerNorm Weight Adjustment

Gemma models add 1 to RMS LayerNorm weights for efficiency. This loader applies the (x + 1) transformation during loading for the following layer norms in each layer:

• input_layernorm.weight
• post_attention_layernorm.weight
• pre_feedforward_layernorm.weight
• post_feedforward_layernorm.weight
• self_attn.k_norm.weight
• self_attn.q_norm.weight

mapping.add_mapping(
    mlc_prefix + mlc_name,
    [hf_prefix + mlc_name],
    functools.partial(
        lambda x, dtype: (x + 1).astype(dtype),
        dtype=named_parameters[mlc_prefix + mlc_name].dtype,
    ),
)

Final Layer Norm

The top-level model.norm.weight also receives the (x + 1) transformation:

mlc_name = "model.norm.weight"
mapping.add_mapping(
    mlc_prefix + mlc_name,
    [hf_prefix + mlc_name],
    functools.partial(
        lambda x, dtype: (x + 1).astype(dtype),
        dtype=named_parameters[mlc_prefix + mlc_name].dtype,
    ),
)

Identity Fallback

Any MLC parameters not yet covered by explicit mappings are mapped with a simple dtype cast (identity transform), adjusting the prefix from MLC to HuggingFace as needed:

for mlc_name, mlc_param in named_parameters.items():
    if mlc_name not in mapping.param_map:
        mapping.add_mapping(
            mlc_name,
            [hf_prefix + mlc_name[len(mlc_prefix):]],
            functools.partial(
                lambda x, dtype: x.astype(dtype),
                dtype=mlc_param.dtype,
            ),
        )

Key Design Decisions

• LayerNorm +1 offset: Gemma models store RMS LayerNorm weights without the +1 bias in HuggingFace format. The loader adds 1 during conversion so the MLC runtime can skip this addition during inference, improving efficiency.
• Prefix adaptation: The dual-prefix scheme allows the same loader to handle both standalone text models and multimodal models that wrap a language model component.
• No QKV fusion: Unlike Llama and Mistral loaders, the Gemma3 loader does not fuse Q/K/V projections. The Gemma3 model architecture keeps these projections separate.

Dependencies

• functools -- for functools.partial to bind dtype arguments
• numpy -- for array concatenation and dtype casting
• mlc_llm.loader.ExternMapping -- the core mapping data structure
• mlc_llm.quantization.Quantization -- quantization configuration
• .gemma3_model.Gemma3Config -- Gemma3 model configuration
• .gemma3_model.Gemma3ForCausalLM -- Gemma3 model class