From ef223519c9dba8334ca13b60b4619e9b61f01373 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Tue, 28 May 2024 11:23:34 -0400
Subject: [PATCH] update deps (#1663) [skip ci]

* update deps and tweak logic so axolotl is pip installable

* use vcs url format

* using dependency_links isn't supported per docs)
---
 .github/workflows/tests.yml |  5 +++++
 requirements.txt            |  8 ++++----
 setup.py                    | 16 +++++++++++-----
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 8f25eddc3..db9173cac 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -84,6 +84,11 @@ jobs:
             python_version: "3.11"
             pytorch: 2.2.2
             num_gpus: 1
+          - cuda: 121
+            cuda_version: 12.1.0
+            python_version: "3.11"
+            pytorch: 2.3.0
+            num_gpus: 1
     steps:
       - name: Checkout
         uses: actions/checkout@v4
diff --git a/requirements.txt b/requirements.txt
index 6723b1798..557c5293f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,7 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
 packaging==23.2
-peft==0.10.0
-transformers==4.40.2
+peft==0.11.1
+transformers==4.41.1
 tokenizers==0.19.1
 bitsandbytes==0.43.1
 accelerate==0.30.1
@@ -16,7 +16,7 @@ flash-attn==2.5.8
 sentencepiece
 wandb
 einops
-xformers==0.0.23.post1
+xformers==0.0.26.post1
 optimum==1.16.2
 hf_transfer
 colorama
@@ -39,6 +39,6 @@ s3fs
 gcsfs
 # adlfs
 
-trl==0.8.5
+trl==0.8.6
 zstandard==0.22.0
 fastcore
diff --git a/setup.py b/setup.py
index 31a6d6716..3d1537edf 100644
--- a/setup.py
+++ b/setup.py
@@ -30,8 +30,11 @@ def parse_requirements():
 
     try:
         if "Darwin" in platform.system():
-            _install_requires.pop(_install_requires.index("xformers==0.0.23.post1"))
+            # don't install xformers on MacOS
+            _install_requires.pop(_install_requires.index("xformers==0.0.26.post1"))
         else:
+            # detect the version of torch already installed
+            # and set it so dependencies don't clobber the torch version
             torch_version = version("torch")
             _install_requires.append(f"torch=={torch_version}")
 
@@ -46,11 +49,14 @@ def parse_requirements():
                 raise ValueError("Invalid version format")
 
             if (major, minor) >= (2, 3):
-                _install_requires.pop(_install_requires.index("xformers==0.0.23.post1"))
-                _install_requires.append("xformers>=0.0.26.post1")
+                pass
             elif (major, minor) >= (2, 2):
-                _install_requires.pop(_install_requires.index("xformers==0.0.23.post1"))
+                _install_requires.pop(_install_requires.index("xformers==0.0.26.post1"))
                 _install_requires.append("xformers>=0.0.25.post1")
+            else:
+                _install_requires.pop(_install_requires.index("xformers==0.0.26.post1"))
+                _install_requires.append("xformers>=0.0.23.post1")
+
     except PackageNotFoundError:
         pass
 
@@ -62,7 +68,7 @@ install_requires, dependency_links = parse_requirements()
 
 setup(
     name="axolotl",
-    version="0.4.0",
+    version="0.4.1",
     description="LLM Trainer",
     long_description="Axolotl is a tool designed to streamline the fine-tuning of various AI models, offering support for multiple configurations and architectures.",
     package_dir={"": "src"},