Add Llama3.1-8B support for vLLM and use KIND_MODEL for vLLM config b…

…y default (#82)
triton-inference-server · Aug 19, 2024 · dda4ca8 · dda4ca8
1 parent a050ec1
commit dda4ca8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/triton_cli/parser.py b/src/triton_cli/parser.py
@@ -60,6 +60,8 @@
     "llama-2-7b-chat": "hf:meta-llama/Llama-2-7b-chat-hf",
     "llama-3-8b": "hf:meta-llama/Meta-Llama-3-8B",
     "llama-3-8b-instruct": "hf:meta-llama/Meta-Llama-3-8B-Instruct",
+    "llama-3.1-8b": "hf:meta-llama/Meta-Llama-3.1-8B",
+    "llama-3.1-8b-instruct": "hf:meta-llama/Meta-Llama-3.1-8B-Instruct",
     # Public
     "gpt2": "hf:gpt2",
     "opt125m": "hf:facebook/opt-125m",

diff --git a/src/triton_cli/repository.py b/src/triton_cli/repository.py
@@ -53,6 +53,7 @@
 # that can be fully autocompleted for a simple deployment.
 MODEL_CONFIG_TEMPLATE = """
 backend: "{backend}"
+instance_group {instance_group}
 """
 
 NGC_CONFIG_TEMPLATE = """
@@ -318,7 +319,10 @@ def __download_hf_model(
 
     def __generate_vllm_model(self, huggingface_id: str):
         backend = "vllm"
-        model_config = MODEL_CONFIG_TEMPLATE.format(backend=backend)
+        instance_group = "[{kind: KIND_MODEL}]"
+        model_config = MODEL_CONFIG_TEMPLATE.format(
+            backend=backend, instance_group=instance_group
+        ).strip()
         model_contents = json.dumps(
             {
                 "model": huggingface_id,