diff --git a/lmquant/dataset/cache/calibration.py b/lmquant/dataset/cache/calibration.py index dfa7856..0eaa05c 100644 --- a/lmquant/dataset/cache/calibration.py +++ b/lmquant/dataset/cache/calibration.py @@ -3,6 +3,8 @@ import functools import gc +import os +import uuid import typing as tp from abc import ABC, abstractmethod @@ -10,7 +12,9 @@ import torch import torch.nn as nn import torch.utils.hooks +from PIL import Image from tqdm import tqdm +from dataclasses import dataclass from tqdm.contrib.logging import logging_redirect_tqdm from ..config import BaseCalibDatasetConfig @@ -20,6 +24,53 @@ __all__ = ["CalibrationCache"] +GLOBAL_CALIB_CACHE: tp.Dict[str, tp.Dict[str, tp.Any]] = {} + + +@dataclass +class CalibSample: + """Calibration sample.""" + + text: str + sample: torch.Tensor + image_srcs: tp.List[str] | None = None + id: str = str(uuid.uuid4()) + + def images(self, model: torch.nn.Module) -> tp.List[Image.Image]: + """Get images.""" + return [Image.open(img).convert("RGB") for img in self.image_srcs] + + def image_tensors(self, model: torch.nn.Module) -> list[torch.Tensor]: + """Get image tensor.""" + assert hasattr(model, "process_images"), "Expecting a VLM with `process_images` method" + return model.process_images(self.images(model), model.config) + + def get_model_inputs(self, model: torch.nn.Module) -> dict[str, torch.Tensor]: + """Process the calibration sample for the model.""" + global GLOBAL_CALIB_CACHE + if GLOBAL_CALIB_CACHE.get(self.id): + model_inputs = GLOBAL_CALIB_CACHE[self.id] + else: + PROCESS_INPUTS_DOC = ( + f" `model.process_inputs` should expect the following arguments:\n" + f" `sample` (torch.Tensor): The input tensor for the model.\n" + f" `images` (tp.List[str] | None): The path to the image file(s). Defaults to `None`.\n" + ) + assert hasattr(model, "process_inputs"), PROCESS_INPUTS_DOC + model_inputs = model.process_inputs(self.text, self.images(model)) + GLOBAL_CALIB_CACHE[self.id] = model_inputs + model_inputs = {k: self.move_to_device(v, next(model.parameters()).device) for k, v in model_inputs.items()} + return model_inputs + + def move_to_device(self, x: tp.Any, device: torch.device) -> None: + """Move the sample to a device.""" + if isinstance(x, torch.Tensor): + return x.to(device=device) + elif isinstance(x, (list, tuple)) and all(isinstance(i, torch.Tensor) for i in x): + return [i.to(device=device) for i in x] + else: + return x + class CalibrationCache(ABC): """Base class for caching calibration dataset.""" @@ -31,7 +82,7 @@ def __init__(self, config: BaseCalibDatasetConfig) -> None: config (BaseCalibrationConfig): Configuration for caching calibration dataset. """ self.config = config - self.cached_samples: list[torch.Tensor] = [] + self.cached_samples: list[CalibSample] = [] @property def num_samples(self) -> int: @@ -43,18 +94,18 @@ def reset(self) -> None: self.cached_samples = [] @abstractmethod - def _iter_samples(self, *args, **kwargs) -> tp.Generator[torch.Tensor, None, None]: + def _iter_samples(self, *args, **kwargs) -> tp.Generator[CalibSample, None, None]: """Iterate over model inputs.""" ... - def iter_samples(self, *args, needs_caching: bool, **kwargs) -> tp.Generator[torch.Tensor, None, None]: + def iter_samples(self, *args, needs_caching: bool, **kwargs) -> tp.Generator[CalibSample, None, None]: """Iterate over model input samples. Args: needs_caching (bool): Whether to cache input samples. Yields: - Generator[torch.Tensor, None, None]: Generator of model input samples. + Generator[CalibSample, None, None]: Generator of model input samples. """ if needs_caching and len(self.cached_samples) > 0: for sample in self.cached_samples: @@ -65,14 +116,14 @@ def iter_samples(self, *args, needs_caching: bool, **kwargs) -> tp.Generator[tor self.cached_samples.append(sample) yield sample - def get_samples(self, *args, needs_caching: bool, **kwargs) -> list[torch.Tensor]: + def get_samples(self, *args, needs_caching: bool, **kwargs) -> list[CalibSample]: """Get model input samples. Args: needs_caching (bool): Whether to cache input samples. Returns: - list[torch.Tensor]: List of model input samples. + list[CalibSample]: List of model input samples. """ if needs_caching: if len(self.cached_samples) == 0: @@ -101,7 +152,8 @@ def _init_cache(self, m: nn.Module, /) -> IOActivationsCache: elif isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): return IOActivationsCache( inputs=ActivationCache( - channels_dim=1, transform=ConvTransformFn(m.kernel_size, m.padding, m.stride, m.dilation) + channels_dim=1, + transform=ConvTransformFn(m.kernel_size, m.padding, m.stride, m.dilation), ), outputs=ActivationCache(channels_dim=1, transform=LinearTransformFn()), ) @@ -114,7 +166,7 @@ def _pre_layer_args_hook( args: tuple[torch.Tensor, ...], args_cache: list[tuple[torch.Tensor]], ) -> None: - assert all(isinstance(x, torch.Tensor) for x in args) + assert all(isinstance(x, torch.Tensor) for x in args), "Not all positional arguments are tensors." args_cache.append(tuple(x.detach().cpu() for x in args)) def _pre_layer_kwargs_hook( @@ -211,7 +263,10 @@ def _iter_layer_activations( # noqa: C901 layer_kwargs_caches[layer_name] = {} layer_hooks.append( module.register_forward_pre_hook( - functools.partial(self._pre_layer_kwargs_hook, kwargs_cache=layer_kwargs_caches[layer_name]), + functools.partial( + self._pre_layer_kwargs_hook, + kwargs_cache=layer_kwargs_caches[layer_name], + ), with_kwargs=True, ) ) @@ -248,7 +303,13 @@ def _iter_layer_activations( # noqa: C901 leave=False, total=self.num_samples, ): - model(sample.to(device=device)) + if isinstance(sample, CalibSample): + if sample.image_srcs is not None: + model(**sample.get_model_inputs(model)) + else: + model(input_ids=sample.sample.to(device=device)) + else: + model(sample.to(device=device)) if psutil.virtual_memory().percent > 90: raise RuntimeError("memory usage > 90%%, aborting") for hook in layer_hooks: @@ -280,7 +341,9 @@ def _iter_layer_activations( # noqa: C901 next_layer_args_cache: list[list[torch.Tensor]] = [] layer_kwargs = layer_kwargs_caches[layer_name] for layer_args in tqdm( - layer_args_cache, desc=f"collecting calibration activations in {layer_name}", leave=False + layer_args_cache, + desc=f"collecting calibration activations in {layer_name}", + leave=False, ): num_args = len(layer_args) layer_args = [arg.to(device=device) for arg in layer_args] diff --git a/lmquant/llm/dataset.py b/lmquant/llm/dataset.py index 787d3b0..546f3d1 100644 --- a/lmquant/llm/dataset.py +++ b/lmquant/llm/dataset.py @@ -8,14 +8,16 @@ import torch import torch.nn as nn -from datasets import load_from_disk +from datasets import load_from_disk, load_dataset from omniconfig import configclass from transformers.cache_utils import DynamicCache from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock +from torch.utils.data import Dataset, DataLoader +from transformers import AutoTokenizer from lmquant.dataset.cache.action import AverageCache, CacheAction, ConcatCache from lmquant.dataset.cache.activation import ActivationCache, IOActivationsCache -from lmquant.dataset.cache.calibration import CalibrationCache +from lmquant.dataset.cache.calibration import CalibrationCache, CalibSample from lmquant.dataset.config import BaseCalibDatasetConfig from lmquant.dataset.transform import LinearTransformFn @@ -24,6 +26,14 @@ __all__ = ["LlmCalibConfig", "LlmCalibrationCache", "LlmConcatCache", "LlmAverageCache"] +def process_sample(sample, is_vlm) -> None: + if is_vlm: + prefix = "<|im_start|>user\n" + assert prefix in sample, "Expecting a VLM sample" + return sample.replace(prefix, prefix + f"\n") + return sample + + @configclass @dataclass(kw_only=True) class LlmCalibConfig(BaseCalibDatasetConfig): @@ -77,7 +87,11 @@ class LlmConcatCache(ConcatCache): """Action for concatenating cached activations.""" def _unpack( - self, name: str, module: nn.Module, args: tuple[torch.Tensor, ...], kwargs: dict[str, tp.Any] | None + self, + name: str, + module: nn.Module, + args: tuple[torch.Tensor, ...], + kwargs: dict[str, tp.Any] | None, ) -> tuple[torch.Tensor, ...]: if len(args) == 0: assert "hidden_states" in kwargs, "hidden_states should be in kwargs if args is empty" @@ -89,7 +103,11 @@ class LlmAverageCache(AverageCache): """Action for averaging cached activations.""" def _unpack( - self, name: str, module: nn.Module, args: tuple[torch.Tensor, ...], kwargs: dict[str, tp.Any] | None + self, + name: str, + module: nn.Module, + args: tuple[torch.Tensor, ...], + kwargs: dict[str, tp.Any] | None, ) -> tuple[torch.Tensor, ...]: if len(args) == 0: assert "hidden_states" in kwargs, "hidden_states should be in kwargs if args is empty" @@ -101,6 +119,7 @@ class LlmCalibrationCache(CalibrationCache): """Cache for collecting calibration dataset for quantizing large language models.""" config: LlmCalibConfig + __is_vlm_dataset: bool = False def __init__(self, config: LlmCalibConfig) -> None: """Initialize LlmCalibrationCache. @@ -148,6 +167,16 @@ def _pre_layer_kwargs_hook( assert cached is None, f"kwargs_cache[{k}] should be None" elif isinstance(v, torch.Tensor): assert v.allclose(cached), f"kwargs_cache[{k}] should be the same as kwargs[{k}]" + elif isinstance(v, tuple): + assert isinstance(cached, tuple), f"kwargs_cache[{k}] should be a tuple" + assert len(v) == len(cached), f"kwargs_cache[{k}] should have the same length as kwargs[{k}]" + for i, (vi, cachedi) in enumerate(zip(v, cached)): + if isinstance(vi, torch.Tensor): + assert vi.allclose( + cachedi + ), f"kwargs_cache[{k}][{i}] should be the same as kwargs[{k}][{i}]" + else: + assert vi == cachedi, f"kwargs_cache[{k}][{i}] should be the same as kwargs[{k}][{i}]" else: assert v == cached, f"kwargs_cache[{k}] should be the same as kwargs[{k}]" else: @@ -167,29 +196,50 @@ def _iter_samples(self, tokenizer: nn.Module) -> tp.Generator[torch.Tensor, None Generator[torch.Tensor, None, None]: Generator for iterating over samples. Each sample is a tensor of shape (1, seq_length). """ - assert tokenizer is not None, "tokenizer is required for pileval dataset" - + assert tokenizer is not None, "tokenizer is required for LLM calibration" # region custom dataset kompress - import os - import json - import logging - kompress_data_path = os.environ["KOMPRESS_DATA_PATH"] - with open(kompress_data_path, "r") as f: - kompress_data_dict = json.load(f) - logging.info(f"Loaded kompress data {kompress_data_dict} from {kompress_data_path}") - split = kompress_data_dict["split"] - dataset_name_or_path = kompress_data_dict["dataset_name_or_path"] - text_column = kompress_data_dict["text_column"] - dataset = load_from_disk(dataset_name_or_path)[split] + try: + import os + import json + import logging + + kompress_data_path = os.environ["KOMPRESS_DATA_PATH"] + with open(kompress_data_path, "r") as f: + kompress_data_dict = json.load(f) + logging.info(f"Loaded kompress data {kompress_data_dict} from {kompress_data_path}") + split = kompress_data_dict["split"] + dataset_name_or_path = kompress_data_dict["dataset_name_or_path"] + text_column = kompress_data_dict["text_column"] + image_column = kompress_data_dict.get("image_column", None) + dataset = load_from_disk(dataset_name_or_path)[split] + self.__is_vlm_dataset = kompress_data_dict.get("is_vlm_dataset", False) # endregion + except Exception as e: + + logging.error(f"[ Couldn't load kompress data. ] Loading from config.dataset_path instead.") + # region dataset pileval + if self.config.data == "pileval": + dataset = load_dataset(self.config.dataset_path, split="validation") + text_column = "text" + elif self.config.data == "vlm_calibration_dataset": + self.__is_vlm_dataset = True + dataset = load_dataset("parquet", data_files=self.config.dataset_path, split="train") + text_column = "conversations" + image_column = "image_url" + else: + raise NotImplementedError(f"Calibration dataset {self.config.data} is not supported") dataset = dataset.shuffle(seed=42) rng = random.Random(42) - samples, num_tokens = [], 0 + samples: list[CalibSample] = [] + num_tokens = 0 for _data in dataset: line = _data[text_column] + imgs = [_data[image_column]] if self.__is_vlm_dataset else None line = line.strip() + line = process_sample(line, self.__is_vlm_dataset) + # line_encoded is a list of token ids line_encoded = tokenizer.encode(line) seq_length = len(line_encoded) @@ -201,15 +251,19 @@ def _iter_samples(self, tokenizer: nn.Module) -> tp.Generator[torch.Tensor, None continue # sample is a tensor of shape (1, seq_length) sample = torch.tensor([line_encoded]) + if seq_length > self.config.seq_length: - tok = rng.randint(0, seq_length - self.config.seq_length) - sample = sample[:, tok : tok + self.config.seq_length] - samples.append(sample) + if self.__is_vlm_dataset: + sample = sample[:, : self.config.seq_length] + else: + tok = rng.randint(0, seq_length - self.config.seq_length) + sample = sample[:, tok : tok + self.config.seq_length] + + samples.append(CalibSample(text=line, sample=sample, image_srcs=imgs)) num_tokens += sample.shape[1] if len(samples) >= self.config.num_samples and num_tokens >= self.config.num_tokens: break - # now concatenate all samples and split according to seq_length - samples = torch.cat(samples, dim=1).split(self.config.seq_length, dim=1) + if num_tokens > self.config.num_tokens: samples = samples[:-1] samples = samples[: self.config.num_samples] @@ -266,36 +320,102 @@ def iter_layer_activations( # noqa: C901 model = model_struct.module else: model_struct = LlmModelStruct.build(model) - backbone_struct = model_struct.backbone_struct - layer_structs = backbone_struct.layer_structs - action = LlmConcatCache("cpu") if action is None else action - for layer_idx, (layer_name, (layer, layer_cache, layer_kwargs)) in enumerate( - self._iter_layer_activations( - model, - *args, - action=action, - layers=backbone_struct.layers, - needs_inputs_fn=needs_inputs_fn, - needs_outputs_fn=needs_outputs_fn, - needs_samples_caching=needs_samples_caching, - **kwargs, - ) - ): - layer_struct = layer_structs[layer_idx] - assert layer_idx == layer_struct.idx - assert layer_name == layer_struct.full_name - assert layer is layer_struct.module - if layer_struct.proj_v_full_name in layer_cache: - cache = layer_cache[layer_struct.proj_v_full_name] - layer_cache[layer_struct.proj_q_full_name] = cache - layer_cache[layer_struct.proj_k_full_name] = cache - if layer_struct.proj_1st_full_names[0] in layer_cache: - for expert_idx in range(layer_struct.config.num_experts): - cache = layer_cache[layer_struct.proj_1st_full_names[expert_idx]] - for name in layer_struct.proj_1st_full_names[expert_idx :: layer_struct.config.num_experts]: - layer_cache[name] = cache - if layer_struct.config.num_experts == 1 and layer_struct.ffn_block_full_name not in layer_cache: - layer_cache[layer_struct.ffn_block_full_name] = layer_cache[layer_struct.proj_1st_full_names[0]] - if layer_struct.config.num_experts > 1 and layer_struct.ffn_block_full_name in layer_cache: - layer_cache[layer_struct.router_full_name] = layer_cache[layer_struct.ffn_block_full_name] - yield layer_name, (layer_struct, layer_cache, layer_kwargs) + backbone_structs = [model_struct.backbone_struct_llm] + if hasattr(model.config, "is_vlm") and model.config.is_vlm: + backbone_structs.insert(0, model_struct.backbone_struct_vit) + print(f"❗️ Quantizing VLM {model.config.is_vlm=}") + + for backbone_struct in backbone_structs: + layer_structs = backbone_struct.layer_structs + action = LlmConcatCache("cpu") if action is None else action + for layer_idx, ( + layer_name, + (layer, layer_cache, layer_kwargs), + ) in enumerate( + self._iter_layer_activations( + model, + *args, + action=action, + layers=backbone_struct.layers, + needs_inputs_fn=needs_inputs_fn, + needs_outputs_fn=needs_outputs_fn, + needs_samples_caching=needs_samples_caching, + **kwargs, + ) + ): + layer_struct = layer_structs[layer_idx] + assert layer_idx == layer_struct.idx + assert layer_name == layer_struct.full_name, f"{layer_name} != {layer_struct.full_name}" + assert layer is layer_struct.module + if layer_struct.proj_v_full_name in layer_cache: + cache = layer_cache[layer_struct.proj_v_full_name] + layer_cache[layer_struct.proj_q_full_name] = cache + layer_cache[layer_struct.proj_k_full_name] = cache + if layer_struct.proj_1st_full_names[0] in layer_cache: + for expert_idx in range(layer_struct.config.num_experts): + cache = layer_cache[layer_struct.proj_1st_full_names[expert_idx]] + for name in layer_struct.proj_1st_full_names[expert_idx :: layer_struct.config.num_experts]: + layer_cache[name] = cache + if layer_struct.config.num_experts == 1 and layer_struct.ffn_block_full_name not in layer_cache: + layer_cache[layer_struct.ffn_block_full_name] = layer_cache[layer_struct.proj_1st_full_names[0]] + if layer_struct.config.num_experts > 1 and layer_struct.ffn_block_full_name in layer_cache: + layer_cache[layer_struct.router_full_name] = layer_cache[layer_struct.ffn_block_full_name] + yield layer_name, (layer_struct, layer_cache, layer_kwargs) + + +class VLMCalibrationDataset(Dataset): + def __init__(self, config, tokenizer, model, device=None, dtype=None): + self.config = config + self.tokenizer = tokenizer + self.model = model + self.device = device if device is not None else model.device + self.dtype = dtype if dtype is not None else model.dtype + self.dataset = load_dataset("parquet", data_files=config.dataset_path, split="train") + self.dataset = self.dataset.shuffle(seed=42) + + def __len__(self): + return min(len(self.dataset), self.config.num_samples) + + def move_sample(self, sample): + if self.device is None and self.dtype is None: + return sample + + if self.device is not None and self.dtype is not None: + sample.to(device=self.device, dtype=self.dtype) + + if self.device is not None: + return sample.to(device=self.device) + elif self.dtype is not None: + return sample.to(dtype=self.dtype) + + def __getitem__(self, idx): + data = self.dataset[idx] + text = data["conversations"] + image_url = data["image_url"] + + # Process the sample + text = text.strip() + text = process_sample(text, is_vlm=True) + + # Tokenize + encoded = self.tokenizer.encode(text, truncation=True, max_length=self.config.seq_length) + input_ids = torch.tensor(encoded) + + # Pad or truncate to seq_length + if len(input_ids) < self.config.seq_length: + input_ids = torch.nn.functional.pad(input_ids, (0, self.config.seq_length - len(input_ids))) + elif len(input_ids) > self.config.seq_length: + input_ids = input_ids[: self.config.seq_length] + + sample: CalibSample = CalibSample(sample=input_ids, images=image_url) + + return { + "input_ids": self.move_sample(sample.sample), + "images": self.move_sample(sample.image_tensors(self.model)), + } + + +def create_vlm_dataloader(config, tokenizer, model): + dataset = VLMCalibrationDataset(config, tokenizer, model) + dataloader = DataLoader(dataset, batch_size=1, shuffle=True) + return dataloader diff --git a/lmquant/llm/model.py b/lmquant/llm/model.py index 9b458cb..71a9ba1 100644 --- a/lmquant/llm/model.py +++ b/lmquant/llm/model.py @@ -3,9 +3,10 @@ from dataclasses import dataclass, field +import os import torch from omniconfig import configclass -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, AutoProcessor from lmquant.model.config import BaseModelConfig @@ -54,11 +55,24 @@ def build( Returns: tuple[AutoModelForCausalLM, AutoTokenizer]: Model and tokenizer. """ - config = AutoConfig.from_pretrained(self.path) - tokenizer = AutoTokenizer.from_pretrained(self.path) + trust_remote_code = eval(os.environ.get("HUGGINGFACE_TRUST_REMOTE_CODE", "False")) + config = AutoConfig.from_pretrained(self.path, trust_remote_code=trust_remote_code) + tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=trust_remote_code) kwargs = {} if cpu else {"device_map": "balanced"} kwargs["torch_dtype"] = dtype - model = AutoModelForCausalLM.from_pretrained(self.path, config=config, **kwargs) + model = AutoModelForCausalLM.from_pretrained( + self.path, config=config, trust_remote_code=trust_remote_code, **kwargs + ) + if self.is_vlm: + assert hasattr(model, "config"), "The model does not have config." + setattr(model.config, "is_vlm", True) + + assert hasattr(model, "process_inputs"), "The model does not have process_inputs." + if hasattr(model, "init_processor"): + model.init_processor( + AutoProcessor.from_pretrained(model.config.image_encoder).image_processor, + tokenizer, + ) patch_attention(model) model.eval() return model, tokenizer diff --git a/lmquant/llm/nn/struct.py b/lmquant/llm/nn/struct.py index e2987e6..57b9973 100644 --- a/lmquant/llm/nn/struct.py +++ b/lmquant/llm/nn/struct.py @@ -44,6 +44,7 @@ ) from transformers.models.qwen2.modeling_qwen2 import ( Qwen2Attention, + Qwen2SdpaAttention, Qwen2Config, Qwen2DecoderLayer, Qwen2ForCausalLM, @@ -52,6 +53,62 @@ Qwen2Model, ) +# LlavaQwen2 +LlavaQwen2Attention = [Qwen2Attention, Qwen2SdpaAttention] +LlavaQwen2DecoderLayer = Qwen2DecoderLayer +LlavaQwen2MLP = Qwen2MLP + +# Siglip +SiglipVisionTransformer = [ + "SigLipVisionTransformer", # NanoLLaVA modeling + "SiglipVisionTransformer", # Dragonfly modeling +] +SiglipEncoderLayer = [ + "SigLipEncoderLayer", # NanoLLaVA modeling + "SiglipEncoderLayer", # Dragonfly modeling +] +SiglipSdpaAttention = [ + "SigLipSdpaAttention", # NanoLLaVA modeling + "SiglipSdpaAttention", # Dragonfly modeling +] +SiglipAttention = [ + "SigLipAttention", # NanoLLaVA modeling + "SiglipAttention", # Dragonfly modeling + *SiglipSdpaAttention, +] +SiglipMLP = [ + "SigLipMLP", # NanoLLaVA modeling + "SiglipMLP", # Dragonfly modeling +] + + +from transformers import PretrainedConfig + + +def object_like(obj: object, *class_or_name: tp.Union[str, type]) -> bool: + try: + if isinstance(class_or_name[0], tp.Tuple): + class_or_name = class_or_name[0] + except IndexError: + assert all(isinstance(x, (str, type)) for x in class_or_name), f"{class_or_name=}" + + class_or_name += tuple(cls.__name__ for cls in list(filter(lambda x: isinstance(x, type), class_or_name))) + + object_is_class_or_name: tp.Callable[[tp.Union[str, type]], bool] = lambda x: ( + x == obj.__class__.__name__ + if isinstance(x, str) + else isinstance(obj, x) or x.__name__ == obj.__class__.__name__ + ) + return any(object_is_class_or_name(x) for x in class_or_name) + + +module_like: tp.Callable[[nn.Module, tp.Union[str, type]], bool] = object_like +"""Check if a module is like the given class or name.""" + +config_like: tp.Callable[[PretrainedConfig, tp.Union[str, type]], bool] = object_like +"""Check if a config is like the given class or name.""" + + __all__ = ["LlmModelStruct", "LlmDecoderLayerStruct", "LlmBackboneStruct"] @@ -122,30 +179,63 @@ class LlmModelStruct: module: nn.Module """the nn.Module of the model""" - backbone: nn.Module + backbone_llm: nn.Module + backbone_vit: nn.Module fc: nn.Linear | None - backbone_name: str + backbone_name_llm: str + backbone_name_vit: str fc_name: str config: LlmConfigStruct - _backbone_struct: tp.Optional["LlmBackboneStruct"] = None + _backbone_struct_llm: tp.Optional["LlmBackboneStruct"] = None + _backbone_struct_vit: tp.Optional["LlmBackboneStruct"] = None + # For backward compatibility @property def backbone_full_name(self) -> str: """Get the backbone full name.""" - return self.backbone_name + print( + f"[WARNING] LlmModelStruct.backbone_full_name will be deprecated. Use backbone_full_name_llm or backbone_full_name_vit instead." + ) + return self.backbone_name_llm + + @property + def backbone_full_name_llm(self) -> str: + """Get the backbone full name.""" + return self.backbone_name_llm + + @property + def backbone_full_name_vit(self) -> str: + """Get the backbone full name.""" + return self.backbone_name_vit @property def fc_full_name(self) -> str: """Get the fc full name.""" return self.fc_name + # For backward compatibility @property def backbone_struct(self) -> "LlmBackboneStruct": """Extract backbone.""" - if self._backbone_struct is None: - self._backbone_struct = extract_llm_backbone(self.backbone, self) - return self._backbone_struct + print( + f"[WARNING] LlmModelStruct.backbone_struct will be deprecated. Use backbone_struct_llm or backbone_struct_vit instead." + ) + return self.backbone_struct_llm + + @property + def backbone_struct_llm(self) -> "LlmBackboneStruct": + """Extract backbone.""" + if self._backbone_struct_llm is None: + self._backbone_struct_llm = extract_llm_backbone(self.backbone_llm, self.backbone_full_name_llm, self) + return self._backbone_struct_llm + + @property + def backbone_struct_vit(self) -> "LlmBackboneStruct": + """Extract backbone.""" + if self._backbone_struct_vit is None: + self._backbone_struct_vit = extract_llm_backbone(self.backbone_vit, self.backbone_full_name_vit, self) + return self._backbone_struct_vit @staticmethod def build(model: nn.Module) -> tp.Optional["LlmModelStruct"]: @@ -176,7 +266,7 @@ class LlmBackboneStruct: layers_name: str final_ln_name: str proj_out_name: str - full_name: str = field(init=False) + full_name: str embedding_full_names: list[str] = field(init=False) proj_in_full_name: str = field(init=False) first_ln_full_name: str = field(init=False) @@ -188,7 +278,6 @@ class LlmBackboneStruct: _layer_structs: list["LlmDecoderLayerStruct"] | None = None def __post_init__(self) -> None: - self.full_name = self.parent.backbone_full_name self.embedding_full_names = [f"{self.full_name}.{name}" for name in self.embedding_names] self.proj_in_full_name = f"{self.full_name}.{self.proj_in_name}" self.first_ln_full_name = f"{self.full_name}.{self.first_ln_name}" @@ -406,10 +495,23 @@ def filter_layer_kwargs_to_attn_kwargs(self, kwargs: dict) -> dict: def extract_llm(model: nn.Module) -> LlmModelStruct | None: """Extract llm into components.""" - if isinstance(model, (OPTForCausalLM, OPTForSequenceClassification, OPTForQuestionAnswering)): - backbone = model.model.decoder - backbone_name = "model.decoder" - elif isinstance( + # region model + if module_like(model, "LlavaQwen2ForCausalLM"): + model.model.vision_tower.load_model() + model.model.vision_tower.to(device=model.model.device) + backbone_vit = model.model.vision_tower.vision_tower.vision_model + backbone_name_vit = "model.vision_tower.vision_tower.vision_model" + backbone_llm = model.model + backbone_name_llm = "model" + elif module_like(model, "DragonflyForCausalLM"): + backbone_vit = model.image_encoder.vision_model + backbone_name_vit = "image_encoder.vision_model" + backbone_llm = model.language_model.model + backbone_name_llm = "language_model.model" + elif module_like(model, (OPTForCausalLM, OPTForSequenceClassification, OPTForQuestionAnswering)): + backbone_llm = model.model.decoder + backbone_name_llm = "model.decoder" + elif module_like( model, ( LlamaForCausalLM, @@ -422,17 +524,30 @@ def extract_llm(model: nn.Module) -> LlmModelStruct | None: Qwen2ForSequenceClassification, ), ): - backbone = model.model - backbone_name = "model" + backbone_llm = model.model + backbone_name_llm = "model" else: raise ValueError(f"Unsupported model type: {type(model)}") - if isinstance(model, (OPTForCausalLM, LlamaForCausalLM, MistralForCausalLM, MixtralForCausalLM, Qwen2ForCausalLM)): + # endregion + + # region fc + if module_like( + model, + ( + OPTForCausalLM, + LlamaForCausalLM, + MistralForCausalLM, + MixtralForCausalLM, + Qwen2ForCausalLM, + "LlavaQwen2ForCausalLM", + ), + ): fc = model.lm_head fc_name = "lm_head" - elif isinstance(model, (OPTForQuestionAnswering)): + elif module_like(model, (OPTForQuestionAnswering)): fc = model.qa_outputs fc_name = "qa_outputs" - elif isinstance( + elif module_like( model, ( OPTForSequenceClassification, @@ -444,10 +559,16 @@ def extract_llm(model: nn.Module) -> LlmModelStruct | None: ): fc = model.score fc_name = "score" + elif module_like(model, "DragonflyForCausalLM"): + fc = model.language_model.lm_head + fc_name = "model.language_model.lm_head" else: raise ValueError(f"Unsupported model type: {type(model)}") + # endregion + + # region config config = model.config - if isinstance(config, OPTConfig): + if config_like(config, OPTConfig): config_struct = LlmConfigStruct( vocab_size=config.vocab_size, hidden_size=config.hidden_size, @@ -461,12 +582,24 @@ def extract_llm(model: nn.Module) -> LlmModelStruct | None: do_norm_before=config.do_layer_norm_before, tie_word_embeddings=config.tie_word_embeddings, ) - elif isinstance(config, (LlamaConfig, MistralConfig, MixtralConfig, Qwen2Config)): + elif config_like( + config, + ( + LlamaConfig, + MistralConfig, + MixtralConfig, + Qwen2Config, + "LlavaQwen2Config", + "DragonflyConfig", + ), + ): + hidden_act_key = "hidden_act" + assert hasattr(config, hidden_act_key), f"{hidden_act_key} not found in {type(config)}" config_struct = LlmConfigStruct( vocab_size=config.vocab_size, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, - intermediate_act=config.hidden_act, + intermediate_act=getattr(config, hidden_act_key), num_hidden_layers=config.num_hidden_layers, num_attention_heads=config.num_attention_heads, num_key_value_heads=config.num_key_value_heads, @@ -477,21 +610,26 @@ def extract_llm(model: nn.Module) -> LlmModelStruct | None: ) else: raise ValueError(f"Unsupported config type: {type(config)}") + # endregion + return LlmModelStruct( module=model, - backbone=backbone, + backbone_llm=backbone_llm, + backbone_vit=backbone_vit, fc=fc, - backbone_name=backbone_name, + backbone_name_llm=backbone_name_llm, + backbone_name_vit=backbone_name_vit, fc_name=fc_name, config=config_struct, ) -def extract_llm_backbone(backbone: nn.Module, parent: LlmModelStruct) -> LlmBackboneStruct | None: +def extract_llm_backbone(backbone: nn.Module, full_name: str, parent: LlmModelStruct) -> LlmBackboneStruct | None: """Extract llm backbone into components.""" - if isinstance(backbone, OPTModel): + # region backbone + if module_like(backbone, OPTModel): backbone = backbone.decoder - if isinstance(backbone, OPTDecoder): + if module_like(backbone, OPTDecoder): embeddings = [backbone.embed_tokens, backbone.embed_positions] layers = backbone.layers first_ln, final_ln = None, backbone.final_layer_norm @@ -500,7 +638,16 @@ def extract_llm_backbone(backbone: nn.Module, parent: LlmModelStruct) -> LlmBack layers_name = "layers" first_ln_name, final_ln_name = "", "final_layer_norm" proj_in_name, proj_out_name = "project_in", "project_out" - elif isinstance(backbone, (LlamaModel, MistralModel, MixtralModel, Qwen2Model)): + elif module_like( + backbone, + ( + LlamaModel, + MistralModel, + MixtralModel, + Qwen2Model, + "LlavaQwen2Model", + ), + ): embeddings = [backbone.embed_tokens] layers = backbone.layers first_ln, final_ln = None, backbone.norm @@ -509,8 +656,19 @@ def extract_llm_backbone(backbone: nn.Module, parent: LlmModelStruct) -> LlmBack layers_name = "layers" first_ln_name, final_ln_name = "", "norm" proj_in_name, proj_out_name = "", "" + elif module_like(backbone, *SiglipVisionTransformer): + embeddings = [backbone.embeddings] + layers = backbone.encoder.layers + first_ln, final_ln = None, backbone.post_layernorm + proj_in, proj_out = None, None + embedding_names = ["embeddings"] + layers_name = "encoder.layers" + first_ln_name, final_ln_name = "", "post_layernorm" + proj_in_name, proj_out_name = "", "" else: raise ValueError(f"Unsupported backbone type: {type(backbone)}") + # endregion + return LlmBackboneStruct( module=backbone, parent=parent, @@ -526,6 +684,7 @@ def extract_llm_backbone(backbone: nn.Module, parent: LlmModelStruct) -> LlmBack layers_name=layers_name, final_ln_name=final_ln_name, proj_out_name=proj_out_name, + full_name=full_name, ) @@ -538,10 +697,43 @@ def extract_llm_layer(layer: nn.Module, layer_idx: int, parent: LlmBackboneStruc Returns: LlmBlockStruct: Block. """ - if isinstance(layer, OPTDecoderLayer): + # region decoder layer + if module_like(layer, *SiglipEncoderLayer): + attn_ln = layer.layer_norm1 + attn_block = layer.self_attn + assert module_like( + attn_block, + *SiglipAttention, + ) + ffn_ln_key = "layer_norm2" + ffn_ln = getattr(layer, ffn_ln_key) + ffn_block = layer.mlp + assert module_like(ffn_block, *SiglipMLP) + proj_qkv = [attn_block.q_proj, attn_block.k_proj, attn_block.v_proj] + proj_out = attn_block.out_proj + proj_1st = [ffn_block.fc1] + proj_2nd = [ffn_block.fc2] + experts = [ffn_block] + router = None + proj_2nd_lowerbound = None + attn_block_kwargs = ( + "attention_mask", + "output_attentions", + ) + attn_ln_name = "layer_norm1" + attn_block_name = "self_attn" + proj_qkv_names = ["q_proj", "k_proj", "v_proj"] + proj_out_name = "out_proj" + ffn_ln_name = ffn_ln_key + ffn_block_name = "mlp" + proj_1st_names = ["fc1"] + proj_2nd_name = "fc2" + experts_name = "" + router_name = "" + elif module_like(layer, OPTDecoderLayer): attn_ln = layer.self_attn_layer_norm attn_block = layer.self_attn - assert isinstance(attn_block, OPTAttention) + assert module_like(attn_block, OPTAttention) ffn_ln = layer.final_layer_norm ffn_block = nn.Sequential(layer.fc1, layer.activation_fn, layer.fc2) proj_qkv = [attn_block.q_proj, attn_block.k_proj, attn_block.v_proj] @@ -568,13 +760,31 @@ def extract_llm_layer(layer: nn.Module, layer_idx: int, parent: LlmBackboneStruc proj_2nd_name = "fc2" experts_name = "" router_name = "" - elif isinstance(layer, (LlamaDecoderLayer, MistralDecoderLayer, Qwen2DecoderLayer)): + elif module_like( + layer, + ( + LlamaDecoderLayer, + MistralDecoderLayer, + Qwen2DecoderLayer, + LlavaQwen2DecoderLayer, + ), + ): attn_ln = layer.input_layernorm attn_block = layer.self_attn - assert isinstance(attn_block, (LlamaAttention, MistralAttention, Qwen2Attention)) - ffn_ln = layer.post_attention_layernorm + assert module_like( + attn_block, + ( + LlamaAttention, + MistralAttention, + Qwen2Attention, + *LlavaQwen2Attention, + ), + ), f"{type(attn_block)=}" + ffn_ln_key = "post_attention_layernorm" + assert hasattr(layer, ffn_ln_key), f"{ffn_ln_key} not found in {type(layer)}" + ffn_ln = getattr(layer, ffn_ln_key) ffn_block = layer.mlp - assert isinstance(ffn_block, (LlamaMLP, MistralMLP, Qwen2MLP)) + assert module_like(ffn_block, (LlamaMLP, MistralMLP, Qwen2MLP, LlavaQwen2MLP)) proj_qkv = [attn_block.q_proj, attn_block.k_proj, attn_block.v_proj] proj_out = attn_block.o_proj proj_1st = [ffn_block.up_proj, ffn_block.gate_proj] @@ -590,25 +800,25 @@ def extract_llm_layer(layer: nn.Module, layer_idx: int, parent: LlmBackboneStruc "use_cache", "cache_position", ) - if not isinstance(layer, LlamaDecoderLayer): + if not module_like(layer, LlamaDecoderLayer): attn_block_kwargs = attn_block_kwargs[:-1] attn_ln_name = "input_layernorm" attn_block_name = "self_attn" proj_qkv_names = ["q_proj", "k_proj", "v_proj"] proj_out_name = "o_proj" - ffn_ln_name = "post_attention_layernorm" + ffn_ln_name = ffn_ln_key ffn_block_name = "mlp" proj_1st_names = ["up_proj", "gate_proj"] proj_2nd_name = "down_proj" experts_name = "" router_name = "" - elif isinstance(layer, MixtralDecoderLayer): + elif module_like(layer, MixtralDecoderLayer): attn_ln = layer.input_layernorm attn_block = layer.self_attn - assert isinstance(attn_block, MixtralAttention) + assert module_like(attn_block, MixtralAttention) ffn_ln = layer.post_attention_layernorm ffn_block = layer.block_sparse_moe - assert isinstance(ffn_block, MixtralSparseMoeBlock) + assert module_like(ffn_block, MixtralSparseMoeBlock) proj_qkv = [attn_block.q_proj, attn_block.k_proj, attn_block.v_proj] proj_out = attn_block.o_proj proj_1st = [expert.w3 for expert in ffn_block.experts] + [expert.w1 for expert in ffn_block.experts] @@ -616,7 +826,13 @@ def extract_llm_layer(layer: nn.Module, layer_idx: int, parent: LlmBackboneStruc experts = [expert for expert in ffn_block.experts] router = ffn_block.gate proj_2nd_lowerbound = None - attn_block_kwargs = ("attention_mask", "position_ids", "past_key_value", "output_attentions", "use_cache") + attn_block_kwargs = ( + "attention_mask", + "position_ids", + "past_key_value", + "output_attentions", + "use_cache", + ) attn_ln_name = "input_layernorm" attn_block_name = "self_attn" proj_qkv_names = ["q_proj", "k_proj", "v_proj"] @@ -629,6 +845,8 @@ def extract_llm_layer(layer: nn.Module, layer_idx: int, parent: LlmBackboneStruc router_name = "gate" else: raise ValueError(f"Unsupported layer type: {type(layer)}") + # endregion + return LlmDecoderLayerStruct( module=layer, parent=parent, diff --git a/lmquant/llm/quant/weight.py b/lmquant/llm/quant/weight.py index 84c136e..3e45abe 100644 --- a/lmquant/llm/quant/weight.py +++ b/lmquant/llm/quant/weight.py @@ -194,14 +194,19 @@ def quantize_llm_weights( quantizers.update(block_quantizers) scale_state_dict.update(block_state_dict) else: - for layer in model.backbone_struct.layer_structs: - block_quantizers, block_state_dict = quantize_llm_decoder_layer_weights( - layer=layer, - config=quant_config, - quant_cache=quant_cache, - return_with_quantizers=return_with_quantizers, - return_with_scale_state_dict=return_with_scale_state_dict, - ) - quantizers.update(block_quantizers) - scale_state_dict.update(block_state_dict) + backbone_structs = [model.backbone_struct_llm] + if hasattr(model.module.config, "is_vlm") and model.module.config.is_vlm: + backbone_structs.insert(0, model.backbone_struct_vit) + print(f"❗️ Quantizing VLM {model.module.config.is_vlm=}") + for backbone_struct in backbone_structs: + for layer in backbone_struct.layer_structs: + block_quantizers, block_state_dict = quantize_llm_decoder_layer_weights( + layer=layer, + config=quant_config, + quant_cache=quant_cache, + return_with_quantizers=return_with_quantizers, + return_with_scale_state_dict=return_with_scale_state_dict, + ) + quantizers.update(block_quantizers) + scale_state_dict.update(block_state_dict) return quant_cache, quantizers, scale_state_dict diff --git a/lmquant/model/config.py b/lmquant/model/config.py index ffa7ce3..99da24d 100644 --- a/lmquant/model/config.py +++ b/lmquant/model/config.py @@ -22,6 +22,7 @@ class BaseModelConfig(ABC): root (str): Root directory path for models. Defaults to ``""``. local_path (str): Local path of the model. Defaults to ``None``. local_root (str): Local root directory path for models. Defaults to ``""``. + is_vlm (bool): Whether the model is a Vision Language Model. Defaults to ``False``. """ name: str @@ -29,6 +30,7 @@ class BaseModelConfig(ABC): root: str = "" local_path: str = None local_root: str = "" + is_vlm: bool = False family: str = field(init=False) def __post_init__(self): diff --git a/pyproject.toml b/pyproject.toml index d58c312..becc2ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ lm_eval = ">= 0.4.2" accelerate = ">= 0.26.0" datasets = ">= 2.16.0" sentencepiece = ">= 0.1.99" -omniconfig = ">= 0.1.5" +omniconfig = "== 0.1.6" protobuf = ">= 5.26.0" [tool.poetry.dependencies.strenum]