From b27dab28acf9ab8a35d2526ee4edeefdc2e8aabd Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 26 Aug 2025 10:06:42 -0400 Subject: [PATCH 1/6] Add support for PyTorch 2.8.0 --- .github/workflows/main.yml | 16 ++++++++++++++++ .github/workflows/tests.yml | 10 ++++++++-- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 3daf39e43f..3f98dd2b42 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -36,6 +36,11 @@ jobs: python_version: "3.11" pytorch: 2.7.1 axolotl_extras: + - cuda: 128 + cuda_version: 12.8.1 + python_version: "3.11" + pytorch: 2.8.0 + axolotl_extras: runs-on: axolotl-gpu-runner steps: - name: Checkout @@ -110,6 +115,11 @@ jobs: python_version: "3.11" pytorch: 2.7.1 axolotl_extras: + - cuda: 128 + cuda_version: 12.8.1 + python_version: "3.11" + pytorch: 2.8.0 + axolotl_extras: runs-on: axolotl-gpu-runner steps: - name: Checkout @@ -169,6 +179,12 @@ jobs: pytorch: 2.7.1 axolotl_extras: vllm is_latest: true + - cuda: 128 + cuda_version: 12.8.1 + python_version: "3.11" + pytorch: 2.8.0 + axolotl_extras: + is_latest: runs-on: axolotl-gpu-runner steps: - name: Checkout diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index fe63aa313a..f2151fd38b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -55,7 +55,7 @@ jobs: fail-fast: false matrix: python_version: ["3.11"] - pytorch_version: ["2.6.0", "2.7.0", "2.7.1"] + pytorch_version: ["2.6.0", "2.7.0", "2.7.1", "2.8.0"] timeout-minutes: 20 steps: @@ -130,7 +130,7 @@ jobs: fail-fast: false matrix: python_version: ["3.11"] - pytorch_version: ["2.6.0", "2.7.0", "2.7.1"] + pytorch_version: ["2.6.0", "2.7.0", "2.7.1", "2.8.0"] timeout-minutes: 20 steps: @@ -298,6 +298,12 @@ jobs: pytorch: 2.7.1 num_gpus: 1 axolotl_extras: + - cuda: 128 + cuda_version: 12.8.1 + python_version: "3.11" + pytorch: 2.8.0 + num_gpus: 1 + axolotl_extras: steps: - name: Checkout uses: actions/checkout@v4 From fb67a67cec506737e0d7b6a1bc9ecdebdd53e9e8 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 26 Aug 2025 10:15:20 -0400 Subject: [PATCH 2/6] loosen triton requirements --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index c51c9d1fe9..5a31fb489f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ # START section of dependencies that don't install on Darwin/MacOS bitsandbytes==0.47.0 # triton 3.4.0 is not compatible with CCE -triton>=3.0.0,<3.4.0 +triton>=3.0.0 mamba-ssm==1.2.0.post1 xformers>=0.0.23.post1 autoawq==0.2.7.post3 From 357709eaf855ce5a425c3858038db46a9c644db7 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 26 Aug 2025 10:42:28 -0400 Subject: [PATCH 3/6] handle torch 2.8.0 in setup.py --- setup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 5aab9d7c08..5bf9ae8403 100644 --- a/setup.py +++ b/setup.py @@ -64,7 +64,9 @@ def parse_requirements(extras_require_map): else: raise ValueError("Invalid version format") - if (major, minor) >= (2, 7): + if (major, minor) >= (2, 8): + pass + elif (major, minor) >= (2, 7): _install_requires.pop(_install_requires.index(xformers_version)) if patch == 0: _install_requires.append("xformers==0.0.30") From 9b807db6edc70158df136b4fd747ec9beedaf917 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 27 Aug 2025 13:10:05 -0400 Subject: [PATCH 4/6] fix versions --- .github/workflows/multi-gpu-e2e.yml | 10 +++++----- .github/workflows/tests.yml | 12 ++++++------ examples/colab-notebooks/colab-axolotl-example.ipynb | 2 +- scripts/cutcrossentropy_install.py | 2 +- src/axolotl/integrations/cut_cross_entropy/README.md | 2 +- .../integrations/cut_cross_entropy/__init__.py | 2 +- 6 files changed, 15 insertions(+), 15 deletions(-) diff --git a/.github/workflows/multi-gpu-e2e.yml b/.github/workflows/multi-gpu-e2e.yml index 3085261511..12a0b06f15 100644 --- a/.github/workflows/multi-gpu-e2e.yml +++ b/.github/workflows/multi-gpu-e2e.yml @@ -36,14 +36,14 @@ jobs: - cuda: 126 cuda_version: 12.6.3 python_version: "3.11" - pytorch: 2.7.0 - axolotl_extras: + pytorch: 2.7.1 + axolotl_extras: vllm num_gpus: 2 nightly_build: "true" - - cuda: 126 - cuda_version: 12.6.3 + - cuda: 128 + cuda_version: 12.8.1 python_version: "3.11" - pytorch: 2.7.1 + pytorch: 2.8.0 axolotl_extras: vllm num_gpus: 2 nightly_build: "true" diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f2151fd38b..59011ee77b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -55,7 +55,7 @@ jobs: fail-fast: false matrix: python_version: ["3.11"] - pytorch_version: ["2.6.0", "2.7.0", "2.7.1", "2.8.0"] + pytorch_version: ["2.6.0", "2.7.1", "2.8.0"] timeout-minutes: 20 steps: @@ -130,7 +130,7 @@ jobs: fail-fast: false matrix: python_version: ["3.11"] - pytorch_version: ["2.6.0", "2.7.0", "2.7.1", "2.8.0"] + pytorch_version: ["2.6.0", "2.7.1", "2.8.0"] timeout-minutes: 20 steps: @@ -240,7 +240,7 @@ jobs: - cuda: 126 cuda_version: 12.6.3 python_version: "3.11" - pytorch: 2.6.0 + pytorch: 2.7.1 num_gpus: 1 axolotl_extras: dockerfile: "Dockerfile-uv.jinja" @@ -340,10 +340,10 @@ jobs: fail-fast: false matrix: include: - - cuda: 124 - cuda_version: 12.4.1 + - cuda: 126 + cuda_version: 12.6.3 python_version: "3.11" - pytorch: 2.6.0 + pytorch: 2.7.1 num_gpus: 1 axolotl_extras: steps: diff --git a/examples/colab-notebooks/colab-axolotl-example.ipynb b/examples/colab-notebooks/colab-axolotl-example.ipynb index 30ef1c3de0..b780a1c48d 100644 --- a/examples/colab-notebooks/colab-axolotl-example.ipynb +++ b/examples/colab-notebooks/colab-axolotl-example.ipynb @@ -40,7 +40,7 @@ "%%capture\n", "# This step can take ~5-10 minutes to install dependencies\n", "!pip install --no-build-isolation axolotl[flash-attn]>=0.9.1\n", - "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0ee9ee8\"" + "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5\"" ] }, { diff --git a/scripts/cutcrossentropy_install.py b/scripts/cutcrossentropy_install.py index b2bb0fcf83..5b49e7427f 100644 --- a/scripts/cutcrossentropy_install.py +++ b/scripts/cutcrossentropy_install.py @@ -29,5 +29,5 @@ print( UNINSTALL_PREFIX - + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0ee9ee8"' + + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5"' ) diff --git a/src/axolotl/integrations/cut_cross_entropy/README.md b/src/axolotl/integrations/cut_cross_entropy/README.md index 02e4e6686c..a64bdd0548 100644 --- a/src/axolotl/integrations/cut_cross_entropy/README.md +++ b/src/axolotl/integrations/cut_cross_entropy/README.md @@ -19,7 +19,7 @@ python scripts/cutcrossentropy_install.py | sh - If you are installing from pip ```bash -pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0ee9ee8" +pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5" ``` ## Usage diff --git a/src/axolotl/integrations/cut_cross_entropy/__init__.py b/src/axolotl/integrations/cut_cross_entropy/__init__.py index 6dd7c97e1f..d0eb1ebdb6 100644 --- a/src/axolotl/integrations/cut_cross_entropy/__init__.py +++ b/src/axolotl/integrations/cut_cross_entropy/__init__.py @@ -35,7 +35,7 @@ _CCE_INSTALL_MESSAGE = ( "Please install Axolotl's fork of cut_cross_entropy with transformers support using " - '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0ee9ee8"`' + '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5"`' ) From 22c3eb42f4fa1b8f59492733a9a2d81e07bfb07a Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 27 Aug 2025 13:30:08 -0400 Subject: [PATCH 5/6] no vllm for torch 2.8.0 --- .github/workflows/multi-gpu-e2e.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/multi-gpu-e2e.yml b/.github/workflows/multi-gpu-e2e.yml index 12a0b06f15..6492e5d3ec 100644 --- a/.github/workflows/multi-gpu-e2e.yml +++ b/.github/workflows/multi-gpu-e2e.yml @@ -44,7 +44,7 @@ jobs: cuda_version: 12.8.1 python_version: "3.11" pytorch: 2.8.0 - axolotl_extras: vllm + axolotl_extras: num_gpus: 2 nightly_build: "true" runs-on: [self-hosted, modal] From cd71594356b32f86cdc1cb09d199cf62b4d352f6 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 28 Aug 2025 09:10:09 -0400 Subject: [PATCH 6/6] remove comment [skip ci] Co-authored-by: NanoCode012 --- requirements.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 5a31fb489f..5accd13ed6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,6 @@ # START section of dependencies that don't install on Darwin/MacOS bitsandbytes==0.47.0 -# triton 3.4.0 is not compatible with CCE triton>=3.0.0 mamba-ssm==1.2.0.post1 xformers>=0.0.23.post1