wehub-resource-sync
86 lines
2.5 KiB
Python
86 lines
2.5 KiB
Python
"""
|
|
This file specifies how MLC's GPT-2 parameter maps from other formats, for example HuggingFace
|
|
PyTorch, HuggingFace safetensors.
|
|
"""
|
|
|
|
import functools
|
|
|
|
from mlc_llm.loader import ExternMapping
|
|
from mlc_llm.quantization import Quantization
|
|
|
|
from .gpt2_model import GPT2Config, GPT2LMHeadModel
|
|
|
|
|
|
def huggingface(model_config: GPT2Config, quantization: Quantization) -> ExternMapping:
|
|
"""Returns a parameter mapping that maps from the names of MLC LLM parameters to
|
|
the names of HuggingFace PyTorch parameters.
|
|
|
|
Parameters
|
|
----------
|
|
model_config : GPT2Config
|
|
The configuration of the GPT-2 model.
|
|
|
|
quantization : Quantization
|
|
The quantization configuration.
|
|
|
|
Returns
|
|
-------
|
|
param_map : ExternMapping
|
|
The parameter mapping from MLC to HuggingFace PyTorch.
|
|
"""
|
|
model = GPT2LMHeadModel(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()
|
|
|
|
mapping.add_mapping(
|
|
"lm_head.weight",
|
|
["wte.weight"],
|
|
functools.partial(
|
|
lambda x, dtype: x.astype(dtype),
|
|
dtype=named_parameters["transformer.wte.weight"].dtype,
|
|
),
|
|
)
|
|
|
|
for i in range(model_config.n_layer):
|
|
mapping.add_unused(f"h.{i}.attn.bias")
|
|
|
|
# Transpose c_attn, c_proj and c_fc weights since GPT-2 uses Conv1D
|
|
for conv1d_weight_name in [
|
|
"attn.c_attn",
|
|
"attn.c_proj",
|
|
"mlp.c_proj",
|
|
"mlp.c_fc",
|
|
]:
|
|
src_name = f"h.{i}.{conv1d_weight_name}.weight"
|
|
mlc_name = f"transformer.{src_name}"
|
|
mapping.add_mapping(
|
|
mlc_name,
|
|
[src_name],
|
|
functools.partial(
|
|
lambda x, dtype: x.transpose().astype(dtype),
|
|
dtype=named_parameters[mlc_name].dtype,
|
|
),
|
|
)
|
|
|
|
for mlc_name, mlc_param in named_parameters.items():
|
|
if mlc_name not in mapping.param_map:
|
|
# transformer.h.0.attn.c_attn.weight --> h.0.attn.c_attn.weight
|
|
source_name = mlc_name.split(".", 1)[1]
|
|
mapping.add_mapping(
|
|
mlc_name,
|
|
[source_name],
|
|
functools.partial(
|
|
lambda x, dtype: x.astype(dtype),
|
|
dtype=mlc_param.dtype,
|
|
),
|
|
)
|
|
|
|
return mapping
|