FEAT: support wizard-v1.1 (#183)

xorbitsai · Jul 14, 2023 · 80003e1 · 80003e1
1 parent 15aaeee
commit 80003e1
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 4 deletions.
diff --git a/xinference/model/llm/__init__.py b/xinference/model/llm/__init__.py
@@ -81,6 +81,36 @@ def install():
         ),
     )
 
+    wizardlm_v1_1_url_generator = lambda model_size, quantization: (
+        f"https://huggingface.co/TheBloke/WizardLM-{model_size}B-V1.1-GGML/resolve/main/"
+        f"wizardlm-{model_size}b-v1.1.ggmlv3.{quantization}.bin"
+    )
+    MODEL_FAMILIES.append(
+        ModelFamily(
+            model_name="wizardlm-v1.1",
+            model_sizes_in_billions=[13],
+            model_format="ggmlv3",
+            quantizations=[
+                "q2_K",
+                "q3_K_L",
+                "q3_K_M",
+                "q3_K_S",
+                "q4_0",
+                "q4_1",
+                "q4_K_M",
+                "q4_K_S",
+                "q5_0",
+                "q5_1",
+                "q5_K_M",
+                "q5_K_S",
+                "q6_K",
+                "q8_0",
+            ],
+            url_generator=wizardlm_v1_1_url_generator,
+            cls=VicunaCensoredGgml,  # according to https://huggingface.co/TheBloke/WizardLM-13B-V1.1-GGML
+        ),
+    )
+
     vicuna_v1_3_url_generator = lambda model_size, quantization: (
         "https://huggingface.co/TheBloke/vicuna-7B-v1.3-GGML/resolve/main/"
         f"vicuna-7b-v1.3.ggmlv3.{quantization}.bin"

diff --git a/xinference/model/llm/core.py b/xinference/model/llm/core.py
@@ -130,18 +130,38 @@ def __init__(
         )
         self._llm = None
 
+    @staticmethod
+    def _is_darwin_and_apple_silicon():
+        return platform.system() == "Darwin" and platform.processor() == "arm"
+
+    @staticmethod
+    def _is_linux():
+        return platform.system() == "Linux"
+
+    def _can_apply_metal(self):
+        return (
+            self.model_spec.quantization == "q4_0"
+            or self.model_spec.quantization == "q4_1"
+        )
+
+    def _can_apply_cublas(self):
+        # TODO: figure out the quantizations supported.
+        return True
+
     def _sanitize_model_config(
         self, llamacpp_model_config: Optional[LlamaCppModelConfig]
     ) -> LlamaCppModelConfig:
         if llamacpp_model_config is None:
             llamacpp_model_config = LlamaCppModelConfig()
         if platform.system() == "Windows":
-            context_length = 512
+            llamacpp_model_config.setdefault("n_ctx", 512)
         else:
-            context_length = 2048
+            llamacpp_model_config.setdefault("n_ctx", 2048)
 
-        llamacpp_model_config.setdefault("n_gpu_layers", self._gpu_layers)
-        llamacpp_model_config.setdefault("n_ctx", context_length)
+        if self._is_darwin_and_apple_silicon() and self._can_apply_metal():
+            llamacpp_model_config.setdefault("n_gpu_layers", 1)
+        elif self._is_linux() and self._can_apply_cublas():
+            llamacpp_model_config.setdefault("n_gpu_layers", self._gpu_layers)
 
         return llamacpp_model_config