[docker] experiment with nvcr pytorch image for torch 2.5.1

.github/workflows/base.yml

@@ -22,36 +22,38 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - cuda: "121"
-            cuda_version: 12.1.1
-            cudnn_version: 8
-            python_version: "3.10"
-            pytorch: 2.3.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "121"
-            cuda_version: 12.1.1
-            cudnn_version: 8
-            python_version: "3.11"
-            pytorch: 2.3.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "124"
-            cuda_version: 12.4.1
-            cudnn_version: ""
-            python_version: "3.10"
-            pytorch: 2.4.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "124"
-            cuda_version: 12.4.1
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: 2.4.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#          - cuda: "121"
+#            cuda_version: 12.1.1
+#            cudnn_version: 8
+#            python_version: "3.10"
+#            pytorch: 2.3.1
+#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#            from_base_img: ""
+#            from_base_tag: ""
+#          - cuda: "121"
+#            cuda_version: 12.1.1
+#            cudnn_version: 8
+#            python_version: "3.11"
+#            pytorch: 2.3.1
+#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#            from_base_img: ""
+#            from_base_tag: ""
+#          - cuda: "124"
+#            cuda_version: 12.4.1
+#            cudnn_version: ""
+#            python_version: "3.11"
+#            pytorch: 2.4.1
+#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+#            from_base_img: ""
+#            from_base_tag: ""
           - cuda: "124"
             cuda_version: 12.4.1
             cudnn_version: ""
             python_version: "3.11"
             pytorch: 2.5.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+            from_base_img: nvcr.io/nvidia/pytorch
+            from_base_tag: 24.10-py3
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -61,7 +63,7 @@ jobs:
         with:
           images: |
             winglian/axolotl-base
-            axolotlai/axolotl-base
+#            axolotlai/axolotl-base
       - name: Login to Docker Hub
         uses: docker/login-action@v2
         with:
@@ -74,7 +76,8 @@ jobs:
         with:
           context: .
           file: ./docker/Dockerfile-base
-          push: ${{ github.event_name != 'pull_request' }}
+          push: true
+#          push: ${{ github.event_name != 'pull_request' }}
           tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
           build-args: |
@@ -84,3 +87,5 @@ jobs:
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
             TORCH_CUDA_ARCH_LIST=${{ matrix.torch_cuda_arch_list }}
+            BASE_IMAGE=${{ matrix.from_base_img || '' }}
+            BASE_TAG=${{ matrix.from_base_tag || '' }}

.github/workflows/tests.yml

@@ -148,63 +148,64 @@ jobs:
         run: |
           find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
 
-  docker-e2e-tests-1st:
-    if: ${{ ! contains(github.event.commits[0].message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' }}
-    # this job needs to be run on self-hosted GPU runners...
-    runs-on: [self-hosted, modal]
-    timeout-minutes: 90
-    needs: [pre-commit, pytest, pytest-sdist]
-
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.4.1
-            num_gpus: 1
-            axolotl_extras:
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-      - name: Install Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-      - name: Install Modal
-        run: |
-          python -m pip install --upgrade pip
-          pip install modal==0.63.64 jinja2
-      - name: Update env vars
-        run: |
-          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
-        run: |
-          modal run cicd.tests
+#  docker-e2e-tests-1st:
+#    if: ${{ ! contains(github.event.commits[0].message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' }}
+#    # this job needs to be run on self-hosted GPU runners...
+#    runs-on: [self-hosted, modal]
+#    timeout-minutes: 90
+#    needs: [pre-commit, pytest, pytest-sdist]
+#
+#    strategy:
+#      fail-fast: false
+#      matrix:
+#        include:
+#          - cuda: 124
+#            cuda_version: 12.4.1
+#            python_version: "3.11"
+#            pytorch: 2.4.1
+#            num_gpus: 1
+#            axolotl_extras:
+#    steps:
+#      - name: Checkout
+#        uses: actions/checkout@v4
+#      - name: Install Python
+#        uses: actions/setup-python@v5
+#        with:
+#          python-version: "3.10"
+#      - name: Install Modal
+#        run: |
+#          python -m pip install --upgrade pip
+#          pip install modal==0.63.64 jinja2
+#      - name: Update env vars
+#        run: |
+#          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
+#        run: |
+#          modal run cicd.tests
 
   docker-e2e-tests:
     if: github.repository_owner == 'axolotl-ai-cloud'
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
     timeout-minutes: 90
-    needs: [pre-commit, pytest, docker-e2e-tests-1st]
+#    needs: [pre-commit, pytest, docker-e2e-tests-1st]
+    needs: [pre-commit, pytest]
 
     strategy:
       fail-fast: false
       matrix:
         include:
-          - cuda: 121
-            cuda_version: 12.1.1
-            python_version: "3.10"
-            pytorch: 2.3.1
-            num_gpus: 1
-            axolotl_extras: mamba-ssm
+#          - cuda: 121
+#            cuda_version: 12.1.1
+#            python_version: "3.10"
+#            pytorch: 2.3.1
+#            num_gpus: 1
+#            axolotl_extras: mamba-ssm
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"
@@ -224,7 +225,7 @@ jobs:
           pip install modal==0.63.64 jinja2
       - name: Update env vars
         run: |
-          echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
+          echo "BASE_TAG=pr-2139-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

cicd/Dockerfile.jinja

@@ -1,4 +1,4 @@
-FROM axolotlai/axolotl-base:{{ BASE_TAG }}
+FROM winglian/axolotl-base:{{ BASE_TAG }}
 
 ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
 ENV AXOLOTL_EXTRAS="{{ AXOLOTL_EXTRAS }}"

docker/Dockerfile

@@ -1,5 +1,6 @@
+ARG BASE_IMAGE=axolotlai/axolotl-base
 ARG BASE_TAG=main-base
-FROM axolotlai/axolotl-base:$BASE_TAG
+FROM $BASE_IMAGE:$BASE_TAG
 
 ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
 ARG AXOLOTL_EXTRAS=""

docker/Dockerfile-base

@@ -3,7 +3,10 @@ ARG CUDNN_VERSION="8"
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
 
-FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
+ARG BASE_IMAGE=nvidia/cuda
+ARG DEFAULT_TAG=${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel-ubuntu${UBUNTU_VERSION}
+ARG BASE_TAG=""
+FROM ${BASE_IMAGE:-nvidia/cuda}:${BASE_TAG:-${DEFAULT_TAG}} AS base-builder
 
 ENV PATH="/root/miniconda3/bin:${PATH}"