add deepspeed-kernels dependency for deepspeed>=0.12.0 (#827)
Browse files- docker/Dockerfile-base +1 -1
docker/Dockerfile-base
CHANGED
@@ -27,7 +27,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
|
|
27 |
WORKDIR /workspace
|
28 |
|
29 |
RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
|
30 |
-
python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA
|
31 |
|
32 |
FROM base-builder AS deepspeed-builder
|
33 |
|
|
|
27 |
WORKDIR /workspace
|
28 |
|
29 |
RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
|
30 |
+
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
|
31 |
|
32 |
FROM base-builder AS deepspeed-builder
|
33 |
|