wehub-resource-sync
118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
"""
|
|
This file specifies how MLC's BERT parameter maps from other formats, for example HuggingFace
|
|
PyTorch, HuggingFace safetensors.
|
|
"""
|
|
|
|
import functools
|
|
from typing import Literal
|
|
|
|
import numpy as np
|
|
|
|
from mlc_llm.loader import ExternMapping
|
|
from mlc_llm.quantization import Quantization
|
|
|
|
from .bert_model import BertConfig, BertModel
|
|
|
|
|
|
def huggingface(
|
|
model_config: BertConfig,
|
|
quantization: Quantization,
|
|
hf_prefix: Literal["", "bert."] = "",
|
|
) -> ExternMapping:
|
|
"""Returns a parameter mapping that maps from the names of MLC LLM parameters to
|
|
the names of HuggingFace PyTorch parameters.
|
|
|
|
Parameters
|
|
----------
|
|
model_config : BertConfig
|
|
The configuration of the BERT model.
|
|
|
|
quantization : Quantization
|
|
The quantization configuration.
|
|
|
|
hf_prefix : Literal["", "bert."]
|
|
Prefix used in HuggingFace weight names. Defaults to "" for standard
|
|
BERT models. Use "bert." for BGE models whose weights are prefixed.
|
|
|
|
Returns
|
|
-------
|
|
param_map : ExternMapping
|
|
The parameter mapping from MLC to HuggingFace PyTorch.
|
|
"""
|
|
model = BertModel(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)
|
|
|
|
mapping = ExternMapping()
|
|
|
|
def to_hf(name: str) -> str:
|
|
return f"{hf_prefix}{name}" if hf_prefix else name
|
|
|
|
for i in range(model_config.num_hidden_layers):
|
|
attn = f"encoder.layer.{i}.attention.self"
|
|
mlc_name = f"{attn}.qkv.weight"
|
|
mlc_param = named_parameters[mlc_name]
|
|
mapping.add_mapping(
|
|
mlc_name,
|
|
[
|
|
to_hf(f"{attn}.query.weight"),
|
|
to_hf(f"{attn}.key.weight"),
|
|
to_hf(f"{attn}.value.weight"),
|
|
],
|
|
functools.partial(
|
|
lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
|
|
dtype=mlc_param.dtype,
|
|
),
|
|
)
|
|
|
|
mlc_name = f"{attn}.qkv.bias"
|
|
mlc_param = named_parameters[mlc_name]
|
|
mapping.add_mapping(
|
|
mlc_name,
|
|
[
|
|
to_hf(f"{attn}.query.bias"),
|
|
to_hf(f"{attn}.key.bias"),
|
|
to_hf(f"{attn}.value.bias"),
|
|
],
|
|
functools.partial(
|
|
lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
|
|
dtype=mlc_param.dtype,
|
|
),
|
|
)
|
|
|
|
for mlc_name, mlc_param in named_parameters.items():
|
|
if mlc_name not in mapping.param_map:
|
|
mapping.add_mapping(
|
|
mlc_name,
|
|
[to_hf(mlc_name)],
|
|
functools.partial(
|
|
lambda x, dtype: x.astype(dtype),
|
|
dtype=mlc_param.dtype,
|
|
),
|
|
)
|
|
|
|
# Mark unused weights that exist in HF but not in MLC
|
|
if hf_prefix:
|
|
mapping.add_unused(f"{hf_prefix}pooler.dense.weight")
|
|
mapping.add_unused(f"{hf_prefix}pooler.dense.bias")
|
|
|
|
return mapping
|
|
|
|
|
|
def huggingface_bge(model_config: BertConfig, quantization: Quantization) -> ExternMapping:
|
|
"""Returns a parameter mapping for BGE models.
|
|
|
|
BGE weights have no prefix but include extra unused weights:
|
|
pooler.dense.weight, pooler.dense.bias, embeddings.position_ids
|
|
"""
|
|
mapping = huggingface(model_config, quantization, "")
|
|
mapping.add_unused("pooler.dense.weight")
|
|
mapping.add_unused("pooler.dense.bias")
|
|
mapping.add_unused("embeddings.position_ids")
|
|
return mapping
|