From 6afba3871d3f6748372afc6289722ce82c5c00c7 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 28 Aug 2025 09:10:40 -0400 Subject: [PATCH] Add support for PyTorch 2.8.0 (#3106) * Add support for PyTorch 2.8.0 * loosen triton requirements * handle torch 2.8.0 in setup.py * fix versions * no vllm for torch 2.8.0 * remove comment Co-authored-by: NanoCode012 --------- Co-authored-by: NanoCode012 --- .github/workflows/main.yml | 16 ++++++++++++++++ .github/workflows/multi-gpu-e2e.yml | 14 +++++++------- .github/workflows/tests.yml | 18 ++++++++++++------ .../colab-axolotl-example.ipynb | 2 +- requirements.txt | 3 +-- scripts/cutcrossentropy_install.py | 2 +- setup.py | 4 +++- .../integrations/cut_cross_entropy/README.md | 2 +- .../integrations/cut_cross_entropy/__init__.py | 2 +- 9 files changed, 43 insertions(+), 20 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 3daf39e43..3f98dd2b4 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/multi-gpu-e2e.yml b/.github/workflows/multi-gpu-e2e.yml index 308526151..6492e5d3e 100644 --- a/.github/workflows/multi-gpu-e2e.yml +++ b/.github/workflows/multi-gpu-e2e.yml @@ -33,13 +33,6 @@ jobs: axolotl_extras: num_gpus: 2 nightly_build: "true" - - cuda: 126 - cuda_version: 12.6.3 - python_version: "3.11" - pytorch: 2.7.0 - axolotl_extras: - num_gpus: 2 - nightly_build: "true" - cuda: 126 cuda_version: 12.6.3 python_version: "3.11" @@ -47,6 +40,13 @@ jobs: axolotl_extras: vllm num_gpus: 2 nightly_build: "true" + - cuda: 128 + cuda_version: 12.8.1 + python_version: "3.11" + pytorch: 2.8.0 + axolotl_extras: + num_gpus: 2 + nightly_build: "true" runs-on: [self-hosted, modal] timeout-minutes: 120 steps: diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index fe63aa313..59011ee77 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.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.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" @@ -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 @@ -334,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 30ef1c3de..b780a1c48 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/requirements.txt b/requirements.txt index c51c9d1fe..5accd13ed 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,8 +2,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 diff --git a/scripts/cutcrossentropy_install.py b/scripts/cutcrossentropy_install.py index b2bb0fcf8..5b49e7427 100644 --- a/scripts/cutcrossentropy_install.py +++ b/scripts/cutcrossentropy_install.py @@ -29,5 +29,5 @@ UV_PREFIX = "uv " if USE_UV else "" 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/setup.py b/setup.py index 5aab9d7c0..5bf9ae840 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") diff --git a/src/axolotl/integrations/cut_cross_entropy/README.md b/src/axolotl/integrations/cut_cross_entropy/README.md index 02e4e6686..a64bdd054 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 6dd7c97e1..d0eb1ebdb 100644 --- a/src/axolotl/integrations/cut_cross_entropy/__init__.py +++ b/src/axolotl/integrations/cut_cross_entropy/__init__.py @@ -35,7 +35,7 @@ LOG = get_logger(__name__) _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"`' )