Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:Mlc ai Mlc llm MiniCPM Loader

From Leeroopedia


Overview

The MiniCPM Loader module (python/mlc_llm/model/minicpm/minicpm_loader.py) defines the parameter mapping for converting MiniCPM model weights from HuggingFace format to MLC LLM's internal representation. It supports both standard dense MLP architectures and Mixture-of-Experts (MoE) configurations, handling the appropriate weight fusion strategies for each.

Location

  • File: python/mlc_llm/model/minicpm/minicpm_loader.py
  • Lines: 150
  • Module: mlc_llm.model.minicpm

Function: huggingface

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

Returns a parameter mapping from MLC LLM parameter names to HuggingFace PyTorch parameter names for the MiniCPM architecture.

Parameters:

Parameter Type Description
model_config MiniCPMConfig The configuration of the MiniCPM model.
quantization Quantization The quantization configuration.

Initialization

model = MiniCPMForCausalLM(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)

Attention QKV Fusion

For each layer, the Q, K, and V projection weights are concatenated along axis 0 into a single wqkv_pack.weight. Note that MiniCPM uses the name wqkv_pack rather than the more common qkv_proj:

attn = f"model.layers.{i}.self_attn"
for weight_type in ["weight"]:
    mlc_name = f"{attn}.wqkv_pack.{weight_type}"
    mapping.add_mapping(
        mlc_name,
        [
            f"{attn}.q_proj.{weight_type}",
            f"{attn}.k_proj.{weight_type}",
            f"{attn}.v_proj.{weight_type}",
        ],
        functools.partial(
            lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
            dtype=mlc_param.dtype,
        ),
    )

Dense MLP Path (num_experts == 0)

When the model has no expert layers, the standard gate/up projection fusion is used:

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

Mixture-of-Experts (MoE) Path

When num_experts > 0, the loader handles three distinct expert weight types:

Expert Gate-Up Fusion (e1_e3)

For each expert, the w1 (gate) and w3 (up) weights are concatenated, then all experts are stacked into a single 3D tensor:

def combine_expert_gate_up(*hf_params, dtype):
    stack = []
    for i in range(0, len(hf_params), 2):
        stack.append(np.concatenate([hf_params[i], hf_params[i + 1]], axis=0))
    return np.stack(stack, axis=0).astype(dtype)

The source parameters are pairs of experts.{expert_id}.w1.weight and experts.{expert_id}.w3.weight for each expert, mapped to a single e1_e3.weight.

Expert Down Projection (e2)

All expert w2.weight tensors are stacked into a single 3D tensor:

mlc_name = f"{mlc_mlp}.e2.weight"
mapping.add_mapping(
    mlc_name,
    [
        f"{mlp}.experts.{expert_id}.w2.weight"
        for expert_id in range(model_config.num_experts)
    ],
    functools.partial(
        lambda *hf_params, dtype: np.stack(hf_params, axis=0).astype(dtype),
        dtype=mlc_param.dtype,
    ),
)

Router Gate Weight

The MoE router gate weight is mapped directly with a dtype cast:

mlc_name = f"{mlc_mlp}.gate.weight"
mapping.add_mapping(
    mlc_name,
    [f"{mlp}.gate.weight"],
    functools.partial(
        lambda x, dtype: x.astype(dtype),
        dtype=mlc_param.dtype,
    ),
)

Tied Embeddings Handling

The identity fallback loop explicitly skips lm_head.weight when tie_word_embeddings is enabled, since the language model head shares the embedding weight:

for mlc_name, mlc_param in named_parameters.items():
    if mlc_name == "lm_head.weight" and model_config.tie_word_embeddings:
        continue
    if mlc_name not in mapping.param_map:
        mapping.add_mapping(
            mlc_name,
            [mlc_name],
            functools.partial(
                lambda x, dtype: x.astype(dtype),
                dtype=mlc_param.dtype,
            ),
        )

Key Design Decisions

  • Dual MLP strategy: The loader conditionally branches between a dense MLP path (gate-up fusion) and a MoE path (expert stacking) based on num_experts.
  • 3D expert stacking: Expert weights are stacked along a new axis 0, creating a 3D tensor of shape (num_experts, ...) for efficient batched expert computation.
  • Tied embedding support: When tie_word_embeddings is set, the lm_head.weight is intentionally excluded from the mapping since it shares the embedding matrix.

Dependencies

  • functools -- for functools.partial and functools.reduce
  • numpy -- for array concatenation, stacking, and dtype casting
  • mlc_llm.loader.ExternMapping -- the core mapping data structure
  • mlc_llm.quantization.Quantization -- quantization configuration
  • .minicpm_model.MiniCPMConfig, .minicpm_model.MiniCPMForCausalLM -- MiniCPM model definitions

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment