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