xinference 0.8.4__py3-none-any.whl → 0.9.0__py3-none-any.whl

xinference/model/llm/pytorch/compression.py

@@ -25,6 +25,8 @@ from torch.nn import functional as F
 from tqdm import tqdm
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
+from ....device_utils import empty_cache
+
 
 @dataclasses.dataclass
 class CompressionConfig:
@@ -153,7 +155,7 @@ def load_compress_model(
             tmp_state_dict[name] = None
             tensor = None
             gc.collect()
-            torch.cuda.empty_cache()
+            empty_cache()
 
     for name in model.state_dict():
         if name not in linear_weights:

xinference/model/llm/pytorch/core.py

@@ -12,10 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
 import os
 from typing import Iterable, Iterator, List, Optional, Union
 
+from ....device_utils import (
+    get_device_preferred_dtype,
+    gpu_count,
+    is_hf_accelerate_supported,
+)
 from ....types import (
     ChatCompletion,
     ChatCompletionChunk,
@@ -115,23 +121,18 @@ class PytorchModel(LLM):
             )
         from .compression import load_compress_model
 
-        cuda_visible_devices_env = os.getenv("CUDA_VISIBLE_DEVICES", None)
-        cuda_visible_devices = (
-            cuda_visible_devices_env.split(",") if cuda_visible_devices_env else []
-        )
-
         quantization = self.quantization
-        num_gpus = len(cuda_visible_devices) if cuda_visible_devices_env != "-1" else 0
+        num_gpus = gpu_count()
         device = self._pytorch_model_config.get("device", "auto")
         self._pytorch_model_config["device"] = select_device(device)
         self._device = self._pytorch_model_config["device"]
 
-        if self._device == "cpu":
-            kwargs = {"torch_dtype": torch.float32}
-        elif self._device == "cuda":
-            kwargs = {"torch_dtype": torch.float16}
-        elif self._device == "mps":
-            kwargs = {"torch_dtype": torch.float16}
+        kwargs = {}
+
+        dtype = get_device_preferred_dtype(self._device)
+
+        if dtype is not None:
+            kwargs["torch_dtype"] = dtype
         else:
             raise ValueError(f"Device {self._device} is not supported in temporary")
 
@@ -142,9 +143,25 @@ class PytorchModel(LLM):
             "trust_remote_code"
         )
         model_format = self.model_spec.model_format
+
+        is_device_map_auto = False
+
+        # This is required for Intel GPU to actually work with accelerate device_map until
+        # https://github.com/intel/intel-extension-for-pytorch/issues/522
+        # is resolved
+        max_memory_env = os.getenv("ACCELERATE_MAX_MEMORY", None)
+
+        if max_memory_env is not None:
+            max_memory_raw = json.loads(max_memory_env)
+            max_memory = {
+                int(k) if k.isdigit() else k: max_memory_raw[k] for k in max_memory_raw
+            }
+            kwargs["max_memory"] = max_memory
+
         if quantization != "none" and model_format == "pytorch":
             if self._device == "cuda" and self._is_linux():
                 kwargs["device_map"] = "auto"
+                is_device_map_auto = True
                 if quantization == "4-bit":
                     kwargs["load_in_4bit"] = True
                     kwargs["bnb_4bit_compute_dtype"] = torch.float16
@@ -178,11 +195,13 @@ class PytorchModel(LLM):
                     logger.debug(f"Model Memory: {self._model.get_memory_footprint()}")
                     return
 
-        if num_gpus > 0 and self._device == "cuda":
+        if num_gpus > 0 and is_hf_accelerate_supported(self._device):
             kwargs.update({"device_map": "auto"})
+            is_device_map_auto = True
+
         self._model, self._tokenizer = self._load_model(**kwargs)
 
-        if self._device == "mps":
+        if not is_device_map_auto:
             self._model.to(self._device)
         logger.debug(f"Model Memory: {self._model.get_memory_footprint()}")
 
