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

default to qlora support, make gptq specific image

.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"
+            axolotl_extras:
-            pytorch: 2.0.0
+          - cuda: "117"
-          - cuda: cu117
             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 }}

.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

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 .` (recommended, supports QLoRA, no gptq/int4 support)
-      - `pip3 install -e .[int4_triton]`
+      - `pip3 install -e .[gptq]` (next best if you don't need QLoRA, but want to use gptq)
-      - `pip3 install -e .`
+      - `pip3 install -e .[gptq_triton]`
 
 ### Dataset
 

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 pip3 install --force-reinstall "peft @ git+https://github.com/huggingface/peft.git@main" \
-RUN python3 -m pip install -U --no-cache-dir pydantic
+            "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

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"
+ARG CUDA="118"
-ENV CUDA_VERSION_BNB=$CUDA_VERSION_BNB
+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" \
+RUN pip3 install awscli && \
-        "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
-        "transformers @ git+https://github.com/huggingface/transformers.git@main" && \
-    pip3 install awscli && \
     # The base image ships with `pydantic==1.8.2` which is not working
     pip3 install -U --no-cache-dir pydantic

scripts/setup-runpod.sh

@@ -1,43 +0,0 @@
-#!/bin/bash
-
-export WANDB_MODE=offline
-export WANDB_CACHE_DIR=/workspace/data/wandb-cache
-mkdir -p $WANDB_CACHE_DIR
-
-mkdir -p /workspace/data/huggingface-cache/{hub,datasets}
-export HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
-export HUGGINGFACE_HUB_CACHE="/workspace/data/huggingface-cache/hub"
-export TRANSFORMERS_CACHE="/workspace/data/huggingface-cache/hub"
-export NCCL_P2P_DISABLE=1
-
-nvidia-smi
-num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
-gpu_indices=$(seq 0 $((num_gpus - 1)) | paste -sd "," -)
-export CUDA_VISIBLE_DEVICES=$gpu_indices
-echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
-
-apt-get update
-apt-get install -y build-essential ninja-build vim git-lfs
-git lfs install
-pip3 install --force-reinstall https://download.pytorch.org/whl/nightly/cu117/torch-2.0.0.dev20230301%2Bcu117-cp38-cp38-linux_x86_64.whl --index-url https://download.pytorch.org/whl/nightly/cu117
-if [ -z "${TORCH_CUDA_ARCH_LIST}" ]; then # only set this if not set yet
-    # this covers most common GPUs that the installed version of pytorch supports
-    # python -c "import torch; print(torch.cuda.get_arch_list())"
-    export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
-fi
-
-# install flash-attn and deepspeed from pre-built wheels for this specific container b/c these take forever to install
-mkdir -p /workspace/wheels
-cd /workspace/wheels
-curl -L -O https://github.com/OpenAccess-AI-Collective/axolotl/raw/wheels/wheels/deepspeed-0.9.2%2B7ddc3b01-cp38-cp38-linux_x86_64.whl
-curl -L -O https://github.com/OpenAccess-AI-Collective/axolotl/raw/wheels/wheels/flash_attn-1.0.4-cp38-cp38-linux_x86_64.whl
-pip install deepspeed-0.9.2%2B7ddc3b01-cp38-cp38-linux_x86_64.whl
-pip install flash_attn-1.0.4-cp38-cp38-linux_x86_64.whl
-pip install "peft @ git+https://github.com/huggingface/peft.git@main" --force-reinstall --no-dependencies
-
-cd /workspace/
-git clone https://github.com/OpenAccess-AI-Collective/axolotl.git
-cd axolotl
-pip install -e .[int4]
-mkdir -p ~/.cache/huggingface/accelerate/
-cp configs/accelerate/default_config.yaml ~/.cache/huggingface/accelerate/default_config.yaml

setup.py

@@ -19,10 +19,10 @@ setup(
     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": [