From 38a4818997b79f7913f897fda284a38982336fe4 Mon Sep 17 00:00:00 2001
From: Augusto de la Torre <augustod@spotify.com>
Date: Fri, 31 Mar 2023 15:59:03 +0200
Subject: [PATCH] Bump libs to test torch 2

---
 docker/Dockerfile               | 20 ++++++++++----------
 docker/requirements-build.txt   |  3 +--
 docker/requirements-runtime.txt |  2 +-
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 96f1002..942e948 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,6 +1,6 @@
 ###################
 # Builder Stage
-FROM nvidia/cuda:11.7.1-devel-ubuntu22.04 AS builder
+FROM nvidia/cuda:11.8.0-devel-ubuntu22.04 AS builder
 
 ARG DEBIAN_FRONTEND=noninteractive
 
@@ -23,19 +23,16 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
 ENV VIRTUAL_ENV=/workspace/venv
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
-
 ADD requirements-build.txt /build
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m venv ${VIRTUAL_ENV} && \
-    pip install -U -I torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url "https://download.pytorch.org/whl/cu117" && \
+    pip install -U -I torch==2.0.0+cu118 torchvision==0.15.1+cu118 --extra-index-url "https://download.pytorch.org/whl/cu118" && \
     pip install -r requirements-build.txt && \
-    export FORCE_CUDA=1 && export TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6" && export CUDA_VISIBLE_DEVICES=0 && \
-    pip install --no-deps git+https://github.com/facebookresearch/xformers.git@e14dba4#egg=xformers
-
+    pip install --no-deps xformers==0.0.18
 
 ###################
 # Runtime Stage
-FROM nvidia/cuda:11.7.1-runtime-ubuntu22.04 as runtime
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04 as runtime
 
 # Use bash shell
 SHELL ["/bin/bash", "-o", "pipefail", "-c"]
@@ -74,18 +71,21 @@ RUN echo "source ${VIRTUAL_ENV}/bin/activate" >> /root/.bashrc
 # Workaround for:
 #   https://github.com/TimDettmers/bitsandbytes/issues/62
 #   https://github.com/TimDettmers/bitsandbytes/issues/73
-ENV LD_LIBRARY_PATH="/usr/local/cuda-11.7/targets/x86_64-linux/lib"
-RUN ln /usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudart.so.11.0 /usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudart.so
+ENV LD_LIBRARY_PATH="/usr/local/cuda-11.8/targets/x86_64-linux/lib/"
+RUN ln /usr/local/cuda/targets/x86_64-linux/lib/libcudart.so.11.8.89 /usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudart.so
+RUN ln /usr/local/cuda/targets/x86_64-linux/lib/libnvrtc.so.11.8.89 /usr/local/cuda-11.8/targets/x86_64-linux/lib/libnvrtc.so
 ADD requirements-runtime.txt /
 RUN pip install --no-cache-dir -r requirements-runtime.txt
 
 WORKDIR /workspace
 RUN git clone https://github.com/victorchall/EveryDream2trainer
+
 WORKDIR /workspace/EveryDream2trainer
+# RUN git checkout torch2
 RUN python utils/get_yamls.py && \
     mkdir -p logs && mkdir -p input
 
 ADD welcome.txt /
 ADD start.sh /
 RUN chmod +x /start.sh
-CMD [ "/start.sh" ]
+CMD [ "/start.sh" ]
\ No newline at end of file
diff --git a/docker/requirements-build.txt b/docker/requirements-build.txt
index 83f3da4..2464c24 100644
--- a/docker/requirements-build.txt
+++ b/docker/requirements-build.txt
@@ -5,5 +5,4 @@ omegaconf==2.2.3
 protobuf==3.20.3
 pyre-extensions==0.0.23
 pytorch-lightning==1.9.2
-transformers==4.25.1
-triton>=2.0.0a2
\ No newline at end of file
+transformers==4.27.1
\ No newline at end of file
diff --git a/docker/requirements-runtime.txt b/docker/requirements-runtime.txt
index 0e7a591..c206e47 100644
--- a/docker/requirements-runtime.txt
+++ b/docker/requirements-runtime.txt
@@ -1,5 +1,5 @@
 aiohttp==3.8.4
-bitsandbytes==0.37.0
+bitsandbytes==0.37.2
 colorama==0.4.6
 ftfy==6.1.1
 ipyevents