Added missing nvidia-ml-py package (#558)

predibase · Jul 26, 2024 · d1a4d09 · d1a4d09
1 parent 15a38d5
commit d1a4d09
Show file tree

Hide file tree

Showing 5 changed files with 1,570 additions and 1,096 deletions.
diff --git a/server/Makefile b/server/Makefile
@@ -19,7 +19,7 @@ gen-server:
 
 install: gen-server
 	pip install pip --upgrade
-	pip install torch==2.2.0
+	pip install torch==2.4.0
 	pip install -r requirements.txt
 	pip install -e ".[bnb, accelerate, quantize, peft, outlines]"
 
@@ -31,7 +31,7 @@ run-dev:
 	# SAFETENSORS_FAST_GPU=1 python -m torch.distributed.run --nproc_per_node=1 lorax_server/cli.py serve flozi00/Mistral-7B-german-assistant-v5-4bit-autogptq --quantize gptq
 
 export-requirements:
-	poetry export -o requirements.txt -E bnb --without-hashes
+	poetry export -o requirements.txt --without-hashes
 
 format:
 	pip install ruff

diff --git a/server/lorax_server/utils/paged_attention.py b/server/lorax_server/utils/paged_attention.py
@@ -12,8 +12,8 @@
         import vllm._custom_ops as ops
     except Exception as e:
         raise ImportError(
-            f"Could not import vllm paged attention. Make sure your installation is correct. Complete error: {e}"
-        )
+            f"Could not import vllm paged attention. Make sure your installation is correct. Error: {e}"
+        ) from e
 
 # TODO(travis): fix for CUDA 8.9 (Lovelace) and 9.0 (Hopper)
 # if torch.cuda.is_available():