chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,101 @@
|
||||
"""
|
||||
This file specifies how MLC's Llava parameter maps from other formats, for example HuggingFace
|
||||
PyTorch, HuggingFace safetensors.
|
||||
"""
|
||||
|
||||
import functools
|
||||
|
||||
import numpy as np
|
||||
|
||||
from mlc_llm.loader import ExternMapping
|
||||
from mlc_llm.loader.standard_loader import make_standard_hf_loader
|
||||
from mlc_llm.quantization import Quantization, make_awq_quant
|
||||
|
||||
from .llava_model import LlavaConfig, LlavaForCausalLM
|
||||
|
||||
awq_quant = make_awq_quant(LlavaForCausalLM)
|
||||
|
||||
|
||||
def _num_layers(config: object) -> int:
|
||||
return config.text_config.num_hidden_layers
|
||||
|
||||
|
||||
huggingface = make_standard_hf_loader(
|
||||
model_cls=LlavaForCausalLM,
|
||||
layer_prefix="language_model.model.layers",
|
||||
add_unused=["rotary_emb.inv_freq"],
|
||||
num_layers_getter=_num_layers,
|
||||
)
|
||||
|
||||
|
||||
def awq(model_config: LlavaConfig, quantization: Quantization) -> ExternMapping:
|
||||
"""Returns a parameter mapping that maps from the names of MLC LLM parameters to
|
||||
the names of AWQ parameters.
|
||||
Parameters
|
||||
----------
|
||||
model_config : LlavaConfig
|
||||
The configuration of the Llava model.
|
||||
|
||||
quantization : Quantization
|
||||
The quantization configuration.
|
||||
|
||||
Returns
|
||||
-------
|
||||
param_map : ExternMapping
|
||||
The parameter mapping from MLC to AWQ.
|
||||
"""
|
||||
model, _ = awq_quant(model_config, quantization)
|
||||
_, _named_params = model.export_tvm(spec=model.get_default_spec())
|
||||
named_parameters = dict(_named_params)
|
||||
|
||||
mapping = ExternMapping()
|
||||
|
||||
for i in range(model_config.text_config.num_hidden_layers):
|
||||
# Add QKV in self attention
|
||||
attn = f"language_model.model.layers.{i}.self_attn"
|
||||
for quantize_suffix in ["qweight", "qzeros", "scales"]:
|
||||
mlc_name = f"{attn}.qkv_proj.{quantize_suffix}"
|
||||
assert mlc_name in named_parameters
|
||||
mlc_param = named_parameters[mlc_name]
|
||||
mapping.add_mapping(
|
||||
mlc_name,
|
||||
[
|
||||
f"{attn}.q_proj.{quantize_suffix}",
|
||||
f"{attn}.k_proj.{quantize_suffix}",
|
||||
f"{attn}.v_proj.{quantize_suffix}",
|
||||
],
|
||||
functools.partial(
|
||||
lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
|
||||
dtype=mlc_param.dtype,
|
||||
),
|
||||
)
|
||||
|
||||
# Concat gate and up in MLP
|
||||
mlp = f"language_model.model.layers.{i}.mlp"
|
||||
for quantize_suffix in ["qweight", "qzeros", "scales"]:
|
||||
mlc_name = f"{mlp}.gate_up_proj.{quantize_suffix}"
|
||||
assert mlc_name in named_parameters
|
||||
mlc_param = named_parameters[mlc_name]
|
||||
mapping.add_mapping(
|
||||
mlc_name,
|
||||
[
|
||||
f"{mlp}.gate_proj.{quantize_suffix}",
|
||||
f"{mlp}.up_proj.{quantize_suffix}",
|
||||
],
|
||||
functools.partial(
|
||||
lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype),
|
||||
dtype=mlc_param.dtype,
|
||||
),
|
||||
)
|
||||
|
||||
# inv_freq is not used in the model
|
||||
mapping.add_unused(f"{attn}.rotary_emb.inv_freq")
|
||||
|
||||
for mlc_name, mlc_param in named_parameters.items():
|
||||
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),
|
||||
)
|
||||
return mapping
|
||||
@@ -0,0 +1,334 @@
|
||||
"""
|
||||
Implementation of LLaVa Model
|
||||
Implements the CLIP Vision Encoder. Uses Llama for the Language Encoder.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import logging
|
||||
from typing import Any, Dict, Optional # noqa: UP035
|
||||
|
||||
from tvm import tirx
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.relax.frontend.nn import Module, Tensor
|
||||
from tvm.relax.frontend.nn.op import permute_dims, reshape, wrap_nested
|
||||
from tvm.relax.op import strided_slice
|
||||
|
||||
from mlc_llm import op as op_ext
|
||||
from mlc_llm.model.model_preset import MODEL_PRESETS
|
||||
from mlc_llm.model.vision import CLIPVisionConfig, CLIPVisionModel, ImageProcessor
|
||||
from mlc_llm.nn import PagedKVCache, RopeMode
|
||||
|
||||
from ...support.config import ConfigBase
|
||||
from ..llama.llama_model import LlamaConfig, LlamaForCausalLM
|
||||
from ..mistral.mistral_model import MistralConfig, MistralForCausalLM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
CONFIG_MAP = {"LlamaForCausalLM": LlamaConfig, "MistralForCausalLM": MistralConfig}
|
||||
ARCHITECTURE_MAP = {
|
||||
"LlamaForCausalLM": LlamaForCausalLM,
|
||||
"MistralForCausalLM": MistralForCausalLM,
|
||||
}
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LlavaConfig(ConfigBase):
|
||||
"""
|
||||
LLaVa Config
|
||||
"""
|
||||
|
||||
image_token_index: int
|
||||
text_config: LlamaConfig
|
||||
vision_config: CLIPVisionConfig
|
||||
vocab_size: int
|
||||
context_window_size: int = -1
|
||||
sliding_window_size: int = -1
|
||||
prefill_chunk_size: int = -1
|
||||
tensor_parallel_shards: int = 1
|
||||
max_batch_size: int = 1
|
||||
text_architecture: str = "LlamaForCausalLM"
|
||||
kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
vision_config_dict: Dict[str, Any] # noqa: UP006
|
||||
if isinstance(self.vision_config, CLIPVisionConfig):
|
||||
vision_config_dict = dataclasses.asdict(self.vision_config)
|
||||
else:
|
||||
vision_config_dict = dict(self.vision_config)
|
||||
|
||||
for k, v in vision_config_dict.pop("kwargs", {}).items():
|
||||
vision_config_dict[k] = v
|
||||
|
||||
self.vision_config = CLIPVisionConfig.from_dict(vision_config_dict)
|
||||
|
||||
text_config_dict: Dict[str, Any] # noqa: UP006
|
||||
if isinstance(self.text_config, ConfigBase):
|
||||
text_config_dict = dataclasses.asdict(self.text_config)
|
||||
else:
|
||||
text_config_dict = dict(self.text_config)
|
||||
|
||||
if "_name_or_path" in text_config_dict:
|
||||
hf_config = self.get_hf_config(text_config_dict)
|
||||
text_config_dict.update(hf_config)
|
||||
architectures = text_config_dict["architectures"]
|
||||
assert len(architectures) == 1
|
||||
self.text_architecture = architectures[0]
|
||||
else:
|
||||
for k, v in text_config_dict.pop("kwargs", {}).items():
|
||||
text_config_dict[k] = v
|
||||
|
||||
self.text_config = CONFIG_MAP[self.text_architecture].from_dict(text_config_dict)
|
||||
|
||||
for k in ["context_window_size", "sliding_window_size", "prefill_chunk_size"]:
|
||||
if getattr(self, k) <= 0:
|
||||
if hasattr(self.text_config, k):
|
||||
setattr(self, k, getattr(self.text_config, k))
|
||||
|
||||
def get_hf_config(self, text_config_dict: Dict[str, Any]) -> Dict[str, Any]: # noqa: UP006
|
||||
"""
|
||||
Get the Hugging Face config of the text model
|
||||
"""
|
||||
|
||||
hf_config: Dict[str, Any] # noqa: UP006
|
||||
try:
|
||||
from transformers import AutoConfig
|
||||
|
||||
hf_config = AutoConfig.from_pretrained(text_config_dict["_name_or_path"]).to_dict()
|
||||
except (ImportError, OSError) as e:
|
||||
# If transformers is not installed, get the config from preset
|
||||
# Llama2 is gated so it throws an OSError. Get the config from preset instead
|
||||
preset_mapping = {
|
||||
"meta-llama/Llama-2-7b-hf": "llama2_7b",
|
||||
"meta-llama/Llama-2-13b-hf": "llama2_13b",
|
||||
"lmsys/vicuna-7b-v1.5": "llama2_7b",
|
||||
"mistralai/Mistral-7B-v0.1": "mistral_7b",
|
||||
}
|
||||
if text_config_dict["_name_or_path"] in preset_mapping:
|
||||
hf_config = MODEL_PRESETS[preset_mapping[text_config_dict["_name_or_path"]]]
|
||||
else:
|
||||
raise ValueError("Unsupported text model") from e
|
||||
|
||||
return hf_config
|
||||
|
||||
|
||||
class LlavaMultiModalProjector(nn.Module):
|
||||
def __init__(self, config: LlavaConfig):
|
||||
super().__init__()
|
||||
|
||||
self.linear_1 = nn.Linear(
|
||||
config.vision_config.hidden_size, config.text_config.hidden_size, bias=True
|
||||
)
|
||||
self.act = nn.GELU()
|
||||
self.linear_2 = nn.Linear(
|
||||
config.text_config.hidden_size, config.text_config.hidden_size, bias=True
|
||||
)
|
||||
|
||||
def forward(self, image_features: Tensor) -> Tensor:
|
||||
hidden_states = self.linear_1(image_features)
|
||||
hidden_states = self.act(hidden_states)
|
||||
hidden_states = self.linear_2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class LlavaForCausalLM(Module):
|
||||
def __init__(self, config: LlavaConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.vision_tower = CLIPVisionModel(config.vision_config)
|
||||
self.image_processor = ImageProcessor()
|
||||
self.multi_modal_projector = LlavaMultiModalProjector(config)
|
||||
self.language_model = ARCHITECTURE_MAP[config.text_architecture](config.text_config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.dtype = "float32"
|
||||
|
||||
def to(self, dtype: Optional[str] = None):
|
||||
super().to(dtype=dtype)
|
||||
self.language_model.to(dtype=dtype)
|
||||
if dtype is not None:
|
||||
self.dtype = dtype
|
||||
|
||||
def embed(self, input_ids: Tensor) -> Tensor:
|
||||
return self.language_model.embed(input_ids)
|
||||
|
||||
def image_preprocess(self, pixel_values: Tensor) -> Tensor:
|
||||
pixel_values = permute_dims(pixel_values, axes=(0, 3, 1, 2)) # NHWC -> NCHW
|
||||
pixel_values = self.image_processor.resize(
|
||||
pixel_values,
|
||||
{
|
||||
"shortest_edge": self.config.vision_config.image_size,
|
||||
},
|
||||
)
|
||||
pixel_values = self.image_processor.crop(
|
||||
pixel_values,
|
||||
{
|
||||
"height": self.config.vision_config.image_size,
|
||||
"width": self.config.vision_config.image_size,
|
||||
},
|
||||
)
|
||||
pixel_values = self.image_processor.rescale(pixel_values)
|
||||
pixel_values = self.image_processor.normalize(pixel_values)
|
||||
return pixel_values
|
||||
|
||||
def image_embed(self, pixel_values: Tensor) -> Tensor:
|
||||
pixel_values = self.image_preprocess(pixel_values)
|
||||
pixel_values = pixel_values.astype(self.dtype)
|
||||
image_features_all = self.vision_tower.forward(pixel_values)
|
||||
image_features = wrap_nested(
|
||||
strided_slice(
|
||||
image_features_all._expr,
|
||||
axes=[1],
|
||||
begin=[1],
|
||||
end=[image_features_all.shape[1]],
|
||||
),
|
||||
name="slice",
|
||||
)
|
||||
image_features = self.multi_modal_projector(image_features)
|
||||
image_features = reshape(image_features, shape=(-1, self.config.text_config.hidden_size))
|
||||
return image_features
|
||||
|
||||
def batch_forward(
|
||||
self,
|
||||
input_embeds: Tensor,
|
||||
paged_kv_cache: PagedKVCache,
|
||||
logit_positions: Optional[Tensor] = None,
|
||||
):
|
||||
op_ext.configure()
|
||||
|
||||
return self.language_model.batch_forward(input_embeds, paged_kv_cache, logit_positions)
|
||||
|
||||
def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
op_ext.configure()
|
||||
|
||||
return self.language_model.prefill(input_embed, paged_kv_cache)
|
||||
|
||||
def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
op_ext.configure()
|
||||
|
||||
return self.language_model.decode(input_embed, paged_kv_cache)
|
||||
|
||||
def batch_prefill(
|
||||
self,
|
||||
input_embeds: Tensor,
|
||||
logit_positions: Tensor,
|
||||
paged_kv_cache: PagedKVCache,
|
||||
):
|
||||
return self.language_model.batch_prefill(input_embeds, logit_positions, paged_kv_cache)
|
||||
|
||||
def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache):
|
||||
return self.language_model.batch_decode(input_embeds, paged_kv_cache)
|
||||
|
||||
def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache):
|
||||
return self.language_model.batch_verify(input_embeds, paged_kv_cache)
|
||||
|
||||
def create_paged_kv_cache(
|
||||
self,
|
||||
max_batch_size: tirx.Var,
|
||||
max_total_seq_len: tirx.Var,
|
||||
prefill_chunk_size: tirx.Var,
|
||||
page_size: tirx.Var,
|
||||
support_sliding_window: tirx.Var,
|
||||
) -> PagedKVCache:
|
||||
return PagedKVCache.create_generic(
|
||||
attn_kind="mha",
|
||||
max_batch_size=max_batch_size,
|
||||
max_total_seq_len=max_total_seq_len,
|
||||
prefill_chunk_size=prefill_chunk_size,
|
||||
page_size=page_size,
|
||||
support_sliding_window=support_sliding_window,
|
||||
num_hidden_layers=self.config.text_config.num_hidden_layers,
|
||||
num_attention_heads=self.config.text_config.num_attention_heads
|
||||
// self.config.tensor_parallel_shards,
|
||||
num_key_value_heads=self.config.text_config.num_key_value_heads
|
||||
// self.config.tensor_parallel_shards,
|
||||
qk_head_dim=self.config.text_config.head_dim,
|
||||
v_head_dim=self.config.text_config.head_dim,
|
||||
rope_mode=RopeMode.NORMAL,
|
||||
rope_scale=1,
|
||||
rope_theta=self.language_model.rope_theta,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
|
||||
def get_default_spec(self):
|
||||
mod_spec = {
|
||||
"embed": {
|
||||
"input_ids": nn.spec.Tensor(["seq_len"], "int32"),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"image_embed": {
|
||||
"pixel_values": nn.spec.Tensor(
|
||||
[1, "image_height", "image_width", 3],
|
||||
"uint8",
|
||||
),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"prefill": {
|
||||
"input_embed": nn.spec.Tensor(
|
||||
[1, "seq_len", self.config.text_config.hidden_size], self.dtype
|
||||
),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"decode": {
|
||||
"input_embed": nn.spec.Tensor(
|
||||
[1, 1, self.config.text_config.hidden_size], self.dtype
|
||||
),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"batch_prefill": {
|
||||
"input_embeds": nn.spec.Tensor(
|
||||
[1, "seq_len", self.config.text_config.hidden_size], self.dtype
|
||||
),
|
||||
"logit_positions": nn.spec.Tensor(["batch_size"], "int32"),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"batch_decode": {
|
||||
"input_embeds": nn.spec.Tensor(
|
||||
["batch_size", 1, self.config.text_config.hidden_size], self.dtype
|
||||
),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"batch_verify": {
|
||||
"input_embeds": nn.spec.Tensor(
|
||||
[1, "seq_len", self.config.text_config.hidden_size], self.dtype
|
||||
),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"create_paged_kv_cache": {
|
||||
"max_batch_size": int,
|
||||
"max_total_seq_len": int,
|
||||
"prefill_chunk_size": int,
|
||||
"page_size": int,
|
||||
"support_sliding_window": int,
|
||||
"$": {
|
||||
"param_mode": "none",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
}
|
||||
return nn.spec.ModuleSpec.from_raw(mod_spec, self)
|
||||
Reference in New Issue
Block a user