ad einops for supporting falcon models (deepjavalibrary#792)

KexinFeng · Aug 16, 2023 · 6585880 · 6585880
1 parent 411f76a
commit 6585880
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/serving/docker/deepspeed.Dockerfile b/serving/docker/deepspeed.Dockerfile
@@ -56,7 +56,7 @@ RUN apt-get update && \
     DEBIAN_FRONTEND=noninteractive apt-get install -yq libaio-dev libopenmpi-dev && \
     pip3 install torch==${torch_version} torchvision==${torch_vision_version} --extra-index-url https://download.pytorch.org/whl/cu117 \
     ${deepspeed_wheel} transformers==${transformers_version} \
-    triton==2.0.0.dev20221202 mpi4py sentencepiece accelerate==${accelerate_version} bitsandbytes==${bitsandbytes_version}\
+    triton==2.0.0.dev20221202 mpi4py sentencepiece einops accelerate==${accelerate_version} bitsandbytes==${bitsandbytes_version}\
     diffusers[torch]==${diffusers_version} opencv-contrib-python-headless && \
     scripts/install_aitemplate.sh && \
     scripts/patch_oss_dlc.sh python && \

diff --git a/serving/docker/fastertransformer.Dockerfile b/serving/docker/fastertransformer.Dockerfile
@@ -60,7 +60,7 @@ RUN apt-get update && apt-get install -y wget git libnuma-dev zlib1g-dev rapidjs
     cd ../../ && rm -rf ompi && \
     scripts/install_python.sh ${python_version} && \
     pip3 install ${torch_wheel} ${ft_wheel} ${tb_wheel} && \
-    pip3 install transformers==${transformers_version} accelerate==${accelerate_version} bitsandbytes==${bitsandbytes_version} && \
+    pip3 install transformers==${transformers_version} accelerate==${accelerate_version} bitsandbytes==${bitsandbytes_version} einops && \
     pip3 install cmake sentencepiece && \
     pip3 cache purge && \
     apt-get clean -y && rm -rf /var/lib/apt/lists/* && \