update test and main/nightly builds (#1797)

* update test and main/nightly builds

* don't install mamba-ssm on 2.4.0 since it has no wheels yet

.github/workflows/main.yml

@@ -13,28 +13,22 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - cuda: 118
+          - cuda: 121
-            cuda_version: 11.8.0
+            cuda_version: 12.1.1
             python_version: "3.10"
-            pytorch: 2.1.2
+            pytorch: 2.3.1
-            axolotl_extras:
+            axolotl_extras: mamba-ssm
-            axolotl_args: "--extra-index-url https://download.pytorch.org/whl/cu118"
           - cuda: 121
-            cuda_version: 12.1.0
+            cuda_version: 12.1.1
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
             python_version: "3.11"
             pytorch: 2.3.1
-            axolotl_extras:
+            axolotl_extras: mamba-ssm
             is_latest: true
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.4.0
+            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -75,27 +69,22 @@ jobs:
     strategy:
       matrix:
         include:
-          - cuda: 118
+          - cuda: 121
-            cuda_version: 11.8.0
+            cuda_version: 12.1.1
             python_version: "3.10"
-            pytorch: 2.1.2
+            pytorch: 2.3.1
             axolotl_extras:
           - cuda: 121
-            cuda_version: 12.1.0
+            cuda_version: 12.1.1
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
             python_version: "3.11"
             pytorch: 2.3.1
             axolotl_extras:
             is_latest: true
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.4.0
+            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -134,7 +123,7 @@ jobs:
       matrix:
         include:
           - cuda: 121
-            cuda_version: 12.1.0
+            cuda_version: 12.1.1
             python_version: "3.11"
             pytorch: 2.3.1
             axolotl_extras:

.github/workflows/nightlies.yml

@@ -12,28 +12,22 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - cuda: 118
+          - cuda: 121
-            cuda_version: 11.8.0
+            cuda_version: 12.1.1
             python_version: "3.10"
-            pytorch: 2.1.2
+            pytorch: 2.3.1
-            axolotl_extras:
-            axolotl_args: "--extra-index-url https://download.pytorch.org/whl/cu118"
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.10"
-            pytorch: 2.1.2
             axolotl_extras:
           - cuda: 121
-            cuda_version: 12.1.0
+            cuda_version: 12.1.1
-            python_version: "3.11"
-            pytorch: 2.2.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
             python_version: "3.11"
             pytorch: 2.3.1
             axolotl_extras:
             is_latest: true
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.4.0
+            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -75,27 +69,22 @@ jobs:
     strategy:
       matrix:
         include:
-          - cuda: 118
+          - cuda: 121
-            cuda_version: 11.8.0
+            cuda_version: 12.1.1
             python_version: "3.10"
-            pytorch: 2.1.2
+            pytorch: 2.3.1
             axolotl_extras:
           - cuda: 121
-            cuda_version: 12.1.0
+            cuda_version: 12.1.1
-            python_version: "3.10"
-            pytorch: 2.1.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
-            python_version: "3.11"
-            pytorch: 2.2.2
-            axolotl_extras:
-          - cuda: 121
-            cuda_version: 12.1.0
             python_version: "3.11"
             pytorch: 2.3.1
             axolotl_extras:
             is_latest: true
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.4.0
+            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout

.github/workflows/tests.yml

@@ -72,27 +72,24 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - cuda: 118
+          - cuda: 121
-            cuda_version: 11.8.0
+            cuda_version: 12.1.1
             python_version: "3.10"
-            pytorch: 2.1.2
+            pytorch: 2.3.1
-            axolotl_args: "--extra-index-url https://download.pytorch.org/whl/cu118"
             num_gpus: 1
+            axolotl_extras: mamba-ssm
           - cuda: 121
-            cuda_version: 12.1.0
+            cuda_version: 12.1.1
-            python_version: "3.10"
-            pytorch: 2.1.2
-            num_gpus: 1
-          - cuda: 121
-            cuda_version: 12.1.0
-            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.1
             num_gpus: 1
+            axolotl_extras: mamba-ssm
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.4.0
+            num_gpus: 1
+            axolotl_extras:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -109,6 +106,7 @@ jobs:
           echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
           echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV
           echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV
+          echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
       - name: Run tests job on Modal

cicd/Dockerfile.jinja

@@ -24,9 +24,9 @@ RUN git fetch origin +$GITHUB_REF && \
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN pip install causal_conv1d
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,optimizers,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,optimizers] $AXOLOTL_ARGS; \
     fi
 
 # So we can test the Docker image

docker/Dockerfile

@@ -22,9 +22,9 @@ WORKDIR /workspace/axolotl
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN pip install causal_conv1d
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,optimizers,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,optimizers] $AXOLOTL_ARGS; \
     fi
 
 # So we can test the Docker image