Merge pull request #16 from foundation-model-stack/pr_xformers

Pr xformers (in main reoo)
foundation-model-stack · Jun 30, 2023 · e94b0bc · e94b0bc
2 parents 459d430 + 2787116
commit e94b0bc
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/nightly-pytorch-compile/Dockerfile b/nightly-pytorch-compile/Dockerfile
@@ -48,9 +48,15 @@ RUN pip install \
     chmod -R g=u /opt/conda
 
 # xformers install (built from source)
-RUN pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers --no-cache-dir && \
+ENV CUDA_HOME=/usr/local/cuda
+ENV MAX_JOBS=2
+ENV TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6"
+RUN apt-get install -y libcublas-11-8 libcublas-dev-11-8 libcusparse-11-8 libcusparse-dev-11-8 libcusolver-11-8 libcusolver-dev-11-8 libcurand-11-8 libcurand-dev-11-8 && \
+    pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers --no-cache-dir && \
     chmod -R g+w /opt/conda && \
-    chmod -R g=u /opt/conda
+    chmod -R g=u /opt/conda && \
+    apt-get remove -y libcublas-11-8 libcublas-dev-11-8 libcusparse-11-8 libcusparse-dev-11-8 libcusolver-11-8 libcusolver-dev-11-8 libcurand-11-8 libcurand-dev-11-8 && \
+    apt-get -y clean
 
 # setup .aim_profile (needed to get aim running)
 RUN mkdir -p ${HOME} && \