From 2734e3f1a282539d09e17acfbfaf6a63f3bc29cd Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Sun, 7 May 2023 10:26:29 -0400
Subject: [PATCH] build base separately

fix arg order for image
fix dockerfile var excaping
move args around
---
 .github/workflows/base.yml | 35 +++++++++++++++++++++++++++++++++++
 .github/workflows/main.yml | 29 +----------------------------
 docker/Dockerfile-base     | 34 +++++++++++++++++++++++++++++++++-
 3 files changed, 69 insertions(+), 29 deletions(-)
 create mode 100644 .github/workflows/base.yml

diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
new file mode 100644
index 000000000..a03b03cce
--- /dev/null
+++ b/.github/workflows/base.yml
@@ -0,0 +1,35 @@
+name: ci-cd
+
+on:
+  push:
+    branches:
+      - "main-base"
+      - "dev-base"
+
+jobs:
+  build-base:
+    if: github.repository_owner == 'OpenAccess-AI-Collective'
+    runs-on: self-hosted
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@v3
+        with:
+          images: winglian/axolotl-base
+      - name: Login to Docker Hub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Build
+        uses: docker/build-push-action@v4
+        with:
+          context: .
+          file: ./docker/Dockerfile-base
+          push: ${{ github.event_name != 'pull_request' }}
+          tags: ${{ steps.metadata.outputs.tags }}
+          labels: ${{ steps.metadata.outputs.labels }}
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index f8d54c031..742080540 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -1,4 +1,4 @@
-name: ci
+name: ci-cd
 
 on:
   push:
@@ -7,34 +7,7 @@ on:
       - "dev"
 
 jobs:
-  build-base:
-    if: github.repository_owner == 'OpenAccess-AI-Collective'
-    runs-on: self-hosted
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-      - name: Docker metadata
-        id: metadata
-        uses: docker/metadata-action@v3
-        with:
-          images: winglian/axolotl-base
-      - name: Login to Docker Hub
-        uses: docker/login-action@v2
-        with:
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v2
-      - name: Build
-        uses: docker/build-push-action@v4
-        with:
-          context: .
-          file: ./docker/Dockerfile-base
-          push: ${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.metadata.outputs.tags }}
-          labels: ${{ steps.metadata.outputs.labels }}
   build-axolotl:
-    needs: build-base
     if: github.repository_owner == 'OpenAccess-AI-Collective'
     runs-on: self-hosted
     steps:
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
index f50c7f734..1f28e4575 100644
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -2,11 +2,14 @@ ARG CUDA_VERSION="11.8.0"
 ARG CUDNN_VERSION="8"
 ARG UBUNTU_VERSION="22.04"
 
-FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION
+FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION as base-builder
 
 ENV PATH="/root/miniconda3/bin:${PATH}"
 
 ARG PYTHON_VERSION="3.9"
+ARG PYTORCH="2.0.0"
+ARG CUDA="cu118"
+
 ENV PYTHON_VERSION=$PYTHON_VERSION
 
 RUN apt-get update
@@ -23,3 +26,32 @@ RUN conda create -n "py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
 ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
 
 WORKDIR /workspace
+
+RUN python3 -m pip install --no-cache-dir -U torch==${PYTORCH} torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/$CUDA
+
+
+FROM base-builder AS flash-attn-builder
+
+WORKDIR /workspace
+
+ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
+
+RUN git clone https://github.com/HazyResearch/flash-attention.git && \
+    cd flash-attention && \
+    python3 setup.py bdist_wheel
+
+FROM base-builder AS deepspeed-builder
+
+WORKDIR /workspace
+
+RUN git clone https://github.com/microsoft/DeepSpeed.git && \
+    cd DeepSpeed && \
+    MAX_CONCURRENCY=8 DS_BUILD_SPARSE_ATTN=0 DS_BUILD_OPS=1 python3 setup.py bdist_wheel
+
+FROM base-builder
+
+RUN mkdir /workspace/wheels
+COPY --from=deepspeed-builder /workspace/DeepSpeed/dist/deepspeed-*.whl wheels
+COPY --from=flash-attn-builder /workspace/flash-attention/dist/flash_attn-*.whl wheels
+
+RUN pip3 install.sh wheels/deepspeed-*.whl wheels/flash_attn-*.whl