wehub-resource-sync
274 lines
10 KiB
Python
274 lines
10 KiB
Python
"""
|
|
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)
|