feat: Re-add deserializing vLLM models

Integrated changes from
ssteel/tensorizer-support branch that allowed
for deserializing vLLM models.

vllm/engine/arg_utils.py

@@ -14,103 +14,7 @@
                          VisionLanguageConfig)
 from vllm.utils import str_to_int_tuple
 
-
-@dataclass
-class TensorizerArgs:
-    download_dir: Union[io.BufferedIOBase, io.RawIOBase, typing.BinaryIO, str,
-                        bytes, os.PathLike, int]
-    device: Optional[Union[torch.device, str]] = None
-    dtype: Optional[torch.dtype] = None
-    ## Commenting out serializer_encryption until I work out how I want to implement it
-    # serializer_encryption: Optional[bool] = False
-    lazy_load: bool = False
-    plaid_mode_buffers: Optional[int] = None
-    verify_hash: bool = False
-    filter_func: Optional[Callable[[str], Union[bool, Any]]] = None
-    deserializer_encryption_key: Optional[str] = None
-
-    def __post_init__(self):
-        self.file_obj = self.tensorizer_uri
-        self.s3_access_key_id = os.environ.get("S3_ACCESS_KEY_ID") or None
-        self.s3_secret_access_key = os.environ.get("S3_SECRET_ACCESS_KEY") or None
-        self.s3_endpoint = os.environ.get("S3_ENDPOINT_URL") or None
-
-        self.credentials = {
-            "s3_access_key_id": self.s3_access_key_id,
-            "s3_secret_access_key": self.s3_secret_access_key,
-            "s3_endpoint": self.s3_endpoint,
-        }
-        self.serializer_params = {
-            # Placeholder for now
-        }
-
-
-        # Omitting self.dtype and self.device as this behaves weirdly
-        self.deserializer_params = {
-            "filter_func": self.filter_func,
-            "lazy_load": self.lazy_load,
-            "plaid_mode": True,
-            "plaid_mode_buffers": self.plaid_mode_buffers,
-            "verify_hash": self.verify_hash,
-            "encryption": self.deserializer_encryption_key,
-            # "dtype":self.dtype,
-            # "device":self.device,
-        }
-
-    @staticmethod
-    def add_cli_args(
-            parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
-        """Tensorizer CLI arguments"""
-        # TODO: Add support for encryption -- CLI args can be base64 encoded
-        #       key/password for --serializer-encryption. Need to revist
-        parser.add_argument(
-            "--serializer-encryption",
-            action='store_true',
-            help="An `EncryptionParams` object holding a password or key"
-            "to use for encryption. If None, no encryption will be used.")
-        parser.add_argument(
-            "--lazy-load",
-            action='store_true',
-            help="If True, tensors will be loaded and cached when keys are"
-            "accessed. If False, all tensors will be loaded into memory up"
-            "front.",
-        )
-        parser.add_argument(
-            "--tensorizer-uri",
-            help="Path to serialized model tensors. Can be a local file path"
-                 "or a S3 URI.",
-        )
-        parser.add_argument(
-            "--plaid-mode-buffers",
-            default=None,
-            help="The number of buffers to use in plaid mode."
-            "This is only used if ``plaid_mode=True``. These buffers"
-            "are used to pipeline the loading and processing of tensors.")
-        parser.add_argument(
-            "--verify-hash",
-            action='store_true',
-            help="If True, the hashes of each tensor will be verified"
-            "against the hashes stored in the metadata. A `HashMismatchError`"
-            "will be raised if any of the hashes do not match.")
-        parser.add_argument(
-            "--deserializer-encryption-key",
-            default=None,
-            help="A `DecryptionParams` object holding a password or key"
-            "to use for decryption. ``None`` (the default) means no decryption."
-        )
-        return parser
-
-    @classmethod
-    def from_cli_args(cls, args: argparse.Namespace) -> 'TensorizerArgs':
-        # Get the list of attributes of this dataclass.
-        attrs = [attr.name for attr in dataclasses.fields(cls)]
-        # Set the attributes from the parsed arguments.
-        tensorizer_args = cls(**{
-            attr: getattr(args, attr)
-            for attr in attrs if hasattr(args, attr)
-        })
-        return tensorizer_args
-
+from vllm.model_executor.tensorizer_loader import TensorizerArgs
 
 @dataclass
 class EngineArgs:

vllm/model_executor/model_loader.py

@@ -78,11 +78,12 @@ def get_model(model_config: ModelConfig, device_config: DeviceConfig,
         # Create a model instance.
         # The weights will be initialized as empty tensors.
         with torch.device(device_config.device):
-            if hasattr(model_class, "supported_lora_modules"):
-                from vllm.model_executor.tensorizer_loader import zero_length_init
-                with zero_length_init():
-                    model = model_class(model_config.hf_config, linear_method,
-                                        lora_config)
+            if model_config.load_format == "tensorizer" and _is_vllm_model(model_config):
+                model = load_with_tensorizer(model_class, model_config)
+                return model.eval()
+            elif hasattr(model_class, "supported_lora_modules"):
+                model = model_class(model_config.hf_config, linear_method,
+                                    lora_config)
             elif lora_config:
                 raise ValueError(
                     f"Model {model_class.__name__} does not support LoRA, "

vllm/model_executor/tensorizer_loader.py

@@ -1,35 +1,228 @@
 import contextlib
 import contextvars
+import dataclasses
 import functools
 import threading
+import time
 import typing
-from types import MethodType
 from typing import Optional
+from typing import Type, Union, Any, Callable
+import io
+import os
+import argparse
+
 
 import torch
+from dataclasses import dataclass
+from tensorizer import TensorDeserializer, stream_io
+from tensorizer.utils import convert_bytes, get_mem_usage, no_init_or_tensor
 from torch import nn
 
-from vllm.model_executor.layers.activation import ScaledActivation
-from vllm.model_executor.layers.linear import ColumnParallelLinear, MergedColumnParallelLinear, RowParallelLinear, \
-    QKVParallelLinear
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from vllm.model_executor.models.mixtral import MixtralMoE
+from vllm.config import ModelConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear
 
 logger = init_logger(__name__)
 
+def load_with_tensorizer(model_cls: Type[nn.Module], model_config: ModelConfig) -> nn.Module:
+    tensorizer = TensorizerAgent(model_cls, model_config)
+    return tensorizer.deserialize()
+
+def _is_vllm_model(model_config: ModelConfig) -> bool:
+    return "vllm" in model_config.tensorizer_args.tensorizer_uri
+
+def _make_model_contiguous(model: nn.Module):
+    # Ensure tensors are saved in memory contiguously
+    for param in model.parameters():
+        param.data = param.data.contiguous()
+
+
+@dataclass
+class TensorizerArgs:
+    tensorizer_uri: Union[
+        io.BufferedIOBase,
+        io.RawIOBase,
+        typing.BinaryIO,
+        str,
+        bytes,
+        os.PathLike,
+        int,
+    ]
+    device: Optional[Union[torch.device, str]] = None
+    dtype: Optional[torch.dtype] = None
+    ## Commenting out serializer_encryption until I work out how I want to implement it
+    # serializer_encryption: Optional[bool] = False
+    lazy_load: bool = False
+    plaid_mode_buffers: Optional[int] = None
+    verify_hash: bool = False
+    filter_func: Optional[Callable[[str], Union[bool, Any]]] = None
+    deserializer_encryption_key: Optional[str] = None
+
+    def __post_init__(self):
+        self.file_obj = self.tensorizer_uri
+        self.s3_access_key_id = os.environ.get("S3_ACCESS_KEY_ID") or None
+        self.s3_secret_access_key = os.environ.get("S3_SECRET_ACCESS_KEY") or None
+        self.s3_endpoint = os.environ.get("S3_ENDPOINT_URL") or None
+
+        self.credentials = {
+            "s3_access_key_id": self.s3_access_key_id,
+            "s3_secret_access_key": self.s3_secret_access_key,
+            "s3_endpoint": self.s3_endpoint,
+        }
+        self.serializer_params = {
+            # Placeholder for now
+        }
+
+
+        # Omitting self.dtype and self.device as this behaves weirdly
+        self.deserializer_params = {
+            "filter_func": self.filter_func,
+            "lazy_load": self.lazy_load,
+            "plaid_mode": True if not self.device == "cpu" else False,
+            "plaid_mode_buffers": self.plaid_mode_buffers,
+            "verify_hash": self.verify_hash,
+            "encryption": self.deserializer_encryption_key,
+            # "dtype":self.dtype,
+            # "device":self.device,
+        }
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        """Tensorizer CLI arguments"""
+        # TODO: Add support for encryption -- CLI args can be base64 encoded
+        #       key/password for --serializer-encryption. Need to revist
+        parser.add_argument(
+            "--serializer-encryption",
+            action="store_true",
+            help="An `EncryptionParams` object holding a password or key"
+            "to use for encryption. If None, no encryption will be used.",
+        )
+        parser.add_argument(
+            "--lazy-load",
+            action="store_true",
+            help="If True, tensors will be loaded and cached when keys are"
+            "accessed. If False, all tensors will be loaded into memory up"
+            "front.",
+        )
+        parser.add_argument(
+            "--tensorizer-uri",
+            help="Path to serialized model tensors. Can be a local file path"
+                 "or a S3 URI.",
+        )
+        parser.add_argument(
+            "--plaid-mode-buffers",
+            default=None,
+            help="The number of buffers to use in plaid mode."
+            "This is only used if ``plaid_mode=True``. These buffers"
+            "are used to pipeline the loading and processing of tensors.",
+        )
+        parser.add_argument(
+            "--verify-hash",
+            action="store_true",
+            help="If True, the hashes of each tensor will be verified"
+            "against the hashes stored in the metadata. A `HashMismatchError`"
+            "will be raised if any of the hashes do not match.",
+        )
+        parser.add_argument(
+            "--deserializer-encryption-key",
+            default=None,
+            help="A `DecryptionParams` object holding a password or key"
+            "to use for decryption. ``None`` (the default) means no decryption.",
+        )
+        return parser
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "TensorizerArgs":
+        # Get the list of attributes of this dataclass.
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        # Set the attributes from the parsed arguments.
+        tensorizer_args = cls(
+            **{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)}
+        )
+        return tensorizer_args
+
+
+
+class TensorizerAgent:
+    def __init__(self, model_cls: Type[nn.Module],
+                 model_config: ModelConfig,
+                 ):
+        self.model_cls = model_cls
+        self.model_config = model_config
+        self.tensorizer_args = self.model_config.tensorizer_args
+        self.serialize_model = not self._verify_path_reachable()
+        self.model = self._init_model()
+
+    def _init_model(self):
+        model_args = self.model_config.hf_config
+        model_args.torch_dtype = self.model_config.dtype
+        model = no_init_or_tensor(lambda: self.model_cls(*[model_args]))
+        return model
+
+    def _verify_path_reachable(self):
+        if not self.tensorizer_args.tensorizer_uri.endswith(".tensors"):
+            raise ValueError(f"download_dir {self.tensorizer_args.tensorizer_uri} must specify a .tensors "
+                             f"file when load_format = tensorizer")
+
+    def deserialize(self):
+        before_mem = get_mem_usage()
+        # Lazy load the tensors from S3 into the model.
+        start = time.time()
+        stream = stream_io.open_stream(self.tensorizer_args.tensorizer_uri, mode="rb", **self.tensorizer_args.credentials)
+        deserializer = TensorDeserializer(stream, **self.deserialize_args)
+        deserializer.load_into_module(self.model)
+        self.model = self.model.to(dtype=self.model_config.dtype)
+        end = time.time()
+
+        # Brag about how fast we are.
+        total_bytes_str = convert_bytes(deserializer.total_tensor_bytes)
+        duration = end - start
+        per_second = convert_bytes(deserializer.total_tensor_bytes / duration)
+        after_mem = get_mem_usage()
+        deserializer.close()
+        logger.info(
+            f"Deserialized {total_bytes_str} in {end - start:0.2f}s, {per_second}/s"
+        )
+        logger.info(f"Memory usage before: {before_mem}")
+        logger.info(f"Memory usage after: {after_mem}")
+
+        return self.model.eval()
+
+    # def serialize(self):
+    #     with torch.device("cuda"):
+    #         model = self.model_cls(self.model_config.hf_config)
+    #     self.model_config.load_format = "auto"
+    #     model.load_weights(
+    #         self.model_config.model,
+    #         self.model_config.download_dir,
+    #         self.model_config.load_format,
+    #         self.model_config.revision,
+    #     )
+    #     _make_model_contiguous(model)
+    #     stream = stream_io.open_stream(self.tensorizer_args.download_dir, "wb", **self.credentials)
+    #     serializer = TensorSerializer(stream, **self.serialize_args)
+    #     logger.info(
+    #         f"Serializing model tensors {self.model_config.model} to {self.tensorizer_args.download_dir}."
+    #     )
+    #     serializer.write_module(model)
+    #     serializer.close()
+    #     logger.info(
+    #         f"Serialization complete. Running the previous command will deserialize the saved model weights."
+    #     )
+    #     return model.eval()
+
 
 ## Monkey patch for Parameter to ensure `requires_grad=False`
 from torch.nn.parameter import Parameter
 
 # Save the original __init__ method for later use
-original_new = Parameter.__new__
+#original_new = Parameter.__new__
 
 def _new(cls, data, requires_grad=False):
     return original_new(cls, data, requires_grad=requires_grad)
 
 # Replace the original __init__ method with our new one
-Parameter.__new__ = _new
+#Parameter.__new__ = _new
 
 def tensorizer_loader(params_dict):
     return _TensorizerWeightsLoaderImpl(params_dict).context_manager()
@@ -168,6 +361,9 @@ def _torch_empty_substitute(*args, **kwargs):
                 args = ((*dimension, 0),)
     return _torch_empty(device = "cuda", requires_grad = False, *args, **kwargs)
 
+
+
+
 # def vpe_weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
 #     param_data = param.data
 #     if self.input_is_parallel:

vllm/model_executor/weight_utils.py

@@ -298,7 +298,7 @@ def hf_model_weights_iterator(
         deserializer_args = tensorizer_args.deserializer_params
         credentials = tensorizer_args.credentials
         stream = open_stream(tensorizer_args.tensorizer_uri, **credentials)
-        with TensorDeserializer(stream, **deserializer_args, device="cuda:0") as state:
+        with TensorDeserializer(stream, **deserializer_args, device="cpu") as state:
             for name, param in state.items():
                 yield name, param
         del state