Vlm support

lmquant/dataset/cache/calibration.py

@@ -3,14 +3,18 @@
 
 import functools
 import gc
+import os
+import uuid
 import typing as tp
 from abc import ABC, abstractmethod
 
 import psutil
 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]