@@ -448,7 +467,7 @@ class PytorchChatModel(PytorchModel, ChatModelMixin):
         generate_config = self._sanitize_generate_config(generate_config)
         # TODO(codingl2k1): qwen hacky to set stop for function call.
         model_family = self.model_family.model_family or self.model_family.model_name
-        if tools and "qwen-chat" == model_family:
+        if tools and model_family in ["qwen-chat", "qwen1.5-chat"]:
             stop = generate_config.get("stop")
             if isinstance(stop, str):
                 generate_config["stop"] = [stop, "Observation:"]

xinference/model/llm/pytorch/qwen_vl.py

@@ -95,9 +95,11 @@ class QwenVLChatModel(PytorchChatModel):
         if not isinstance(content, str):
             # TODO(codingl2k1): Optimize _ensure_url
             content = [
-                {"image": _ensure_url(c["image_url"]["url"]), "type": "image"}
-                if c.get("type") == "image_url"
-                else c
+                (
+                    {"image": _ensure_url(c["image_url"]["url"]), "type": "image"}
+                    if c.get("type") == "image_url"
+                    else c
+                )
                 for c in content
             ]
             content = sorted(content, key=operator.itemgetter("type"))

xinference/model/llm/pytorch/spec_decoding_utils.py

@@ -17,6 +17,8 @@ import time
 import uuid
 from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple
 
+from ....device_utils import empty_cache
+
 try:
     import torch
     from torch.nn import functional as F
@@ -526,4 +528,4 @@ def speculative_generate_stream(
     del kv_cache
     del draft_kv_cache
     gc.collect()
-    torch.cuda.empty_cache()
+    empty_cache()

xinference/model/llm/pytorch/spec_model.py

@@ -13,9 +13,13 @@
 # limitations under the License.
 
 import logging
-import os
 from typing import Iterator, List, Optional, Union
 
+from ....device_utils import (
+    get_device_preferred_dtype,
+    gpu_count,
+    is_hf_accelerate_supported,
+)
 from ....types import Completion, CompletionChunk, Embedding
 from ...utils import select_device
 from .. import LLMFamilyV1, LLMSpecV1
@@ -73,30 +77,26 @@ class SpeculativeModel(PytorchChatModel):
 
     def load(self):
         try:
-            import torch
+            import torch  # noqa: F401
         except ImportError:
             raise ImportError(
                 f"Failed to import module 'torch'. Please make sure 'torch' is installed.\n\n"
             )
 
-        cuda_visible_devices_env = os.getenv("CUDA_VISIBLE_DEVICES", None)
-        cuda_visible_devices = (
-            cuda_visible_devices_env.split(",") if cuda_visible_devices_env else []
-        )
-
-        num_gpus = len(cuda_visible_devices) if cuda_visible_devices_env != "-1" else 0
+        num_gpus = gpu_count()
         device = self._pytorch_model_config.get("device", "auto")
         self._pytorch_model_config["device"] = select_device(device)
         self._device = self._pytorch_model_config["device"]
 
-        if self._device == "cpu":
-            kwargs = {"torch_dtype": torch.float32}
-        elif self._device == "cuda":
-            kwargs = {"torch_dtype": torch.float16}
-        elif self._device == "mps":
-            kwargs = {"torch_dtype": torch.float16}
+        kwargs = {}
+
+        dtype = get_device_preferred_dtype(self._device)
+
+        if dtype is not None:
+            kwargs["torch_dtype"] = dtype
         else:
             raise ValueError(f"Device {self._device} is not supported in temporary")
+
         kwargs["trust_remote_code"] = self._pytorch_model_config.get(
             "trust_remote_code"
         )
@@ -106,15 +106,18 @@ class SpeculativeModel(PytorchChatModel):
                 "Quantization is not supported by speculative decoding yet"
             )
 
-        if num_gpus > 0 and self._device == "cuda":
+        is_device_map_auto = False
+
+        if num_gpus > 0 and is_hf_accelerate_supported(self._device):
             kwargs.update({"device_map": "auto"})
