Merge pull request #108 from OpenAccess-AI-Collective/docker-gptq

default to qlora support, make gptq specific image
axolotl-ai-cloud · May 30, 2023 · bbc5bc5 · bbc5bc5
2 parents 4df9da7 + 48612f8
commit bbc5bc5
Show file tree

Hide file tree

Showing 7 changed files with 48 additions and 72 deletions.
diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
@@ -14,14 +14,17 @@ jobs:
     strategy:
       matrix:
         include:
-          - cuda: cu118
+          - cuda: "118"
             cuda_version: 11.8.0
-            cuda_version_bnb: "118"
-            pytorch: 2.0.0
-          - cuda: cu117
+            axolotl_extras:
+          - cuda: "117"
             cuda_version: 11.7.0
-            cuda_version_bnb: "117"
             pytorch: 1.13.1
+            axolotl_extras:
+          - cuda: "118"
+            cuda_version: 11.8.0
+            pytorch: 2.0.0
+            axolotl_extras: gptq
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -43,12 +46,12 @@ jobs:
           context: .
           file: ./docker/Dockerfile-base
           push: ${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }}
+          tags: ${{ steps.metadata.outputs.tags }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
           build-args: |
             CUDA_VERSION=${{ matrix.cuda_version }}
-            CUDA_VERSION_BNB=${{ matrix.cuda_version_bnb }}
             CUDA=${{ matrix.cuda }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
+            AXOLOTL_EXTRAS=${{ matrix.axolotl_extras }}
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -16,9 +16,15 @@ jobs:
           - cuda: cu118
             cuda_version: 11.8.0
             pytorch: 2.0.0
+            axolotl_extras:
+          - cuda: cu118
+            cuda_version: 11.8.0
+            pytorch: 2.0.0
+            axolotl_extras: gptq
           - cuda: cu117
             cuda_version: 11.7.0
             pytorch: 1.13.1
+            axolotl_extras:
     runs-on: self-hosted
     steps:
       - name: Checkout
@@ -40,10 +46,10 @@ jobs:
         with:
           context: .
           build-args: |
-            BASE_TAG=${{ github.ref_name }}-base-${{ matrix.cuda }}-${{ matrix.pytorch }}
+            BASE_TAG=${{ github.ref_name }}-base-${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           file: ./docker/Dockerfile
           push: ${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }}
+          tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
@@ -57,9 +63,15 @@ jobs:
           - cuda: cu118
             cuda_version: 11.8.0
             pytorch: 2.0.0
+            axolotl_extras:
+          - cuda: cu118
+            cuda_version: 11.8.0
+            pytorch: 2.0.0
+            axolotl_extras: gptq
           - cuda: cu117
             cuda_version: 11.7.0
             pytorch: 1.13.1
+            axolotl_extras:
     runs-on: self-hosted
     steps:
       - name: Checkout
@@ -81,10 +93,10 @@ jobs:
         with:
           context: .
           build-args: |
-            BASE_TAG=${{ github.ref_name }}-${{ matrix.cuda }}-${{ matrix.pytorch }}
+            BASE_TAG=${{ github.ref_name }}-${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           file: ./docker/Dockerfile-runpod
           push: ${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }}
+          tags: ${{ steps.metadata.outputs.tags }}-${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
diff --git a/README.md b/README.md
@@ -32,7 +32,7 @@
 ```bash
 git clone https://github.com/OpenAccess-AI-Collective/axolotl
 
-pip3 install -e .[int4]
+pip3 install -e .
 
 accelerate config
 
@@ -59,9 +59,9 @@ accelerate launch scripts/finetune.py examples/lora-openllama-3b/config.yml \
   1. Install python **3.9**
 
   2. Install python dependencies with ONE of the following:
-      - `pip3 install -e .[int4]` (recommended)
-      - `pip3 install -e .[int4_triton]`
-      - `pip3 install -e .`
+      - `pip3 install -e .` (recommended, supports QLoRA, no gptq/int4 support)
+      - `pip3 install -e .[gptq]` (next best if you don't need QLoRA, but want to use gptq)
+      - `pip3 install -e .[gptq_triton]`
 
 ### Dataset
 

diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -2,19 +2,26 @@ ARG BASE_TAG=main-base
 FROM winglian/axolotl-base:$BASE_TAG
 
 ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
+ARG AXOLOTL_EXTRAS=""
 
 RUN apt-get update && \
     apt-get install -y vim curl
 
 WORKDIR /workspace
 
-# The base image ships with `pydantic==1.8.2` which is not working
-RUN python3 -m pip install -U --no-cache-dir pydantic
+RUN pip3 install --force-reinstall "peft @ git+https://github.com/huggingface/peft.git@main" \
+            "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
+            "transformers @ git+https://github.com/huggingface/transformers.git@main"
 
 RUN mkdir axolotl
 COPY . axolotl/
+# If AXOLOTL_EXTRAS is set, append it in brackets
 RUN cd axolotl && \
-    pip install -e .[int4]
+    if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
+        pip install -e .[$AXOLOTL_EXTRAS]; \
+    else \
+        pip install -e .; \
+    fi
 
 # helper for huggingface-login cli
 RUN git config --global credential.helper store
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
@@ -9,7 +9,7 @@ ENV PATH="/root/miniconda3/bin:${PATH}"
 
 ARG PYTHON_VERSION="3.9"
 ARG PYTORCH="2.0.0"
-ARG CUDA="cu118"
+ARG CUDA="118"
 
 ENV PYTHON_VERSION=$PYTHON_VERSION
 
@@ -29,7 +29,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
 WORKDIR /workspace
 
 RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
-    python3 -m pip install --no-cache-dir -U torch==${PYTORCH} torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/$CUDA
+    python3 -m pip install --no-cache-dir -U torch==${PYTORCH} torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu$CUDA
 
 
 FROM base-builder AS flash-attn-builder
@@ -61,12 +61,12 @@ RUN git clone https://github.com/microsoft/DeepSpeed.git && \
 FROM base-builder AS bnb-builder
 
 WORKDIR /workspace
-ARG CUDA_VERSION_BNB="118"
-ENV CUDA_VERSION_BNB=$CUDA_VERSION_BNB
+ARG CUDA="118"
+ENV CUDA=$CUDA
 
 RUN git clone https://github.com/TimDettmers/bitsandbytes.git && \
     cd bitsandbytes && \
-    CUDA_VERSION=$CUDA_VERSION_BNB make cuda11x && \
+    CUDA_VERSION=$CUDA make cuda11x && \
     python setup.py bdist_wheel
 
 FROM base-builder
@@ -93,9 +93,6 @@ COPY --from=flash-attn-builder /workspace/flash-attention/csrc/layer_norm/dist/d
 RUN pip3 install wheels/deepspeed-*.whl wheels/flash_attn-*.whl wheels/fused_dense_lib-*.whl wheels/xentropy_cuda_lib-*.whl wheels/rotary_emb-*.whl wheels/dropout_layer_norm-*.whl
 RUN cd /workspace/builds/bitsandbytes && python3 setup.py install
 RUN git lfs install --skip-repo
-RUN pip3 install "peft @ git+https://github.com/huggingface/peft.git@main" \
-        "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
-        "transformers @ git+https://github.com/huggingface/transformers.git@main" && \
-    pip3 install awscli && \
+RUN pip3 install awscli && \
     # The base image ships with `pydantic==1.8.2` which is not working
     pip3 install -U --no-cache-dir pydantic
diff --git a/scripts/setup-runpod.sh b/scripts/setup-runpod.sh
diff --git a/setup.py b/setup.py
@@ -19,10 +19,10 @@
     packages=find_packages(),
     install_requires=install_requires,
     extras_require={
-        "int4": [
+        "gptq": [
             "alpaca_lora_4bit @ git+https://github.com/winglian/alpaca_lora_4bit.git@setup_pip",
         ],
-        "int4_triton": [
+        "gptq_triton": [
             "alpaca_lora_4bit[triton] @ git+https://github.com/winglian/alpaca_lora_4bit.git@setup_pip",
         ],
         "extras": [