wip ao upgrade

2026-01-05 18:23:33 +00:00
316 changed files with 1203 additions and 34096 deletions
--- a/.github/CONTRIBUTING.md
+++ b/.github/CONTRIBUTING.md
@@ -68,12 +68,7 @@ You can skip certain CI checks by including specific keywords in your commit mes

 ### Code Style

-axolotl uses [Ruff](https://docs.astral.sh/ruff/) as its code style guide. Please ensure that your code follows these guidelines.
-
-Use the pre-commit linter to ensure that your code is formatted consistently.
-```bash
-pre-commit run --all-files
-```
+axolotl uses [{codestyle}]({URLofCodestyle}) as its code style guide. Please ensure that your code follows these guidelines.

 ### Commit Messages

@@ -83,6 +78,6 @@ Write clear and concise commit messages that briefly describe the changes made i

 - [GitHub Help](https://help.github.com/)
 - [GitHub Pull Request Documentation](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests)
- [Ruff](https://docs.astral.sh/ruff/)
+- [{codestyle}]({URLofCodestyle})

 Thank you once again for your interest in contributing to axolotl. We look forward to collaborating with you and creating an even better project together!
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -15,11 +15,6 @@
 <!--- Include details of your testing environment, tests ran to see how -->
 <!--- your change affects other areas of the code, etc. -->

-## AI Usage Disclaimer
-
-<!--- Was AI (e.g., ChatGPT, Claude, Copilot) used to generate or assist with this PR? -->
-<!--- Please indicate: No / Yes (specify which tool and to what extent) -->
-
 ## Screenshots (if appropriate)

 ## Types of changes
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -15,17 +15,12 @@ on:
      - '.github/workflows/base.yml'
  workflow_dispatch:

-permissions:
-  contents: read
-
 jobs:
  build-base:
    if: ${{ github.repository_owner == 'axolotl-ai-cloud' && (github.event_name != 'pull_request' || !github.event.pull_request.draft) }}
    timeout-minutes: 480
    # this job needs to be run on self-hosted GPU runners...
    runs-on: ubuntu-latest-m
-    env:
-      HAS_DOCKERHUB_CREDS: ${{ secrets.DOCKERHUB_USERNAME != '' && secrets.DOCKERHUB_TOKEN != '' }}
    strategy:
      fail-fast: false
      matrix:
@@ -37,7 +32,6 @@ jobs:
            pytorch: 2.8.0
            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64"
          - cuda: "128"
            cuda_version: 12.8.1
            cudnn_version: ""
@@ -45,7 +39,6 @@ jobs:
            pytorch: 2.9.0
            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
          - cuda: "128"
            cuda_version: 12.8.1
            cudnn_version: ""
@@ -53,31 +46,6 @@ jobs:
            pytorch: 2.9.1
            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "128"
-            cuda_version: 12.8.1
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: 2.10.0
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "128"
-            cuda_version: 12.8.1
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.10.0
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
-#          - cuda: "129"
-#            cuda_version: 12.9.1
-#            cudnn_version: ""
-#            python_version: "3.12"
-#            pytorch: 2.9.1
-#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-#            dockerfile: "Dockerfile-base"
-#            platforms: "linux/amd64,linux/arm64"
          - cuda: "130"
            cuda_version: 13.0.0
            cudnn_version: ""
@@ -85,23 +53,6 @@ jobs:
            pytorch: 2.9.1
            torch_cuda_arch_list: "9.0+PTX"
            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "130"
-            cuda_version: 13.0.0
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.9.1
-            torch_cuda_arch_list: "9.0+PTX"
-            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "130"
-            cuda_version: 13.0.0
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.10.0
-            torch_cuda_arch_list: "9.0+PTX"
-            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
 #          - cuda: "128"
 #            cuda_version: 12.8.1
 #            cudnn_version: ""
@@ -127,19 +78,17 @@ jobs:
          images: |
            axolotlai/axolotl-base
      - name: Login to Docker Hub
-        uses: docker/login-action@v3
-        if: ${{ github.event_name != 'pull_request' && env.HAS_DOCKERHUB_CREDS == 'true' }}
+        uses: docker/login-action@v2
        with:
          username: ${{ secrets.DOCKERHUB_USERNAME }}
          password: ${{ secrets.DOCKERHUB_TOKEN }}
      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3
      - name: Build
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@v4
        with:
          context: .
          file: ./docker/${{ matrix.dockerfile }}
-          platforms: ${{ matrix.platforms }}
          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 }}
@@ -154,8 +103,6 @@ jobs:
    if: ${{ github.repository_owner == 'axolotl-ai-cloud' && (github.event_name != 'pull_request' || !github.event.pull_request.draft) }}
    timeout-minutes: 480
    runs-on: ubuntu-latest-m
-    env:
-      HAS_DOCKERHUB_CREDS: ${{ secrets.DOCKERHUB_USERNAME != '' && secrets.DOCKERHUB_TOKEN != '' }}
    strategy:
      fail-fast: false
      matrix:
@@ -167,7 +114,6 @@ jobs:
            pytorch: 2.8.0
            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64"
          - cuda: "128"
            cuda_version: 12.8.1
            cudnn_version: ""
@@ -175,15 +121,6 @@ jobs:
            pytorch: 2.9.1
            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "128"
-            cuda_version: 12.8.1
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.9.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
          - cuda: "128"
            cuda_version: 12.8.1
            cudnn_version: ""
@@ -191,31 +128,6 @@ jobs:
            pytorch: 2.9.0
            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "128"
-            cuda_version: 12.8.1
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: 2.10.0
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "128"
-            cuda_version: 12.8.1
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.10.0
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-#          - cuda: "129"
-#            cuda_version: 12.9.1
-#            cudnn_version: ""
-#            python_version: "3.12"
-#            pytorch: 2.9.1
-#            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-#            dockerfile: "Dockerfile-uv-base"
-#            platforms: "linux/amd64,linux/arm64"
          - cuda: "130"
            cuda_version: 13.0.0
            cudnn_version: ""
@@ -223,23 +135,6 @@ jobs:
            pytorch: 2.9.1
            torch_cuda_arch_list: "9.0+PTX"
            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "130"
-            cuda_version: 13.0.0
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.9.1
-            torch_cuda_arch_list: "9.0+PTX"
-            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "130"
-            cuda_version: 13.0.0
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.10.0
-            torch_cuda_arch_list: "9.0+PTX"
-            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
    steps:
      - name: Checkout
        uses: actions/checkout@v4
@@ -250,19 +145,17 @@ jobs:
          images: |
            axolotlai/axolotl-base-uv
      - name: Login to Docker Hub
-        uses: docker/login-action@v3
-        if: ${{ github.event_name != 'pull_request' && env.HAS_DOCKERHUB_CREDS == 'true' }}
+        uses: docker/login-action@v2
        with:
          username: ${{ secrets.DOCKERHUB_USERNAME }}
          password: ${{ secrets.DOCKERHUB_TOKEN }}
      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3
      - name: Build
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@v4
        with:
          context: .
          file: ./docker/${{ matrix.dockerfile }}
-          platforms: ${{ matrix.platforms }}
          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 }}
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -13,9 +13,6 @@ on:
       - ".pre-commit-config.yaml"
  workflow_dispatch:

-permissions:
-  contents: read
-
 jobs:
  pre-commit:
    name: pre-commit
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -8,9 +8,6 @@ on:
      - "v*"
  workflow_dispatch:

-permissions:
-  contents: read
-
 jobs:
  build-axolotl:
    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
@@ -23,44 +20,22 @@ jobs:
            python_version: "3.11"
            pytorch: 2.8.0
            axolotl_extras:
-            platforms: "linux/amd64"
+            is_latest: true
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
            pytorch: 2.9.0
            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
            pytorch: 2.9.1
            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-            is_latest: true
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-#          - cuda: 129
-#            cuda_version: 12.9.1
-#            python_version: "3.12"
+#          - cuda: 130
+#            cuda_version: 13.0.0
+#            python_version: "3.11"
 #            pytorch: 2.9.1
 #            axolotl_extras:
-#            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.11"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
    runs-on: axolotl-gpu-runner
    steps:
      - name: Checkout
@@ -86,7 +61,6 @@ jobs:
        uses: docker/build-push-action@v5
        with:
          context: .
-          platforms: ${{ matrix.platforms }}
          build-args: |
            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
            CUDA=${{ matrix.cuda }}
@@ -101,89 +75,11 @@ jobs:
            ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
          labels: ${{ steps.metadata.outputs.labels }}

-  build-axolotl-uv:
-    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.11"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.12"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-            is_latest: true
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.11"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-    runs-on: axolotl-gpu-runner
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-      - name: Docker metadata
-        id: metadata
-        uses: docker/metadata-action@v5
-        with:
-          images: |
-            axolotlai/axolotl-uv
-          tags: |
-            type=ref,event=branch
-            type=pep440,pattern={{version}}
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-      - name: Login to Docker Hub
-        uses: docker/login-action@v3
-        with:
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-      # guidance for testing before pushing: https://docs.docker.com/build/ci/github-actions/test-before-push/
-      - name: Build and export to Docker
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          platforms: ${{ matrix.platforms }}
-          build-args: |
-            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
-            CUDA=${{ matrix.cuda }}
-            PYTORCH_VERSION=${{ matrix.pytorch }}
-            AXOLOTL_ARGS=${{ matrix.axolotl_args }}
-            AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}
-          file: ./docker/Dockerfile-uv
-          push: ${{ github.event_name != 'pull_request' }}
-          tags: |
-            ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
-            ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
-            ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
-          labels: ${{ steps.metadata.outputs.labels }}
-
  build-axolotl-cloud:
    needs: build-axolotl
    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
    # this job needs to be run on self-hosted GPU runners...
    strategy:
-      fail-fast: false
      matrix:
        include:
          - cuda: 128
@@ -191,44 +87,22 @@ jobs:
            python_version: "3.11"
            pytorch: 2.8.0
            axolotl_extras:
-            platforms: "linux/amd64"
+            is_latest: true
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
            pytorch: 2.9.0
            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
            pytorch: 2.9.1
            axolotl_extras:
-            is_latest: true
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-#          - cuda: 129
-#            cuda_version: 12.9.1
-#            python_version: "3.12"
+#          - cuda: 130
+#            cuda_version: 13.0.0
+#            python_version: "3.11"
 #            pytorch: 2.9.1
 #            axolotl_extras:
-#            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.11"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
    runs-on: axolotl-gpu-runner
    steps:
      - name: Checkout
@@ -253,7 +127,6 @@ jobs:
        uses: docker/build-push-action@v5
        with:
          context: .
-          platforms: ${{ matrix.platforms }}
          build-args: |
            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
            CUDA=${{ matrix.cuda }}
@@ -264,96 +137,21 @@ jobs:
             ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
          labels: ${{ steps.metadata.outputs.labels }}

-  build-axolotl-cloud-uv:
-    needs: build-axolotl-uv
-    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
-    # this job needs to be run on self-hosted GPU runners...
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.11"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.12"
-            pytorch: 2.9.1
-            axolotl_extras:
-            is_latest: true
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.11"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.12"
-            pytorch: 2.10.0
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-    runs-on: axolotl-gpu-runner
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-      - name: Docker metadata
-        id: metadata
-        uses: docker/metadata-action@v5
-        with:
-          images: |
-            axolotlai/axolotl-cloud-uv
-          tags: |
-            type=ref,event=branch
-            type=pep440,pattern={{version}}
-      - name: Login to Docker Hub
-        uses: docker/login-action@v3
-        with:
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-      - name: Build
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          platforms: ${{ matrix.platforms }}
-          build-args: |
-            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
-            CUDA=${{ matrix.cuda }}
-          file: ./docker/Dockerfile-cloud-uv
-          push: ${{ github.event_name != 'pull_request' }}
-          tags: |
-             ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
-             ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
-          labels: ${{ steps.metadata.outputs.labels }}
-
  build-axolotl-cloud-no-tmux:
    needs: build-axolotl
    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
    # this job needs to be run on self-hosted GPU runners...
    strategy:
-      fail-fast: false
      matrix:
        include:
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
-            pytorch: 2.9.1
+            pytorch: 2.8.0
            axolotl_extras:
-            is_latest: true
-          - cuda: 130
-            cuda_version: 13.0.0
+            is_latest:
+          - cuda: 128
+            cuda_version: 12.8.1
            python_version: "3.11"
            pytorch: 2.9.1
            axolotl_extras:
@@ -382,7 +180,6 @@ jobs:
        uses: docker/build-push-action@v5
        with:
          context: .
-          platforms: linux/amd64,linux/arm64
          build-args: |
            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
            CUDA=${{ matrix.cuda }}
--- a/.github/workflows/multi-gpu-e2e.yml
+++ b/.github/workflows/multi-gpu-e2e.yml
@@ -8,7 +8,6 @@ on:
      - 'setup.py'
      - 'pyproject.toml'
      - '.github/workflows/multi-gpu-e2e.yml'
-      - 'scripts/cutcrossentropy_install.py'
      - 'src/axolotl/core/trainers/mixins/sequence_parallel.py'
      - 'src/axolotl/utils/distributed.py'
  workflow_dispatch:
@@ -20,9 +19,6 @@ concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}

-permissions:
-  contents: read
-
 env:
  MODAL_IMAGE_BUILDER_VERSION: "2025.06"

@@ -39,27 +35,14 @@ jobs:
            pytorch: 2.8.0
            axolotl_extras: fbgemm-gpu
            num_gpus: 2
-#          - cuda: 129
-#            cuda_version: 12.9.1
-#            python_version: "3.12"
-#            pytorch: 2.9.1
-#            axolotl_extras: "fbgemm-gpu"
-#            num_gpus: 2
-#            dockerfile: "Dockerfile-uv.jinja"
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.11"
-            pytorch: 2.9.1
-            axolotl_extras:
-#            axolotl_extras: fbgemm-gpu
-            num_gpus: 2
+            nightly_build: "true"
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
-            pytorch: 2.10.0
-            axolotl_extras: "fbgemm-gpu"
+            pytorch: 2.9.1
+            axolotl_extras: fbgemm-gpu
            num_gpus: 2
-            dockerfile: "Dockerfile-uv.jinja"
+            nightly_build: "true"
    runs-on: [self-hosted, modal]
    timeout-minutes: 120
    steps:
@@ -81,9 +64,8 @@ jobs:
          echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
-          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
+          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
+          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
      - name: Run tests job on Modal
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
        run: |
          modal run -m cicd.multigpu
--- a/.github/workflows/nightlies.yml
+++ b/.github/workflows/nightlies.yml
@@ -5,9 +5,6 @@ on:
  schedule:
    - cron: '0 0 * * *'  # Runs at 00:00 UTC every day

-permissions:
-  contents: read
-
 jobs:
  build-axolotl:
    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
--- a/.github/workflows/precommit-autoupdate.yml
+++ b/.github/workflows/precommit-autoupdate.yml
@@ -5,8 +5,6 @@ on:
    - cron: '0 0 1 * *'  # Run monthly
  workflow_dispatch:  # Manual kickoff

-permissions: {}
-
 jobs:
  auto-update:
    runs-on: ubuntu-latest
--- a/.github/workflows/preview-docs.yml
+++ b/.github/workflows/preview-docs.yml
@@ -14,8 +14,14 @@ on:
      - .github/workflows/preview-docs.yml

 permissions:
-  contents: read
+  checks: write
+  contents: write
+  deployments: write
+  issues: write
+  discussions: write
+  pages: write
  pull-requests: write
+  statuses: write

 jobs:
  preview:
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -3,11 +3,9 @@ name: publish pypi
 on:
  push:
    tags:
-      - "v*"
+      - 'v*'
  workflow_dispatch:

-permissions: {}
-
 jobs:
  setup_release:
    name: Create Release
@@ -30,8 +28,7 @@ jobs:
      name: pypi
      url: https://pypi.org/p/axolotl
    permissions:
-      contents: read
-      id-token: write # IMPORTANT: this permission is mandatory for trusted publishing
+      id-token: write  # IMPORTANT: this permission is mandatory for trusted publishing
    steps:
      - name: Check out repository code
        uses: actions/checkout@v4
@@ -43,17 +40,17 @@ jobs:

      - name: Install dependencies
        run: |
-          pip3 install wheel packaging==26.0
+          pip3 install wheel packaging==23.2
          pip3 install --no-build-isolation -e .
          pip3 install -r requirements-dev.txt -r requirements-tests.txt

      - name: Extract tag name
        id: tag
-        run: echo "TAG_NAME=$(echo $GITHUB_REF | cut -d / -f 3)" >> "$GITHUB_OUTPUT"
+        run: echo ::set-output name=TAG_NAME::$(echo $GITHUB_REF | cut -d / -f 3)

-      - name: Update version in VERSION file
+      - name: Update version in setup.py
        run: |
-          echo "${{ steps.tag.outputs.TAG_NAME }}" | sed 's/^v//' > VERSION
+          sed -i -E 's/version="([0-9.]+)",/version="${{ steps.tag.outputs.TAG_NAME }}",/g' setup.py

      - name: Build a source dist
        run: |
--- a/.github/workflows/tests-nightly.yml
+++ b/.github/workflows/tests-nightly.yml
@@ -3,13 +3,6 @@ on:
  workflow_dispatch:
  schedule:
    - cron: '0 0 * * *'  # Runs at 00:00 UTC every day
-  pull_request:
-    types: [opened, synchronize, reopened, ready_for_review]
-    paths:
-      - '.github/workflows/tests-nightly.yml'
-
-permissions:
-  contents: read

 jobs:
  pre-commit:
@@ -25,26 +18,15 @@ jobs:
        env:
          SKIP: no-commit-to-branch

-  prime-cdn-s3-cache:
-    name: Prefetch S3 once to prime the CDN cache
-    runs-on: ubuntu-latest
-    if: ${{ !github.event.pull_request.draft }}
-    timeout-minutes: 10
-    steps:
-      - name: Restore Cache from S3
-        id: hf-cache-restore-s3
-        run: |
-          curl -v -H "Range: bytes=0-1023" -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst > /dev/null
-
  pytest:
    name: PyTest
    runs-on: ubuntu-latest
-    needs: [prime-cdn-s3-cache]
    strategy:
      fail-fast: false
+      max-parallel: 2
      matrix:
-        python_version: ["3.12"]  # TODO include py3.14 once https://github.com/mistralai/mistral-common/pull/194 is merged
-        pytorch_version: ["2.8.0", "2.9.1", "2.10.0"]
+        python_version: ["3.11"]
+        pytorch_version: ["2.8.0", "2.9.0", "2.9.1"]
    timeout-minutes: 20

    steps:
@@ -55,7 +37,7 @@ jobs:
        id: hf-cache-restore-s3
        run: |
          mkdir -p /home/runner/.cache/huggingface/hub
-          curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xf - -C /home/runner/.cache/huggingface/hub/  --use-compress-program unzstd
+          curl -L https://d1dttdx32dkk5p.cloudfront.net/hf-cache.tar.zst | tar -xf - -C /home/runner/.cache/huggingface/hub/  --use-compress-program unzstd

      - name: Setup Python
        uses: actions/setup-python@v5
@@ -66,7 +48,7 @@ jobs:
      - name: upgrade pip
        run: |
          pip3 install --upgrade pip
-          pip3 install --upgrade packaging==26.0 setuptools==78.1.1 wheel
+          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 wheel

      - name: Install PyTorch
        run: |
@@ -120,23 +102,16 @@ jobs:
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
-            pytorch: 2.9.1
+            pytorch: 2.8.0
            num_gpus: 1
            axolotl_extras:
            nightly_build: "true"
          - cuda: 128
            cuda_version: 12.8.1
            python_version: "3.11"
-            pytorch: 2.10.0
-            num_gpus: 1
-            axolotl_extras:
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.12"
            pytorch: 2.9.1
            num_gpus: 1
            axolotl_extras:
-            dockerfile: "Dockerfile-uv.jinja"
            nightly_build: "true"
    steps:
      - name: Checkout
@@ -157,11 +132,9 @@ jobs:
          echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
-          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
+          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
      - name: Run tests job on Modal
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
        run: |
          modal run cicd.e2e_tests
  docker-e2e-multigpu-tests:
@@ -202,8 +175,7 @@ jobs:
          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
+          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
      - name: Run tests job on Modal
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
        run: |
          modal run cicd.multigpu
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -28,9 +28,6 @@ concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}

-permissions:
-  contents: read
-
 env:
  TRANSFORMERS_IS_CI: "yes"

@@ -49,32 +46,16 @@ jobs:
        env:
          SKIP: no-commit-to-branch

-  prime-cdn-s3-cache:
-    name: Prefetch S3 once to prime the CDN cache
-    runs-on: ubuntu-latest
-    if: ${{ !github.event.pull_request.draft }}
-    timeout-minutes: 10
-    steps:
-      - name: Restore Cache from S3
-        id: hf-cache-restore-s3
-        run: |
-          curl -v -H "Range: bytes=0-1023" -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst > /dev/null
-
  pytest:
    name: PyTest
    runs-on: ubuntu-latest
    if: ${{ !github.event.pull_request.draft }}
-    needs: [prime-cdn-s3-cache]
+#    needs: [preload-cache]
    strategy:
      fail-fast: false
      matrix:
-        python_version: ["3.12"]  # TODO include py3.14 once https://github.com/mistralai/mistral-common/pull/194 is merged
-        pytorch_version: ["2.8.0", "2.9.1", "2.10.0"]
-#        exclude:
-#          - python_version: "3.14"
-#            pytorch_version: "2.8.0"
-#          - python_version: "3.14"
-#            pytorch_version: "2.9.1"
+        python_version: ["3.11"]
+        pytorch_version: ["2.8.0", "2.9.0", "2.9.1"]
    timeout-minutes: 20

    steps:
@@ -89,7 +70,7 @@ jobs:
        id: hf-cache-restore-s3
        run: |
          mkdir -p ~/.cache/huggingface/hub
-          curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xpf - -C ~/.cache/huggingface/hub/  --use-compress-program unzstd --strip-components=1
+          curl -L https://d1dttdx32dkk5p.cloudfront.net/hf-cache.tar.zst | tar -xpf - -C ~/.cache/huggingface/hub/  --use-compress-program unzstd --strip-components=1
          ls -ltr ~/.cache/huggingface/hub/

      - name: Setup Python
@@ -101,7 +82,7 @@ jobs:
      - name: upgrade pip
        run: |
          pip3 install --upgrade pip
-          pip3 install --upgrade packaging==26.0 setuptools==75.8.0 wheel
+          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 wheel

      - name: Install PyTorch
        run: |
@@ -129,10 +110,10 @@ jobs:

      - name: Pre-Download dataset fixture
        run: |
-          hf download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
+          huggingface-cli download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures

      - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan

      - name: Run tests
        run: |
@@ -146,7 +127,7 @@ jobs:
          pytest -v --durations=10 tests/cli/ --cov=axolotl --cov-append --cov-report=xml

      - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan

      - name: Upload coverage to Codecov
        uses: codecov/codecov-action@v5
@@ -160,18 +141,12 @@ jobs:
    name: PyTest from Source Dist
    runs-on: ubuntu-latest
    if: ${{ !github.event.pull_request.draft }}
-    needs: [prime-cdn-s3-cache]
    strategy:
      fail-fast: false
      matrix:
-        python_version: ["3.12"]  # TODO include py3.14 once https://github.com/mistralai/mistral-common/pull/194 is merged
-        pytorch_version: ["2.8.0", "2.9.1", "2.10.0"]
-#        exclude:
-#          - python_version: "3.14"
-#            pytorch_version: "2.8.0"
-#          - python_version: "3.14"
-#            pytorch_version: "2.9.1"
-    timeout-minutes: 30
+        python_version: ["3.11"]
+        pytorch_version: ["2.8.0", "2.9.0", "2.9.1"]
+    timeout-minutes: 20

    steps:
      - name: cleanup node
@@ -185,7 +160,7 @@ jobs:
        id: hf-cache-restore-s3
        run: |
          mkdir -p ~/.cache/huggingface/hub
-          curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xpf - -C ~/.cache/huggingface/hub/  --use-compress-program unzstd --strip-components=1
+          curl -L https://d1dttdx32dkk5p.cloudfront.net/hf-cache.tar.zst | tar -xpf - -C ~/.cache/huggingface/hub/  --use-compress-program unzstd --strip-components=1
          ls -ltr ~/.cache/huggingface/hub/

      - name: Setup Python
