feat: add cpu optional deps (#43)

Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · May 26, 2023 · 8e95e34 · 8e95e34
1 parent ef5d369
commit 8e95e34
Show file tree

Hide file tree

Showing 6 changed files with 10 additions and 7 deletions.
diff --git a/images/base/Dockerfile b/images/base/Dockerfile
@@ -62,7 +62,7 @@ RUN mkdir -p /workspace
 
 COPY . workspace/
 WORKDIR /workspace
-RUN pip install -e .
+RUN pip install -e .[cpu]
 
 RUN modelz-llm --dry-run
 

diff --git a/images/chatglm-6b/Dockerfile b/images/chatglm-6b/Dockerfile
@@ -62,7 +62,7 @@ RUN mkdir -p /workspace
 
 COPY . workspace/
 WORKDIR /workspace
-RUN pip install -e .
+RUN pip install -e .[cpu]
 
 # RUN modelz-llm --dry-run --model THUDM/chatglm-6b
 

diff --git a/images/fastchat-t5-3b/Dockerfile b/images/fastchat-t5-3b/Dockerfile
@@ -62,7 +62,7 @@ RUN mkdir -p /workspace
 
 COPY . workspace/
 WORKDIR /workspace
-RUN pip install -e .
+RUN pip install -e .[cpu]
 
 RUN modelz-llm --dry-run --model lmsys/fastchat-t5-3b-v1.0
 

diff --git a/images/llama-7b/Dockerfile b/images/llama-7b/Dockerfile
@@ -63,7 +63,7 @@ RUN mkdir -p /workspace
 
 COPY . workspace/
 WORKDIR /workspace
-RUN pip install -e .
+RUN pip install -e .[cpu]
 
 # RUN modelz-llm --dry-run --model decapoda-research/llama-7b-hf
 

diff --git a/images/vicuna-7b/Dockerfile b/images/vicuna-7b/Dockerfile
@@ -63,7 +63,7 @@ RUN mkdir -p /workspace
 
 COPY . workspace/
 WORKDIR /workspace
-RUN pip install -e .
+RUN pip install -e .[cpu]
 
 # RUN modelz-llm --dry-run --model lmsys/vicuna-7b-delta-v1.1
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -17,13 +17,16 @@ dependencies = [
 ]
 [project.optional-dependencies]
 cpu = [
-    "torch --extra-index-url https://download.pytorch.org/whl/cpu",
+    # only include torch CPU for Linux x86_64
+    "torch @ https://download.pytorch.org/whl/cpu/torch-2.0.1%2Bcpu-cp39-cp39-linux_x86_64.whl ; python_version == '3.9'",
+    "torch @ https://download.pytorch.org/whl/cpu/torch-2.0.1%2Bcpu-cp310-cp310-linux_x86_64.whl ; python_version == '3.10'",
+    "torch @ https://download.pytorch.org/whl/cpu/torch-2.0.1%2Bcpu-cp311-cp311-linux_x86_64.whl ; python_version == '3.11'",
     "transformers",
     "sentencepiece",
     "accelerate",
 ]
 gpu = [
-    "torch --extra-index-url https://download.pytorch.org/whl/cu116",
+    "torch",  # CUDA 11.7, ref https://pytorch.org/get-started/locally/
     "transformers",
     "sentencepiece",
     "accelerate",