move vllm to optional-dependencies

microsoft · Aug 23, 2024 · 9d72763 · 9d72763
1 parent 1d6949a
commit 9d72763
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 4 deletions.
diff --git a/mttl/__init__.py b/mttl/__init__.py
@@ -15,7 +15,9 @@
 except ImportError:
     from mttl.logging import logger
 
-    logger.info("Flash Attention not available")
+    logger.info(
+        'Flash Attention not available. You can install it with `pip install -e ".[flash_attn]`".'
+    )
 
 import torch
 

diff --git a/mttl/vllm_engines/engines.py b/mttl/vllm_engines/engines.py
@@ -16,7 +16,9 @@
     from vllm.model_executor.parallel_utils.parallel_state import destroy_model_parallel
 except ImportError:
     LLM = object
-    logger.warning("VLLM is not installed. Please install it to use LLMEngine.")
+    logger.warning(
+        'VLLM is not installed. Please install it with `pip install -e ".[vllm]"` to use LLMEngine.'
+    )
 
 
 def save_merged_model(model, model_path, hf_path="/tmp/merged"):

diff --git a/pyproject.toml b/pyproject.toml
@@ -33,10 +33,14 @@ dev = [
     "isort",
 ]
 
-flash-att = [
+flash-attn = [
     "flash-attn>=2.6.0",
 ]
 
+vllm = [
+    "vllm",
+]
+
 [project.urls]
 "Homepage" = "https://github.com/microsoft/mttl"
 "Bug Tracker" = "https://github.com/microsoft/mttl/issues"

diff --git a/requirements.txt b/requirements.txt
@@ -19,7 +19,6 @@ matplotlib
 openai
 ray
 nevergrad
-vllm
 evaluate
 seaborn
 azure-storage-blob