@@ -197,7 +172,7 @@ jobs:
      - name: upgrade pip
        run: |
          pip3 install --upgrade pip
-          pip3 install --upgrade packaging==26.0 setuptools==75.8.0 setuptools_scm build wheel psutil
+          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 setuptools_scm build wheel psutil

      - name: Install PyTorch
        run: |
@@ -225,7 +200,7 @@ jobs:
          axolotl --help

      - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan

      - name: Run tests
        run: |
@@ -234,10 +209,10 @@ jobs:
          pytest -v --durations=10 tests/cli/

      - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan

  gate-skip-e2e:
-    needs: [pre-commit]
+    needs: [pre-commit, pytest, pytest-sdist]
    runs-on: ubuntu-latest
    outputs:
      skip: ${{ steps.compute.outputs.skip }}
@@ -273,16 +248,16 @@ jobs:
    # this job needs to be run on self-hosted GPU runners...
    runs-on: [self-hosted, modal]
    timeout-minutes: 120
-    needs: [pre-commit, pytest]
+    needs: [pre-commit, pytest, pytest-sdist, gate-skip-e2e]

    strategy:
      fail-fast: false
      matrix:
        include:
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.12"
-            pytorch: 2.9.1
+          - cuda: 128
+            cuda_version: 12.8.1
+            python_version: "3.11"
+            pytorch: 2.8.0
            num_gpus: 1
            axolotl_extras:
            dockerfile: "Dockerfile-uv.jinja"
@@ -306,10 +281,9 @@ jobs:
          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
          echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
      - name: Run tests job on Modal
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
        run: |
          modal run cicd.e2e_tests

@@ -342,18 +316,6 @@ jobs:
            pytorch: 2.9.1
            num_gpus: 1
            axolotl_extras:
-          - cuda: 128
-            cuda_version: 12.8.1
-            python_version: "3.11"
-            pytorch: 2.10.0
-            num_gpus: 1
-            axolotl_extras:
-          - cuda: 130
-            cuda_version: 13.0.0
-            python_version: "3.11"
-            pytorch: 2.9.1
-            num_gpus: 1
-            axolotl_extras:
    steps:
      - name: Checkout
        uses: actions/checkout@v4
@@ -375,10 +337,9 @@ jobs:
          echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
          echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV
+          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
      - name: Run tests job on Modal
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
        run: |
          modal run cicd.e2e_tests

@@ -418,6 +379,7 @@ jobs:
          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
          echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
      - name: Run tests job on Modal
        run: |
          modal run cicd.cleanup
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -11,7 +11,7 @@ repos:
    -   id: no-commit-to-branch
        args: ['--branch', 'main']
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.4
+    rev: v0.14.10
    hooks:
    -   id: ruff
        args: [--fix]
@@ -26,7 +26,7 @@ repos:
            'pydantic>=2.5.3',
        ]
 -   repo: https://github.com/PyCQA/bandit
