From 34c0a86a1152430f4189278a5f8f1da375f5f36b Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Fri, 8 Sep 2023 11:58:54 -0400
Subject: [PATCH] update readme to point to direct link to runpod template,
 cleanup install instrucitons (#532)

* update readme to point to direct link to runpod template, cleanup install instrucitons

* default install flash-attn and auto-gptq now too

* update readme w flash-attn extra

* fix version in setup
---
 .github/workflows/tests.yml |  4 ++--
 README.md                   | 20 ++++----------------
 docker/Dockerfile           |  4 ++--
 requirements.txt            |  2 +-
 setup.py                    |  9 ++-------
 5 files changed, 11 insertions(+), 28 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index d5184def6..9ff08db07 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -24,8 +24,8 @@ jobs:
 
       - name: Install dependencies
         run: |
-          pip install -e .
-          pip install -r requirements-tests.txt
+          pip3 install -e .
+          pip3 install -r requirements-tests.txt
 
       - name: Run tests
         run: |
diff --git a/README.md b/README.md
index 90dd96e8b..775592efe 100644
--- a/README.md
+++ b/README.md
@@ -90,8 +90,7 @@ accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml \
   ```bash
   docker run --gpus '"all"' --rm -it winglian/axolotl:main-py3.10-cu118-2.0.1
   ```
-  - `winglian/axolotl-runpod:main-py3.10-cu118-2.0.1`: for runpod
-  - `winglian/axolotl-runpod:main-py3.9-cu118-2.0.1-gptq`: for gptq
+  - `winglian/axolotl-runpod:main-latest`: for runpod or use this [direct link](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
 
   Or run on the current files for development:
 
@@ -104,19 +103,9 @@ accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml \
 
   2. Install pytorch stable https://pytorch.org/get-started/locally/
 
-  3. Install python dependencies with ONE of the following:
-      - Recommended, supports QLoRA, NO gptq/int4 support
+  3. Install axolotl along with python dependencies
         ```bash
-        pip3 install -e .
-        pip3 install -U git+https://github.com/huggingface/peft.git
-        ```
-      - gptq/int4 support, NO QLoRA
-        ```bash
-        pip3 install -e .[gptq]
-        ```
-      - same as above but not recommended
-        ```bash
-        pip3 install -e .[gptq_triton]
+        pip3 install -e .[flash-attn]
         ```
 
 - LambdaLabs
@@ -151,10 +140,9 @@ accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml \
   git clone https://github.com/OpenAccess-AI-Collective/axolotl
   cd axolotl
 
-  pip3 install -e . # change depend on needs
+  pip3 install -e .
   pip3 install protobuf==3.20.3
   pip3 install -U --ignore-installed requests Pillow psutil scipy
-  pip3 install git+https://github.com/huggingface/peft.git # not for gptq
   ```
 
   5. Set path
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 683ca75ff..8608e2348 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -15,9 +15,9 @@ RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN cd axolotl && \
     if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[flash-attn,gptq,$AXOLOTL_EXTRAS]; \
+        pip install -e .[flash-attn,$AXOLOTL_EXTRAS]; \
     else \
-        pip install -e .[flash-attn,gptq]; \
+        pip install -e .[flash-attn]; \
     fi
 
 # fix so that git fetch/pull from remote works
diff --git a/requirements.txt b/requirements.txt
index 1c8e97dff..6d33f5728 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,7 +12,7 @@ evaluate
 fire
 PyYAML>=6.0
 datasets
-flash-attn>=2.0.8
+flash-attn>=2.2.1
 sentencepiece
 wandb
 einops
diff --git a/setup.py b/setup.py
index 973d656cd..a53603f8b 100644
--- a/setup.py
+++ b/setup.py
@@ -7,9 +7,7 @@ def parse_requirements():
     _install_requires = []
     _dependency_links = []
     with open("./requirements.txt", encoding="utf-8") as requirements_file:
-        lines = [
-            r.strip() for r in requirements_file.readlines() if "auto-gptq" not in r
-        ]
+        lines = [r.strip() for r in requirements_file.readlines()]
         for line in lines:
             if line.startswith("--extra-index-url"):
                 # Handle custom index URLs
@@ -33,11 +31,8 @@ setup(
     install_requires=install_requires,
     dependency_links=dependency_links,
     extras_require={
-        "gptq": [
-            "auto-gptq",
-        ],
         "flash-attn": [
-            "flash-attn==2.0.8",
+            "flash-attn>=2.2.1",
         ],
         "extras": [
             "deepspeed",