+            is_device_map_auto = True
 
         self._model, self._tokenizer = self._load_model(
             model_path=self.model_path,
             revision=self.model_spec.model_revision,
             **kwargs,
         )
-        if self._device == "mps":
+        if not is_device_map_auto:
             self._model.to(self._device)
         logger.debug(
             f"Model {self.model_uid} memory footprint: {self._model.get_memory_footprint()}"
@@ -125,7 +128,7 @@ class SpeculativeModel(PytorchChatModel):
             revision=self._draft_model_spec.model_revision,
             **kwargs,
         )
-        if self._device == "mps":
+        if not is_device_map_auto:
             self._model.to(self._device)
         logger.debug(
             f"Draft model {self.model_uid} memory footprint: {self._model.get_memory_footprint()}"

xinference/model/llm/pytorch/utils.py

@@ -29,6 +29,7 @@ from transformers.generation.logits_process import (
     TopPLogitsWarper,
 )
 
+from ....device_utils import empty_cache
 from ....types import (
     CompletionChoice,
     CompletionChunk,
@@ -122,7 +123,7 @@ def generate_stream(
         temperature, repetition_penalty, top_p, top_k
     )
 
-    if "qwen" in str(type(model)).lower():
+    if ".modeling_qwen." in str(type(model)).lower():
         # TODO: hacky
         input_ids = tokenizer(prompt, allowed_special="all").input_ids
     else:
@@ -335,7 +336,7 @@ def generate_stream(
     # clean
     del past_key_values, out
     gc.collect()
-    torch.cuda.empty_cache()
+    empty_cache()
 
 
 @torch.inference_mode()
@@ -489,4 +490,4 @@ def generate_stream_falcon(
 
     # clean
     gc.collect()
-    torch.cuda.empty_cache()
+    empty_cache()

xinference/model/llm/pytorch/yi_vl.py

@@ -57,16 +57,18 @@ class YiVLChatModel(PytorchChatModel):
         from ....thirdparty.llava.mm_utils import load_pretrained_model
         from ....thirdparty.llava.model.constants import key_info
 
-        device = self._pytorch_model_config.get("device", "auto")
-        device = select_device(device)
+        self._device = self._pytorch_model_config.get("device", "auto")
+        self._device = select_device(self._device)
 
         key_info["model_path"] = self.model_path
+        # Default device_map is auto, it can loads model to multiple cards.
+        # If the device_map is set to cuda, then only 1 card can be used.
         (
             self._tokenizer,
             self._model,
             self._image_processor,
             _,
-        ) = load_pretrained_model(self.model_path, device_map=device)
+        ) = load_pretrained_model(self.model_path, device_map=self._device)
 
     @staticmethod
     def _message_content_to_yi(content) -> Union[str, tuple]:
@@ -187,7 +189,7 @@ class YiVLChatModel(PytorchChatModel):
                 prompt, self._tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
             )
             .unsqueeze(0)
-            .cuda()
+            .to(self._device)
         )
 
         images = state.get_images(return_pil=True)
@@ -210,7 +212,9 @@ class YiVLChatModel(PytorchChatModel):
         max_new_tokens = generate_config.get("max_tokens", 512)
         generate_kwargs = {
             "input_ids": input_ids,
-            "images": image_tensor.unsqueeze(0).to(dtype=torch.bfloat16).cuda(),
+            "images": image_tensor.unsqueeze(0)
+            .to(dtype=torch.bfloat16)
+            .to(self._device),
             "streamer": streamer,
             "do_sample": True,
             "top_p": float(top_p),

xinference/model/llm/utils.py

@@ -402,6 +402,15 @@ Begin!"""
                     else:
                         ret += role + ": </s>"
             return ret
+        elif prompt_style.style_name == "gemma":
+            ret = ""
+            for message in chat_history:
+                content = message["content"]
+                role = get_role(message["role"])
+                ret += "<start_of_turn>" + role + "\n"
+                if content:
+                    ret += content + "<end_of_turn>\n"
+            return ret
         else:
             raise ValueError(f"Invalid prompt style: {prompt_style.style_name}")
 
@@ -556,7 +565,7 @@ Begin!"""
             content, func, args = cls._eval_gorilla_openfunctions_arguments(c, tools)
         elif "chatglm3" == family:
             content, func, args = cls._eval_chatglm3_arguments(c, tools)
-        elif "qwen-chat" == family:
+        elif family in ["qwen-chat", "qwen1.5-chat"]:
             content, func, args = cls._eval_qwen_chat_arguments(c, tools)
         else:
             raise Exception(

xinference/model/llm/vllm/core.py

@@ -56,6 +56,7 @@ class VLLMModelConfig(TypedDict, total=False):
     max_num_batched_tokens: int
     max_num_seqs: int
     quantization: Optional[str]
+    max_model_len: Optional[int]
 
 
 class VLLMGenerateConfig(TypedDict, total=False):
@@ -98,6 +99,8 @@ VLLM_SUPPORTED_CHAT_MODELS = [
     "mixtral-instruct-v0.1",
     "chatglm3",
 ]
+if VLLM_INSTALLED and vllm.__version__ >= "0.3.0":
+    VLLM_SUPPORTED_CHAT_MODELS.append("qwen1.5-chat")
 
 
 class VLLMModel(LLM):
@@ -151,6 +154,7 @@ class VLLMModel(LLM):
         model_config.setdefault("gpu_memory_utilization", 0.90)
         model_config.setdefault("max_num_seqs", 256)
         model_config.setdefault("quantization", None)
+        model_config.setdefault("max_model_len", 4096)
 
         return model_config
 

xinference/model/rerank/core.py

@@ -14,28 +14,22 @@
 
 import logging
 import os
-import shutil
 import uuid
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple
 
 import numpy as np
-from pydantic import BaseModel
 
 from ...constants import XINFERENCE_CACHE_DIR
 from ...types import Document, DocumentObj, Rerank
-from ..core import ModelDescription
-from ..utils import is_model_cached, valid_model_revision
+from ..core import CacheableModelSpec, ModelDescription
+from ..utils import is_model_cached
 
 logger = logging.getLogger(__name__)
 
 # Used for check whether the model is cached.
 # Init when registering all the builtin models.
 MODEL_NAME_TO_REVISION: Dict[str, List[str]] = defaultdict(list)
-
-SUPPORTED_SCHEMES = ["s3"]
-
-
 RERANK_MODEL_DESCRIPTIONS: Dict[str, List[Dict]] = defaultdict(list)
 
 
@@ -45,7 +39,7 @@ def get_rerank_model_descriptions():
     return copy.deepcopy(RERANK_MODEL_DESCRIPTIONS)
 
 
-class RerankModelSpec(BaseModel):
+class RerankModelSpec(CacheableModelSpec):
     model_name: str
     language: List[str]
     model_id: str
@@ -180,135 +174,10 @@ def get_cache_status(
     return is_model_cached(model_spec, MODEL_NAME_TO_REVISION)
 
 
-def cache_from_uri(
-    model_spec: RerankModelSpec,
-    self_hosted_storage: bool = False,
-) -> str:
-    from fsspec import AbstractFileSystem, filesystem
-
-    from ..utils import copy_from_src_to_dst, parse_uri
-
-    cache_dir = get_cache_dir(model_spec)
-    if os.path.exists(cache_dir):
-        logger.info(f"Rerank cache {cache_dir} exists")
-        return cache_dir
-
-    assert model_spec.model_uri is not None
-    src_scheme, src_root = parse_uri(model_spec.model_uri)
-    if src_root.endswith("/"):
-        # remove trailing path separator.
-        src_root = src_root[:-1]
-
-    if src_scheme == "file":
-        if not os.path.isabs(src_root):
-            raise ValueError(
-                f"Model URI cannot be a relative path: {model_spec.model_uri}"
-            )
-        os.makedirs(XINFERENCE_CACHE_DIR, exist_ok=True)
-        os.symlink(src_root, cache_dir, target_is_directory=True)
-        return cache_dir
-    elif src_scheme in SUPPORTED_SCHEMES:
-        # use anonymous connection for self-hosted storage.
-        src_fs: AbstractFileSystem = filesystem(src_scheme, anon=self_hosted_storage)
-        local_fs: AbstractFileSystem = filesystem("file")
-
-        files_to_download = []
-        os.makedirs(cache_dir, exist_ok=True)
-
-        for path, _, files in src_fs.walk(model_spec.model_uri):
-            for file in files:
-                src_path = f"{path}/{file}"
-                local_path = src_path.replace(src_root, cache_dir)
-                files_to_download.append((src_path, local_path))
-
-        from concurrent.futures import ThreadPoolExecutor
-
-        failed = False
-        with ThreadPoolExecutor(max_workers=min(len(files_to_download), 4)) as executor:
-            futures = [
-                (
-                    src_path,
-                    executor.submit(
-                        copy_from_src_to_dst, src_fs, src_path, local_fs, local_path
-                    ),
-                )
-                for src_path, local_path in files_to_download
-            ]
-            for src_path, future in futures:
-                if failed:
-                    future.cancel()
-                else:
-                    try:
-                        future.result()
-                    except:
-                        logger.error(f"Download {src_path} failed", exc_info=True)
-                        failed = True
-
-        if failed:
-            logger.warning(f"Removing cache directory: {cache_dir}")
-            shutil.rmtree(cache_dir, ignore_errors=True)
-            raise RuntimeError(
-                f"Failed to download rerank model '{model_spec.model_name}' "
-            )
-        return cache_dir
-    else:
-        raise ValueError(f"Unsupported URL scheme: {src_scheme}")
-
-
 def cache(model_spec: RerankModelSpec):
-    from huggingface_hub import snapshot_download as hf_download
-    from modelscope.hub.snapshot_download import snapshot_download as ms_download
-
-    from ..utils import retry_download, symlink_local_file
-
-    if (
-        hasattr(model_spec, "model_uri")
-        and getattr(model_spec, "model_uri", None) is not None
-    ):
-        logger.info(f"Rerank model caching from URI: {model_spec.model_uri}")
-        return cache_from_uri(model_spec=model_spec)
-
-    cache_dir = get_cache_dir(model_spec)
-    if not os.path.exists(cache_dir):
-        os.makedirs(cache_dir, exist_ok=True)
-    meta_path = os.path.join(cache_dir, "__valid_download")
-    if valid_model_revision(meta_path, model_spec.model_revision):
-        return cache_dir
-
-    if model_spec.model_hub == "modelscope":
-        logger.info(
-            f"Download {model_spec.model_name} from modelscope {model_spec.model_id}"
-        )
-        download_dir = retry_download(
-            ms_download,
-            model_spec.model_name,
-            None,
-            model_spec.model_id,
-            revision=model_spec.model_revision,
-        )
-        for subdir, dirs, files in os.walk(download_dir):
-            for file in files:
-                relpath = os.path.relpath(os.path.join(subdir, file), download_dir)
-                symlink_local_file(os.path.join(subdir, file), cache_dir, relpath)
-    else:
-        logger.info(
-            f"Download {model_spec.model_name} from huggingface {model_spec.model_id}"
-        )
-        retry_download(
-            hf_download,
-            model_spec.model_name,
-            None,
-            model_spec.model_id,
-            revision=model_spec.model_revision,
-            local_dir=cache_dir,
-            local_dir_use_symlinks=True,
-        )
-    with open(meta_path, "w") as f:
-        import json
+    from ..utils import cache
 
-        desc = RerankModelDescription(None, None, model_spec)
-        json.dump(desc.to_dict(), f)
-    return cache_dir
+    return cache(model_spec, RerankModelDescription)
 
 
 def create_rerank_model_instance(