-    rev: 1.9.4
+    rev: 1.9.2
    hooks:
    -   id: bandit
        args: [
--- a/.runpod/README.md
+++ b/.runpod/README.md
@@ -123,7 +123,7 @@ datasets:
 | --------------------------------- | -------------------------- | ----------------------------------- |
 | `dataset_prepared_path`           | `"data/last_run_prepared"` | Path for prepared dataset           |
 | `push_dataset_to_hub`             | `""`                       | Push dataset to HF hub              |
-| `dataset_num_proc`                | `4`                        | Number of preprocessing processes   |
+| `dataset_processes`               | `4`                        | Number of preprocessing processes   |
 | `dataset_keep_in_memory`          | `false`                    | Keep dataset in memory              |
 | `shuffle_merged_datasets`         | `true`                     | Shuffle merged datasets             |
 | `shuffle_before_merging_datasets` | `false`                    | Shuffle each dataset before merging |
--- a/.runpod/src/config/config.yaml
+++ b/.runpod/src/config/config.yaml
@@ -39,6 +39,7 @@
 #     type: # linear | dynamic
 #     factor: # float

+
 # # Whether you are training a 4-bit GPTQ quantized model
 # gptq: true
 # gptq_groupsize: 128 # group size
@@ -106,7 +107,7 @@
 # push_dataset_to_hub: # repo path
 # # The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()`
 # # if not set.
-# dataset_num_proc: # defaults to os.cpu_count() if not set
+# dataset_processes: # defaults to os.cpu_count() if not set
 # # push checkpoints to hub
 # hub_model_id: # repo path to push finetuned model
 # # how to push checkpoints to hub
@@ -223,6 +224,9 @@
 # eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
 # eval_table_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128

+# # Save model as safetensors (require safetensors package)
+# save_safetensors:
+
 # # Whether to mask out or include the human's prompt from the training labels
 # train_on_inputs: false
 # # Group similarly sized data to minimize padding.
@@ -348,6 +352,8 @@
 # # Allow overwrite yml config using from cli
 # strict:

+
+
 base_model: ${BASE_MODEL}
 base_model_ignore_patterns: ${BASE_MODEL_IGNORE_PATTERNS}
 base_model_config: ${BASE_MODEL_CONFIG}
@@ -406,7 +412,7 @@ chat_template_jinja: ${CHAT_TEMPLATE_JINJA}
 default_system_message: ${DEFAULT_SYSTEM_MESSAGE}
 dataset_prepared_path: ${DATASET_PREPARED_PATH}
 push_dataset_to_hub: ${PUSH_DATASET_TO_HUB}
-dataset_num_proc: ${DATASET_NUM_PROC}
+dataset_processes: ${DATASET_PROCESSES}
 dataset_keep_in_memory: ${DATASET_KEEP_IN_MEMORY}
 hub_model_id: ${HUB_MODEL_ID}
 hub_strategy: ${HUB_STRATEGY}
@@ -506,6 +512,7 @@ profiler_steps: ${PROFILER_STEPS}
 loss_watchdog_threshold: ${LOSS_WATCHDOG_THRESHOLD}
 loss_watchdog_patience: ${LOSS_WATCHDOG_PATIENCE}

+save_safetensors: ${SAVE_SAFETENSORS}
 train_on_inputs: ${TRAIN_ON_INPUTS}
 group_by_length: ${GROUP_BY_LENGTH}
 gradient_checkpointing: ${GRADIENT_CHECKPOINTING}
--- a/README.md
+++ b/README.md
@@ -29,23 +29,8 @@

 ## 🎉 Latest Updates

- 2026/03:
-  - New model support has been added in Axolotl for [Mistral Small 4](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/mistral4), [Qwen3.5, Qwen3.5 MoE](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/qwen3.5), [GLM-4.7-Flash](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/glm47-flash), [GLM-4.6V](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/glm46v), and [GLM-4.5-Air](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/glm45).
-  - [MoE expert quantization](https://docs.axolotl.ai/docs/expert_quantization.html) support (via `quantize_moe_experts: true`) greatly reduces VRAM when training MoE models (FSDP2 compat).
- 2026/02:
-  - [ScatterMoE LoRA](https://github.com/axolotl-ai-cloud/axolotl/pull/3410) support. LoRA fine-tuning directly on MoE expert weights using custom Triton kernels.
-  - Axolotl now has support for [SageAttention](https://github.com/axolotl-ai-cloud/axolotl/pull/2823) and [GDPO](https://github.com/axolotl-ai-cloud/axolotl/pull/3353) (Generalized DPO).
- 2026/01:
-  - New integration for [EAFT](https://github.com/axolotl-ai-cloud/axolotl/pull/3366) (Entropy-Aware Focal Training), weights loss by entropy of the top-k logit distribution, and [Scalable Softmax](https://github.com/axolotl-ai-cloud/axolotl/pull/3338), improves long context in attention.
- 2025/12:
-  - Axolotl now includes support for [Kimi-Linear](https://docs.axolotl.ai/docs/models/kimi-linear.html), [Plano-Orchestrator](https://docs.axolotl.ai/docs/models/plano.html), [MiMo](https://docs.axolotl.ai/docs/models/mimo.html), [InternVL 3.5](https://docs.axolotl.ai/docs/models/internvl3_5.html), [Olmo3](https://docs.axolotl.ai/docs/models/olmo3.html), [Trinity](https://docs.axolotl.ai/docs/models/trinity.html), and [Ministral3](https://docs.axolotl.ai/docs/models/ministral3.html).
-  - [Distributed Muon Optimizer](https://github.com/axolotl-ai-cloud/axolotl/pull/3264) support has been added for FSDP2 pretraining.
+- 2025/12: Axolotl now includes support for [Kimi-Linear](https://docs.axolotl.ai/docs/models/kimi-linear.html), [Plano-Orchestrator](https://docs.axolotl.ai/docs/models/plano.html), [MiMo](https://docs.axolotl.ai/docs/models/mimo.html), [InternVL 3.5](https://docs.axolotl.ai/docs/models/internvl3_5.html), [Olmo3](https://docs.axolotl.ai/docs/models/olmo3.html), [Trinity](https://docs.axolotl.ai/docs/models/trinity.html), and [Ministral3](https://docs.axolotl.ai/docs/models/ministral3.html).
 - 2025/10: New model support has been added in Axolotl for: [Qwen3 Next](https://docs.axolotl.ai/docs/models/qwen3-next.html), [Qwen2.5-vl, Qwen3-vl](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/qwen2_5-vl), [Qwen3, Qwen3MoE](https://docs.axolotl.ai/docs/models/qwen3.html), [Granite 4](https://docs.axolotl.ai/docs/models/granite4.html), [HunYuan](https://docs.axolotl.ai/docs/models/hunyuan.html), [Magistral 2509](https://docs.axolotl.ai/docs/models/magistral/vision.html), [Apertus](https://docs.axolotl.ai/docs/models/apertus.html), and [Seed-OSS](https://docs.axolotl.ai/docs/models/seed-oss.html).
-
-<details>
-
-<summary>Expand older updates</summary>
-
 - 2025/09: Axolotl now has text diffusion training. Read more [here](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/diffusion).
 - 2025/08: QAT has been updated to include NVFP4 support. See [PR](https://github.com/axolotl-ai-cloud/axolotl/pull/3107).
 - 2025/07:
@@ -54,10 +39,15 @@
  - FP8 finetuning with fp8 gather op is now possible in Axolotl via `torchao`. Get started [here](https://docs.axolotl.ai/docs/mixed_precision.html#sec-fp8)!
  - [Voxtral](https://docs.axolotl.ai/docs/models/voxtral.html), [Magistral 1.1](https://docs.axolotl.ai/docs/models/magistral.html), and [Devstral](https://docs.axolotl.ai/docs/models/devstral.html) with mistral-common tokenizer support has been integrated in Axolotl!
  - TiledMLP support for single-GPU to multi-GPU training with DDP, DeepSpeed and FSDP support has been added to support Arctic Long Sequence Training. (ALST). See [examples](https://github.com/axolotl-ai-cloud/axolotl/tree/main/examples/alst) for using ALST with Axolotl!
- 2025/06: Magistral with mistral-common tokenizer support has been added to Axolotl. See [docs](https://docs.axolotl.ai/docs/models/magistral.html) to start training your own Magistral models with Axolotl!
 - 2025/05: Quantization Aware Training (QAT) support has been added to Axolotl. Explore the [docs](https://docs.axolotl.ai/docs/qat.html) to learn more!
- 2025/04: Llama 4 support has been added in Axolotl. See [docs](https://docs.axolotl.ai/docs/models/llama-4.html) to start training your own Llama 4 models with Axolotl's linearized version!
+
+<details>
+
+<summary>Expand older updates</summary>
+
 - 2025/03: Axolotl has implemented Sequence Parallelism (SP) support. Read the [blog](https://huggingface.co/blog/axolotl-ai-co/long-context-with-sequence-parallelism-in-axolotl) and [docs](https://docs.axolotl.ai/docs/sequence_parallelism.html) to learn how to scale your context length when fine-tuning.
+- 2025/06: Magistral with mistral-common tokenizer support has been added to Axolotl. See [docs](https://docs.axolotl.ai/docs/models/magistral.html) to start training your own Magistral models with Axolotl!
+- 2025/04: Llama 4 support has been added in Axolotl. See [docs](https://docs.axolotl.ai/docs/models/llama-4.html) to start training your own Llama 4 models with Axolotl's linearized version!
 - 2025/03: (Beta) Fine-tuning Multimodal models is now supported in Axolotl. Check out the [docs](https://docs.axolotl.ai/docs/multimodal.html) to fine-tune your own!
 - 2025/02: Axolotl has added LoRA optimizations to reduce memory usage and improve training speed for LoRA and QLoRA in single GPU and multi-GPU training (DDP and DeepSpeed). Jump into the [docs](https://docs.axolotl.ai/docs/lora_optims.html) to give it a try.
 - 2025/02: Axolotl has added GRPO support. Dive into our [blog](https://huggingface.co/blog/axolotl-ai-co/training-llms-w-interpreter-feedback-wasm) and [GRPO example](https://github.com/axolotl-ai-cloud/grpo_code) and have some fun!
@@ -72,10 +62,10 @@ Axolotl is a free and open-source tool designed to streamline post-training and
 Features:

 - **Multiple Model Support**: Train various models like GPT-OSS, LLaMA, Mistral, Mixtral, Pythia, and many more models available on the Hugging Face Hub.
- **Multimodal Training**: Fine-tune vision-language models (VLMs) including LLaMA-Vision, Qwen2-VL, Pixtral, LLaVA, SmolVLM2, GLM-4.6V, InternVL 3.5, Gemma 3n, and audio models like Voxtral with image, video, and audio support.
- **Training Methods**: Full fine-tuning, LoRA, QLoRA, GPTQ, QAT, Preference Tuning (DPO, IPO, KTO, ORPO), RL (GRPO, GDPO), and Reward Modelling (RM) / Process Reward Modelling (PRM).
+- **Multimodal Training**: Fine-tune vision-language models (VLMs) including LLaMA-Vision, Qwen2-VL, Pixtral, LLaVA, SmolVLM2, and audio models like Voxtral with image, video, and audio support.
+- **Training Methods**: Full fine-tuning, LoRA, QLoRA, GPTQ, QAT, Preference Tuning (DPO, IPO, KTO, ORPO), RL (GRPO), and Reward Modelling (RM) / Process Reward Modelling (PRM).
 - **Easy Configuration**: Re-use a single YAML configuration file across the full fine-tuning pipeline: dataset preprocessing, training, evaluation, quantization, and inference.
- **Performance Optimizations**: [Multipacking](https://docs.axolotl.ai/docs/multipack.html), [Flash Attention 2/3/4](https://docs.axolotl.ai/docs/attention.html#flash-attention), [Xformers](https://docs.axolotl.ai/docs/attention.html#xformers), [Flex Attention](https://docs.axolotl.ai/docs/attention.html#flex-attention), [SageAttention](https://docs.axolotl.ai/docs/attention.html#sageattention), [Liger Kernel](https://docs.axolotl.ai/docs/custom_integrations.html#liger-kernels), [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy), [ScatterMoE](https://docs.axolotl.ai/docs/custom_integrations.html#kernels-integration), [Sequence Parallelism (SP)](https://docs.axolotl.ai/docs/sequence_parallelism.html), [LoRA optimizations](https://docs.axolotl.ai/docs/lora_optims.html), [Multi-GPU training (FSDP1, FSDP2, DeepSpeed)](https://docs.axolotl.ai/docs/multi-gpu.html), [Multi-node training (Torchrun, Ray)](https://docs.axolotl.ai/docs/multi-node.html), and many more!
+- **Performance Optimizations**: [Multipacking](https://docs.axolotl.ai/docs/multipack.html), [Flash Attention](https://github.com/Dao-AILab/flash-attention), [Xformers](https://github.com/facebookresearch/xformers), [Flex Attention](https://pytorch.org/blog/flexattention/), [Liger Kernel](https://github.com/linkedin/Liger-Kernel), [Cut Cross Entropy](https://github.com/apple/ml-cross-entropy/tree/main), [Sequence Parallelism (SP)](https://docs.axolotl.ai/docs/sequence_parallelism.html), [LoRA optimizations](https://docs.axolotl.ai/docs/lora_optims.html), [Multi-GPU training (FSDP1, FSDP2, DeepSpeed)](https://docs.axolotl.ai/docs/multi-gpu.html), [Multi-node training (Torchrun, Ray)](https://docs.axolotl.ai/docs/multi-node.html), and many more!
 - **Flexible Dataset Handling**: Load from local, HuggingFace, and cloud (S3, Azure, GCP, OCI) datasets.
 - **Cloud Ready**: We ship [Docker images](https://hub.docker.com/u/axolotlai) and also [PyPI packages](https://pypi.org/project/axolotl/) for use on cloud platforms and local hardware.

@@ -98,7 +88,7 @@ Features:
 #### Using pip

 ```bash
-pip3 install -U packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install -U packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]

 # Download example axolotl configs, deepspeed configs
--- a/1
+++ b/1
@@ -1 +0,0 @@
-0.16.0.dev0
--- a/_quarto.yml
+++ b/_quarto.yml
@@ -251,6 +251,7 @@ website:
                - docs/models/olmo3.qmd
                - docs/models/trinity.qmd
                - docs/models/arcee.qmd
+                - docs/models/mistral.qmd
                - section: "Ministral3"
                  contents:
                    - docs/models/ministral3.qmd
@@ -265,7 +266,6 @@ website:
                - docs/models/mistral-small.qmd
                - docs/models/voxtral.qmd
                - docs/models/devstral.qmd
-                - docs/models/mistral.qmd
                - docs/models/llama-4.qmd
                - docs/models/llama-2.qmd
                - docs/models/qwen3-next.qmd
@@ -320,7 +320,6 @@ website:
            - docs/multipack.qmd
            - docs/mixed_precision.qmd
            - docs/optimizers.qmd
-            - docs/attention.qmd

        - section: "Advanced Features"
          contents:
@@ -331,7 +330,6 @@ website:
            - docs/sequence_parallelism.qmd
            - docs/gradient_checkpointing.qmd
            - docs/nd_parallelism.qmd
-            - docs/expert_quantization.qmd

        - section: "Troubleshooting"
          contents:
--- a/benchmarks/bench_entropy.py
+++ b/benchmarks/bench_entropy.py
@@ -1,208 +0,0 @@
-"""Benchmark for entropy_from_logits Triton kernel vs original chunked implementation.
-
-Usage: CUDA_VISIBLE_DEVICES=0 python benchmarks/bench_entropy.py
-"""
-
-import gc
-import statistics
-
-import torch
-import torch.nn.functional as F
-
-from axolotl.monkeypatch.trainer.utils import entropy_from_logits
-
-V = 151936  # Qwen vocab
-WARMUP = 5
-BENCH_ITERS = 20
-MEM_ITERS = 10
-
-
-def entropy_from_logits_original(logits: torch.Tensor, chunk_size: int = 128):
-    """Original chunked implementation (reference)."""
-    original_shape = logits.shape[:-1]
-    num_classes = logits.shape[-1]
-    flat_logits = logits.reshape(-1, num_classes)
-    entropies = []
-    for chunk in flat_logits.split(chunk_size, dim=0):
-        logps = F.log_softmax(chunk, dim=-1)
-        chunk_entropy = -(torch.exp(logps) * logps).sum(-1)
-        entropies.append(chunk_entropy)
-    return torch.cat(entropies, dim=0).reshape(original_shape)
-
-
-def _clean_gpu():
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    torch.cuda.reset_accumulated_memory_stats()
-    torch.cuda.synchronize()
-
-
-def profile_time(fn, logits, n_iters=BENCH_ITERS):
-    for _ in range(WARMUP):
-        out = fn(logits, chunk_size=128)
-        del out
-    torch.cuda.synchronize()
-
-    times = []
-    for _ in range(n_iters):
-        s = torch.cuda.Event(enable_timing=True)
-        e = torch.cuda.Event(enable_timing=True)
-        s.record()
-        out = fn(logits, chunk_size=128)
-        e.record()
-        torch.cuda.synchronize()
-        times.append(s.elapsed_time(e))
-        del out
-    return times
-
-
-def profile_memory(fn, logits, n_iters=MEM_ITERS):
-    for _ in range(WARMUP):
-        out = fn(logits, chunk_size=128)
-        del out
-    torch.cuda.synchronize()
-
-    peaks = []
-    for _ in range(n_iters):
-        _clean_gpu()
-        base = torch.cuda.max_memory_allocated()
-        out = fn(logits, chunk_size=128)
-        torch.cuda.synchronize()
-        peaks.append(torch.cuda.max_memory_allocated() - base)
-        del out
-    return [p / 1e6 for p in peaks]
-
-
-def fmt(values, unit=""):
-    mean = statistics.mean(values)
-    std = statistics.stdev(values) if len(values) > 1 else 0.0
-    return f"{mean:8.2f} ± {std:5.2f} {unit}  [min={min(values):.2f}, max={max(values):.2f}]"
-
-
-def benchmark_contiguous():
-    print("=" * 60)
-    print(
-        f"CONTIGUOUS BENCHMARK  (warmup={WARMUP}, time={BENCH_ITERS}, mem={MEM_ITERS})"
-    )
-    print("=" * 60)
-
-    configs = [
-        (1, 2048),
-        (1, 8192),
-        (1, 16384),
-        (4, 4096),
-        (8, 2048),
-        (16, 2048),
-        (16, 4096),
-    ]
-
-    for B, L in configs:
-        mem_gb = B * L * V * 2 / 1e9
-        if mem_gb > 28:
-            print(f"\n  skip B={B}, L={L} ({mem_gb:.1f} GB)")
-            continue
-
-        N = B * L
-        print(f"\n{'─' * 60}")
-        print(f"B={B:2d}, L={L:5d}  ({N:6d} rows, logits {mem_gb:.2f} GB)")
-        print(f"{'─' * 60}")
-
-        torch.manual_seed(42)
-        logits = torch.randn(B, L, V, device="cuda", dtype=torch.bfloat16)
-
-        t_orig = profile_time(entropy_from_logits_original, logits)
-        t_triton = profile_time(entropy_from_logits, logits)
-        orig_mean = statistics.mean(t_orig)
-        triton_mean = statistics.mean(t_triton)
-
-        print("  TIME (ms):")
-        print(f"    original: {fmt(t_orig, 'ms')}")
-        print(f"    triton:   {fmt(t_triton, 'ms')}")
-        print(f"    speedup:  {orig_mean / triton_mean:.2f}x")
-
-        m_orig = profile_memory(entropy_from_logits_original, logits)
-        m_triton = profile_memory(entropy_from_logits, logits)
-        orig_peak = statistics.mean(m_orig)
-        triton_peak = statistics.mean(m_triton)
-
-        print("  MEMORY (peak overhead):")
-        print(f"    original: {fmt(m_orig, 'MB')}")
-        print(f"    triton:   {fmt(m_triton, 'MB')}")
-        print(f"    saved:    {orig_peak - triton_peak:.1f} MB")
-
-        del logits
-        _clean_gpu()
-
-
-def benchmark_noncontiguous():
-    print("\n" + "=" * 60)
-    print(
-        f"NON-CONTIGUOUS BENCHMARK  (warmup={WARMUP}, time={BENCH_ITERS}, mem={MEM_ITERS})"
-    )
-    print("=" * 60)
-
-    configs = [
-        (4, 2048, "transpose"),
-        (4, 8192, "transpose"),
-        (8, 2048, "transpose"),
-        (4, 4096, "slice_batch"),
-    ]
-
-    for B, L, method in configs:
-        torch.manual_seed(42)
-
-        if method == "transpose":
-            raw = torch.randn(L, B, V, device="cuda", dtype=torch.bfloat16)
-            logits_nc = raw.transpose(0, 1)
-            raw_gb = L * B * V * 2 / 1e9
-        elif method == "slice_batch":
-            raw = torch.randn(B * 2, L, V, device="cuda", dtype=torch.bfloat16)
-            logits_nc = raw[::2]
-            raw_gb = B * 2 * L * V * 2 / 1e9
-        else:
-            continue
-
-        if raw_gb > 28:
-            print(f"\n  skip B={B}, L={L}, {method} ({raw_gb:.1f} GB)")
-            del raw, logits_nc
-            torch.cuda.empty_cache()
-            continue
-
-        N = B * L
-        print(f"\n{'─' * 60}")
-        print(f"B={B}, L={L}  {method}  ({N} rows, raw {raw_gb:.2f} GB)")
-        print(f"{'─' * 60}")
-
-        def original_with_copy(logits, chunk_size=128):
-            return entropy_from_logits_original(
-                logits.contiguous(), chunk_size=chunk_size
-            )
-
-        t_orig = profile_time(original_with_copy, logits_nc)
-        t_triton = profile_time(entropy_from_logits, logits_nc)
-        orig_mean = statistics.mean(t_orig)
-        triton_mean = statistics.mean(t_triton)
-
-        print("  TIME (ms):")
-        print(f"    orig+copy:     {fmt(t_orig, 'ms')}")
-        print(f"    triton-strided:{fmt(t_triton, 'ms')}")
-        print(f"    speedup:       {orig_mean / triton_mean:.2f}x")
-
-        m_orig = profile_memory(original_with_copy, logits_nc)
-        m_triton = profile_memory(entropy_from_logits, logits_nc)
-        orig_peak = statistics.mean(m_orig)
-        triton_peak = statistics.mean(m_triton)
-
-        print("  MEMORY (peak overhead):")
-        print(f"    orig+copy:     {fmt(m_orig, 'MB')}")
-        print(f"    triton-strided:{fmt(m_triton, 'MB')}")
-        print(f"    saved:         {orig_peak - triton_peak:.1f} MB")
-
-        del raw, logits_nc
-        _clean_gpu()
-
-
-if __name__ == "__main__":
-    benchmark_contiguous()
-    benchmark_noncontiguous()
--- a/benchmarks/bench_scattermoe_lora.py
+++ b/benchmarks/bench_scattermoe_lora.py
@@ -1,284 +0,0 @@
-"""Benchmark for ScatterMoE LoRA Triton kernels.
-
-Measures forward, backward dX, and backward dA/dB kernels at common MoE
-model shapes. Reports per-kernel timings, LoRA overhead vs base scatter2scatter,
-and full fwd+bwd autograd throughput.
-
-Usage:
-  CUDA_VISIBLE_DEVICES=0 python benchmarks/bench_scattermoe_lora.py
-  CUDA_VISIBLE_DEVICES=0 python benchmarks/bench_scattermoe_lora.py --ranks 16 64
-  CUDA_VISIBLE_DEVICES=0 python benchmarks/bench_scattermoe_lora.py --models Qwen/Qwen3.5-35B-A3B
-"""
-
-import argparse
-import gc
-import time
-from functools import partial
-
-import torch
-
-from axolotl.integrations.kernels.libs.scattermoe_lora.kernels import (
-    lora_ops,
-    ops as base_ops,
-)
-from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_experts import (
-    flatten_sort_count,
-)
-from axolotl.integrations.kernels.libs.scattermoe_lora.parallel_linear_lora import (
-    ScatterMoELoRA,
-)
-
-DEVICE = "cuda"
-DTYPE = torch.bfloat16
-WARMUP = 5
-ITERS = 20
-
-# ─── Model configs ──────────────────────────────────────────────────────────
-
-BUILTIN_CONFIGS = {
-    "Qwen3.5-35B-A3B": (256, 2048, 512, 8),  # E, H, I, k
-    "Qwen3-30B-A3B": (128, 2048, 768, 8),
-    "OLMoE-1B-7B": (64, 2048, 1024, 8),
-    "Mixtral-8x7B": (8, 4096, 14336, 2),
-}
-
-
-def _resolve_config(spec):
-    """Resolve a model spec to (E, H, I, k). Accepts builtin names or HF IDs."""
-    key = spec.lower().replace("/", "-")
-    for name, cfg in BUILTIN_CONFIGS.items():
-        if key in name.lower() or name.lower() in key:
-            return name, cfg
-
-    from transformers import AutoConfig
-
-    hf_cfg = AutoConfig.from_pretrained(spec, trust_remote_code=True)
-    if callable(getattr(hf_cfg, "get_text_config", None)):
-        tc = hf_cfg.get_text_config()
-        if hasattr(tc, "model_type") and tc.model_type != hf_cfg.model_type:
-            hf_cfg = tc
-    hidden = hf_cfg.hidden_size
-    inter = getattr(hf_cfg, "moe_intermediate_size", None) or hf_cfg.intermediate_size
-    experts = (
-        getattr(hf_cfg, "num_experts", None)
-        or getattr(hf_cfg, "num_local_experts", None)
-        or getattr(hf_cfg, "n_routed_experts", None)
-    )
-    top_k = (
-        getattr(hf_cfg, "num_experts_per_tok", None)
-        or getattr(hf_cfg, "num_experts_per_token", None)
-        or 2
-    )
-    name = spec.split("/")[-1]
-    return name, (experts, hidden, inter, top_k)
-
-
-# ─── Benchmark helpers ──────────────────────────────────────────────────────
-
-
-def _clean():
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.synchronize()
-
-
-def _bench(fn, warmup=WARMUP, iters=ITERS):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    times = []
-    for _ in range(iters):
-        torch.cuda.synchronize()
-        t0 = time.perf_counter()
-        fn()
-        torch.cuda.synchronize()
-        times.append((time.perf_counter() - t0) * 1000)
-    times.sort()
-    return times[len(times) // 2]
-
-
-def _setup(num_experts, K, N, T, top_k, R):
-    torch.manual_seed(42)
-    x = torch.randn(T, K, device=DEVICE, dtype=DTYPE)
-    W = torch.randn(num_experts, K, N, device=DEVICE, dtype=DTYPE) * 0.02
-    lora_A = torch.randn(R * num_experts, K, device=DEVICE, dtype=DTYPE) * 0.01
-    lora_B = torch.randn(N, R * num_experts, device=DEVICE, dtype=DTYPE) * 0.01
-    logits = torch.randn(T, num_experts, device=DEVICE)
-    _, top_idx = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
-    sei, ssi, eo = flatten_sort_count(top_idx, num_experts)
-    gx = base_ops.group(x, ssi, fan_out=top_k)
-    dy = torch.randn(gx.size(0), N, device=DEVICE, dtype=DTYPE)
-    return x, W, lora_A, lora_B, sei, ssi, eo, gx, dy
-
-
-# ─── Kernel wrappers (avoid B023 loop-variable capture) ──────────────────────
-
-
-def _call_fwd(x, W, sei, ssi, top_k, lA, lB):
-    return lora_ops.scatter2scatter_lora(
-        X=x,
-        W=W,
-        sorted_expert_idxs=sei,
-        sorted_scattered_idxs=ssi,
-        k=top_k,
-        lora_A=lA,
-        lora_B=lB,
-        scaling=2.0,
-    )
-
-
-def _call_base(x, W, sei, ssi, top_k):
-    return base_ops.scatter2scatter(
-        X=x,
-        W=W,
-        sorted_expert_idxs=sei,
-        sorted_scattered_idxs=ssi,
-        k=top_k,
-    )
-
-
-def _call_dx(dy, W, sei, ssi, lA, lB):
-    return lora_ops.scatter2scatter_lora_dX(
-        DY=dy,
-        W=W,
-        sorted_expert_idxs=sei,
-        sorted_scattered_idxs=ssi,
-        k=1,
-        lora_A=lA,
-        lora_B=lB,
-        scaling=2.0,
-        dy_grouped=True,
-        dx_grouped=False,
-    )
-
-
-def _call_bwd(dy, gx, lA, lB, eo, num_experts):
-    return lora_ops.group_bwd_lora(
-        DY=dy,
-        X=gx,
-        lora_A=lA,
-        lora_B=lB,
-        expert_offsets=eo,
-        E=num_experts,
-        scaling=2.0,
-    )
-
-
-# ─── Main ────────────────────────────────────────────────────────────────────
-
-
-def main():
-    parser = argparse.ArgumentParser(description="ScatterMoE LoRA kernel benchmark")
-    parser.add_argument(
-        "--models",
-        "-m",
-        nargs="+",
-        help="Model names or HF IDs (default: all builtins)",
-    )
-    parser.add_argument("--ranks", "-r", nargs="+", type=int, default=[16, 32, 64])
-    parser.add_argument("--seq-len", "-T", type=int, default=2048)
-    args = parser.parse_args()
-
-    T = args.seq_len
-    print(f"GPU: {torch.cuda.get_device_name()}")
-    print(f"T={T}, ranks={args.ranks}\n")
-
-    if args.models:
-        configs = [_resolve_config(m) for m in args.models]
-    else:
-        configs = list(BUILTIN_CONFIGS.items())
-
-    for model_name, (num_experts, hidden, inter, top_k) in configs:
-        print(f"{'=' * 70}")
-        print(f"  {model_name}: E={num_experts}, H={hidden}, I={inter}, k={top_k}")
-        print(f"{'=' * 70}")
-
-        for R in args.ranks:
-            for proj, K, N in [("gate_up", hidden, 2 * inter), ("down", inter, hidden)]:
-                _clean()
-                x, W, lA, lB, sei, ssi, eo, gx, dy = _setup(
-                    num_experts, K, N, T, top_k, R
-                )
-
-                # Forward with LoRA (auto-dispatched: fused or split)
-                dispatch = (
-                    "split"
-                    if (
-                        num_experts <= lora_ops._SPLIT_LORA_FWD_MAX_EXPERTS
-                        and K * N >= lora_ops._SPLIT_LORA_FWD_THRESHOLD
-                    )
-                    else "fused"
-                )
-                t_fwd = _bench(partial(_call_fwd, x, W, sei, ssi, top_k, lA, lB))
-                t_base = _bench(partial(_call_base, x, W, sei, ssi, top_k))
-                t_dx = _bench(partial(_call_dx, dy, W, sei, ssi, lA, lB))
-                t_bwd = _bench(partial(_call_bwd, dy, gx, lA, lB, eo, num_experts))
-
-                total = t_fwd + t_dx + t_bwd
-                overhead = t_fwd / t_base - 1 if t_base > 0 else 0
-
-                print(
-                    f"  R={R:>2} {proj:<8}  "
-                    f"fwd={t_fwd:>6.2f}ms [{dispatch}]  "
-                    f"base={t_base:>6.2f}ms "
-                    f"(+{overhead * 100:.0f}%)  "
-                    f"dx={t_dx:>6.2f}ms  bwd={t_bwd:>6.2f}ms  "
-                    f"total={total:>6.2f}ms"
-                )
-
-                # Full autograd fwd+bwd with memory measurement
-                x_ag = x.clone().requires_grad_(True)
-                lA_ag = lA.clone().requires_grad_(True)
-                lB_ag = lB.clone().requires_grad_(True)
-
-                def _run_autograd(
-                    _x=x_ag,
-                    _W=W,
-                    _k=top_k,
-                    _sei=sei,
-                    _ssi=ssi,
-                    _eo=eo,
-                    _lA=lA_ag,
-                    _lB=lB_ag,
-                ):
-                    out = ScatterMoELoRA.apply(
-                        _x,
-                        _W,
-                        _k,
-                        _sei,
-                        _ssi,
-                        _eo,
-                        _lA,
-                        _lB,
-                        2.0,
-                        None,
-                        None,
-                        False,
-                        False,
-                        True,
-                        False,
-                    )
-                    out.sum().backward()
-                    _x.grad = None
-                    _lA.grad = None
-                    _lB.grad = None
-
-                t_full = _bench(_run_autograd)
-
-                _clean()
-                torch.cuda.reset_peak_memory_stats()
-                mem_before = torch.cuda.memory_allocated()
-                _run_autograd()
-                torch.cuda.synchronize()
-                mem_peak = torch.cuda.max_memory_allocated() - mem_before
-
-                print(
-                    f"         full_fwd_bwd={t_full:>6.2f}ms  "
-                    f"peak_delta={mem_peak / 1e6:>6.1f}MB"
-                )
-
-        print()
-
-
-if __name__ == "__main__":
-    main()
--- a/benchmarks/bench_selective_logsoftmax.py
+++ b/benchmarks/bench_selective_logsoftmax.py
@@ -1,191 +0,0 @@
-"""Benchmark for selective_log_softmax Triton kernel vs original implementation.
-
-Usage: CUDA_VISIBLE_DEVICES=0 python benchmarks/bench_selective_logsoftmax.py
-"""
-
-import gc
-import statistics
-
-import torch
-
-from axolotl.monkeypatch.trainer.utils import (
-    selective_log_softmax,
-    selective_log_softmax_original,
-)
-
-V = 151936  # Qwen vocab
-WARMUP = 5
-BENCH_ITERS = 20
-MEM_ITERS = 10
-
-
-def _clean_gpu():
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    torch.cuda.reset_accumulated_memory_stats()
-    torch.cuda.synchronize()
-
-
-def profile_time(fn, args, n_iters=BENCH_ITERS):
-    for _ in range(WARMUP):
-        fn(*args)
-    torch.cuda.synchronize()
-
-    times = []
-    for _ in range(n_iters):
-        s = torch.cuda.Event(enable_timing=True)
-        e = torch.cuda.Event(enable_timing=True)
-        s.record()
-        fn(*args)
-        e.record()
-        torch.cuda.synchronize()
-        times.append(s.elapsed_time(e))
-    return times
-
-
-def profile_memory(fn, args, n_iters=MEM_ITERS):
-    for _ in range(WARMUP):
-        out = fn(*args)
-        del out
-    torch.cuda.synchronize()
-
-    peaks = []
-    for _ in range(n_iters):
-        _clean_gpu()
-        base = torch.cuda.max_memory_allocated()
-        out = fn(*args)
-        torch.cuda.synchronize()
-        peaks.append(torch.cuda.max_memory_allocated() - base)
-        del out
-    return [p / 1e6 for p in peaks]
-
-
-def fmt(values, unit=""):
-    mean = statistics.mean(values)
-    std = statistics.stdev(values) if len(values) > 1 else 0.0
-    return f"{mean:8.2f} ± {std:5.2f} {unit}  [min={min(values):.2f}, max={max(values):.2f}]"
-
-
-def benchmark_forward():
-    print("=" * 60)
-    print(f"FORWARD BENCHMARK  (warmup={WARMUP}, time={BENCH_ITERS}, mem={MEM_ITERS})")
-    print("=" * 60)
-
-    configs = [
-        (1, 2048),
-        (1, 8192),
-        (4, 4096),
-        (8, 2048),
-        (16, 2048),
-        (16, 4096),
-    ]
-
-    for B, L in configs:
-        mem_gb = B * L * V * 2 / 1e9
-        if mem_gb > 28:
-            print(f"\n  skip B={B}, L={L} ({mem_gb:.1f} GB)")
-            continue
-
-        N = B * L
-        print(f"\n{'─' * 60}")
-        print(f"B={B:2d}, L={L:5d}  ({N:6d} rows, logits {mem_gb:.2f} GB)")
-        print(f"{'─' * 60}")
-
-        torch.manual_seed(42)
-        logits = torch.randn(B, L, V, device="cuda", dtype=torch.bfloat16)
-        index = torch.randint(0, V, (B, L), device="cuda")
-
-        t_orig = profile_time(selective_log_softmax_original, (logits, index))
-        t_triton = profile_time(selective_log_softmax, (logits, index))
-        orig_mean = statistics.mean(t_orig)
-        triton_mean = statistics.mean(t_triton)
-
-        print("  TIME (ms):")
-        print(f"    original: {fmt(t_orig, 'ms')}")
-        print(f"    triton:   {fmt(t_triton, 'ms')}")
-        print(f"    speedup:  {orig_mean / triton_mean:.2f}x")
-
-        m_orig = profile_memory(selective_log_softmax_original, (logits, index))
-        m_triton = profile_memory(selective_log_softmax, (logits, index))
-        orig_peak = statistics.mean(m_orig)
-        triton_peak = statistics.mean(m_triton)
-
-        print("  MEMORY (peak overhead):")
-        print(f"    original: {fmt(m_orig, 'MB')}")
-        print(f"    triton:   {fmt(m_triton, 'MB')}")
-        print(f"    saved:    {orig_peak - triton_peak:.1f} MB")
-
-        del logits, index
-        _clean_gpu()
-
-
-def benchmark_backward():
-    print("\n" + "=" * 60)
-    print(f"FWD+BWD BENCHMARK  (warmup={WARMUP}, time={BENCH_ITERS}, mem={MEM_ITERS})")
-    print("=" * 60)
-
-    configs = [
-        (1, 2048),
-        (1, 8192),
-        (4, 4096),
-        (8, 2048),
-        (16, 2048),
-        (16, 4096),
-    ]
-
-    def fwd_bwd_original(logits, index):
-        logits.grad = None
-        out = selective_log_softmax_original(logits, index)
-        out.sum().backward()
-
-    def fwd_bwd_triton(logits, index):
-        logits.grad = None
-        out = selective_log_softmax(logits, index)
-        out.sum().backward()
-
-    for B, L in configs:
-        mem_gb = B * L * V * 2 / 1e9
-        if mem_gb > 20:
-            print(f"\n  skip B={B}, L={L} ({mem_gb:.1f} GB, need room for grads)")
-            continue
-
-        N = B * L
-        print(f"\n{'─' * 60}")
-        print(f"B={B:2d}, L={L:5d}  ({N:6d} rows, logits {mem_gb:.2f} GB)")
-        print(f"{'─' * 60}")
-
-        torch.manual_seed(42)
-        logits_orig = torch.randn(
-            B, L, V, device="cuda", dtype=torch.bfloat16, requires_grad=True
-        )
-        logits_tri = logits_orig.detach().clone().requires_grad_(True)
-        index = torch.randint(0, V, (B, L), device="cuda")
-
-        t_orig = profile_time(fwd_bwd_original, (logits_orig, index))
-        t_triton = profile_time(fwd_bwd_triton, (logits_tri, index))
-        orig_mean = statistics.mean(t_orig)
-        triton_mean = statistics.mean(t_triton)
-
-        print("  FWD+BWD TIME (ms):")
-        print(f"    original: {fmt(t_orig, 'ms')}")
-        print(f"    triton:   {fmt(t_triton, 'ms')}")
-        print(f"    speedup:  {orig_mean / triton_mean:.2f}x")
-
-        m_orig = profile_memory(fwd_bwd_original, (logits_orig, index))
-        m_triton = profile_memory(fwd_bwd_triton, (logits_tri, index))
-        orig_peak = statistics.mean(m_orig)
-        triton_peak = statistics.mean(m_triton)
-
-        print("  FWD+BWD MEMORY (peak overhead):")
-        print(f"    original: {fmt(m_orig, 'MB')}")
-        print(f"    triton:   {fmt(m_triton, 'MB')}")
-        print(f"    saved:    {orig_peak - triton_peak:.1f} MB")
-
-        del logits_orig, logits_tri, index
-        _clean_gpu()
-
-
-if __name__ == "__main__":
-    benchmark_forward()
-    benchmark_backward()
--- a/cicd/Dockerfile-uv.jinja
+++ b/cicd/Dockerfile-uv.jinja
@@ -11,7 +11,7 @@ ENV NIGHTLY_BUILD="{{ NIGHTLY_BUILD }}"
 ENV HF_HOME="{{ HF_HOME }}"

 RUN apt-get update && \
-    apt-get install -y --allow-change-held-packages vim curl nano zstd libnccl2 libnccl-dev ibverbs-providers ibverbs-utils infiniband-diags librdmacm-dev librdmacm1 rdmacm-utils slurm-wlm
+    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev ibverbs-providers ibverbs-utils infiniband-diags librdmacm-dev librdmacm1 rdmacm-utils slurm-wlm

 WORKDIR /workspace

@@ -31,9 +31,8 @@ RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
        sed -i 's#^datasets.*#datasets @ git+https://github.com/huggingface/datasets.git@main#' requirements.txt; \
    fi

-RUN uv pip install packaging==26.0 setuptools==78.1.1
+RUN uv pip install packaging==23.2 setuptools==75.8.0
 RUN uv pip install torchvision
-RUN uv pip uninstall causal_conv1d
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
        uv pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
    else \
--- a/cicd/Dockerfile.jinja
+++ b/cicd/Dockerfile.jinja
@@ -12,7 +12,7 @@ ENV HF_HOME="{{ HF_HOME }}"
 ENV AXOLOTL_DATASET_NUM_PROC="8"

 RUN apt-get update && \
-    apt-get install -y --allow-change-held-packages vim curl nano zstd libnccl2 libnccl-dev ibverbs-providers ibverbs-utils infiniband-diags librdmacm-dev librdmacm1 rdmacm-utils slurm-wlm
+    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev ibverbs-providers ibverbs-utils infiniband-diags librdmacm-dev librdmacm1 rdmacm-utils slurm-wlm

 WORKDIR /workspace

@@ -32,8 +32,7 @@ RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
        sed -i 's#^datasets.*#datasets @ git+https://github.com/huggingface/datasets.git@main#' requirements.txt; \
    fi

-RUN pip install packaging==26.0 setuptools==78.1.1 psutil
-RUN pip uninstall -y causal_conv1d
+RUN pip install packaging==23.2 setuptools==75.8.0 psutil
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
    else \
--- a/cicd/cicd.sh
+++ b/cicd/cicd.sh
@@ -3,13 +3,6 @@ set -e

 python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__"

-curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xpf - -C "${HF_HOME}/hub/"  --use-compress-program unzstd --strip-components=1
-# hf download "NousResearch/Meta-Llama-3-8B"
-# hf download "NousResearch/Meta-Llama-3-8B-Instruct"
-# hf download "microsoft/Phi-4-reasoning"
-# hf download "microsoft/Phi-3.5-mini-instruct"
-# hf download "microsoft/Phi-3-medium-128k-instruct"
-
 # Run unit tests with initial coverage report
 pytest -v --durations=10 -n8 \
  --ignore=tests/e2e/ \
--- a/cicd/multigpu.py
+++ b/cicd/multigpu.py
@@ -17,8 +17,7 @@ template_loader = jinja2.FileSystemLoader(searchpath=cicd_path)
 template_env = jinja2.Environment(
    loader=template_loader, autoescape=select_autoescape()
 )
-dockerfile = os.environ.get("E2E_DOCKERFILE", "Dockerfile.jinja")
-df_template = template_env.get_template(dockerfile)
+df_template = template_env.get_template("Dockerfile.jinja")

 df_args = {
    "AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
@@ -28,11 +27,8 @@ df_args = {
    "CUDA": os.environ.get("CUDA", "126"),
    "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
    "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
-    "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
    "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
    "HF_HOME": "/workspace/data/huggingface-cache/hub",
-    "PYTHONUNBUFFERED": os.environ.get("PYTHONUNBUFFERED", "1"),
-    "DEEPSPEED_LOG_LEVEL": os.environ.get("DEEPSPEED_LOG_LEVEL", "WARNING"),
 }

 dockerfile_contents = df_template.render(**df_args)
--- a/cicd/multigpu.sh
+++ b/cicd/multigpu.sh
@@ -2,7 +2,7 @@
 set -e

 # Only run two tests at a time to avoid OOM on GPU (with coverage collection)
-pytest -v --durations=10 -n2 --maxfail=3 \
+pytest -v --durations=10 -n2 --maxfail=4 \
  --ignore=/workspace/axolotl/tests/e2e/multigpu/solo/ \
  --ignore=/workspace/axolotl/tests/e2e/multigpu/patched/ \
  /workspace/axolotl/tests/e2e/multigpu/ \
--- a/cicd/single_gpu.py
+++ b/cicd/single_gpu.py
@@ -68,6 +68,10 @@ def run_cmd(cmd: str, run_folder: str):
    sp_env["AXOLOTL_DATASET_NUM_PROC"] = "8"

    # Propagate errors from subprocess.
-    exit_code = subprocess.call(cmd.split(), cwd=run_folder, env=sp_env)  # nosec
-    if exit_code:
-        raise RuntimeError(f"Command '{cmd}' failed with exit code {exit_code}")
+    try:
+        exit_code = subprocess.call(cmd.split(), cwd=run_folder, env=sp_env)  # nosec
+        if exit_code:
+            print(f"Command '{cmd}' failed with exit code {exit_code}")
+            return exit_code
+    except Exception as e:  # pylint: disable=broad-except
+        print(f"Command '{cmd}' failed with exception {e}")
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -6,7 +6,6 @@ ARG AXOLOTL_EXTRAS=""
 ARG AXOLOTL_ARGS=""
 ARG CUDA="118"
 ARG PYTORCH_VERSION="2.1.2"
-ARG TARGETARCH

 ENV PYTORCH_VERSION=$PYTORCH_VERSION

@@ -21,18 +20,13 @@ RUN git clone --depth=1 https://github.com/axolotl-ai-cloud/axolotl.git

 WORKDIR /workspace/axolotl

-# If AXOLOTL_EXTRAS is set, append it in brackets; don't install deepspeed with arm64
-RUN pip uninstall -y causal_conv1d
-RUN if [ "$TARGETARCH" = "arm64" ]; then \
-        BASE_EXTRAS="flash-attn,ring-flash-attn,optimizers,ray"; \
+# If AXOLOTL_EXTRAS is set, append it in brackets
+RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
    else \
-        BASE_EXTRAS="deepspeed,flash-attn,ring-flash-attn,optimizers,ray"; \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
    fi && \
-    if [ "$AXOLOTL_EXTRAS" != "" ]; then \
-        pip install --no-build-isolation -e .[$BASE_EXTRAS,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
-    else \
-        pip install --no-build-isolation -e .[$BASE_EXTRAS] $AXOLOTL_ARGS; \
-    fi && \    python scripts/unsloth_install.py | sh && \
+    python scripts/unsloth_install.py | sh && \
    python scripts/cutcrossentropy_install.py | sh && \
    pip install pytest && \
    pip cache purge
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -2,16 +2,14 @@ ARG CUDA_VERSION="11.8.0"
 ARG CUDNN_VERSION="8"
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
-ARG TARGETARCH

 FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder

 ENV PATH="/root/miniconda3/bin:${PATH}"

-ARG TARGETARCH
-ARG PYTHON_VERSION="3.11"
+ARG PYTHON_VERSION="3.10"
 ARG PYTORCH_VERSION="2.1.2"
-ARG CUDA="128"
+ARG CUDA="118"
 ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"

 ENV PYTHON_VERSION=$PYTHON_VERSION
@@ -24,17 +22,11 @@ RUN apt-get update \
        librdmacm-dev librdmacm1 rdmacm-utils slurm-wlm \
    && rm -rf /var/cache/apt/archives \
    && rm -rf /var/lib/apt/lists/* \
-    && if [ "$TARGETARCH" = "amd64" ]; then \
-        MINICONDA_ARCH="x86_64"; \
-    elif [ "$TARGETARCH" = "arm64" ]; then \
-        MINICONDA_ARCH="aarch64"; \
-    else \
-        echo "Unsupported architecture: $TARGETARCH"; exit 1; \
-    fi \
-    && wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh \
+    && wget \
+    https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
    && mkdir /root/.conda \
-    && bash Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh -b \
-    && rm -f Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh \
+    && bash Miniconda3-latest-Linux-x86_64.sh -b \
+    && rm -f Miniconda3-latest-Linux-x86_64.sh \
    && conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main \
    && conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r \
    && conda create -n "py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
@@ -43,7 +35,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"

 WORKDIR /workspace

-RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==26.0 setuptools==75.8.0 wheel psutil && \
+RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==23.2 setuptools==75.8.0 wheel psutil && \
    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} torchvision --extra-index-url https://download.pytorch.org/whl/cu$CUDA && \
    python3 -m pip cache purge

@@ -59,18 +51,8 @@ RUN git lfs install --skip-repo && \
    pip3 install -U --no-cache-dir pydantic==1.10.10 && \
    pip3 cache purge

-# Map Python version (e.g., 3.12 -> cp312)
-RUN PYTHON_CP="cp$(echo $PYTHON_VERSION | tr -d '.')" && \
-    # Map PyTorch version (e.g., 2.9.1 -> torch2.9, 2.10.0 -> torch2.10)
-    TORCH_TAG="torch$(echo $PYTORCH_VERSION | grep -oP '^\d+\.\d+')" && \
-    # Map architecture
-    case "$TARGETARCH" in \
-        amd64) ARCH_TAG="x86_64" ;; \
-        arm64) ARCH_TAG="aarch64" ;; \
-        *) echo "Unsupported architecture: $TARGETARCH"; exit 1 ;; \
-    esac && \
-    WHL_VERSION="v0.7.16" && \
-    WHL_FILE="flash_attn-2.8.3+cu${CUDA}${TORCH_TAG}-${PYTHON_CP}-${PYTHON_CP}-linux_${ARCH_TAG}.whl" && \
-    wget -nv "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/${WHL_VERSION}/${WHL_FILE}" && \
-    pip3 install --no-cache-dir "${WHL_FILE}" && \
-    rm "${WHL_FILE}"
+RUN if [ "$PYTORCH_VERSION" =~ ^2\.9\.[0-9]+$ ] && [ "$CUDA" = "128" ] ; then \
+        wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+        pip3 install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+        rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+    fi
--- a/docker/Dockerfile-base-nightly
+++ b/docker/Dockerfile-base-nightly
@@ -30,7 +30,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"

 WORKDIR /workspace

-RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==26.0 setuptools==75.8.0 wheel && \
+RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==23.2 setuptools==75.8.0 wheel && \
    python3 -m pip install --no-cache-dir -U torch --extra-index-url https://download.pytorch.org/whl/nightly/cu$CUDA && \
    python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \
    python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" && \
--- a/docker/Dockerfile-cloud-uv
+++ b/docker/Dockerfile-cloud-uv
@@ -1,30 +0,0 @@
-ARG BASE_TAG=main
-FROM axolotlai/axolotl-uv:$BASE_TAG
-
-ENV HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
-ENV HF_HUB_CACHE="/workspace/data/huggingface-cache/hub"
-ENV HF_HOME="/workspace/data/huggingface-cache/hub"
-ENV HF_HUB_ENABLE_HF_TRANSFER="1"
-
-EXPOSE 8888
-EXPOSE 22
-
-COPY scripts/cloud-entrypoint.sh /root/cloud-entrypoint.sh
-COPY scripts/motd /etc/motd
-
-RUN uv pip install jupyterlab notebook ipywidgets && \
-    jupyter lab clean
-RUN apt update && \
-    apt install --yes --no-install-recommends openssh-server tmux iproute2 nvtop && \
-    rm -rf /var/cache/apt/archives && \
-    rm -rf /var/lib/apt/lists/* && \
-    mkdir -p ~/.ssh && \
-    chmod 700 ~/.ssh && \
-    printf "\n[[ -z \"\$TMUX\"  ]] && { tmux attach-session -t ssh_tmux || tmux new-session -s ssh_tmux; exit; }\n" >> ~/.bashrc && \
-    printf "[ ! -z \"\$TERM\" -a -r /etc/motd ] && cat /etc/motd\n" >> ~/.bashrc && \
-    chmod +x /workspace/axolotl/scripts/cloud-entrypoint.sh && \
-    chmod +x /root/cloud-entrypoint.sh && \
-    echo 'set-option -g history-limit 5000' >> ~/.tmux.conf
-
-ENTRYPOINT ["/root/cloud-entrypoint.sh"]
-CMD ["sleep", "infinity"]
--- a/docker/Dockerfile-uv
+++ b/docker/Dockerfile-uv
@@ -1,48 +0,0 @@
-ARG BASE_TAG=main-base
-FROM axolotlai/axolotl-base-uv:$BASE_TAG
-
-ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
-ARG AXOLOTL_EXTRAS=""
-ARG AXOLOTL_ARGS=""
-ARG CUDA="118"
-ARG PYTORCH_VERSION="2.1.2"
-ARG TARGETARCH
-
-ENV PYTORCH_VERSION=$PYTORCH_VERSION
-
-RUN apt-get update && \
-    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev rsync s3fs && \
-    rm -rf /var/cache/apt/archives && \
-    rm -rf /var/lib/apt/lists/*
-
-WORKDIR /workspace
-
-RUN git clone --depth=1 https://github.com/axolotl-ai-cloud/axolotl.git
-
-WORKDIR /workspace/axolotl
-
-# If AXOLOTL_EXTRAS is set, append it in brackets; don't install deepspeed with arm64
-RUN uv pip uninstall causal_conv1d
-RUN if [ "$TARGETARCH" = "arm64" ]; then \
-        BASE_EXTRAS="flash-attn,ring-flash-attn,optimizers,ray"; \
-    else \
-        BASE_EXTRAS="deepspeed,flash-attn,ring-flash-attn,optimizers,ray"; \
-    fi && \
-    if [ "$AXOLOTL_EXTRAS" != "" ]; then \
-        uv pip install --no-build-isolation -e .[$BASE_EXTRAS,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
-    else \
-        uv pip install --no-build-isolation -e .[$BASE_EXTRAS] $AXOLOTL_ARGS; \
-    fi && \
-    python scripts/unsloth_install.py --uv | sh && \
-    python scripts/cutcrossentropy_install.py --uv | sh && \
-    uv pip install pytest && \
-    uv cache clean
-
-# fix so that git fetch/pull from remote works with shallow clone
-RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
-    git config --get remote.origin.fetch && \
-    git config --global credential.helper store
-
-COPY .axolotl-complete.bash /root/.axolotl-complete.bash
-RUN chmod +x /root/.axolotl-complete.bash && \
-    echo 'source /root/.axolotl-complete.bash' >> ~/.bashrc
--- a/docker/Dockerfile-uv-base
+++ b/docker/Dockerfile-uv-base
@@ -2,11 +2,9 @@ ARG CUDA_VERSION="12.6.3"
 ARG CUDNN_VERSION=""
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
-ARG TARGETARCH

 FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder

-ARG TARGETARCH
 ARG PYTHON_VERSION="3.11"
 ARG PYTORCH_VERSION="2.6.0"
 ARG CUDA="126"
@@ -33,25 +31,12 @@ ENV PATH="/workspace/axolotl-venv/bin:${PATH}"

 RUN uv pip install packaging setuptools wheel psutil \
    && uv pip install torch==${PYTORCH_VERSION} torchvision \
+    && uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" \
+    && uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" \
    && uv pip install awscli pydantic

-RUN if [ "$TARGETARCH" = "amd64" ]; then \
-        uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main"; \
-        uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main"; \
+RUN if [ "$PYTORCH_VERSION" = "2.9.0" ] && [ "$CUDA" = "128" ] ; then \
+        wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+        uv pip install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+        rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
    fi
-
-# Map Python version (e.g., 3.12 -> cp312)
-RUN PYTHON_CP="cp$(echo $PYTHON_VERSION | tr -d '.')" && \
-    # Map PyTorch version (e.g., 2.9.1 -> torch2.9, 2.10.0 -> torch2.10)
-    TORCH_TAG="torch$(echo $PYTORCH_VERSION | grep -oP '^\d+\.\d+')" && \
-    # Map architecture
-    case "$TARGETARCH" in \
-        amd64) ARCH_TAG="x86_64" ;; \
-        arm64) ARCH_TAG="aarch64" ;; \
-        *) echo "Unsupported architecture: $TARGETARCH"; exit 1 ;; \
-    esac && \
-    WHL_VERSION="v0.7.16" && \
-    WHL_FILE="flash_attn-2.8.3+cu${CUDA}${TORCH_TAG}-${PYTHON_CP}-${PYTHON_CP}-linux_${ARCH_TAG}.whl" && \
-    wget -nv "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/${WHL_VERSION}/${WHL_FILE}" && \
-    uv pip install --no-cache-dir "${WHL_FILE}" && \
-    rm "${WHL_FILE}"
--- a/docs/amd_hpc.qmd
+++ b/docs/amd_hpc.qmd
@@ -86,7 +86,7 @@ export HF_DATASETS_OFFLINE=1
 Download a base model using the Hugging Face CLI:

 ```bash
-hf download meta-llama/Meta-Llama-3.1-8B --local-dir ~/hfdata/llama3.1-8B
+huggingface-cli download meta-llama/Meta-Llama-3.1-8B --local-dir ~/hfdata/llama3.1-8B
 ```

 ### 10. Create Axolotl Configuration
--- a/docs/attention.qmd
+++ b/docs/attention.qmd
@@ -1,178 +0,0 @@
---
-title: Attention
-description: Supported attention modules in Axolotl
---
-
-## SDP Attention
-
-This is the default built-in attention in PyTorch.
-
-```yaml
-sdp_attention: true
-```
-
-For more details: [PyTorch docs](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
-
-## Flash Attention
-
-Axolotl supports Flash Attention 2, 3, and 4. The best available version is used automatically
-based on your installed packages and GPU.
-
-```yaml
-flash_attention: true
-```
-
-For more details: [Flash Attention](https://github.com/Dao-AILab/flash-attention/)
-
-### Flash Attention 2
-
-Requirements: Ampere, Ada, or Hopper GPUs (Turing or lower not supported)
-
-```bash
-pip install flash-attn --no-build-isolation
-```
-
-::: {.callout-tip}
-
-If you get `undefined symbol` while training, ensure you installed PyTorch prior to Axolotl.
-Alternatively, try reinstall or downgrade a version.
-
-:::
-
-### Flash Attention 3
-
-Requirements: Hopper only and CUDA 12.8 (recommended)
-
-```bash
-git clone https://github.com/Dao-AILab/flash-attention.git
-cd flash-attention/hopper
-
-python setup.py install
-```
-
-### Flash Attention 4
-
-Requirements: Hopper or Blackwell GPUs
-
-```bash
-pip install flash-attn-4
-```
-
-Or from source:
-
-```bash
-git clone https://github.com/Dao-AILab/flash-attention.git
-cd flash-attention/flash_attn/cute
-
-pip install -e .
-
-# FA2's flash_attn package includes a cute/ stub that shadows FA4.
-# Remove it so Python can find the real FA4 module:
-rm -r $(python -c "import flash_attn; print(flash_attn.__path__[0])")/cute
-```
-
-::: {.callout-note}
-
-**Hopper (SM90) users**: The backward kernel is not yet included in the pip package. To use FA4
-for training on Hopper, install from source using the instructions above.
-
-:::
-
-::: {.callout-warning}
-
-FA4 only supports head dimensions up to 128 (`d ≤ 128`). The DeepSeek shape `(192, 128)` is
-also supported but only on Blackwell. Axolotl automatically detects incompatible head dimensions
-and falls back to FA2/3.
-
-:::
-
-For more details: [flash-attention/flash_attn/cute](https://github.com/Dao-AILab/flash-attention/tree/main/flash_attn/cute)
-
-### AMD
-
-Requirements: ROCm 6.0 and above.
-
-See [Flash Attention AMD docs](https://github.com/Dao-AILab/flash-attention/tree/main?tab=readme-ov-file#amd-rocm-support).
-
-## Flex Attention
-
-A flexible PyTorch API for attention used in combination with `torch.compile`.
-
-```yaml
-flex_attention: true
-
-# recommended
-torch_compile: true
-```
-
-::: {.callout-note}
-
-We recommend using latest stable version of PyTorch for best performance.
-
-:::
-
-For more details: [PyTorch docs](https://pytorch.org/blog/flexattention/)
-
-## SageAttention
-
-Attention kernels with QK Int8 and PV FP16 accumulator.
-
-```yaml
-sage_attention: true
-```
-
-Requirements: Ampere, Ada, or Hopper GPUs
-
-```bash
-pip install sageattention==2.2.0 --no-build-isolation
-```
-
-::: {.callout-warning}
-
-Only LoRA/QLoRA recommended at the moment. We found loss drop to 0 for full finetuning. See [GitHub Issue](https://github.com/thu-ml/SageAttention/issues/198).
-
-:::
-
-For more details: [Sage Attention](https://github.com/thu-ml/SageAttention)
-
-::: {.callout-note}
-
-We do not support SageAttention 3 at the moment. If you are interested on adding this or improving SageAttention implementation, please make an Issue.
-
-:::
-
-
-## xFormers
-
-```yaml
-xformers_attention: true
-```
-
-::: {.callout-tip}
-
-We recommend using with Turing GPUs or below (such as on Colab).
-
-:::
-
-For more details: [xFormers](https://github.com/facebookresearch/xformers)
-
-## Shifted Sparse Attention
-
-::: {.callout-warning}
-
-We plan to deprecate this! If you use this feature, we recommend switching to methods above.
-
-:::
-
-Requirements: LLaMA model architecture
-
-```yaml
-flash_attention: true
-s2_attention: true
-```
-
-::: {.callout-tip}
-
-No sample packing support!
-
-:::
--- a/docs/cli.qmd
+++ b/docs/cli.qmd
@@ -210,8 +210,6 @@ axolotl lm-eval config.yml
 Configuration options:

 ```yaml
-lm_eval_model: # model to evaluate (local or hf path)
-
 # List of tasks to evaluate
 lm_eval_tasks:
  - arc_challenge
@@ -220,7 +218,7 @@ lm_eval_batch_size: # Batch size for evaluation
 output_dir: # Directory to save evaluation results
 ```

-See [LM Eval Harness integration docs](https://docs.axolotl.ai/docs/custom_integrations.html#language-model-evaluation-harness-lm-eval) for full configuration details.
+See [LM Eval Harness](https://github.com/EleutherAI/lm-evaluation-harness) for more details.

 ### delinearize-llama4

--- a/docs/expert_quantization.qmd
+++ b/docs/expert_quantization.qmd
@@ -1,67 +0,0 @@
---
-title: "MoE Expert Quantization"
-description: "Reduce VRAM usage when training MoE model adapters by quantizing expert weights on load"
---
-
-Transformers v5 changed MoE expert layers from `nn.Linear` to fused `nn.Parameter` (3D+ tensors).
-This means `bitsandbytes` can no longer quantize them during model loading, resulting in all expert
-weights being loaded in full bf16 precision and causing massive VRAM usage.
-
-`quantize_moe_experts` solves this by quantizing expert weights during model loading.
-It intercepts the weight loading process, quantizes each expert tensor on the fly, and
-immediately frees the original bf16 tensor from VRAM. This dramatically reduces peak memory.
-For example, GLM-4.7-Flash QLoRA drops from ~127GiB to ~23GiB reserved memory.
-
-## Usage
-
-Enable expert quantization in your Axolotl config:
-
-```yaml
-quantize_moe_experts: true
-```
-
-This works with both 4-bit (QLoRA) and 8-bit (LoRA) quantization.
-
-### Expert LoRA targeting
-
-You can optionally apply LoRA adapters directly to expert weights using `lora_target_parameters`:
-
-```yaml
-lora_target_parameters:
-  - mlp.experts.gate_up_proj
-  - mlp.experts.down_proj
-  # - mlp.gate.weight  # router
-```
-
-::: {.callout-note}
-`lora_dropout` must be `0` when using `lora_target_parameters`.
-:::
-
-## Requirements
-
- Requires (`adapter: lora` and `load_in_8bit: true`) or (`adapter: qlora` and `load_in_4bit: true`)
- CUDA GPUs only (not tested with ROCm or other backends)
- FSDP2 compatible for distributed training
-
-## Limitations
-
- `lora_target_linear` is not compatible with `quantize_moe_experts`. See [Expert LoRA targeting](#expert-lora-targeting) instead.
- `cpu_ram_efficient_loading` hangs / takes long time with FSDP2 + QLoRA.
- Total model parameter count may display incorrectly (trainable param count is correct).
- FSDP LoRA (8-bit) may have a large initial VRAM spike at the first 1-2 steps, which then drops. QLoRA does not exhibit this.
- FSDP2 may use more VRAM per GPU than single GPU training due to not all layers being properly sharded across ranks.
- Model loading takes longer due to on-demand quantization, even on consecutive runs.
- DeepSpeed has not been tested.
-
-## Implementation details
-
-The quantization is applied by patching transformers to intercept weight loading.
-When a 3D+ CUDA tensor with "expert" in its name is detected:
-
- **4-bit mode:** Uses bitsandbytes NF4 parametrization (configurable via `bnb_4bit_quant_type`).
- **8-bit mode:** Uses a custom row-wise int8 parametrization with bitsandbytes dequantization.
-
-The original bf16 tensor is freed immediately after quantization. Multiple sub-patches are applied to
-transformers, PEFT and accelerate FSDP2 to support these parametrized expert modules.
-
-For full implementation details, see [PR #3439](https://github.com/axolotl-ai-cloud/axolotl/pull/3439).
--- a/docs/installation.qmd
+++ b/docs/installation.qmd
@@ -165,7 +165,7 @@ We recommend using WSL2 (Windows Subsystem for Linux) or Docker.
   ```
 4. (Optional) Login to Hugging Face:
   ```{.bash}
-   hf auth login
+   huggingface-cli login
   ```

 ## Troubleshooting {#sec-troubleshooting}
--- a/docs/lora_optims.qmd
+++ b/docs/lora_optims.qmd
@@ -89,10 +89,6 @@ lora_o_kernel: true
 Currently, LoRA kernels are not supported for RLHF training, only SFT.
 :::

-::: {.callout-warning}
-LoRA kernels do not support remote modeling code.
-:::
-
 ## Requirements

 - One or more NVIDIA or AMD GPUs (in order to use the Triton kernels)
--- a/docs/multimodal.qmd
+++ b/docs/multimodal.qmd
@@ -13,15 +13,12 @@ format:
 - [Pixtral](#sec-pixtral)
 - [Llava-1.5](#sec-llava-15)
 - [Mistral-Small-3.1](#sec-mistral-small-31)
- [Mistral-Small-4](#sec-mistral-small-4)
 - [Magistral-Small-2509](#sec-magistral-small-2509)
 - [Voxtral](#sec-voxtral)
 - [Gemma-3](#sec-gemma-3)
 - [Gemma-3n](#sec-gemma-3n)
 - [Qwen2-VL](#sec-qwen2-vl)
 - [Qwen2.5-VL](#sec-qwen25-vl)
- [Qwen3.5](#sec-qwen3-5)
- [GLM-4.6V](#sec-glm-4-6v)
 - [SmolVLM2](#sec-smolvlm2)
 - [LFM2-VL](#sec-lfm2-vl)
 - [Intern-VL](#sec-intern-vl)
@@ -110,12 +107,6 @@ Please make sure to install vision lib via `pip install 'mistral-common[opencv]=
 base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
 ```

-### Mistral-Small-4 {#sec-mistral-small-4}
-
-```yaml
-base_model: mistralai/Mistral-Small-4-119B-2603
-```
-
 ### Magistral-Small-2509 {#sec-magistral-small-2509}

 ::: {.callout-tip}
@@ -192,26 +183,6 @@ base_model: Qwen/Qwen3-VL-4B-Instruct
 chat_template: qwen2_vl  # same as qwen2-vl
 ```

-### Qwen3.5 {#sec-qwen3-5}
-
-```yaml
-base_model: Qwen/Qwen3.5-9B
-
-chat_template: qwen3_5
-```
-
-### GLM-4.6V {#sec-glm-4-6v}
-
-Both GLM-4.6V (106B MoE) and GLM-4.6V-Flash (9B) are supported.
-
-```yaml
-# GLM-4.6V (106B MoE version)
-base_model: zai-org/GLM-4.6V
-
-# OR GLM-4.6V-Flash (9B version)
-base_model: zai-org/GLM-4.6V-Flash
-```
-
 ### SmolVLM2 {#sec-smolvlm2}

 ::: {.callout-tip}
--- a/docs/optimizations.qmd
+++ b/docs/optimizations.qmd
@@ -66,15 +66,6 @@ Provides efficient Triton kernels to improve training speed and reduce memory us

 - **Learn more:** [Custom Integrations - Liger Kernels](custom_integrations.qmd#liger-kernels)

-### Expert Kernels
-
-Optimized kernel implementations for Mixture of Experts (MoE) model training.
-
- **ScatterMoE**: Triton-based MoE kernels with fused LoRA support.
- **SonicMoE**: CUTLASS-based MoE kernels for NVIDIA Hopper and Blackwell GPUs.
-
- **Learn more:** [Custom Integrations - Kernels Integration](custom_integrations.qmd#kernels-integration)
-
 ## Long Context Models

 Techniques to train models on sequences longer than their original context window.
@@ -140,10 +131,3 @@ Simulates quantization effects during training, helping the model adapt and pote
 Allows you to finetune LoRA adapters on top of a model that has already been quantized using the GPTQ method.

 - **Example:** [GPTQ LoRA Example](https://github.com/axolotl-ai-cloud/axolotl/blob/main/examples/llama-2/gptq-lora.yml)
-
-### MoE Expert Quantization
-
-Quantizes MoE expert weights on load to reduce VRAM when training MoE models with adapters. Required for Transformers v5+ MoE models where experts use fused `nn.Parameter` tensors.
-
- **Config:** `quantize_moe_experts: true`
- **Learn more:** [MoE Expert Quantization](expert_quantization.qmd)
--- a/docs/rlhf.qmd
+++ b/docs/rlhf.qmd
@@ -17,7 +17,6 @@ feedback. Various methods include, but not limited to:
 - [Kahneman-Tversky Optimization (KTO)](#kto)
 - [Odds Ratio Preference Optimization (ORPO)](#orpo)
 - [Group Relative Policy Optimization (GRPO)](#grpo)
- [Group Reward-Decoupled Policy Optimization (GDPO)](#gdpo)


 ## RLHF using Axolotl
@@ -721,309 +720,6 @@ trl:

 For more information, see [GRPO docs](https://huggingface.co/docs/trl/v0.17.0/en/grpo_trainer#loss-types).

-#### Async GRPO
-
-Async GRPO overlaps vLLM generation with training by producing rollouts in a background thread. While the model trains on the current batch, the next batch is already being generated. This can significantly reduce wall-clock time per step.
-
-```yaml
-trl:
-  use_data_producer: true     # Enable data producer protocol
-  use_vllm: true
-  async_prefetch: true         # Generate rollouts in background thread
-  prefetch_depth: 1            # Number of rollouts to prefetch
-  vllm_sync_interval: 2        # Sync weights to vLLM every N steps
-```
-
-::: {.callout-note}
-Because the background thread generates completions with slightly stale model weights, async GRPO uses importance sampling correction to account for the distribution shift. This is controlled by `vllm_importance_sampling_correction: true` (default when async is enabled).
-:::
-
-##### vLLM LoRA Sync
-
-By default, weight sync to vLLM merges the LoRA adapter into the base model and broadcasts all parameters via NCCL. LoRA sync is a faster alternative that saves only the adapter weights to the filesystem and has vLLM load them natively using Punica kernels.
-
-```yaml
-adapter: lora
-lora_r: 32
-lora_alpha: 64
-lora_target_linear: true
-
-trl:
-  vllm_lora_sync: true         # Enable native LoRA sync
-```
-
-When `vllm_lora_sync: true` is set, axolotl automatically selects the LoRA-aware vLLM serve module. Start vLLM as usual:
-
-```bash
-CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
-```
-
-Then start training on a separate GPU:
-
-```bash
-CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml
-```
-
-::: {.callout-tip}
-LoRA sync is especially beneficial with multi-GPU training (FSDP/DeepSpeed), where NCCL merge-sync can cause GPU contention with vLLM generation.
-:::
-
-##### Streaming Partial Batch
-
-Instead of scoring the entire batch at once, streaming mode scores one prompt group at a time. This enables finer-grained zero-advantage skipping and reduces peak memory usage during scoring.
-
-```yaml
-trl:
-  streaming_partial_batch: true
-```
-
-##### Importance Sampling Correction
-
-When using async prefetch, completions are generated from a slightly older version of the model. Importance sampling (IS) correction adjusts the policy gradient to account for this distribution shift.
-
-```yaml
-trl:
-  vllm_importance_sampling_correction: true   # Enable IS correction
-  importance_sampling_level: token             # 'token' or 'sequence'
-  off_policy_mask_threshold: 0.5              # Mask sequences with IS ratio below this
-```
-
- `importance_sampling_level: token` applies per-token IS ratios (recommended with Liger kernel)
- `importance_sampling_level: sequence` applies per-sequence IS ratios
- `off_policy_mask_threshold` masks out sequences where the IS ratio indicates they are too far off-policy
-
-##### Replay Buffer
-
-The replay buffer caches rollout groups that had learning signal (non-zero reward variance) and uses them to replace zero-signal groups in later batches.
-
-```yaml
-trl:
-  replay_buffer_size: 100       # Max cached groups (0 = disabled)
-  replay_recompute_logps: true  # Recompute log-probs for replayed data (recommended)
-```
-
-::: {.callout-note}
-When `replay_recompute_logps: true` (default), old log-probabilities are recomputed using the current model weights. This fixes the IS mismatch that would otherwise occur when replaying stale data.
-:::
-
-##### Deferred Re-rolling
-
-Failed prompts (where the model produces zero reward for all generations) are buffered and re-injected into later batches when the model may be better equipped to solve them.
-
-```yaml
-trl:
-  reroll_start_fraction: 0.5    # Start re-rolling after 50% of training
-  reroll_max_groups: 1          # Max groups to replace per batch
-```
-
-##### Zero-Advantage Batch Skipping
-
-When all advantages in a micro-batch are zero (no learning signal), the forward/backward pass is skipped entirely. This is enabled by default and logged as `skipped_zero_adv_batches=1`.
-
-```yaml
-trl:
-  skip_zero_advantage_batches: true   # default
-```
-
-##### Parallel Reward Workers
-
-Reward functions that use `signal.alarm()` (e.g., `math_verify`) must run in the main thread. Parallel reward workers use subprocesses to work around this limitation while enabling concurrent reward computation.
-
-```yaml
-trl:
-  reward_num_workers: 4         # Number of subprocess workers (1 = no parallelism)
-```
-
-##### Full Async GRPO Example
-
-```yaml
-base_model: Qwen/Qwen2.5-1.5B-Instruct
-
-vllm:
-    host: 0.0.0.0
-    port: 8000
-    gpu_memory_utilization: 0.35
-    dtype: auto
-
-adapter: lora
-lora_r: 32
-lora_alpha: 64
-lora_target_linear: true
-
-rl: grpo
-trl:
-  use_data_producer: true
-  use_vllm: true
-  async_prefetch: true
-  prefetch_depth: 1
-  vllm_sync_interval: 2
-  vllm_lora_sync: true
-  streaming_partial_batch: true
-  vllm_importance_sampling_correction: true
-  off_policy_mask_threshold: 0.5
-  importance_sampling_level: token
-  num_generations: 8
-  max_completion_length: 512
-  reward_funcs:
-    - rewards.accuracy_reward
-  reroll_start_fraction: 0.5
-  replay_buffer_size: 100
-  reward_num_workers: 4
-  skip_zero_advantage_batches: true
-
-datasets:
-  - path: AI-MO/NuminaMath-TIR
-    type: rewards.prompt_transform
-    split: train
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-max_steps: 500
-learning_rate: 1e-5
-bf16: true
-gradient_checkpointing: true
-```
-
-```bash
-# Terminal 1: Start vLLM on GPU 0
-CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
-
-# Terminal 2: Train on GPU 1
-CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml
-```
-
-##### Multi-GPU Async GRPO
-
-Async GRPO supports FSDP and DeepSpeed ZeRO-3 for multi-GPU training. vLLM runs on one GPU while training is distributed across the remaining GPUs.
-
-**FSDP:**
-
-```yaml
-fsdp:
-  - full_shard
-  - auto_wrap
-fsdp_config:
-  fsdp_transformer_layer_cls_to_wrap: Qwen2DecoderLayer
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-```
-
-**DeepSpeed ZeRO-3:**
-
-```yaml
-deepspeed: deepspeed_configs/zero3_bf16.json
-gradient_checkpointing_kwargs:
-  use_reentrant: true   # Required for ZeRO-3
-```
-
-```bash
-# Terminal 1: Start vLLM on GPU 0
-CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml
-
-# Terminal 2: Train on GPUs 0,1
-CUDA_VISIBLE_DEVICES=0,1 accelerate launch --num_processes 2 -m axolotl.cli.train config.yaml
-```
-
-::: {.callout-important}
-With multi-GPU async prefetch, only rank 0 generates completions in the background thread. Results are broadcast to all ranks on the main thread. This avoids FSDP/DeepSpeed collective deadlocks from unsynchronized background threads.
-:::
-
-### GDPO
-
-GDPO (Group Reward-Decoupled Policy Optimization) extends GRPO for multi-reward training. It addresses the **reward advantage collapse** problem by normalizing each reward function independently before combining them.
-
-::: {.callout-tip}
-Use GDPO when training with multiple reward functions. For single reward, GRPO and GDPO produce equivalent results.
-:::
-
-Paper: [https://arxiv.org/pdf/2501.05242](https://arxiv.org/pdf/2501.05242)
-
-GDPO uses TRL's native `multi_objective_aggregation` parameter under the hood. When you set `rl: gdpo`, axolotl automatically configures TRL to use `normalize_then_sum` aggregation.
-
-```yaml
-base_model: Qwen/Qwen2.5-1.5B-Instruct
-
-vllm:
-    host: 0.0.0.0
-    port: 8000
-    tensor_parallel_size: 2
-    gpu_memory_utilization: 0.85
-
-rl: gdpo
-
-trl:
-    beta: 0.001
-    max_completion_length: 256
-    use_vllm: true
-    num_generations: 4
-    reward_funcs:
-        - rewards.format_reward
-        - rewards.correctness_reward
-    reward_weights: [1.0, 2.0]
-
-datasets:
-    - path: openai/gsm8k
-      name: main
-      type: rewards.oai_gsm8k_transform
-```
-
-You can also use GRPO with explicit aggregation control:
-
-```yaml
-rl: grpo
-trl:
-    multi_objective_aggregation: normalize_then_sum  # GDPO behavior
-    # or: sum_then_normalize  # Default GRPO behavior
-```
-
-#### GDPO vs GRPO
-
-| Aspect | GRPO | GDPO |
-|--------|------|------|
-| **Aggregation** | `sum_then_normalize` | `normalize_then_sum` |
-| **Multi-reward** | May collapse advantages | Preserves reward signals |
-| **Single reward** | Standard behavior | Equivalent to GRPO |
-
-#### Why GDPO?
-
-When using multiple rewards with GRPO, different reward combinations can produce identical advantages:
-
-```
-# Example: format + correctness rewards
-[format=0, correct=3] → sum=3
-[format=1, correct=2] → sum=3  ← GRPO sees these as equal!
-[format=2, correct=1] → sum=3
-[format=3, correct=0] → sum=3
-```
-
-GDPO normalizes each reward independently, preserving their relative differences.
-
-#### Reward Functions
-
-GDPO uses the same reward function format as GRPO:
-
-```python
-# rewards.py
-def format_reward(completions, **kwargs) -> list[float]:
-    return [1.0 if len(c) > 10 else 0.0 for c in completions]
-
-def correctness_reward(completions, answers, **kwargs) -> list[float]:
-    rewards = []
-    for completion, answer in zip(completions, answers):
-        # Your scoring logic here
-        rewards.append(score)
-    return rewards
-```
-
-#### Sequence Parallelism
-
-GDPO supports sequence parallelism for long-context training:
-
-```yaml
-rl: gdpo
-context_parallel_size: 2
-```
-
 ### SimPO

 SimPO uses [CPOTrainer](https://huggingface.co/docs/trl/main/en/cpo_trainer) but with alternative loss function.
--- a/examples/apertus/README.md
+++ b/examples/apertus/README.md
@@ -15,7 +15,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl

-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'

 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
--- a/examples/arcee/README.md
+++ b/examples/arcee/README.md
@@ -17,7 +17,7 @@ Thanks to the team at Arcee.ai for using Axolotl in supervised fine-tuning the A
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl

-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'

 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
--- 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@63b15e6\""
+    "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@318b7e2\""
   ]
  },
  {
--- a/examples/devstral/README.md
+++ b/examples/devstral/README.md
@@ -16,7 +16,7 @@ Thanks to the team at MistralAI for giving us early access to prepare for this r

 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```

--- a/examples/devstral/devstral-small-qlora.yml
+++ b/examples/devstral/devstral-small-qlora.yml
@@ -52,7 +52,6 @@ gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
 flash_attention: true
-scaling_softmax: true

 loss_watchdog_threshold: 5.0
 loss_watchdog_patience: 3
--- a/examples/eaft/eaft-example.yml
+++ b/examples/eaft/eaft-example.yml
@@ -1,77 +0,0 @@
-base_model: google/gemma-3-1b-it
-
-model_type: Gemma3ForCausalLM
-cls_model_config: Gemma3TextConfig
-
-# gemma3 doesn't seem to play nice with ddp
-ddp_find_unused_parameters: true
-
-chat_template: gemma3
-eot_tokens:
-  - <end_of_turn>
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
-
-datasets:
-  - path: cgato/SlimOrcaDedupCleaned
-    type: chat_template
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-dataset_prepared_path:
-val_set_size: 0
-output_dir: ./outputs/eaft-gemma-3-1b
-
-use_eaft: true
-eaft_alpha: 1.0
-eaft_k: 20
-
-sequence_len: 1024
-sample_packing: false
-
-adapter:
-lora_model_dir:
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-eval_batch_size: 1
-max_steps: 1000
-evaluation_strategy: "no"
-optimizer: adamw_torch_fused
-lr_scheduler: cosine
-learning_rate: 5e-5
-
-train_on_inputs: false
-group_by_length: false
-bf16: auto
-fp16:
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-
-early_stopping_patience:
-resume_from_checkpoint:
-local_rank:
-logging_steps: 1
-xformers_attention:
-flash_attention: true
-
-warmup_ratio: 0.1
-weight_decay: 0.0
-debug:
-deepspeed:
-fsdp:
-fsdp_config:
-special_tokens:
--- a/examples/gemma3/gemma-3-1b-qlora.yml
+++ b/examples/gemma3/gemma-3-1b-qlora.yml
@@ -1,5 +1,7 @@
 base_model: google/gemma-3-1b-it

+model_type: Gemma3ForCausalLM
+
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name

@@ -24,15 +26,10 @@ datasets:
 val_set_size: 0.0
 output_dir: ./outputs/out

-# Freeze vision tower
-unfrozen_parameters:
-  - ^model\.language_model\..*
-  - ^lm_head\..*
-
 adapter: qlora
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0
+lora_dropout: 0.05
 lora_target_linear: true

 sequence_len: 2048
--- a/examples/gemma3/gemma-3-270m-qlora.yml
+++ b/examples/gemma3/gemma-3-270m-qlora.yml
@@ -1,5 +1,7 @@
 base_model: google/gemma-3-270m-it

+model_type: Gemma3ForCausalLM
+
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name

@@ -24,15 +26,10 @@ datasets:
 val_set_size: 0.0
 output_dir: ./outputs/out

-# Freeze vision tower
-unfrozen_parameters:
-  - ^model\.language_model\..*
-  - ^lm_head\..*
-
 adapter: qlora
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0
+lora_dropout: 0.05
 lora_target_linear: true

 sequence_len: 2048
--- a/examples/gemma3/gemma-3-4b-qlora.yml
+++ b/examples/gemma3/gemma-3-4b-qlora.yml
@@ -1,5 +1,8 @@
 base_model: google/gemma-3-4b-it

+# Need to set else transformers tries to load vision too
+model_type: Gemma3ForCausalLM
+
 load_in_4bit: true

 # gemma3 doesn't seem to play nice with ddp
@@ -20,11 +23,6 @@ dataset_prepared_path: last_run_prepared
 val_set_size: 0.01
 output_dir: ./outputs/out

-# Freeze vision tower
-unfrozen_parameters:
-  - ^model\.language_model\..*
-  - ^lm_head\..*
-
 adapter: qlora
 lora_model_dir:

@@ -34,8 +32,8 @@ sample_packing: true

 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0
-lora_target_linear: true
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'

 wandb_project:
 wandb_entity:
--- a/examples/gemma3/gemma-3-4b-vision-qlora.yml
+++ b/examples/gemma3/gemma-3-4b-vision-qlora.yml
@@ -31,7 +31,7 @@ pad_to_sequence_len: false

 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0
+lora_dropout: 0.05
 lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'

 wandb_project:
--- a/examples/gemma3n/README.md
+++ b/examples/gemma3n/README.md
@@ -10,7 +10,7 @@ Gemma-3n is a family of multimodal models from Google found on [HuggingFace](htt

 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```

--- a/examples/glm45/README.md
+++ b/examples/glm45/README.md
@@ -1,72 +0,0 @@
-# Finetune Z.ai's GLM-4.5-Air with Axolotl
-
-[GLM-4.5-Air](https://huggingface.co/zai-org/GLM-4.5-Air) is a MoE model by Z.ai.
-
-This guide shows how to fine-tune it with Axolotl.
-
-## Getting started
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
-
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
-
-3. Run the finetuning example:
-
-```bash
-# QLoRA (1x80GB @ ~63.4GiB/GPU)
-axolotl train examples/glm45/glm-45-air-qlora.yaml
-```
-
-### Dataset
-
-In addition to the standard OpenAI Messages format, GLM-4.5 supports an extra parameter for thinking in the assistant section.
-
-```json
-{
-    "role": "assistant",
-    "reasoning_content": "...",  // or have </think>...</think> in `content`
-    "content": "..."
-}
-```
-
-Make sure you set the below extra attributes if needed:
-
-```yaml
-datasets:
-  - path: ...
-    type: chat_template
-    message_property_mappings:
-      role: role
-      content: content
-
-    #   tool_calls: tool_calls  # uncomment if using tools
-    #   reasoning_content: reasoning_content  # uncomment if have reasoning
-
-# Uncomment if training on tool role (you would rarely if ever need this)
-# eot_tokens:
-#   - <|observation|>
-```
-
-### Tips
-
- The role name for tools in this template is `tool`.
- You will see this Axolotl WARNING — this is expected as the template does not use EOS:
-  ```
-  EOS token '<|endoftext|>' not found in chat_template. Please check if your template/EOS token is correct.
-  ```
- You can run a full finetuning by removing `adapter: qlora`, `load_in_4bit: true`, and `quantize_moe_experts: true` from the config.
- **LoRA kernels**: Incompatible with this model. Must be explicitly disabled (`lora_*_kernel: false`).
- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-
-## Optimization Guides
-
-Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
-
-## Related Resources
-
- [GLM-4.5-Air on HuggingFace](https://huggingface.co/zai-org/GLM-4.5-Air)
- [GLM-4.5 Blog](https://z.ai/blog/glm-4.5)
- [Axolotl Docs](https://docs.axolotl.ai)
- [Axolotl Website](https://axolotl.ai)
- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
--- a/examples/glm45/glm-45-air-qlora.yaml
+++ b/examples/glm45/glm-45-air-qlora.yaml
@@ -1,64 +0,0 @@
-base_model: zai-org/GLM-4.5-Air
-
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_8bit: false
-load_in_4bit: true
-
-quantize_moe_experts: true # important
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/lora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 16
-lora_alpha: 8
-lora_dropout: 0
-lora_target_modules:
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
-gradient_accumulation_steps: 2
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
--- a/examples/glm46v/README.md
+++ b/examples/glm46v/README.md
@@ -1,44 +0,0 @@
-# Finetune GLM-4.6V with Axolotl
-
-GLM-4.6V is a family of vision-language models from ZhipuAI found on [HuggingFace](https://huggingface.co/zai-org/GLM-4.6V). This guide shows how to fine-tune it with Axolotl for vision-language tasks.
-
-
-
-## Getting started
-
-1. Install Axolotl from source following the [installation guide](https://docs.axolotl.ai/docs/installation.html#sec-edge-build).
-
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
-
-
-3. Run the fine-tuning:
-
-    glm-4-6v-flash(9B)
-    ```bash
-    axolotl train examples/glm46v/glm-4-6v-flash-qlora.yaml
-    ```
-
-Let us know how it goes. Happy finetuning! 🚀
-
-## Tips
-
- Vision datasets should follow the format described in the [multimodal docs](https://docs.axolotl.ai/docs/multimodal.html#dataset-format)
- You can run a **full finetuning** by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
- Read more on how to load your own dataset in the [dataset loading docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-
-## Supported Models
-
- **GLM-4.6V**: Full vision-language model (`zai-org/GLM-4.6V`)
- **GLM-4.6V-Flash**: Faster variant (`zai-org/GLM-4.6V-Flash`)
-
-## Optimization Guides
-
-Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
-
-## Related Resources
-
- [ZhipuAI GLM-4.6V](https://huggingface.co/zai-org/GLM-4.6V)
- [Axolotl Docs](https://docs.axolotl.ai)
- [Axolotl Website](https://axolotl.ai)
- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
--- a/examples/glm46v/glm-4-6v-flash-ddp.yaml
+++ b/examples/glm46v/glm-4-6v-flash-ddp.yaml
@@ -1,53 +0,0 @@
-base_model: zai-org/GLM-4.6V-Flash
-trust_remote_code: true
-
-processor_type: AutoProcessor
-load_in_4bit: true
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-ddp_find_unused_parameters: true
-
-output_dir: ./outputs/glm-4-6v-flash-qlora
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-sequence_len: 2048
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-logging_steps: 1
-sdp_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 0
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/glm46v/glm-4-6v-flash-qlora.yaml
+++ b/examples/glm46v/glm-4-6v-flash-qlora.yaml
@@ -1,50 +0,0 @@
-base_model: zai-org/GLM-4.6V-Flash
-trust_remote_code: true
-
-processor_type: AutoProcessor
-load_in_4bit: true
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-output_dir: ./outputs/glm-4-6v-flash-qlora
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-sequence_len: 2048
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-logging_steps: 1
-sdp_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 0
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/glm47-flash/README.md
+++ b/examples/glm47-flash/README.md
@@ -1,65 +0,0 @@
-# Finetune Z.ai's GLM-4.7-Flash with Axolotl
-
-[GLM-4.7-Flash](https://huggingface.co/zai-org/GLM-4.7-Flash) is a 30B-A3B MoE model by Z.ai.
-
-This guide shows how to fine-tune it with Axolotl.
-
-## Getting started
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
-
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
-
-3. Run the finetuning example:
-
-```bash
-# QLoRA
-# - no target experts (1x48GB @ ~24GiB/GPU)
-# - target experts (1x48GB @ ~34GiB/GPU)
-axolotl train examples/glm47-flash/qlora.yaml
-
-# QLoRA FSDP2 no target experts (2x48GB @ ~29GiB/GPU)
-axolotl train examples/glm47-flash/qlora_fsdp.yaml
-```
-
-```bash
-# LoRA
-# - no target experts (1x48GB @ ~35GiB/GPU)
-# - target experts (1x48GB @ OOM. Projected ~45-50GiB/GPU)
-axolotl train examples/glm47-flash/lora.yaml
-
-# LoRA FSDP2 no target experts (2x48GB @ ~43GiB/GPU)
-axolotl train examples/glm47-flash/lora_fsdp.yaml
-```
-
-### MoE Expert Quantization & Expert LoRA
-
-This model quantize expert weights on load. To learn about expert quantization, expert LoRA targeting, and related limitations, see the [MoE Expert Quantization](https://docs.axolotl.ai/docs/expert_quantization.html) docs.
-
-## Limitations
-
- **lora_target_linear**: Incompatible for this model.
- **LoRA kernels**: Incompatible with this model due to non-standard attention projections (DSA). Must be explicitly disabled (`lora_*_kernel: false`).
-
-
-### TIPS
-
- For inference, the official Z.ai team recommends these default settings (most tasks):
-  - `temperature: 1.0`
-  - `top_p: 0.95`
-  - `max_new_tokens: 131072`
- You can run a full finetuning by removing `adapter: qlora`, `load_in_4bit: true`, and `quantize_moe_experts: true` from the config. This is heavy, so we have not tested this.
- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-
-## Optimization Guides
-
-Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
-
-## Related Resources
-
- [GLM-4.7-Flash on HuggingFace](https://huggingface.co/zai-org/GLM-4.7-Flash)
- [GLM-4.7 Blog](https://z.ai/blog/glm-4.7)
- [Axolotl Docs](https://docs.axolotl.ai)
- [Axolotl Website](https://axolotl.ai)
- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
--- a/examples/glm47-flash/lora.yaml
+++ b/examples/glm47-flash/lora.yaml
@@ -1,65 +0,0 @@
-base_model: zai-org/GLM-4.7-Flash
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_8bit: true
-quantize_moe_experts: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/glm4.7-flash-lora-8bit-out
-
-adapter: lora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0
-lora_target_modules:
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-# Uncomment to also target MoE expert weights:
-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-# LoRA kernels incompatible with DSA attention
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_torch_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
--- a/examples/glm47-flash/lora_fsdp.yaml
+++ b/examples/glm47-flash/lora_fsdp.yaml
@@ -1,75 +0,0 @@
-base_model: zai-org/GLM-4.7-Flash
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_8bit: true
-quantize_moe_experts: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/glm4.7-flash-lora-8bit-fsdp-out
-
-adapter: lora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0
-lora_target_modules:
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-# Uncomment to also target MoE expert weights:
-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-# LoRA kernels incompatible with DSA attention
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_torch_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-fsdp_config:
-  fsdp_version: 2
-  offload_params: false
-  cpu_ram_efficient_loading: false
-  auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  transformer_layer_cls_to_wrap: Glm4MoeLiteDecoderLayer
-  state_dict_type: FULL_STATE_DICT
-  sharding_strategy: FULL_SHARD
-  reshard_after_forward: true
-  activation_checkpointing: true
--- a/examples/glm47-flash/qlora.yaml
+++ b/examples/glm47-flash/qlora.yaml
@@ -1,65 +0,0 @@
-base_model: zai-org/GLM-4.7-Flash
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_4bit: true
-quantize_moe_experts: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/glm4.7-flash-qlora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0
-lora_target_modules:
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-# Uncomment to also target MoE expert weights:
-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-# LoRA kernels incompatible with DSA attention
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_torch_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
--- a/examples/glm47-flash/qlora_fsdp.yaml
+++ b/examples/glm47-flash/qlora_fsdp.yaml
@@ -1,75 +0,0 @@
-base_model: zai-org/GLM-4.7-Flash
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_4bit: true
-quantize_moe_experts: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/glm4.7-flash-qlora-fsdp-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0
-lora_target_modules:
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-# Uncomment to also target MoE expert weights:
-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-# LoRA kernels incompatible with DSA attention
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_torch_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-fsdp_config:
-  fsdp_version: 2
-  offload_params: false
-  cpu_ram_efficient_loading: false
-  auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  transformer_layer_cls_to_wrap: Glm4MoeLiteDecoderLayer
-  state_dict_type: FULL_STATE_DICT
-  sharding_strategy: FULL_SHARD
-  reshard_after_forward: true
-  activation_checkpointing: true
--- a/examples/gpt-oss/README.md
+++ b/examples/gpt-oss/README.md
@@ -14,7 +14,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations

 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```

--- a/examples/granite4/README.md
+++ b/examples/granite4/README.md
@@ -15,7 +15,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl

-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'

 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
--- a/examples/hunyuan/README.md
+++ b/examples/hunyuan/README.md
@@ -13,7 +13,7 @@ Tencent released a family of opensource models called HunYuan with varying param
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl

-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'

 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
--- a/examples/jamba/qlora_fsdp_large.yaml
+++ b/examples/jamba/qlora_fsdp_large.yaml
@@ -19,6 +19,7 @@ datasets:
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
 output_dir: jamba-large-fsdp-qlora-ft
+save_safetensors: true
 adapter: qlora
 sequence_len: 2048
 sample_packing: true
--- a/examples/llama-3/3b-qat-mxfp4.yaml
+++ b/examples/llama-3/3b-qat-mxfp4.yaml
@@ -1,65 +0,0 @@
-base_model: meta-llama/Llama-3.2-3B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
-
-plugins:
-  - axolotl.integrations.liger.LigerPlugin
-
-liger_rope: true
-liger_rms_norm: true
-liger_glu_activation: true
-liger_layer_norm: true
-liger_fused_linear_cross_entropy: true
-
-datasets:
-  - path: yahma/alpaca-cleaned
-    type: alpaca
-    split: train[:95%]
-
-output_dir: ./outputs/qat_out/
-dataset_prepared_path: ./outputs/dataset_prepared
-
-sequence_len: 2048
-flash_attention: true
-
-qat:
-  activation_dtype: mxfp4
-  weight_dtype: mxfp4
-  group_size: 32
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_checkpointing: true
-activation_offloading: true
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_8bit
-
-cosine_constant_lr_ratio: 0
-cosine_min_lr_ratio: 1.0
-learning_rate: 2e-5
-save_only_model: true
-bf16: true
-
-resume_from_checkpoint:
-logging_steps: 1
-
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-warmup_ratio: 0.1
-weight_decay: 0.0
-
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config
--- a/examples/llama-3/qlora-1b-gdpo.yaml
+++ b/examples/llama-3/qlora-1b-gdpo.yaml
@@ -1,68 +0,0 @@
-base_model: meta-llama/Llama-3.2-1B-Instruct
-
-chat_template: llama3
-
-rl: gdpo
-
-trl:
-  beta: 0.001
-  max_completion_length: 128
-  num_generations: 2
-  temperature: 0.7
-  top_p: 0.95
-
-  use_vllm: false
-
-
-  multi_objective_aggregation: normalize_then_sum
-
-  reward_funcs:
-    - rwd.format_reward
-    - rwd.correctness_reward
-  reward_weights: [1.0, 2.0]
-
-  log_completions: true
-  num_completions_to_print: 3
-  scale_rewards: true
-
-datasets:
-  - path: openai/gsm8k
-    name: main
-    split: train[:1000]
-    type: rwd.gsm8k_transform
-
-val_set_size: 0.0
-output_dir: ./outputs/llama3-gdpo-out
-
-sequence_len: 512
-sample_packing: false
-pad_to_sequence_len: false
-
-gradient_accumulation_steps: 8
-micro_batch_size: 1
-num_epochs: 1
-max_steps: 100
-
-optimizer: adamw_torch_fused
-lr_scheduler: cosine
-learning_rate: 5e-5
-weight_decay: 0.01
-warmup_steps: 10
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-
-flash_attention: true
-logging_steps: 1
-save_steps: 50
-save_safetensors: true
-
-special_tokens:
-  pad_token: "<|end_of_text|>"
-
-
-seed: 42
--- a/examples/llama-3/qlora-fsdp-405b.yaml
+++ b/examples/llama-3/qlora-fsdp-405b.yaml
@@ -12,6 +12,7 @@ datasets:
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
 output_dir: ./outputs/out/qlora-llama3_1-405b
+save_safetensors: true

 adapter: qlora

--- a/examples/magistral/README.md
+++ b/examples/magistral/README.md
@@ -14,7 +14,7 @@ Thanks to the team at MistralAI for giving us early access to prepare for these

 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.7.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```

--- a/examples/mamba/config.yml
+++ b/examples/mamba/config.yml
@@ -47,5 +47,6 @@ saves_per_epoch: 1
 weight_decay: 0.0
 special_tokens:
 tokens:
+save_safetensors: False

 # save_first_step: true  # uncomment this to validate checkpoint saving works with your config
--- a/examples/ministral3/ministral3-3b-qlora.yaml
+++ b/examples/ministral3/ministral3-3b-qlora.yaml
@@ -59,7 +59,6 @@ gradient_checkpointing: true
 resume_from_checkpoint:
 logging_steps: 1
 flash_attention: true
-scaling_softmax: true

 warmup_ratio: 0.1
 evals_per_epoch: 1
--- a/examples/mistral4/README.md
+++ b/examples/mistral4/README.md
@@ -1,85 +0,0 @@
-# Finetune Mistral Small 4 with Axolotl
-
-Mistral Small 4 is a 119B parameter (6.5B active) multimodal MoE model from MistralAI that unifies instruct, reasoning, and coding capabilities into a single model. It is available on HuggingFace at [Mistral-Small-4-119B-2603](https://huggingface.co/mistralai/Mistral-Small-4-119B-2603).
-
-Thanks to the team at MistralAI for giving us early access to prepare for this release.
-
-## Getting started
-
-Note: Training this model requires weights in BF16 which we will link to later.
-Users interested in training can convert / descale the existing FP8 weights.
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
-
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage
-
-3. Install transformers from main
-
-  ```bash
-  pip install git+https://github.com/huggingface/transformers.git
-  ```
-
-4. Run one of the example configs:
-
-  ```bash
-  # text-only
-  axolotl train examples/mistral4/qlora-text.yml  # no experts ~69 GiB, experts ~93 GiB
-  axolotl train examples/mistral4/fft-text.yml
-
-  # text + vision
-  # run: wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
-  axolotl train examples/mistral4/qlora-vision.yml  # no experts ~68 GiB
-  axolotl train examples/mistral4/fft-vision.yml
-  ```
-
-Note: FFT configs provided as reference. Please adjust hyperparameters as needed.
-
-## Reasoning Effort
-
-The chat template supports a `reasoning_effort` variable to control the model's reasoning depth:
-
- `"none"` — instruct mode (default)
- `"high"` — reasoning mode with explicit thinking steps
-
-Pass it via `chat_template_kwargs` under your dataset config:
-
-```yaml
-datasets:
-  - path: your/dataset
-    type: chat_template
-    chat_template_kwargs:
-      reasoning_effort: high
-```
-
-## Thinking Support
-
-The chat template supports a `thinking` content type in assistant messages for training on reasoning traces (rendered as `[THINK]...[/THINK]` blocks).
-
-To use thinking datasets, add the `thinking` mapping via `message_property_mappings`:
-
-```yaml
-datasets:
-  - path: your/thinking-dataset
-    type: chat_template
-    message_property_mappings:
-      role: role
-      content: content
-      thinking: thinking
-    chat_template_kwargs:
-      reasoning_effort: high
-```
-
-See the [Magistral thinking guide](../magistral/think/README.md) for dataset format details.
-
-## Tips
-
- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
- The text dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
- The vision model requires multi-modal dataset format as documented [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
-
-## Related Resources
-
- [MistralAI Mistral Small 4 Blog](https://mistral.ai/news/mistral-small-4)
- [Axolotl Docs](https://docs.axolotl.ai)
- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
--- a/examples/mistral4/fft-text.yml
+++ b/examples/mistral4/fft-text.yml
@@ -1,58 +0,0 @@
-base_model: mistralai/Mistral-Small-4-119B-2603
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-  - axolotl.integrations.kernels.KernelsPlugin
-use_kernels: true
-use_sonicmoe: true
-
-# only train language model layers, freeze vision tower
-unfrozen_parameters:
-  - model.language_model.*
-  - lm_head
-  - embed_tokens
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-output_dir: ./outputs/out
-
-sequence_len: 2048
-sample_packing: true
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 2e-5
-
-bf16: true
-tf32: true
-
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
-
-fsdp_version: 2
-fsdp_config:
-  offload_params: false
-  cpu_ram_efficient_loading: false
-  state_dict_type: FULL_STATE_DICT
-  auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  transformer_layer_cls_to_wrap: Mistral4DecoderLayer
-  reshard_after_forward: true
-  activation_checkpointing: true
--- a/examples/mistral4/fft-vision.yml
+++ b/examples/mistral4/fft-vision.yml
@@ -1,57 +0,0 @@
-base_model: mistralai/Mistral-Small-4-119B-2603
-processor_type: AutoProcessor
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-  - axolotl.integrations.kernels.KernelsPlugin
-use_kernels: true
-use_sonicmoe: true
-
-# vision requirements
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-datasets:
-  - path: Nanobit/text-vision-2k-test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-output_dir: ./outputs/out
-
-sequence_len: 2048
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 2e-5
-
-bf16: true
-tf32: true
-
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
-
-fsdp_version: 2
-fsdp_config:
-  offload_params: false
-  cpu_ram_efficient_loading: false
-  state_dict_type: FULL_STATE_DICT
-  auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  transformer_layer_cls_to_wrap: Mistral4DecoderLayer
-  reshard_after_forward: true
-  activation_checkpointing: true
--- a/examples/mistral4/qlora-text.yml
+++ b/examples/mistral4/qlora-text.yml
@@ -1,58 +0,0 @@
-base_model: mistralai/Mistral-Small-4-119B-2603
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_4bit: true
-quantize_moe_experts: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-output_dir: ./outputs/out
-
-adapter: qlora
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
-
-# uncomment to train on expert layers
-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-# lora_mlp_kernel: false
-# lora_qkv_kernel: false
-# lora_o_kernel: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: true
-tf32: true
-
-gradient_checkpointing: true
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/mistral4/qlora-vision.yml
+++ b/examples/mistral4/qlora-vision.yml
@@ -1,63 +0,0 @@
-base_model: mistralai/Mistral-Small-4-119B-2603
-processor_type: AutoProcessor
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_4bit: true
-quantize_moe_experts: true
-
-# vision chat template requirements
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-datasets:
-  - path: Nanobit/text-vision-2k-test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-output_dir: ./outputs/out
-
-adapter: qlora
-
-sequence_len: 2048
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
-
-# uncomment to train on expert layers
-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-# lora_mlp_kernel: false
-# lora_qkv_kernel: false
-# lora_o_kernel: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: true
-tf32: true
-
-gradient_checkpointing: true
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/nemotron/nemotron-mini-4b-qlora.yaml
+++ b/examples/nemotron/nemotron-mini-4b-qlora.yaml
@@ -1,57 +0,0 @@
-base_model: nvidia/Nemotron-Mini-4B-Instruct
-
-load_in_8bit: false
-load_in_4bit: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/nemotron-mini-4b-qlora
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 4096
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-lora_target_modules:
-  - q_proj
-  - k_proj
-  - v_proj
-  - o_proj
-  - up_proj
-  - down_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-special_tokens:
--- a/examples/qwen3-next/README.md
+++ b/examples/qwen3-next/README.md
@@ -6,13 +6,30 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations

 ## Getting started

-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as Qwen3-Next is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).

-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
+    Here is an example of how to install from main for pip:
+
+```bash
+# Ensure you have Pytorch installed (Pytorch 2.6.0 min)
+git clone https://github.com/axolotl-ai-cloud/axolotl.git
+cd axolotl
+
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
+pip3 install --no-build-isolation -e '.[flash-attn]'
+
+# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
+python scripts/cutcrossentropy_install.py | sh
+```
+
+2. Install Qwen3-Next transformers commit
+```bash
+pip3 uninstall -y transformers && pip3 install "git+https://github.com/huggingface/transformers.git@b9282355bea846b54ed850a066901496b19da654"
+```

 3. Install FLA for improved performance
 ```bash
-pip3 uninstall -y causal-conv1d && pip3 install flash-linear-attention==0.4.1
+pip3 uninstall -y causal-conv1d && pip3 install flash-linear-attention==0.3.2
 ```

 4. Run the finetuning example:
@@ -21,7 +38,7 @@ pip3 uninstall -y causal-conv1d && pip3 install flash-linear-attention==0.4.1
 axolotl train examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml
 ```

-This config uses about ~47 GiB (no target experts) and ~71GiB (target experts) VRAM.
+This config uses about 45.62 GiB VRAM.

 Let us know how it goes. Happy finetuning! 🚀

--- a/examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml
+++ b/examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml
@@ -9,8 +9,6 @@ plugins:
 load_in_8bit: false
 load_in_4bit: true

-quantize_moe_experts: true
-
 datasets:
  - path: fozziethebeat/alpaca_messages_2k_test
    type: chat_template
@@ -27,7 +25,7 @@ sample_packing: true

 lora_r: 16
 lora_alpha: 8
-lora_dropout: 0
+lora_dropout: 0.05
 lora_target_modules:
  - linear_attn.in_proj_ba
  - linear_attn.in_proj_qkvz
@@ -36,19 +34,12 @@ lora_target_modules:
  - shared_expert.down_proj
  - shared_expert.gate_proj
  - shared_expert_gate
+  - mlp.gate
  - q_proj
  - v_proj
  - k_proj
  - o_proj

-# lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
 wandb_project:
 wandb_entity:
 wandb_watch:
--- a/examples/qwen3.5/122b-a10b-moe-qlora.yaml
+++ b/examples/qwen3.5/122b-a10b-moe-qlora.yaml
@@ -1,71 +0,0 @@
-base_model: Qwen/Qwen3.5-122B-A10B
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-strict: false
-
-chat_template: qwen3_5
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-val_set_size: 0.0
-output_dir: ./outputs/out
-dataset_prepared_path: last_run_prepared
-
-sequence_len: 2048
-sample_packing: true
-
-load_in_4bit: true
-quantize_moe_experts: true
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0
-lora_target_modules:
-  - q_proj
-  - k_proj
-  - v_proj
-  - o_proj
-
-#lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_4bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 4
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
--- a/examples/qwen3.5/27b-fft.yaml
+++ b/examples/qwen3.5/27b-fft.yaml
@@ -1,59 +0,0 @@
-base_model: Qwen/Qwen3.5-27B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-# Full fine-tune (FFT) of the text-only path of Qwen3.5-27B.
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-strict: false
-
-chat_template: qwen3_5
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-val_set_size: 0.0
-output_dir: ./outputs/out
-dataset_prepared_path: last_run_prepared
-
-sequence_len: 2048
-sample_packing: true
-
-# Freeze vision encoder
-unfrozen_parameters:
-  - model\.language_model\..*
-  - lm_head\..*
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 4
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
--- a/examples/qwen3.5/27b-qlora.yaml
+++ b/examples/qwen3.5/27b-qlora.yaml
@@ -1,72 +0,0 @@
-base_model: Qwen/Qwen3.5-27B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-# Note: Qwen3.5 is an early-fusion VLM (image+text). This config fine-tunes
-# the text-only path. For multimodal (image+text) fine-tuning, add image
-# columns to your dataset following axolotl's multimodal dataset format.
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-strict: false
-
-chat_template: qwen3_5
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-val_set_size: 0.0
-output_dir: ./outputs/out
-dataset_prepared_path: last_run_prepared
-
-sequence_len: 2048
-sample_packing: true
-
-load_in_4bit: true
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_target_modules:
-  - q_proj
-  - k_proj
-  - v_proj
-  - o_proj
-  - down_proj
-  - up_proj
-  # Uncomment below to also target the linear attention projections.
-  # These use separate in_proj_qkv / in_proj_z / out_proj (Qwen3.5-specific).
-  # - linear_attn.in_proj_qkv
-  # - linear_attn.in_proj_z
-  # - linear_attn.out_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_4bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 4
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
--- a/examples/qwen3.5/35b-a3b-moe-qlora.yaml
+++ b/examples/qwen3.5/35b-a3b-moe-qlora.yaml
@@ -1,70 +0,0 @@
-base_model: Qwen/Qwen3.5-35B-A3B
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-strict: false
-
-chat_template: qwen3_5
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-val_set_size: 0.0
-output_dir: ./outputs/out
-dataset_prepared_path: last_run_prepared
-
-sequence_len: 2048
-sample_packing: true
-
-load_in_4bit: true
-quantize_moe_experts: true
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0
-lora_target_modules:
-  - q_proj
-  - k_proj
-  - v_proj
-  - o_proj
-
-#lora_target_parameters:
-#   - mlp.experts.gate_up_proj
-#   - mlp.experts.down_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_4bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-lora_mlp_kernel: false
-lora_qkv_kernel: false
-lora_o_kernel: false
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 4
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
--- a/examples/qwen3.5/9b-fft-vision.yaml
+++ b/examples/qwen3.5/9b-fft-vision.yaml
@@ -1,49 +0,0 @@
-base_model: Qwen/Qwen3.5-9B
-processor_type: AutoProcessor
-
-# Required for multimodal training
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-chat_template: qwen3_5
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-sequence_len: 4096
-pad_to_sequence_len: false
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
--- a/examples/qwen3.5/9b-lora-vision.yaml
+++ b/examples/qwen3.5/9b-lora-vision.yaml
@@ -1,68 +0,0 @@
-base_model: Qwen/Qwen3.5-9B
-processor_type: AutoProcessor
-
-# These 3 lines are required for vision/multimodal training
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-chat_template: qwen3_5
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-adapter: lora
-lora_model_dir:
-
-sequence_len: 8192
-pad_to_sequence_len: false
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-# Targets the language model attention and MLP layers.
-# Qwen3.5 is early-fusion: all layers (including those seeing vision tokens) share
-# the same transformer stack, so standard attention targets work for both modalities.
-lora_target_modules:
-  - q_proj
-  - k_proj
-  - v_proj
-  - o_proj
-  - down_proj
-  - up_proj
-  # Uncomment to also target the linear attention (GatedDeltaNet) projections:
-  # - linear_attn.in_proj_qkv
-  # - linear_attn.in_proj_z
-  # - linear_attn.out_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: true
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/qwen3.5/README.md
+++ b/examples/qwen3.5/README.md
@@ -1,74 +0,0 @@
-# Finetune Qwen3.5 with Axolotl
-
-[Qwen3.5](https://huggingface.co/collections/Qwen/qwen35) is a hybrid architecture model series combining Gated DeltaNet linear attention with standard Transformer attention. All Qwen3.5 models are early-fusion vision-language models: dense variants use `Qwen3_5ForConditionalGeneration` and MoE variants use `Qwen3_5MoeForConditionalGeneration`.
-
-Vision and text tokens are processed through the same transformer stack. The configs below train on text-only data unless noted otherwise. See `9b-lora-vision.yaml` for a multimodal example.
-
-Available configs:
-
-| Config | Model | Type | Peak VRAM |
-|---|---|---|---|
-| `27b-qlora.yaml` | Qwen3.5-27B | Dense VLM, text-only QLoRA | ~47 GiB |
-| `27b-fft.yaml` | Qwen3.5-27B | Dense VLM, text-only FFT (vision frozen) | ~53 GiB |
-| `35b-a3b-moe-qlora.yaml` | Qwen3.5-35B-A3B | MoE, text-only QLoRA | — |
-| `122b-a10b-moe-qlora.yaml` | Qwen3.5-122B-A10B | MoE, text-only QLoRA | — |
-| `9b-lora-vision.yaml` | Qwen3.5-9B | Vision+text LoRA, single GPU | — |
-| `9b-fft-vision.yaml` | Qwen3.5-9B | Vision+text FFT, single GPU | ~61 GiB |
-
-
-## Getting started
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
-
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
-
-3. Install FLA for sample packing support with the Gated DeltaNet linear attention layers:
-```bash
-pip3 uninstall -y causal-conv1d && pip3 install flash-linear-attention==0.4.1
-```
-> FLA is required when `sample_packing: true`. Without it, training raises a `RuntimeError` on packed sequences. Vision configs use `sample_packing: false` so FLA is optional there.
-
-4. Run a finetuning example:
-
-```bash
-# Dense 27B text-only (QLoRA, ~47 GiB VRAM with sample packing)
-axolotl train examples/qwen3.5/27b-qlora.yaml
-
-# Dense 27B text-only FFT with vision encoder frozen (~53 GiB, single 80 GiB GPU)
-axolotl train examples/qwen3.5/27b-fft.yaml
-
-# MoE 35B-A3B text-only (QLoRA)
-axolotl train examples/qwen3.5/35b-a3b-moe-qlora.yaml
-
-# MoE 122B-A10B text-only (QLoRA)
-axolotl train examples/qwen3.5/122b-a10b-moe-qlora.yaml
-
-# 9B vision+text (LoRA, multimodal dataset)
-axolotl train examples/qwen3.5/9b-lora-vision.yaml
-
-# 9B vision+text FFT, single 80 GiB GPU (~61 GiB peak)
-axolotl train examples/qwen3.5/9b-fft-vision.yaml
-
-```
-
-### TIPS
-
- For inference, you can experiment with `temperature: 0.7`, `top_p: 0.8`, `top_k: 20`, and `min_p: 0`.
- For **text-only FFT** on 27B, use `27b-fft.yaml` which sets `unfrozen_parameters` to freeze the vision encoder (`model.visual.*`) — this avoids wasting optimizer state on parameters that receive no gradient from text-only data.
- You can run a full finetuning of smaller configs by removing `adapter: qlora` and `load_in_4bit: true`. See [Multi-GPU](#optimization-guides) below.
- Read more on loading your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
- For **multimodal** finetuning, set `processor_type: AutoProcessor`, `skip_prepare_dataset: true`, and `remove_unused_columns: false` as shown in `9b-lora-vision.yaml`.
- The Gated DeltaNet linear attention layers (`linear_attn.*`) can optionally be added to `lora_target_modules` — they are commented out by default.
-
-## Optimization Guides
-
- [Optimizations Guide](https://docs.axolotl.ai/docs/optimizations.html)
-
-## Related Resources
-
- [Qwen3.5 Blog](https://qwenlm.github.io/blog/qwen3.5/)
- [Axolotl Docs](https://docs.axolotl.ai)
- [Axolotl Website](https://axolotl.ai)
- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)
--- a/examples/swanlab/README.md
+++ b/examples/swanlab/README.md
@@ -1,285 +0,0 @@
-# SwanLab Integration Examples
-
-This directory contains example configurations demonstrating SwanLab integration with Axolotl.
-
-## Examples Overview
-
-### 1. DPO with Completion Logging
-**File**: `dpo-swanlab-completions.yml`
-
-Demonstrates DPO (Direct Preference Optimization) training with RLHF completion table logging.
-
-**Features**:
- Basic SwanLab experiment tracking
- Completion table logging (prompts, chosen/rejected responses, rewards)
- Memory-bounded buffer for long training runs
- Cloud sync configuration
-
-**Best for**: RLHF practitioners who want to analyze model outputs qualitatively
-
-**Quick start**:
-```bash
-export SWANLAB_API_KEY=your-api-key
-accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-completions.yml
-```
-
---
-
-### 2. LoRA with Performance Profiling
-**File**: `lora-swanlab-profiling.yml`
-
-Demonstrates standard LoRA fine-tuning with performance profiling enabled.
-
-**Features**:
- SwanLab experiment tracking
- Automatic profiling of trainer methods
- Profiling metrics visualization
- Performance optimization guidance
-
-**Best for**: Engineers optimizing training performance and comparing different configurations
-
-**Quick start**:
-```bash
-export SWANLAB_API_KEY=your-api-key
-accelerate launch -m axolotl.cli.train examples/swanlab/lora-swanlab-profiling.yml
-```
-
---
-
-### 3. Full-Featured DPO Production Setup
-**File**: `dpo-swanlab-full-featured.yml`
-
-Comprehensive production-ready configuration with ALL SwanLab features enabled.
-
-**Features**:
- Experiment tracking with team workspace
- RLHF completion logging
- Performance profiling
- Lark (Feishu) team notifications
- Private deployment support
- Production checklist and troubleshooting
-
-**Best for**: Production RLHF training with team collaboration
-
-**Quick start**:
-```bash
-export SWANLAB_API_KEY=your-api-key
-export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
-export SWANLAB_LARK_SECRET=your-webhook-secret
-accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-full-featured.yml
-```
-
---
-
-### 4. Custom Trainer Profiling (Python)
-**File**: `custom_trainer_profiling.py`
-
-Python code examples showing how to add SwanLab profiling to custom trainers.
-
-**Features**:
- `@swanlab_profile` decorator examples
- Context manager profiling for fine-grained timing
- `ProfilingConfig` for advanced filtering and throttling
- Multiple profiling patterns and best practices
-
-**Best for**: Advanced users creating custom trainers
-
-**Usage**:
-```python
-from custom_trainer_profiling import CustomTrainerWithProfiling
-# See file for detailed examples and patterns
-```
-
---
-
-## Feature Matrix
-
-| Example | Tracking | Completion Logging | Profiling | Lark Notifications | Team Workspace |
-|---------|----------|-------------------|-----------|-------------------|----------------|
-| dpo-swanlab-completions.yml | ✅ | ✅ | ✅ (auto) | ➖ (commented) | ➖ (commented) |
-| lora-swanlab-profiling.yml | ✅ | ➖ (disabled) | ✅ (auto) | ➖ (commented) | ➖ (commented) |
-| dpo-swanlab-full-featured.yml | ✅ | ✅ | ✅ (auto) | ✅ | ✅ |
-| custom_trainer_profiling.py | N/A | N/A | ✅ (manual) | N/A | N/A |
-
---
-
-## Configuration Quick Reference
-
-### Basic SwanLab Setup
-```yaml
-plugins:
-  - axolotl.integrations.swanlab.SwanLabPlugin
-
-use_swanlab: true
-swanlab_project: my-project
-swanlab_experiment_name: my-experiment
-swanlab_mode: cloud  # cloud, local, offline, disabled
-```
-
-### RLHF Completion Logging
-```yaml
-swanlab_log_completions: true
-swanlab_completion_log_interval: 100  # Log every 100 steps
-swanlab_completion_max_buffer: 128    # Memory-bounded buffer
-```
-
-### Lark Team Notifications
-```yaml
-swanlab_lark_webhook_url: https://open.feishu.cn/...
-swanlab_lark_secret: your-webhook-secret  # Required for production
-```
-
-### Team Workspace
-```yaml
-swanlab_workspace: my-research-team
-```
-
-### Private Deployment
-```yaml
-swanlab_web_host: https://swanlab.yourcompany.com
-swanlab_api_host: https://api.swanlab.yourcompany.com
-```
-
---
-
-## Authentication
-
-### Recommended: Environment Variable
-```bash
-export SWANLAB_API_KEY=your-api-key
-export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
-export SWANLAB_LARK_SECRET=your-webhook-secret
-```
-
-### Alternative: Config File (less secure)
-```yaml
-swanlab_api_key: your-api-key
-swanlab_lark_webhook_url: https://open.feishu.cn/...
-swanlab_lark_secret: your-webhook-secret
-```
-
---
-
-## Common Use Cases
-
-### Use Case 1: Migrate from WandB to SwanLab
-Start with `lora-swanlab-profiling.yml`, add your model/dataset config, disable WandB:
-```yaml
-use_swanlab: true
-use_wandb: false
-```
-
-### Use Case 2: Analyze DPO Model Outputs
-Use `dpo-swanlab-completions.yml`, adjust completion logging interval based on your training length:
-```yaml
-swanlab_completion_log_interval: 50   # More frequent for short training
-swanlab_completion_log_interval: 200  # Less frequent for long training
-```
-
-### Use Case 3: Optimize Training Performance
-Use `lora-swanlab-profiling.yml`, run multiple experiments with different optimizations:
- Baseline: `flash_attention: false, gradient_checkpointing: false`
- Flash Attention: `flash_attention: true`
- Gradient Checkpointing: `gradient_checkpointing: true`
- Both: `flash_attention: true, gradient_checkpointing: true`
-
-Compare profiling metrics in SwanLab dashboard.
-
-### Use Case 4: Production RLHF with Team Collaboration
-Use `dpo-swanlab-full-featured.yml`, set up team workspace and Lark notifications:
-```yaml
-swanlab_workspace: ml-team
-swanlab_lark_webhook_url: ...
-swanlab_lark_secret: ...
-```
-
---
-
-## Viewing Your Experiments
-
-### Cloud Mode
-Visit [https://swanlab.cn](https://swanlab.cn) and navigate to your project.
-
-**Dashboard sections**:
- **Metrics**: Training loss, learning rate, profiling metrics
- **Tables**: RLHF completions (for DPO/KTO/ORPO/GRPO)
- **Config**: Hyperparameters and configuration
- **System**: Resource usage (GPU, memory, CPU)
- **Files**: Logged artifacts
-
-### Local Mode
-```bash
-swanlab watch ./swanlog
-# Open browser to http://localhost:5092
-```
-
---
-
-## Troubleshooting
-
-### SwanLab not initializing
-```bash
-# Check API key
-echo $SWANLAB_API_KEY
-
-# Verify SwanLab is installed
-pip show swanlab
-
-# Check config
-grep -A 5 "use_swanlab" your-config.yml
-```
-
-### Completions not appearing
- Verify you're using an RLHF trainer (DPO/KTO/ORPO/GRPO)
- Check `swanlab_log_completions: true`
- Wait for `swanlab_completion_log_interval` steps
- Look for "Registered SwanLab RLHF completion logging" in logs
-
-### Lark notifications not working
- Test webhook manually: `curl -X POST "$SWANLAB_LARK_WEBHOOK_URL" ...`
- Verify `SWANLAB_LARK_SECRET` is set correctly
- Check bot is added to Lark group chat
- Look for "Registered Lark notification callback" in logs
-
-### Profiling metrics not appearing
- Verify `use_swanlab: true`
- Check SwanLab is initialized (look for init log message)
- Profiling metrics are under "profiling/" namespace
- Profiling auto-enabled when SwanLab is enabled
-
---
-
-## Performance Notes
-
-### Overhead Comparison
-
-| Feature | Overhead per Step | Memory Usage |
-|---------|------------------|--------------|
-| Basic tracking | < 0.1% | ~10 MB |
-| Completion logging | < 0.5% | ~64 KB (buffer=128) |
-| Profiling | < 0.1% | ~1 KB |
-| **Total** | **< 0.7%** | **~10 MB** |
-
-### Best Practices
-1. Use ONE logging tool in production (disable WandB/MLflow when using SwanLab)
-2. Adjust completion log interval based on training length (100-200 steps)
-3. Keep completion buffer size reasonable (128-512)
-4. Profile critical path methods first (training_step, compute_loss)
-5. Use ProfilingConfig to throttle high-frequency operations
-
---
-
-## Further Reading
-
- **Full Documentation**: [src/axolotl/integrations/swanlab/README.md](../../src/axolotl/integrations/swanlab/README.md)
- **SwanLab Docs**: [https://docs.swanlab.cn](https://docs.swanlab.cn)
- **Axolotl Docs**: [https://axolotl-ai-cloud.github.io/axolotl/](https://axolotl-ai-cloud.github.io/axolotl/)
- **DPO Paper**: [Direct Preference Optimization](https://arxiv.org/abs/2305.18290)
-
---
-
-## Contributing
-
-Found an issue or have an improvement? Please submit a PR or open an issue:
- [Axolotl Issues](https://github.com/axolotl-ai-cloud/axolotl/issues)
- [SwanLab Issues](https://github.com/SwanHubX/SwanLab/issues)
--- a/examples/swanlab/custom_trainer_profiling.py
+++ b/examples/swanlab/custom_trainer_profiling.py
@@ -1,299 +0,0 @@
-"""Example: Custom Trainer with SwanLab Profiling
-
-This example demonstrates how to add SwanLab profiling to your custom trainer.
-
-Features:
- @swanlab_profile decorator for automatic profiling
- swanlab_profiling_context for fine-grained profiling
- ProfilingConfig for advanced filtering and throttling
-
-Usage:
-    1. Create your custom trainer extending AxolotlTrainer
-    2. Add @swanlab_profile decorators to methods you want to profile
-    3. Use swanlab_profiling_context for fine-grained profiling within methods
-    4. Enable SwanLab in your config (use_swanlab: true)
-
-See also:
-    - examples/swanlab/lora-swanlab-profiling.yml for config
-    - src/axolotl/integrations/swanlab/profiling.py for implementation
-"""
-
-from axolotl.core.trainers.base import AxolotlTrainer
-from axolotl.integrations.swanlab.profiling import (
-    ProfilingConfig,
-    swanlab_profile,
-    swanlab_profiling_context,
-    swanlab_profiling_context_advanced,
-)
-
-
-class CustomTrainerWithProfiling(AxolotlTrainer):
-    """Custom trainer with SwanLab profiling enabled.
-
-    This trainer demonstrates three profiling patterns:
-    1. Decorator-based profiling (@swanlab_profile)
-    2. Context manager profiling (swanlab_profiling_context)
-    3. Advanced profiling with filtering (ProfilingConfig)
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        # Create custom profiling config for high-frequency operations
-        self.fast_op_config = ProfilingConfig(
-            enabled=True,
-            min_duration_ms=0.5,  # Only log if duration > 0.5ms
-            log_interval=50,  # Log every 50th call
-        )
-
-    # ========================================================================
-    # Pattern 1: Decorator-based Profiling
-    # ========================================================================
-    # Best for: Methods you always want to profile
-    # Overhead: ~2-5 microseconds per call (negligible)
-
-    @swanlab_profile
-    def training_step(self, model, inputs):
-        """Main training step - always profile.
-
-        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.training_step
-        """
-        return super().training_step(model, inputs)
-
-    @swanlab_profile
-    def compute_loss(self, model, inputs, return_outputs=False):
-        """Loss computation - always profile.
-
-        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.compute_loss
-        """
-        return super().compute_loss(model, inputs, return_outputs)
-
-    @swanlab_profile
-    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
-        """Prediction step - always profile.
-
-        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.prediction_step
-        """
-        return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys)
-
-    # ========================================================================
-    # Pattern 2: Fine-grained Context Manager Profiling
-    # ========================================================================
-    # Best for: Profiling specific code blocks within a method
-    # Use case: When you want to profile forward vs backward separately
-
-    def complex_training_step(self, model, inputs):
-        """Training step with fine-grained profiling.
-
-        Profiling metrics:
-        - profiling/Time taken: CustomTrainerWithProfiling.forward_pass
-        - profiling/Time taken: CustomTrainerWithProfiling.backward_pass
-        - profiling/Time taken: CustomTrainerWithProfiling.optimizer_step
-        """
-        # Profile just the forward pass
-        with swanlab_profiling_context(self, "forward_pass"):
-            outputs = model(**inputs)
-            loss = outputs.loss
-
-        # Profile just the backward pass
-        with swanlab_profiling_context(self, "backward_pass"):
-            loss.backward()
-
-        # Profile optimizer step
-        with swanlab_profiling_context(self, "optimizer_step"):
-            self.optimizer.step()
-            self.optimizer.zero_grad()
-
-        return outputs
-
-    # ========================================================================
-    # Pattern 3: Advanced Profiling with Filtering
-    # ========================================================================
-    # Best for: High-frequency operations where you want to throttle logging
-    # Use case: Methods called 100+ times per step
-
-    def _prepare_inputs(self, inputs):
-        """Prepare inputs - throttled profiling.
-
-        This method is called frequently (once per batch), so we throttle
-        profiling to reduce overhead:
-        - Only log if duration > 0.5ms (skip very fast operations)
-        - Only log every 50th call (reduce logging frequency)
-
-        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.prepare_inputs
-        """
-        with swanlab_profiling_context_advanced(
-            self, "prepare_inputs", config=self.fast_op_config
-        ):
-            return super()._prepare_inputs(inputs)
-
-    def _prepare_input_for_model(self, input_ids):
-        """Another high-frequency operation - throttled profiling.
-
-        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.prepare_input_for_model
-        """
-        with swanlab_profiling_context_advanced(
-            self, "prepare_input_for_model", config=self.fast_op_config
-        ):
-            # Your custom input preparation logic
-            return input_ids
-
-    # ========================================================================
-    # Pattern 4: Exception-safe Profiling
-    # ========================================================================
-    # Profiling is exception-safe: duration is logged even if method raises
-
-    @swanlab_profile
-    def potentially_failing_method(self):
-        """This method may raise an exception.
-
-        SwanLab profiling will still log the duration before re-raising.
-        Profiling metric: profiling/Time taken: CustomTrainerWithProfiling.potentially_failing_method
-        """
-        # Do some work
-        result = self._do_risky_computation()
-
-        # If this raises, profiling duration is still logged
-        if result < 0:
-            raise ValueError("Invalid result")
-
-        return result
-
-    def _do_risky_computation(self):
-        """Placeholder for risky computation."""
-        return 42
-
-
-# ============================================================================
-# Advanced Example: Custom ProfilingConfig Per Method
-# ============================================================================
-
-
-class AdvancedProfilingTrainer(AxolotlTrainer):
-    """Trainer with method-specific profiling configurations."""
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        # Different profiling configs for different method types
-        self.critical_path_config = ProfilingConfig(
-            enabled=True,
-            min_duration_ms=0.0,  # Log everything on critical path
-            log_interval=1,  # Log every call
-        )
-
-        self.fast_path_config = ProfilingConfig(
-            enabled=True,
-            min_duration_ms=1.0,  # Only log if > 1ms
-            log_interval=100,  # Log every 100th call
-        )
-
-        self.debug_config = ProfilingConfig(
-            enabled=True,
-            min_duration_ms=0.0,  # Log everything
-            log_interval=1,  # Log every call
-        )
-
-    def training_step(self, model, inputs):
-        """Critical path - log everything."""
-        with swanlab_profiling_context_advanced(
-            self, "training_step", config=self.critical_path_config
-        ):
-            return super().training_step(model, inputs)
-
-    def _prepare_inputs(self, inputs):
-        """Fast path - throttle logging."""
-        with swanlab_profiling_context_advanced(
-            self, "prepare_inputs", config=self.fast_path_config
-        ):
-            return super()._prepare_inputs(inputs)
-
-    def _debug_method(self, data):
-        """Debug-only method - verbose logging."""
-        with swanlab_profiling_context_advanced(
-            self, "debug_method", config=self.debug_config
-        ):
-            # Your debug logic
-            pass
-
-
-# ============================================================================
-# How to Use This Custom Trainer
-# ============================================================================
-
-"""
-To use this custom trainer:
-
-1. Save this file to your project (e.g., my_custom_trainer.py)
-
-2. Create a config file that uses your custom trainer:
-
-    # config.yml
-    base_model: NousResearch/Llama-3.2-1B
-
-    # ... other config ...
-
-    plugins:
-      - axolotl.integrations.swanlab.SwanLabPlugin
-
-    use_swanlab: true
-    swanlab_project: my-profiling-experiment
-
-    # Optional: Specify custom trainer
-    # (Or modify axolotl to use your custom trainer class)
-
-3. Run training:
-
-    export SWANLAB_API_KEY=your-api-key
-    accelerate launch -m axolotl.cli.train config.yml
-
-4. View profiling metrics in SwanLab dashboard:
-   - profiling/Time taken: CustomTrainerWithProfiling.training_step
-   - profiling/Time taken: CustomTrainerWithProfiling.forward_pass
-   - profiling/Time taken: CustomTrainerWithProfiling.backward_pass
-   - etc.
-
-5. Compare profiling metrics across runs:
-   - Run baseline without optimizations
-   - Run with flash_attention enabled
-   - Run with gradient_checkpointing enabled
-   - Compare profiling metrics to see performance impact
-"""
-
-# ============================================================================
-# Tips for Effective Profiling
-# ============================================================================
-
-"""
-1. Profile the critical path first:
-   - training_step, compute_loss, prediction_step
-   - These methods are called most frequently and have biggest impact
-
-2. Use throttling for high-frequency operations:
-   - Methods called 100+ times per step
-   - Use log_interval=50 or log_interval=100
-   - Reduces profiling overhead and dashboard clutter
-
-3. Filter noise with min_duration_ms:
-   - Set min_duration_ms=1.0 to skip very fast operations
-   - Focus on operations that actually take time
-
-4. Compare across runs:
-   - Run same config multiple times to check consistency
-   - Compare different optimization strategies
-   - Track profiling trends over time
-
-5. Monitor distributed training:
-   - Check for per-rank timing differences
-   - Look for stragglers (slower ranks)
-   - Identify synchronization bottlenecks
-
-6. Disable profiling in production:
-   - from axolotl.integrations.swanlab.profiling import DEFAULT_PROFILING_CONFIG
-   - DEFAULT_PROFILING_CONFIG.enabled = False
-
-7. Exception handling:
-   - Profiling is exception-safe
-   - Duration logged even if method raises
-   - Useful for debugging methods that fail intermittently
-"""
--- a/examples/swanlab/dpo-swanlab-completions.yml
+++ b/examples/swanlab/dpo-swanlab-completions.yml
@@ -1,168 +0,0 @@
-# SwanLab DPO Training Example with Completion Logging
-#
-# This example demonstrates DPO (Direct Preference Optimization) training
-# with SwanLab integration for experiment tracking and completion table logging.
-#
-# Features enabled:
-# - SwanLab experiment tracking
-# - RLHF completion table logging (prompts, chosen/rejected responses, rewards)
-# - Lark (Feishu) team notifications (optional)
-#
-# To run:
-#   export SWANLAB_API_KEY=your-api-key
-#   accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-completions.yml
-
-# Model Configuration
-base_model: meta-llama/Meta-Llama-3-8B-Instruct
-model_type: LlamaForCausalLM
-tokenizer_type: AutoTokenizer
-
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-  eos_token: <|eot_id|>
-
-# Quantization
-load_in_8bit: true
-load_in_4bit: false
-
-# LoRA Configuration
-adapter: lora
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-
-# DPO Configuration
-chat_template: llama3
-rl: dpo
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_dpo_test
-    type: chat_template.default
-    field_messages: conversation
-    field_chosen: chosen
-    field_rejected: rejected
-    message_property_mappings:
-      role: role
-      content: content
-    roles:
-      system:
-        - system
-      user:
-        - user
-      assistant:
-        - assistant
-
-# Dataset and Output
-dataset_prepared_path:
-val_set_size: 0.05
-output_dir: ./outputs/dpo-swanlab-out
-
-# Training Configuration
-sequence_len: 4096
-sample_packing: false
-micro_batch_size: 2
-gradient_accumulation_steps: 4
-num_epochs: 4
-
-# Optimization
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-warmup_ratio: 0.1
-weight_decay: 0.0
-
-# Precision
-bf16: auto
-tf32: false
-
-# Performance
-gradient_checkpointing: true
-flash_attention: true
-
-# Checkpointing and Logging
-logging_steps: 1
-evals_per_epoch: 4
-saves_per_epoch: 1
-
-# ============================================================================
-# SwanLab Integration
-# ============================================================================
-
-plugins:
-  - axolotl.integrations.swanlab.SwanLabPlugin
-
-# Basic SwanLab Configuration
-use_swanlab: true
-swanlab_project: dpo-training
-swanlab_experiment_name: llama-3-dpo-completions-demo
-swanlab_description: "DPO training with completion table logging"
-swanlab_mode: cloud  # Options: cloud, local, offline, disabled
-
-# SwanLab Authentication
-# Recommended: Set via environment variable
-#   export SWANLAB_API_KEY=your-api-key
-# Or set in config (less secure):
-# swanlab_api_key: your-api-key
-
-# Optional: Team workspace
-# swanlab_workspace: my-research-team
-
-# ============================================================================
-# RLHF Completion Table Logging
-# ============================================================================
-#
-# Automatically logs model completions to SwanLab for qualitative analysis:
-# - Prompts from your DPO dataset
-# - Chosen responses (preferred)
-# - Rejected responses (non-preferred)
-# - Reward differences
-#
-# View the table in SwanLab dashboard under "rlhf_completions"
-
-swanlab_log_completions: true
-swanlab_completion_log_interval: 100  # Log every 100 training steps
-swanlab_completion_max_buffer: 128    # Keep last 128 completions in memory
-
-# Memory Usage Notes:
-# - Buffer size 128: ~64 KB (default, recommended)
-# - Buffer size 512: ~256 KB (for more historical completions)
-# - Buffer size 1024: ~512 KB (maximum for very long training runs)
-
-# Performance Notes:
-# - Completion logging overhead: < 0.5% per training step
-# - Only logs every N steps to minimize impact
-# - Memory-bounded buffer prevents memory leaks
-
-# ============================================================================
-# Optional: Lark (Feishu) Team Notifications
-# ============================================================================
-#
-# Get real-time training notifications in your team chat
-# Uncomment to enable:
-
-# swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
-# swanlab_lark_secret: your-webhook-secret  # Recommended for production
-
-# Notifications sent for:
-# - Training start
-# - Training completion
-# - Training errors
-# - Metric milestones (if configured)
-
-# ============================================================================
-# Optional: Private SwanLab Deployment
-# ============================================================================
-#
-# For enterprise users with private SwanLab deployment:
-
-# swanlab_web_host: https://swanlab.yourcompany.com
-# swanlab_api_host: https://api.swanlab.yourcompany.com
-
-# ============================================================================
-# Disable WandB if you're migrating from it
-# ============================================================================
-
-# wandb_project:
-# wandb_entity:
-# use_wandb: false
--- a/examples/swanlab/dpo-swanlab-full-featured.yml
+++ b/examples/swanlab/dpo-swanlab-full-featured.yml
@@ -1,329 +0,0 @@
-# SwanLab Full-Featured DPO Training Example
-#
-# This example demonstrates ALL SwanLab integration features:
-# - Experiment tracking with cloud sync
-# - RLHF completion table logging
-# - Performance profiling
-# - Lark (Feishu) team notifications
-# - Team workspace collaboration
-#
-# Use this as a reference for production RLHF training setups.
-#
-# To run:
-#   export SWANLAB_API_KEY=your-api-key
-#   export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
-#   export SWANLAB_LARK_SECRET=your-webhook-secret
-#   accelerate launch -m axolotl.cli.train examples/swanlab/dpo-swanlab-full-featured.yml
-
-# ============================================================================
-# Model Configuration
-# ============================================================================
-
-base_model: meta-llama/Meta-Llama-3-8B-Instruct
-model_type: LlamaForCausalLM
-tokenizer_type: AutoTokenizer
-
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-  eos_token: <|eot_id|>
-
-# Quantization for efficient training
-load_in_8bit: true
-load_in_4bit: false
-
-# ============================================================================
-# LoRA Configuration
-# ============================================================================
-
-adapter: lora
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true  # Target all linear layers
-
-# ============================================================================
-# DPO (Direct Preference Optimization) Configuration
-# ============================================================================
-
-chat_template: llama3
-rl: dpo  # Enable DPO trainer
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_dpo_test
-    type: chat_template.default
-    field_messages: conversation
-    field_chosen: chosen
-    field_rejected: rejected
-    message_property_mappings:
-      role: role
-      content: content
-    roles:
-      system:
-        - system
-      user:
-        - user
-      assistant:
-        - assistant
-
-# ============================================================================
-# Dataset and Output Configuration
-# ============================================================================
-
-dataset_prepared_path:
-val_set_size: 0.05
-output_dir: ./outputs/dpo-swanlab-full-featured-out
-
-# ============================================================================
-# Training Configuration
-# ============================================================================
-
-sequence_len: 4096
-sample_packing: false
-
-micro_batch_size: 2
-gradient_accumulation_steps: 4
-num_epochs: 4
-
-# ============================================================================
-# Optimization
-# ============================================================================
-
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-warmup_ratio: 0.1
-weight_decay: 0.0
-
-# ============================================================================
-# Precision and Performance
-# ============================================================================
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-flash_attention: true
-
-# ============================================================================
-# Checkpointing and Logging
-# ============================================================================
-
-logging_steps: 1
-evals_per_epoch: 4
-saves_per_epoch: 1
-
-# ============================================================================
-# SwanLab Integration - Full Configuration
-# ============================================================================
-
-plugins:
-  - axolotl.integrations.swanlab.SwanLabPlugin
-
-# ------------------------------------------------------------------------------
-# Basic SwanLab Configuration
-# ------------------------------------------------------------------------------
-
-use_swanlab: true
-swanlab_project: dpo-production
-swanlab_experiment_name: llama-3-dpo-full-featured-v1
-swanlab_description: |
-  Production DPO training with all SwanLab features enabled:
-  - Completion table logging for qualitative analysis
-  - Performance profiling for optimization
-  - Lark notifications for team collaboration
-
-swanlab_mode: cloud  # Options: cloud, local, offline, disabled
-
-# ------------------------------------------------------------------------------
-# Team Collaboration
-# ------------------------------------------------------------------------------
-
-# Workspace for team collaboration (shared experiments)
-swanlab_workspace: ml-research-team
-
-# Authentication (recommended: use environment variable)
-#   export SWANLAB_API_KEY=your-api-key
-# Or set in config (less secure):
-# swanlab_api_key: your-api-key
-
-# ------------------------------------------------------------------------------
-# RLHF Completion Table Logging
-# ------------------------------------------------------------------------------
-# Automatically logs model completions for qualitative analysis:
-# - Prompts from your DPO dataset
-# - Chosen responses (preferred)
-# - Rejected responses (non-preferred)
-# - Reward differences
-#
-# View in SwanLab dashboard under "rlhf_completions" table
-
-swanlab_log_completions: true
-swanlab_completion_log_interval: 100  # Log every 100 steps
-swanlab_completion_max_buffer: 256    # Larger buffer for long training runs
-
-# Buffer size recommendations:
-# - 128: Default, ~64 KB memory (recommended for most cases)
-# - 256: ~128 KB memory (this config, good for longer training)
-# - 512: ~256 KB memory (maximum for very long runs)
-
-# ------------------------------------------------------------------------------
-# Lark (Feishu) Team Notifications
-# ------------------------------------------------------------------------------
-# Get real-time training notifications in your team chat
-#
-# Notifications sent for:
-# - Training start
-# - Training completion
-# - Training errors
-# - Metric milestones (if configured)
-
-# Recommended: Set via environment variables
-#   export SWANLAB_LARK_WEBHOOK_URL=https://open.feishu.cn/...
-#   export SWANLAB_LARK_SECRET=your-webhook-secret
-
-# Or set in config (less secure):
-# swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
-# swanlab_lark_secret: your-webhook-secret  # REQUIRED for production
-
-# Security note: ALWAYS use swanlab_lark_secret in production to prevent
-# unauthorized parties from sending fake notifications to your team chat.
-
-# ------------------------------------------------------------------------------
-# Performance Profiling
-# ------------------------------------------------------------------------------
-# Profiling is automatically enabled when SwanLab is enabled.
-# Metrics logged to SwanLab under "profiling/" namespace:
-#   profiling/Time taken: AxolotlTrainer.training_step
-#   profiling/Time taken: AxolotlTrainer.compute_loss
-#   profiling/Time taken: AxolotlTrainer.prediction_step
-#
-# Use these metrics to:
-# - Identify bottlenecks in training loop
-# - Compare performance across different configurations
-# - Monitor performance regressions over time
-# - Debug unexpected slowdowns
-
-# For custom profiling in your own trainer, see:
-#   examples/swanlab/custom_trainer_profiling.py
-
-# ------------------------------------------------------------------------------
-# Optional: Private SwanLab Deployment
-# ------------------------------------------------------------------------------
-# For enterprise users with private SwanLab deployment:
-
-# swanlab_web_host: https://swanlab.yourcompany.com
-# swanlab_api_host: https://api.swanlab.yourcompany.com
-
-# ------------------------------------------------------------------------------
-# Optional: Model Checkpointing to SwanLab
-# ------------------------------------------------------------------------------
-# Log model checkpoints to SwanLab (coming soon)
-
-swanlab_log_model: false
-
-# ============================================================================
-# Disable Other Logging Tools (Recommended)
-# ============================================================================
-# Using multiple logging tools simultaneously can impact performance:
-# - Expected overhead: ~1-2% per logger
-# - Potential config/callback conflicts
-#
-# For production training, use ONLY SwanLab:
-
-# wandb_project:
-# use_wandb: false
-#
-# use_mlflow: false
-#
-# use_comet: false
-
-# ============================================================================
-# Expected Training Behavior
-# ============================================================================
-
-# With this configuration, you should see:
-#
-# 1. SwanLab Initialization (rank 0 only):
-#    INFO: SwanLab initialized for project: dpo-production
-#    INFO: SwanLab experiment: llama-3-dpo-full-featured-v1
-#    INFO: SwanLab mode: cloud
-#    INFO: SwanLab workspace: ml-research-team
-#
-# 2. Completion Logging (rank 0 only):
-#    INFO: Registered SwanLab RLHF completion logging callback for DPOTrainer
-#          (log_interval=100, max_buffer=256)
-#
-# 3. Lark Notifications (rank 0 only):
-#    INFO: Registered Lark notification callback with HMAC authentication
-#
-# 4. Distributed Training Detection (if multi-GPU):
-#    INFO: Distributed training detected (world_size=N)
-#    INFO: Only rank 0 will initialize SwanLab
-#    INFO: Other ranks will skip SwanLab to avoid conflicts
-#
-# 5. Training Start Notification (Lark):
-#    Your team chat receives: "Training started: llama-3-dpo-full-featured-v1"
-#
-# 6. Periodic Completion Logging:
-#    Every 100 steps, completion table is updated in SwanLab dashboard
-#
-# 7. Training Complete Notification (Lark):
-#    Your team chat receives: "Training completed: llama-3-dpo-full-featured-v1"
-#    With link to SwanLab dashboard and final metrics
-#
-# 8. SwanLab Dashboard Shows:
-#    - Training metrics (loss, learning rate, etc.)
-#    - Completion table (rlhf_completions)
-#    - Profiling metrics (profiling/Time taken: ...)
-#    - Hyperparameters and configuration
-#    - System resource usage
-
-# ============================================================================
-# Production Checklist
-# ============================================================================
-
-# Before deploying to production, verify:
-# ✅ SwanLab API key is set via environment variable (not in config)
-# ✅ Lark webhook secret is set (required for HMAC authentication)
-# ✅ Workspace is set to your team's workspace
-# ✅ Experiment name is descriptive and unique
-# ✅ Only SwanLab is enabled (other loggers disabled)
-# ✅ Completion logging buffer size is appropriate for your training duration
-# ✅ Private deployment hosts are set (if using enterprise SwanLab)
-# ✅ Test run completes successfully and shows up in SwanLab dashboard
-# ✅ Lark notifications are received in team chat
-# ✅ Profiling metrics are logged correctly
-
-# ============================================================================
-# Troubleshooting
-# ============================================================================
-
-# If SwanLab initialization fails:
-# 1. Check SWANLAB_API_KEY environment variable is set
-# 2. Verify swanlab_project is set in config
-# 3. Check swanlab_mode is valid (cloud/local/offline/disabled)
-# 4. Verify internet connectivity (for cloud mode)
-
-# If Lark notifications not received:
-# 1. Check SWANLAB_LARK_WEBHOOK_URL is set correctly
-# 2. Verify SWANLAB_LARK_SECRET matches your Lark bot settings
-# 3. Test webhook manually: curl -X POST "$SWANLAB_LARK_WEBHOOK_URL" ...
-# 4. Check training logs for "Registered Lark notification callback"
-# 5. Verify bot is added to the target Lark group chat
-
-# If completions not appearing in SwanLab:
-# 1. Verify you're using an RLHF trainer (DPO/KTO/ORPO/GRPO)
-# 2. Check swanlab_log_completions is true
-# 3. Wait for log_interval steps (default: 100)
-# 4. Check training logs for "Registered SwanLab RLHF completion logging"
-
-# If profiling metrics not appearing:
-# 1. Verify use_swanlab is true
-# 2. Check SwanLab is initialized (check logs)
-# 3. Look under "profiling/" namespace in dashboard
-# 4. Profiling may be disabled if DEFAULT_PROFILING_CONFIG.enabled = False
-
-# For more help:
-# - SwanLab docs: https://docs.swanlab.cn
-# - Axolotl SwanLab integration: src/axolotl/integrations/swanlab/README.md
-# - GitHub issues: https://github.com/axolotl-ai-cloud/axolotl/issues
--- a/examples/swanlab/lora-swanlab-profiling.yml
+++ b/examples/swanlab/lora-swanlab-profiling.yml
@@ -1,178 +0,0 @@
-# SwanLab LoRA Training Example with Performance Profiling
-#
-# This example demonstrates standard LoRA fine-tuning with SwanLab integration
-# for performance profiling and optimization.
-#
-# Features enabled:
-# - SwanLab experiment tracking
-# - Performance profiling (training step, forward/backward pass timing)
-# - Real-time metrics visualization
-#
-# To run:
-#   export SWANLAB_API_KEY=your-api-key
-#   accelerate launch -m axolotl.cli.train examples/swanlab/lora-swanlab-profiling.yml
-
-# Model Configuration
-base_model: NousResearch/Llama-3.2-1B
-
-# Dataset Configuration
-datasets:
-  - path: teknium/GPT4-LLM-Cleaned
-    type: alpaca
-
-val_set_size: 0.1
-output_dir: ./outputs/lora-swanlab-profiling-out
-
-# LoRA Configuration
-adapter: lora
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-# Training Configuration
-sequence_len: 2048
-sample_packing: true
-eval_sample_packing: true
-
-micro_batch_size: 2
-gradient_accumulation_steps: 2
-num_epochs: 1
-
-# Optimization
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-warmup_ratio: 0.1
-weight_decay: 0.0
-
-# Precision
-bf16: auto
-tf32: false
-
-# Performance
-gradient_checkpointing: true
-flash_attention: true
-
-# Checkpointing and Logging
-logging_steps: 1
-evals_per_epoch: 4
-saves_per_epoch: 1
-
-# Loss Monitoring
-loss_watchdog_threshold: 5.0
-loss_watchdog_patience: 3
-
-special_tokens:
-  pad_token: "<|end_of_text|>"
-
-# ============================================================================
-# SwanLab Integration
-# ============================================================================
-
-plugins:
-  - axolotl.integrations.swanlab.SwanLabPlugin
-
-# Basic SwanLab Configuration
-use_swanlab: true
-swanlab_project: lora-profiling
-swanlab_experiment_name: llama-3.2-1b-profiling-demo
-swanlab_description: "LoRA fine-tuning with performance profiling"
-swanlab_mode: cloud  # Options: cloud, local, offline, disabled
-
-# SwanLab Authentication
-# Recommended: Set via environment variable
-#   export SWANLAB_API_KEY=your-api-key
-# Or set in config (less secure):
-# swanlab_api_key: your-api-key
-
-# Optional: Team workspace
-# swanlab_workspace: my-ml-team
-
-# ============================================================================
-# Performance Profiling
-# ============================================================================
-#
-# SwanLab automatically profiles trainer methods when enabled.
-# Profiling metrics appear in SwanLab dashboard under "profiling/" namespace.
-#
-# Built-in profiling:
-# - Minimal overhead (< 0.1% per step)
-# - High-precision timing (microsecond accuracy)
-# - Exception-safe (logs duration even if method fails)
-#
-# View profiling metrics in SwanLab dashboard:
-#   profiling/Time taken: AxolotlTrainer.training_step
-#   profiling/Time taken: AxolotlTrainer.compute_loss
-#   profiling/Time taken: AxolotlTrainer.prediction_step
-#
-# For custom profiling in your own trainer, see:
-#   examples/swanlab/custom_trainer_profiling.py
-
-# Completion logging is disabled for non-RLHF trainers
-swanlab_log_completions: false  # Only works with DPO/KTO/ORPO/GRPO
-
-# ============================================================================
-# Optional: Compare with Multiple Runs
-# ============================================================================
-#
-# To compare profiling metrics across different configurations:
-#
-# 1. Run baseline without flash attention:
-#    swanlab_experiment_name: llama-3.2-1b-no-flash-attn
-#    flash_attention: false
-#
-# 2. Run with gradient checkpointing:
-#    swanlab_experiment_name: llama-3.2-1b-grad-checkpoint
-#    gradient_checkpointing: true
-#
-# 3. Run with both:
-#    swanlab_experiment_name: llama-3.2-1b-optimized
-#    flash_attention: true
-#    gradient_checkpointing: true
-#
-# Then compare profiling metrics in SwanLab dashboard to see performance impact
-
-# ============================================================================
-# Optional: Lark (Feishu) Team Notifications
-# ============================================================================
-#
-# Get notified when profiling experiments complete:
-
-# swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
-# swanlab_lark_secret: your-webhook-secret
-
-# ============================================================================
-# Profiling Best Practices
-# ============================================================================
-#
-# 1. Run multiple epochs to see profiling trends over time
-# 2. Ignore first ~10 steps (warmup period, slower)
-# 3. Look for outliers (steps that take significantly longer)
-# 4. Compare profiling metrics before/after optimization changes
-# 5. Monitor per-rank profiling in distributed training
-#
-# Common bottlenecks to profile:
-# - training_step: Overall step time (should be consistent)
-# - compute_loss: Loss computation (scales with sequence length)
-# - prediction_step: Evaluation time (can be slow for large val sets)
-#
-# If you see inconsistent timing:
-# - Check for data loading bottlenecks
-# - Monitor GPU utilization (may be CPU-bound)
-# - Check for gradient accumulation effects
-# - Verify CUDA kernel synchronization
-
-# ============================================================================
-# Disable WandB if you're migrating from it
-# ============================================================================
-
-# wandb_project:
-# use_wandb: false
--- a/examples/trinity/README.md
+++ b/examples/trinity/README.md
@@ -8,15 +8,13 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations

 1. Install Axolotl following the main from the [installation guide](https://docs.axolotl.ai/docs/installation.html#sec-edge-build).

-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
-
-3. Run the finetuning example:
+2. Run the finetuning example:

    ```bash
    axolotl train examples/trinity/trinity-nano-preview-qlora.yaml
    ```

-This config uses about 24.9 GiB VRAM (w/o CCE).
+This config uses about 24.9 GiB VRAM.

 Let us know how it goes. Happy finetuning! 🚀

@@ -31,6 +29,10 @@ Let us know how it goes. Happy finetuning! 🚀

 Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).

+## Limitations
+
+**Cut Cross Entropy (CCE)**: Currently not supported. We plan to include CCE support for Trinity in the near future.
+
 ## Related Resources

 - [Trinity Blog](https://www.arcee.ai/blog/the-trinity-manifesto)
--- a/examples/trinity/trinity-nano-preview-qlora.yaml
+++ b/examples/trinity/trinity-nano-preview-qlora.yaml
@@ -1,4 +1,5 @@
 base_model: arcee-ai/Trinity-Nano-Preview
+trust_remote_code: true
 revision_of_model: 2ee94b0

 # Automatically upload checkpoint and final model to HF
--- a/examples/voxtral/README.md
+++ b/examples/voxtral/README.md
@@ -12,7 +12,7 @@ Thanks to the team at MistralAI for giving us early access to prepare for this r

 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```

--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=64", "wheel", "setuptools_scm>=8", "packaging==26.0"]
+requires = ["setuptools>=64", "wheel", "setuptools_scm>=8", "packaging==23.2"]
 build-backend = "setuptools.build_meta"

 [project]
@@ -24,9 +24,6 @@ Repository = "https://github.com/axolotl-ai-cloud/axolotl.git"
 py-modules = ["setuptools_axolotl_dynamic_dependencies"]
 include-package-data = true

-[tool.setuptools.dynamic]
-version = { file = "VERSION" }
-
 [tool.setuptools.cmdclass]
 build_py = "setuptools_axolotl_dynamic_dependencies.BuildPyCommand"

@@ -60,6 +57,3 @@ indent-style = "space"
 skip-magic-trailing-comma = false
 line-ending = "auto"
 docstring-code-format = false
-
-[tool.uv.extra-build-dependencies]
-axolotl = ["huggingface_hub"]
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,29 +1,26 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/

 # START section of dependencies that don't install on Darwin/MacOS
-bitsandbytes==0.49.1
-triton>=3.4.0
+bitsandbytes==0.48.2
+triton>=3.0.0
 mamba-ssm==1.2.0.post1
 xformers>=0.0.23.post1
-liger-kernel==0.7.0
+liger-kernel==0.6.4
 # END section

-packaging==26.0
-huggingface_hub>=1.1.7
-peft>=0.18.1
+packaging==23.2
+
+huggingface_hub>=0.36.0
+peft>=0.18.0
 tokenizers>=0.22.1
-transformers==5.3.0
-accelerate==1.13.0
-datasets==4.5.0
-deepspeed>=0.18.6,<0.19.0
-trl==0.29.0
-hf_xet==1.3.2
-kernels==0.12.2
-
-fla-core==0.4.1
-flash-linear-attention==0.4.1
-
-trackio>=0.16.1
+transformers==4.57.1
+accelerate==1.12.0
+datasets==4.4.2
+deepspeed>=0.18.3
+trl==0.25.1
+hf_xet==1.2.0
+kernels==0.11.5
+trackio>=0.13.0
 typing-extensions>=4.15.0

 optimum==1.16.2
@@ -66,7 +63,7 @@ langdetect==1.0.9
 immutabledict==4.2.0
 antlr4-python3-runtime==4.13.2

-torchao==0.16.0
+torchao==0.15.0
 openenv-core==0.1.0
 schedulefree==1.4.1

@@ -75,4 +72,4 @@ axolotl-contribs-mit==0.0.6
 # telemetry
 posthog==6.7.11

-mistral-common==1.10.0
+mistral-common==1.8.6
--- 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@63b15e6"'
+    + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@318b7e2"'
 )
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,6 @@
 """setup.py for axolotl"""

+import ast
 import os
 import platform
 import re
@@ -25,19 +26,6 @@ def parse_requirements(extras_require_map):
                _install_requires.append(line)
    try:
        xformers_version = [req for req in _install_requires if "xformers" in req][0]
-        install_xformers = platform.machine() != "aarch64"
-        if platform.machine() == "aarch64":
-            # skip on ARM64
-            skip_packages = [
-                "torchao",
-                "fla-core",
-                "flash-linear-attention",
-            ]
-            _install_requires = [
-                req
-                for req in _install_requires
-                if re.split(r"[>=<]", req)[0].strip() not in skip_packages
-            ]
        if "Darwin" in platform.system():
            # skip packages not compatible with OSX
            skip_packages = [
@@ -74,68 +62,44 @@ def parse_requirements(extras_require_map):
            else:
                raise ValueError("Invalid version format")

-            torch_parts = torch_version.split("+")
-            if len(torch_parts) == 2:
-                torch_cuda_version = torch_parts[1]
-                _dependency_links.append(
-                    f"https://download.pytorch.org/whl/{torch_cuda_version}"
-                )
-
            if (major, minor) >= (2, 9):
                extras_require_map.pop("fbgemm-gpu")
-                extras_require_map["fbgemm-gpu"] = [
-                    "fbgemm-gpu==1.4.0",
-                    "fbgemm-gpu-genai==1.4.2",
-                ]
+                extras_require_map["fbgemm-gpu"] = ["fbgemm-gpu-genai==1.4.1"]
                extras_require_map["vllm"] = ["vllm==0.11.1"]
-                if not install_xformers:
-                    _install_requires.pop(_install_requires.index(xformers_version))
-                extras_require_map["vllm"] = ["vllm==0.13.0"]
-                if patch == 0:
-                    extras_require_map["vllm"] = ["vllm==0.13.0"]
-                else:
-                    extras_require_map["vllm"] = ["vllm==0.14.0"]
            elif (major, minor) >= (2, 8):
                extras_require_map.pop("fbgemm-gpu")
                extras_require_map["fbgemm-gpu"] = ["fbgemm-gpu-genai==1.3.0"]
                extras_require_map["vllm"] = ["vllm==0.11.0"]
-                if not install_xformers:
-                    _install_requires.pop(_install_requires.index(xformers_version))
            elif (major, minor) >= (2, 7):
                _install_requires.pop(_install_requires.index(xformers_version))
                if patch == 0:
-                    if install_xformers:
-                        _install_requires.append("xformers==0.0.30")
+                    _install_requires.append("xformers==0.0.30")
                    # vllm 0.9.x is incompatible with latest transformers
                    extras_require_map.pop("vllm")
                else:
-                    if install_xformers:
-                        _install_requires.append("xformers==0.0.31")
+                    _install_requires.append("xformers==0.0.31")
                    extras_require_map["vllm"] = ["vllm==0.10.1"]
            elif (major, minor) >= (2, 6):
                _install_requires.pop(_install_requires.index(xformers_version))
-                if install_xformers:
-                    _install_requires.append("xformers==0.0.29.post3")
+                _install_requires.append("xformers==0.0.29.post3")
                # since we only support 2.6.0+cu126
                _dependency_links.append("https://download.pytorch.org/whl/cu126")
                extras_require_map.pop("vllm")
            elif (major, minor) >= (2, 5):
                _install_requires.pop(_install_requires.index(xformers_version))
-                if install_xformers:
-                    if patch == 0:
-                        _install_requires.append("xformers==0.0.28.post2")
-                    else:
-                        _install_requires.append("xformers>=0.0.28.post3")
+                if patch == 0:
+                    _install_requires.append("xformers==0.0.28.post2")
+                else:
+                    _install_requires.append("xformers>=0.0.28.post3")
                extras_require_map.pop("vllm")
            elif (major, minor) >= (2, 4):
                extras_require_map.pop("vllm")
-                if install_xformers:
-                    if patch == 0:
-                        _install_requires.pop(_install_requires.index(xformers_version))
-                        _install_requires.append("xformers>=0.0.27")
-                    else:
-                        _install_requires.pop(_install_requires.index(xformers_version))
-                        _install_requires.append("xformers==0.0.28.post1")
+                if patch == 0:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                    _install_requires.append("xformers>=0.0.27")
+                else:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                    _install_requires.append("xformers==0.0.28.post1")
            else:
                raise ValueError("axolotl requires torch>=2.4")

@@ -146,11 +110,15 @@ def parse_requirements(extras_require_map):

 def get_package_version():
    with open(
-        Path(os.path.dirname(os.path.abspath(__file__))) / "VERSION",
+        Path(os.path.dirname(os.path.abspath(__file__)))
+        / "src"
+        / "axolotl"
+        / "__init__.py",
        "r",
        encoding="utf-8",
    ) as fin:
-        version_ = fin.read().strip()
+        version_match = re.search(r"^__version__\s*=\s*(.*)$", fin.read(), re.MULTILINE)
+    version_ = ast.literal_eval(version_match.group(1))
    return version_


--- a/Show More
+++ b/Show More