chore: import upstream snapshot with attribution

python/mlc_llm/model/bert/bert_loader.py

@@ -0,0 +1,117 @@
+"""
+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

python/mlc_llm/model/bert/bert_model.py

@@ -0,0 +1,273 @@
+"""
+Implementation for BERT architecture.
+"""
+
+import dataclasses
+from functools import partial
+from typing import Any, Dict, Optional  # noqa: UP035
+
+from tvm import te, tirx
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import Tensor, op
+
+from mlc_llm import op as op_ext
+from mlc_llm.support import logging
+from mlc_llm.support.config import ConfigBase
+from mlc_llm.support.style import bold
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class BertConfig(ConfigBase):
+    """Configuration of the BERT model."""
+
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    intermediate_size: int
+    hidden_act: str
+    layer_norm_eps: float
+    context_window_size: int = 0
+    prefill_chunk_size: int = 0
+    tensor_parallel_shards: int = 1
+    type_vocab_size: int = 2
+    pad_token_id: int = 0
+    position_offset: int = 0
+    head_dim: int = 0
+    max_batch_size: int = 1
+    kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)  # noqa: UP006
+
+    def __post_init__(self):
+        if self.intermediate_size is None or self.intermediate_size == -1:
+            self.intermediate_size = 4 * self.hidden_size
+        if self.context_window_size == 0:
+            for name in ["max_position_embeddings", "max_sequence_length"]:
+                if name in self.kwargs:
+                    self.context_window_size = self.kwargs.pop(name)
+                    logger.info(
+                        "%s not found in config.json. Falling back to %s (%d)",
+                        bold("context_window_size"),
+                        bold(name),
+                        self.context_window_size,
+                    )
+                    break
+            else:
+                raise ValueError(
+                    "Unable to determine the maximum sequence length, because none of "
+                    "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is "
+                    "provided in `config.json`."
+                )
+        if self.head_dim == 0:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+        assert self.head_dim * self.num_attention_heads == self.hidden_size
+        if self.prefill_chunk_size == 0:
+            logger.info(
+                "%s defaults to %s (%d)",
+                bold("prefill_chunk_size"),
+                bold("context_window_size"),
+                self.context_window_size,
+            )
+            self.prefill_chunk_size = self.context_window_size
+        elif self.prefill_chunk_size > self.context_window_size:
+            logger.info(
+                "Overriding %s from %d to %d (%s)",
+                bold("prefill_chunk_size"),
+                self.prefill_chunk_size,
+                self.context_window_size,
+                bold("context_window_size"),
+            )
+            self.prefill_chunk_size = self.context_window_size
+
+
+class BertSelfAttention(nn.Module):
+    def __init__(self, config: BertConfig):
+        if config.num_attention_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_attention_heads} attention heads"
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
+        self.num_heads = config.num_attention_heads // config.tensor_parallel_shards
+        self.head_dim = config.head_dim
+
+        self.qkv = nn.Linear(
+            in_features=config.hidden_size,
+            out_features=3 * self.num_heads * self.head_dim,
+            bias=True,
+        )
+
+    def forward(self, hidden_states: Tensor, attention_mask: Tensor):
+        d, h = self.head_dim, self.num_heads
+        b, s, _ = hidden_states.shape
+
+        qkv = self.qkv(hidden_states)
+        qkv = op.reshape(qkv, (b, s, 3 * h, d))
+        q, k, v = op.split(qkv, 3, axis=2)
+
+        # Attention
+        output = op_ext.attention(q, k, v, attention_mask)
+        return output
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: Tensor, input_tensor: Tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertAttention(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.self = BertSelfAttention(config)
+        self.output = BertSelfOutput(config)
+
+    def forward(self, hidden_states: Tensor, attention_mask: Tensor):
+        self_output = self.self(hidden_states, attention_mask)
+        attention_output = self.output(self_output, hidden_states)
+        return attention_output
+
+
+ACT2FN = {
+    "gelu": partial(nn.gelu, approximate=False),
+    "relu": nn.relu,
+    "silu": nn.silu,
+    "swish": nn.silu,
+    "gelu_new": partial(nn.gelu, approximate=True),
+}
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.intermediate_act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states: Tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class BertOutput(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: Tensor, input_tensor: Tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertLayer(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.attention = BertAttention(config)
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config)
+
+    def forward(self, hidden_states: Tensor, attention_mask: Tensor):
+        attention_output = self.attention(hidden_states, attention_mask)
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(self, hidden_states: Tensor, attention_mask: Tensor):
+        for layer in self.layer:
+            hidden_states = layer(hidden_states, attention_mask)
+        return hidden_states
+
+
+class BertEmbeddings(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, dtype="float32")
+        self.position_embeddings = nn.Embedding(
+            config.context_window_size, config.hidden_size, dtype="float32"
+        )
+        self.token_type_embeddings = nn.Embedding(
+            config.type_vocab_size, config.hidden_size, dtype="float32"
+        )
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, input_ids: Tensor, token_type_ids: Tensor, position_ids: Tensor):
+        words_embeddings = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = words_embeddings + position_embeddings + token_type_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        return embeddings
+
+
+class BertModel(nn.Module):
+    def __init__(self, config: BertConfig):
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        self.dtype = "float32"
+
+    def to(self, dtype: Optional[str] = None):
+        super().to(dtype=dtype)
+        if dtype is not None:
+            self.dtype = dtype
+
+    def forward(self, inputs: Tensor, attention_mask: Tensor):
+        # TODO: XLM-RoBERTa models use position indices starting from pad_token_id + 1
+        # (e.g., [2, 3, 4, ...] when pad_token_id=1), while this implementation uses
+        # [0, 1, 2, ...]. For XLM-RoBERTa models (e.g., bge-m3), the position_embeddings
+        # weights need to be shifted during weight conversion to compensate.
+        def _input_positions(inputs: te.Tensor):
+            b, s = inputs.shape
+            return te.compute((b, s), lambda _, j: j.astype("int32"), name="input_positions")
+
+        input_positions = op.tensor_expr_op(
+            _input_positions,
+            name_hint="input_positions",
+            args=[inputs],
+        )
+
+        token_type_ids = op.zeros(inputs.shape, dtype="int32")
+
+        embeddings = self.embeddings(inputs, token_type_ids, input_positions)
+        encoder_output = self.encoder(embeddings, attention_mask)
+        return encoder_output
+
+    def prefill(self, inputs: Tensor, attention_mask: Tensor):
+        def _attention_mask(mask: te.Tensor, zero, batch_size, seq_len):
+            return te.compute(
+                (batch_size, 1, seq_len, seq_len),
+                lambda b, _, i, j: tirx.if_then_else(
+                    tirx.any(mask[b, i] == zero, mask[b, j] == zero),
+                    tirx.min_value(self.dtype),
+                    tirx.max_value(self.dtype),
+                ),
+                name="attention_mask_prefill",
+            )
+
+        batch_size, seq_len = inputs.shape
+        attention_mask_2d = op.tensor_expr_op(
+            _attention_mask,
+            name_hint="attention_mask_prefill",
+            args=[attention_mask, tirx.IntImm("int32", 0), batch_size, seq_len],
+        )
+        return self.forward(inputs, attention_mask_2d)
+
+    def get_default_spec(self):
+        mod_spec = {
+            "prefill": {
+                "inputs": nn.spec.Tensor(["batch_size", "seq_len"], "int32"),
+                "attention_mask": nn.spec.Tensor(["batch_size", "seq_len"], "int32"),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+        }
+        return nn.spec.ModuleSpec.from_raw(mod_spec, self)