move deepspeed-kernels install to setup.py

axolotl-ai-cloud · Jan 24, 2024 · 4ce5990 · 4ce5990
1 parent f3a03b7
commit 4ce5990
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
@@ -29,7 +29,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
 WORKDIR /workspace
 
 RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
-    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} deepspeed-kernels --extra-index-url https://download.pytorch.org/whl/cu$CUDA
+    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA
 
 RUN git lfs install --skip-repo && \
     pip3 install awscli && \

diff --git a/setup.py b/setup.py
@@ -57,6 +57,7 @@ def parse_requirements():
         ],
         "deepspeed": [
             "deepspeed>=0.13.1",
+            "deepspeed-kernels",
         ],
         "mamba-ssm": [
             "mamba-ssm==1.0.1",