fix test dims

.github/CONTRIBUTING.md

@@ -31,10 +31,7 @@ PRs are **greatly welcome**!
 
 Please run below to setup env
 ```bash
-# Install axolotl + dev and test dependencies from lockfile
+pip3 install -r requirements-dev.txt -r requirements-tests.txt
-export UV_TORCH_BACKEND=cu128  # or cu130
-uv sync --extra flash-attn --extra deepspeed --group dev --group test
-source .venv/bin/activate
 pre-commit install
 
 # test

.github/workflows/lint.yml

@@ -6,7 +6,7 @@ on:
       types: [opened, synchronize, reopened, ready_for_review]
       paths:
        - '**.py'
-       - 'pyproject.toml'
+       - 'requirements.txt'
        - '.github/workflows/*.yml'
        - "*.[q]md"
        - "examples/**/*.y[a]?ml"

.github/workflows/multi-gpu-e2e.yml

@@ -3,15 +3,17 @@ name: docker-multigpu-tests-biweekly
 on:
   pull_request:
     paths:
-      - "tests/e2e/multigpu/**.py"
+      - 'tests/e2e/multigpu/**.py'
-      - "pyproject.toml"
+      - 'requirements.txt'
-      - ".github/workflows/multi-gpu-e2e.yml"
+      - 'setup.py'
-      - "scripts/cutcrossentropy_install.py"
+      - 'pyproject.toml'
-      - "src/axolotl/core/trainers/mixins/sequence_parallel.py"
+      - '.github/workflows/multi-gpu-e2e.yml'
-      - "src/axolotl/utils/distributed.py"
+      - 'scripts/cutcrossentropy_install.py'
+      - 'src/axolotl/core/trainers/mixins/sequence_parallel.py'
+      - 'src/axolotl/utils/distributed.py'
   workflow_dispatch:
   schedule:
-    - cron: "0 0 * * 1,4" # Runs at 00:00 UTC every monday & thursday
+    - cron: '0 0 * * 1,4'  # Runs at 00:00 UTC every monday & thursday
 
 # Cancel jobs on the same ref if a new one is triggered
 concurrency:
@@ -31,19 +33,19 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          #          - cuda: 129
+#          - cuda: 129
-          #            cuda_version: 12.9.1
+#            cuda_version: 12.9.1
-          #            python_version: "3.12"
+#            python_version: "3.12"
-          #            pytorch: 2.9.1
+#            pytorch: 2.9.1
-          #            axolotl_extras: "fbgemm-gpu"
+#            axolotl_extras: "fbgemm-gpu"
-          #            num_gpus: 2
+#            num_gpus: 2
-          #            dockerfile: "Dockerfile-uv.jinja"
+#            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
+#            axolotl_extras: fbgemm-gpu
             num_gpus: 2
           - cuda: 128
             cuda_version: 12.8.1
@@ -51,6 +53,7 @@ jobs:
             pytorch: 2.10.0
             axolotl_extras: "fbgemm-gpu"
             num_gpus: 2
+            dockerfile: "Dockerfile-uv.jinja"
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
     steps:
@@ -72,7 +75,7 @@ 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-uv.jinja'}}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

.github/workflows/pypi.yml

@@ -8,9 +8,6 @@ on:
 
 permissions: {}
 
-env:
-  UV_SYSTEM_PYTHON: "1"
-
 jobs:
   setup_release:
     name: Create Release
@@ -44,15 +41,11 @@ jobs:
         with:
           python-version: "3.11"
 
-      - name: Install uv
-        uses: astral-sh/setup-uv@v7
-
       - name: Install dependencies
         run: |
-          uv pip install wheel packaging
+          pip3 install wheel packaging==26.0
-          uv pip install --no-build-isolation -e .
+          pip3 install --no-build-isolation -e .
-          uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+          pip3 install -r requirements-dev.txt -r requirements-tests.txt
-            codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
 
       - name: Extract tag name
         id: tag

.github/workflows/tests-nightly.yml

@@ -2,18 +2,15 @@ name: Tests Nightly against upstream main
 on:
   workflow_dispatch:
   schedule:
-    - cron: "0 0 * * *" # Runs at 00:00 UTC every day
+    - 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"
+      - '.github/workflows/tests-nightly.yml'
 
 permissions:
   contents: read
 
-env:
-  UV_SYSTEM_PYTHON: "1"
-
 jobs:
   pre-commit:
     name: pre-commit
@@ -23,7 +20,7 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: "3.11"
-          cache: "pip" # caching pip dependencies
+          cache: 'pip' # caching pip dependencies
       - uses: pre-commit/action@v3.0.1
         env:
           SKIP: no-commit-to-branch
@@ -46,7 +43,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ["3.12"] # TODO include py3.14 once https://github.com/mistralai/mistral-common/pull/194 is merged
+        python_version: ["3.12"]  # TODO include py3.14 once https://github.com/mistralai/mistral-common/pull/194 is merged
         pytorch_version: ["2.9.1", "2.10.0"]
     timeout-minutes: 20
 
@@ -64,34 +61,36 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python_version }}
+          cache: 'pip' # caching pip dependencies
 
-      - name: Install uv
+      - name: upgrade pip
-        uses: astral-sh/setup-uv@v7
+        run: |
+          pip3 install --upgrade pip
+          pip3 install --upgrade packaging==26.0 setuptools==78.1.1 wheel
 
       - name: Install PyTorch
         run: |
-          uv pip install torch==${{ matrix.pytorch_version }} torchvision
+          pip3 install torch==${{ matrix.pytorch_version }} torchvision
-          uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt
+
+      - name: Update requirements.txt
+        run: |
+          sed -i 's#^transformers.*#transformers @ git+https://github.com/huggingface/transformers.git@main#' requirements.txt
+          sed -i 's#^peft.*#peft @ git+https://github.com/huggingface/peft.git@main#' requirements.txt
+          sed -i 's#^accelerate.*#accelerate @ git+https://github.com/huggingface/accelerate.git@main#' requirements.txt
+          sed -i 's#^trl.*#trl @ git+https://github.com/huggingface/trl.git@main#' requirements.txt
+          sed -i 's#^datasets.*#datasets @ git+https://github.com/huggingface/datasets.git@main#' requirements.txt
 
       - name: Install dependencies
         run: |
-          uv pip install --no-build-isolation -e . --override /tmp/torch-pin.txt
+          pip3 show torch
-          python scripts/cutcrossentropy_install.py --uv | sh
+          pip3 install --no-build-isolation -U -e .
-          uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+          python scripts/unsloth_install.py | sh
-            codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
+          python scripts/cutcrossentropy_install.py | sh
-
+          pip3 install -r requirements-dev.txt -r requirements-tests.txt
-      - name: Override with nightly HF packages
-        run: |
-          uv pip install --no-deps \
-            "transformers @ git+https://github.com/huggingface/transformers.git@main" \
-            "peft @ git+https://github.com/huggingface/peft.git@main" \
-            "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
-            "trl @ git+https://github.com/huggingface/trl.git@main" \
-            "datasets @ git+https://github.com/huggingface/datasets.git@main"
 
       - name: Make sure PyTorch version wasn't clobbered
         run: |
-          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__, f'Expected torch ${{ matrix.pytorch_version }} but got {torch.__version__}'"
+          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__"
 
       - name: Ensure axolotl CLI was installed
         run: |
@@ -103,6 +102,9 @@ jobs:
           pytest -v --durations=10 tests/patched/
           pytest -v --durations=10 tests/cli/
 
+      - name: cleanup pip cache
+        run: |
+          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
 
   docker-e2e-tests:
     if: github.repository_owner == 'axolotl-ai-cloud'
@@ -134,6 +136,7 @@ jobs:
             pytorch: 2.9.1
             num_gpus: 1
             axolotl_extras:
+            dockerfile: "Dockerfile-uv.jinja"
             nightly_build: "true"
     steps:
       - name: Checkout
@@ -154,7 +157,7 @@ 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-uv.jinja'}}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
           echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         env:

.github/workflows/tests.yml

@@ -6,19 +6,21 @@ on:
     branches:
       - "main"
     paths:
-      - "**.py"
+      - '**.py'
-      - "pyproject.toml"
+      - 'requirements.txt'
-      - ".github/workflows/*.yml"
+      - '.github/workflows/*.yml'
-      - "cicd/cicd.sh"
+      - 'requirements-tests.txt'
-      - "cicd/Dockerfile-uv.jinja"
+      - 'cicd/cicd.sh'
+      - 'cicd/Dockerfile.jinja'
   pull_request:
-    types: [opened, synchronize, reopened, ready_for_review]
+      types: [opened, synchronize, reopened, ready_for_review]
-    paths:
+      paths:
-      - "**.py"
+       - '**.py'
-      - "pyproject.toml"
+       - 'requirements.txt'
-      - ".github/workflows/*.yml"
+       - '.github/workflows/*.yml'
-      - "cicd/cicd.sh"
+       - 'requirements-tests.txt'
-      - "cicd/Dockerfile-uv.jinja"
+       - 'cicd/cicd.sh'
+       - 'cicd/Dockerfile.jinja'
   workflow_dispatch:
 
 # Cancel jobs on the same ref if a new one is triggered
@@ -31,7 +33,6 @@ permissions:
 
 env:
   TRANSFORMERS_IS_CI: "yes"
-  UV_SYSTEM_PYTHON: "1"
 
 jobs:
   pre-commit:
@@ -43,7 +44,7 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: "3.11"
-          cache: "pip" # caching pip dependencies
+          cache: 'pip' # caching pip dependencies
       - uses: pre-commit/action@v3.0.1
         env:
           SKIP: no-commit-to-branch
@@ -93,25 +94,32 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python_version }}
+          cache: 'pip' # caching pip dependencies
 
-      - name: Install uv
+      - name: upgrade pip
-        uses: astral-sh/setup-uv@v7
+        run: |
+          pip3 install --upgrade pip
+          pip3 install --upgrade packaging==26.0 setuptools==75.8.0 wheel
 
       - name: Install PyTorch
         run: |
-          uv pip install torch==${{ matrix.pytorch_version }} torchvision
+          pip3 install --no-cache-dir torch==${{ matrix.pytorch_version }} torchvision
-          uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt
 
       - name: Install dependencies
         run: |
-          uv pip install --no-build-isolation -e . --override /tmp/torch-pin.txt
+          pip3 show torch
-          python scripts/cutcrossentropy_install.py --uv | sh
+          pip3 install --no-cache-dir --no-build-isolation -U -e .
-          uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+          python scripts/unsloth_install.py | sh
-            codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
+          python scripts/cutcrossentropy_install.py | sh
+          pip3 install -r requirements-dev.txt -r requirements-tests.txt
+
+      - name: cleanup pip cache
+        run: |
+          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
 
       - name: Make sure PyTorch version wasn't clobbered
         run: |
-          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__, f'Expected torch ${{ matrix.pytorch_version }} but got {torch.__version__}'"
+          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__"
 
       - name: Ensure axolotl CLI was installed
         run: |
@@ -180,27 +188,33 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python_version }}
+          cache: 'pip' # caching pip dependencies
 
-      - name: Install uv
+      - name: upgrade pip
-        uses: astral-sh/setup-uv@v7
+        run: |
+          pip3 install --upgrade pip
+          pip3 install --upgrade packaging==26.0 setuptools==75.8.0 setuptools_scm build wheel psutil
 
       - name: Install PyTorch
         run: |
-          uv pip install torch==${{ matrix.pytorch_version }} torchvision
+          pip3 install --no-cache-dir torch==${{ matrix.pytorch_version }} torchvision
-          uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt
 
       - name: Install dependencies
         run: |
-          uv pip install packaging setuptools_scm build wheel psutil
+          pip3 show torch
           python -m build --no-isolation --sdist
-          uv pip install --no-build-isolation dist/axolotl*.tar.gz --override /tmp/torch-pin.txt
+          pip3 install --no-cache-dir --no-build-isolation dist/axolotl*.tar.gz
-          python scripts/cutcrossentropy_install.py --uv | sh
+          python scripts/unsloth_install.py | sh
-          uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+          python scripts/cutcrossentropy_install.py | sh
-            codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
+          pip3 install -r requirements-dev.txt -r requirements-tests.txt
+
+      - name: cleanup pip cache
+        run: |
+          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
 
       - name: Make sure PyTorch version wasn't clobbered
         run: |
-          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__, f'Expected torch ${{ matrix.pytorch_version }} but got {torch.__version__}'"
+          python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__"
 
       - name: Ensure axolotl CLI was installed
         run: |
@@ -277,6 +291,7 @@ jobs:
             pytorch: 2.9.1
             num_gpus: 1
             axolotl_extras:
+            dockerfile: "Dockerfile-uv.jinja"
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -297,7 +312,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 "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
@@ -359,7 +374,7 @@ 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 "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV
+          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

MANIFEST.in

@@ -1,6 +1,7 @@
+include requirements.txt
 include README.md
 include LICENSE
-include VERSION
+include src/setuptools_axolotl_dynamic_dependencies.py
 include src/axolotl/utils/chat_templates/templates/*.jinja
 include AGENTS.md
 recursive-include docs/agents *.md

README.md

@@ -95,11 +95,14 @@ Features:
 
 ### Installation
 
-```bash
+#### Using uv (recommended)
-# install uv if you don't already have it installed (restart shell after)
-curl -LsSf https://astral.sh/uv/install.sh | sh
 
-# change depending on system
+```bash
+# install uv if you don't already have it installed
+curl -LsSf https://astral.sh/uv/install.sh | sh
+source $HOME/.local/bin/env
+
+# CUDA 12.8.1 tends to have better package compatibility
 export UV_TORCH_BACKEND=cu128
 
 # create a new virtual environment
@@ -109,6 +112,23 @@ source .venv/bin/activate
 uv pip install torch==2.10.0 torchvision
 uv pip install --no-build-isolation axolotl[deepspeed]
 
+# recommended - install cut-cross-entropy
+uv pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@main"
+
+# (optional) - prefetch flash-attn2 and causal-conv1d kernels
+uv run --python 3.12 python -c "from kernels import get_kernel; get_kernel('kernels-community/flash-attn2'); get_kernel('kernels-community/causal-conv1d')"
+
+# Download example axolotl configs, deepspeed configs
+axolotl fetch examples
+axolotl fetch deepspeed_configs  # OPTIONAL
+```
+
+#### Using pip
+
+```bash
+pip3 install -U packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]
+
 # Download example axolotl configs, deepspeed configs
 axolotl fetch examples
 axolotl fetch deepspeed_configs  # OPTIONAL
@@ -118,7 +138,7 @@ axolotl fetch deepspeed_configs  # OPTIONAL
 
 Installing with Docker can be less error prone than installing in your own environment.
 ```bash
-docker run --gpus '"all"' --ipc=host --rm -it axolotlai/axolotl:main-latest
+docker run --gpus '"all"' --rm -it axolotlai/axolotl:main-latest
 ```
 
 Other installation approaches are described [here](https://docs.axolotl.ai/docs/installation.html).

_quarto.yml

@@ -134,6 +134,7 @@ quartodoc:
         - monkeypatch.stablelm_attn_hijack_flash
         - monkeypatch.trainer_fsdp_optim
         - monkeypatch.transformers_fa_utils
+        - monkeypatch.unsloth_
         - monkeypatch.data.batch_dataset_fetcher
         - monkeypatch.mixtral
         - monkeypatch.gradient_checkpointing.offload_cpu
@@ -326,6 +327,7 @@ website:
         - section: "Advanced Features"
           contents:
             - docs/fsdp_qlora.qmd
+            - docs/unsloth.qmd
             - docs/torchao.qmd
             - docs/custom_integrations.qmd
             - docs/sequence_parallelism.qmd

cicd/Dockerfile-uv.jinja

@@ -22,6 +22,15 @@ WORKDIR /workspace/axolotl
 RUN git fetch origin +$GITHUB_REF && \
     git checkout FETCH_HEAD
 
+# If AXOLOTL_EXTRAS is set, append it in brackets
+RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
+        sed -i 's#^transformers.*#transformers @ git+https://github.com/huggingface/transformers.git@main#' requirements.txt; \
+        sed -i 's#^peft.*#peft @ git+https://github.com/huggingface/peft.git@main#' requirements.txt; \
+        sed -i 's#^accelerate.*#accelerate @ git+https://github.com/huggingface/accelerate.git@main#' requirements.txt; \
+        sed -i 's#^trl.*#trl @ git+https://github.com/huggingface/trl.git@main#' requirements.txt; \
+        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 torchvision
 RUN uv pip uninstall causal_conv1d
@@ -31,21 +40,11 @@ RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
         uv pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
     fi
 
-# Override with nightly HF packages for nightly builds
+RUN python scripts/unsloth_install.py --uv | sh
-RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
-        uv pip install --no-deps \
-            "transformers @ git+https://github.com/huggingface/transformers.git@main" \
-            "peft @ git+https://github.com/huggingface/peft.git@main" \
-            "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
-            "trl @ git+https://github.com/huggingface/trl.git@main" \
-            "datasets @ git+https://github.com/huggingface/datasets.git@main"; \
-    fi
-
 RUN python scripts/cutcrossentropy_install.py --uv | sh
 
 # So we can test the Docker image
-RUN uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \
+RUN uv pip install -r requirements-dev.txt -r requirements-tests.txt
-    codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse
 
 # fix so that git fetch/pull from remote works
 RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \

cicd/Dockerfile.jinja

@@ -0,0 +1,54 @@
+FROM axolotlai/axolotl-base:{{ BASE_TAG }}
+
+ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+ENV AXOLOTL_EXTRAS="{{ AXOLOTL_EXTRAS }}"
+ENV AXOLOTL_ARGS="{{ AXOLOTL_ARGS }}"
+ENV CUDA="{{ CUDA }}"
+ENV PYTORCH_VERSION="{{ PYTORCH_VERSION }}"
+ENV GITHUB_REF="{{ GITHUB_REF }}"
+ENV GITHUB_SHA="{{ GITHUB_SHA }}"
+ENV NIGHTLY_BUILD="{{ NIGHTLY_BUILD }}"
+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
+
+WORKDIR /workspace
+
+RUN git clone --depth=1 https://github.com/axolotl-ai-cloud/axolotl.git
+
+WORKDIR /workspace/axolotl
+
+RUN git fetch origin +$GITHUB_REF && \
+    git checkout FETCH_HEAD
+
+# If AXOLOTL_EXTRAS is set, append it in brackets
+RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
+        sed -i 's#^transformers.*#transformers @ git+https://github.com/huggingface/transformers.git@main#' requirements.txt; \
+        sed -i 's#^peft.*#peft @ git+https://github.com/huggingface/peft.git@main#' requirements.txt; \
+        sed -i 's#^accelerate.*#accelerate @ git+https://github.com/huggingface/accelerate.git@main#' requirements.txt; \
+        sed -i 's#^trl.*#trl @ git+https://github.com/huggingface/trl.git@main#' requirements.txt; \
+        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 if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+    else \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
+    fi
+
+RUN python scripts/unsloth_install.py | sh
+RUN python scripts/cutcrossentropy_install.py | sh
+
+# So we can test the Docker image
+RUN pip install -r requirements-dev.txt -r requirements-tests.txt
+
+# fix so that git fetch/pull from remote works
+RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
+    git config --get remote.origin.fetch
+
+# helper for huggingface-login cli
+RUN git config --global credential.helper store

cicd/cicd.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e
 
-python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__, f'Expected torch $PYTORCH_VERSION but got {torch.__version__}'"
+python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__"
 
 set -o pipefail
 for i in 1 2 3; do

cicd/multigpu.py

@@ -17,7 +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-uv.jinja")
+dockerfile = os.environ.get("E2E_DOCKERFILE", "Dockerfile.jinja")
 df_template = template_env.get_template(dockerfile)
 
 df_args = {

cicd/single_gpu.py

@@ -16,7 +16,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-uv.jinja")
+dockerfile = os.environ.get("E2E_DOCKERFILE", "Dockerfile.jinja")
 df_template = template_env.get_template(dockerfile)
 
 df_args = {

docker/Dockerfile

@@ -32,7 +32,7 @@ RUN if [ "$TARGETARCH" = "arm64" ]; 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 && \
+    fi && \    python scripts/unsloth_install.py | sh && \
     python scripts/cutcrossentropy_install.py | sh && \
     pip install pytest && \
     pip cache purge

docker/Dockerfile-uv

@@ -33,6 +33,7 @@ RUN if [ "$TARGETARCH" = "arm64" ]; then \
     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

docs/debugging.qmd

@@ -76,9 +76,8 @@ datasets:
 Make sure you have an [editable install](https://setuptools.pypa.io/en/latest/userguide/development_mode.html) of Axolotl, which ensures that changes you make to the code are reflected at runtime.  Run the following commands from the root of this project:
 
 ```bash
-export UV_TORCH_BACKEND=cu128  # or cu130
+pip3 install packaging
-uv sync --extra flash-attn --extra deepspeed --group dev --group test
+pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'
-source .venv/bin/activate
 ```
 
 #### Remote Hosts
@@ -209,17 +208,17 @@ cd axolotl
 Next, run the desired docker image and mount the current directory. Below is a docker command you can run to do this:[^2]
 
 ```bash
-docker run --privileged --gpus '"all"' --shm-size 10g --rm -it --name axolotl --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --mount type=bind,src="${PWD}",target=/workspace/axolotl -v ${HOME}/.cache/huggingface:/root/.cache/huggingface axolotlai/axolotl-uv:main-latest
+docker run --privileged --gpus '"all"' --shm-size 10g --rm -it --name axolotl --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --mount type=bind,src="${PWD}",target=/workspace/axolotl -v ${HOME}/.cache/huggingface:/root/.cache/huggingface axolotlai/axolotl:main-py3.10-cu118-2.0.1
 ```
 
 >[!Tip]
 > To understand which containers are available, see the [Docker section of the README](../README.md#docker) and the [DockerHub repo](https://hub.docker.com/r/axolotlai/axolotl/tags).  For details of how the Docker containers are built, see axolotl's [Docker CI builds](../.github/workflows/main.yml).
 
-You will now be in the container.  Next, install Axolotl with dev dependencies:
+You will now be in the container.  Next, perform an editable install of Axolotl:
 
 ```bash
-uv sync --extra flash-attn --extra deepspeed --group dev --group test
+pip3 install packaging
-source .venv/bin/activate
+pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'
 ```
 
 ### Attach To Container

docs/docker.qmd

@@ -6,30 +6,23 @@ format:
     toc-depth: 4
 ---
 
-This section describes the different Docker images that are released by AxolotlAI at
+This section describes the different Docker images that are released by AxolotlAI at [Docker Hub](https://hub.docker.com/u/axolotlai).
-[Docker Hub](https://hub.docker.com/u/axolotlai).
 
 ::: {.callout-important}
-For Blackwell GPUs, please use the tags with PyTorch 2.9.1 and CUDA 12.8.
+For Blackwell GPUs, please use the tags with PyTorch 2.7.1 and CUDA 12.8.
-:::
-
-::: {.callout-tip}
-Each image below is available in a **uv variant** that uses [uv](https://docs.astral.sh/uv/) with
-a relocatable venv (`/workspace/axolotl-venv`) instead of Miniconda + pip. Append `-uv` to the image name
-(e.g. `axolotlai/axolotl-base-uv`). Tags follow the same format. We recommend the uv images for new deployments.
 :::
 
 ## Base
 
-The base image is the most minimal image that can install Axolotl. It is based on the `nvidia/cuda` image.
+The base image is the most minimal image that can install Axolotl. It is based on the `nvidia/cuda` image. It includes python, torch, git, git-lfs, awscli, pydantic, and more.
-It includes python, torch, git, git-lfs, awscli, pydantic, and more.
 
 #### Image
 
-| Variant | Image | Docker Hub |
+```
-|---------|-------|------------|
+axolotlai/axolotl-base
-| pip | `axolotlai/axolotl-base` | [Link](https://hub.docker.com/r/axolotlai/axolotl-base) |
+```
-| uv | `axolotlai/axolotl-base-uv` | [Link](https://hub.docker.com/r/axolotlai/axolotl-base-uv) |
+
+Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl-base)
 
 #### Tags format
 
@@ -39,10 +32,8 @@ main-base-py{python_version}-cu{cuda_version}-{pytorch_version}
 
 Tags examples:
 
+- `main-base-py3.11-cu128-2.8.0`
 - `main-base-py3.11-cu128-2.9.1`
-- `main-base-py3.12-cu128-2.10.0`
-- `main-base-py3.12-cu130-2.9.1`
-- `main-base-py3.12-cu130-2.10.0`
 
 ## Main
 
@@ -50,10 +41,11 @@ The main image is the image that is used to run Axolotl. It is based on the `axo
 
 #### Image
 
-| Variant | Image | Docker Hub |
+```
-|---------|-------|------------|
+axolotlai/axolotl
-| pip | `axolotlai/axolotl` | [Link](https://hub.docker.com/r/axolotlai/axolotl) |
+```
-| uv | `axolotlai/axolotl-uv` | [Link](https://hub.docker.com/r/axolotlai/axolotl-uv) |
+
+Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl)
 
 #### Tags format {#sec-main-tags}
 
@@ -61,7 +53,7 @@ The main image is the image that is used to run Axolotl. It is based on the `axo
 # on push to main
 main-py{python_version}-cu{cuda_version}-{pytorch_version}
 
-# latest main (currently torch 2.9.1, python 3.11, cuda 12.8)
+# latest main (currently torch 2.6.0, python 3.11, cuda 12.4)
 main-latest
 
 # nightly build
@@ -79,12 +71,11 @@ There may be some extra tags appended to the image, like `-vllm` which installs
 
 Tags examples:
 
+- `main-py3.11-cu128-2.8.0`
 - `main-py3.11-cu128-2.9.1`
-- `main-py3.12-cu128-2.10.0`
-- `main-py3.12-cu130-2.9.1`
-- `main-py3.12-cu130-2.10.0`
 - `main-latest`
-- `main-20260315-py3.11-cu128-2.9.1`
+- `main-20250303-py3.11-cu124-2.6.0`
+- `main-20250303-py3.11-cu126-2.6.0`
 - `0.12.0`
 
 ## Cloud
@@ -99,10 +90,11 @@ Jupyter lab is run by default. Set `JUPYTER_DISABLE=1` in the environment variab
 
 #### Image
 
-| Variant | Image | Docker Hub |
+```
-|---------|-------|------------|
+axolotlai/axolotl-cloud
-| pip | `axolotlai/axolotl-cloud` | [Link](https://hub.docker.com/r/axolotlai/axolotl-cloud) |
+```
-| uv | `axolotlai/axolotl-cloud-uv` | [Link](https://hub.docker.com/r/axolotlai/axolotl-cloud-uv) |
+
+Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl-cloud)
 
 #### Tags format
 

docs/installation.qmd

@@ -15,30 +15,64 @@ This guide covers all the ways you can install and set up Axolotl for your envir
 
 - NVIDIA GPU (Ampere architecture or newer for `bf16` and Flash Attention) or AMD GPU
 - Python ≥3.11
-- PyTorch ≥2.9.0
+- PyTorch ≥2.6.0
 
-## Installation {#sec-installation}
+## Installation Methods {#sec-installation-methods}
+
+::: {.callout-important}
+Please make sure to have Pytorch installed before installing Axolotl in your local environment.
+
+Follow the instructions at: [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/)
+:::
 
 ::: {.callout-important}
 For Blackwell GPUs, please use Pytorch 2.9.1 and CUDA 12.8.
 :::
 
-### Quick Install {#sec-uv}
+### PyPI Installation (Recommended) {#sec-pypi}
 
-Axolotl uses [uv](https://docs.astral.sh/uv/) as its package manager. uv is a fast, reliable Python package installer and resolver built in Rust.
+```{.bash}
+pip3 install -U packaging setuptools wheel ninja
+pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]
+```
 
-Install uv if not already installed:
+We use `--no-build-isolation` in order to detect the installed PyTorch version (if
+installed) in order not to clobber it, and so that we set the correct version of
+dependencies that are specific to the PyTorch version or other installed
+co-dependencies.
+
+### uv Installation {#sec-uv}
+
+uv is a fast, reliable Python package installer and resolver built in Rust. It offers significant performance improvements over pip and provides better dependency resolution, making it an excellent choice for complex environments.
+
+Install uv if not already installed
 ```{.bash}
 curl -LsSf https://astral.sh/uv/install.sh | sh
 source $HOME/.local/bin/env
 ```
 
-Choose your CUDA version (e.g. `cu128`, `cu130`), create a venv, and install:
+Choose your CUDA version to use with PyTorch; e.g. `cu124`, `cu126`, `cu128`,
+then create the venv and activate
 ```{.bash}
-export UV_TORCH_BACKEND=cu128  # or cu130
+export UV_TORCH_BACKEND=cu126
 uv venv --no-project --relocatable
 source .venv/bin/activate
-uv pip install --no-build-isolation axolotl[flash-attn,deepspeed]
+```
+
+Install PyTorch
+- PyTorch 2.6.0 recommended
+```{.bash}
+uv pip install packaging setuptools wheel
+uv pip install torch==2.6.0
+uv pip install awscli pydantic
+```
+
+Install axolotl from PyPi
+```{.bash}
+uv pip install --no-build-isolation axolotl[deepspeed,flash-attn]
+
+# optionally install with vLLM if you're using torch==2.6.0 and want to train w/ GRPO
+uv pip install --no-build-isolation axolotl[deepspeed,flash-attn,vllm]
 ```
 
 ### Edge/Development Build {#sec-edge-build}
@@ -48,17 +82,14 @@ For the latest features between releases:
 ```{.bash}
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
-export UV_TORCH_BACKEND=cu128  # or cu130
+pip3 install -U packaging setuptools wheel ninja
-uv sync --extra flash-attn --extra deepspeed
+pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'
-source .venv/bin/activate
 ```
 
-`uv sync` creates a `.venv`, installs exact pinned versions from `uv.lock`, and sets up an editable install automatically.
-
 ### Docker {#sec-docker}
 
 ```{.bash}
-docker run --gpus '"all"' --rm -it --ipc=host axolotlai/axolotl-uv:main-latest
+docker run --gpus '"all"' --rm -it axolotlai/axolotl:main-latest
 ```
 
 For development with Docker:
@@ -75,12 +106,12 @@ docker run --privileged --gpus '"all"' --shm-size 10g --rm -it \
   --ulimit memlock=-1 --ulimit stack=67108864 \
   --mount type=bind,src="${PWD}",target=/workspace/axolotl \
   -v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
-  axolotlai/axolotl-uv:main-latest
+  axolotlai/axolotl:main-latest
 ```
 :::
 
 ::: {.callout-important}
-For Blackwell GPUs, please use `axolotlai/axolotl-uv:main-py3.11-cu128-2.9.1` or the cloud variant `axolotlai/axolotl-cloud-uv:main-py3.11-cu128-2.9.1`.
+For Blackwell GPUs, please use `axolotlai/axolotl:main-py3.11-cu128-2.9.1` or the cloud variant `axolotlai/axolotl-cloud:main-py3.11-cu128-2.9.1`.
 :::
 
 Please refer to the [Docker documentation](docker.qmd) for more information on the different Docker images that are available.
@@ -91,7 +122,7 @@ Please refer to the [Docker documentation](docker.qmd) for more information on t
 
 For providers supporting Docker:
 
-- Use `axolotlai/axolotl-cloud-uv:main-latest`
+- Use `axolotlai/axolotl-cloud:main-latest`
 - Available on:
     - [RunPod](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
     - [Vast.ai](https://cloud.vast.ai?ref_id=62897&template_id=bdd4a49fa8bce926defc99471864cace&utm_source=axolotl&utm_medium=partner&utm_campaign=template_launch_july2025&utm_content=docs_link)
@@ -110,7 +141,7 @@ For providers supporting Docker:
 ### macOS {#sec-macos}
 
 ```{.bash}
-uv pip install --no-build-isolation -e '.'
+pip3 install --no-build-isolation -e '.'
 ```
 
 See @sec-troubleshooting for Mac-specific issues.
@@ -121,44 +152,21 @@ See @sec-troubleshooting for Mac-specific issues.
 We recommend using WSL2 (Windows Subsystem for Linux) or Docker.
 :::
 
-## Migrating from pip to uv {#sec-migrating}
+## Environment Managers {#sec-env-managers}
 
-If you have an existing pip-based Axolotl installation, you can migrate to uv:
+### Conda/Pip venv {#sec-conda}
 
-```{.bash}
+1. Install Python ≥3.11
-# Install uv
+2. Install PyTorch: https://pytorch.org/get-started/locally/
-curl -LsSf https://astral.sh/uv/install.sh | sh
+3. Install Axolotl:
-source $HOME/.local/bin/env
+   ```{.bash}
-
+   pip3 install -U packaging setuptools wheel ninja
-# Create a fresh venv (recommended for a clean start)
+   pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'
-export UV_TORCH_BACKEND=cu128  # or cu130
+   ```
-uv venv --no-project --relocatable
+4. (Optional) Login to Hugging Face:
-source .venv/bin/activate
+   ```{.bash}
-
+   hf auth login
-# Reinstall axolotl
+   ```
-uv pip install --no-build-isolation axolotl[flash-attn,deepspeed]
-```
-
-## Using pip (Alternative) {#sec-pip}
-
-If you are unable to install uv, you can still use pip directly.
-
-::: {.callout-important}
-Please make sure to have PyTorch installed before installing Axolotl with pip.
-
-Follow the instructions at: [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/)
-:::
-
-```{.bash}
-pip3 install -U packaging setuptools wheel ninja
-pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]
-```
-
-For editable/development installs:
-```{.bash}
-pip3 install -U packaging setuptools wheel ninja
-pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'
-```
 
 ## Troubleshooting {#sec-troubleshooting}
 

docs/unsloth.qmd

@@ -0,0 +1,53 @@
+---
+title: "Unsloth"
+description: "Hyper-optimized QLoRA finetuning for single GPUs"
+---
+
+### Overview
+
+Unsloth provides hand-written optimized kernels for LLM finetuning that slightly improve speed and VRAM over
+standard industry baselines.
+
+::: {.callout-important}
+Due to breaking changes in transformers `v4.48.0`, users will need to downgrade to `<=v4.47.1` to use this patch.
+
+This will later be deprecated in favor of [LoRA Optimizations](lora_optims.qmd).
+:::
+
+
+### Installation
+
+The following will install the correct unsloth and extras from source.
+
+```bash
+python scripts/unsloth_install.py | sh
+```
+
+### Usage
+
+Axolotl exposes a few configuration options to try out unsloth and get most of the performance gains.
+
+Our unsloth integration is currently limited to the following model architectures:
+ - llama
+
+These options are specific to LoRA finetuning and cannot be used for multi-GPU finetuning
+```yaml
+unsloth_lora_mlp: true
+unsloth_lora_qkv: true
+unsloth_lora_o: true
+```
+
+These options are composable and can be used with multi-gpu finetuning
+```yaml
+unsloth_cross_entropy_loss: true
+unsloth_rms_norm: true
+unsloth_rope: true
+```
+
+### Limitations
+
+- Single GPU only; e.g. no multi-gpu support
+- No deepspeed or FSDP support (requires multi-gpu)
+- LoRA + QLoRA support only. No full fine tunes or fp8 support.
+- Limited model architecture support. Llama, Phi, Gemma, Mistral only
+- No MoE support.

examples/LiquidAI/README.md

@@ -15,7 +15,8 @@ Thanks to the team at LiquidAI for giving us early access to prepare for these r
     Here is an example of how to install from pip:
     ```bash
     # Ensure you have a compatible version of Pytorch installed
-    uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+    pip3 install packaging setuptools wheel ninja
+    pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
     ```
 
 2.  Run one of the finetuning examples below.
@@ -34,7 +35,7 @@ Thanks to the team at LiquidAI for giving us early access to prepare for these r
 
     **LFM2-MoE**
     ```bash
-    uv pip install git+https://github.com/huggingface/transformers.git@0c9a72e4576fe4c84077f066e585129c97bfd4e6
+    pip install git+https://github.com/huggingface/transformers.git@0c9a72e4576fe4c84077f066e585129c97bfd4e6
 
     # LoRA SFT (1x48GB @ 16.2GiB)
     axolotl train examples/LiquidAI/lfm2-8b-a1b-lora.yaml
@@ -44,7 +45,7 @@ Thanks to the team at LiquidAI for giving us early access to prepare for these r
 
 - **Installation Error**: If you encounter `ImportError: ... undefined symbol ...` or `ModuleNotFoundError: No module named 'causal_conv1d_cuda'`, the `causal-conv1d` package may have been installed incorrectly. Try uninstalling it:
   ```bash
-  uv pip uninstall causal-conv1d
+  pip uninstall -y causal-conv1d
   ```
 
 - **Dataset Loading**: Read more on how to load your own dataset in our [documentation](https://docs.axolotl.ai/docs/dataset_loading.html).

examples/apertus/README.md

@@ -15,7 +15,8 @@ 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
 
-uv pip install --no-build-isolation -e '.[flash-attn]'
+pip3 install packaging==26.0 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
@@ -30,7 +31,7 @@ python scripts/cutcrossentropy_install.py | sh
 # For those using our Docker image, use the below path.
 export CUDA_HOME=/usr/local/cuda
 
-uv pip install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
+pip3 install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
 ```
 
 For any installation errors, see [XIELU Installation Issues](#xielu-installation-issues)
@@ -66,7 +67,7 @@ If those didn't help, please try the below solutions:
 1. Pass env for CMAKE and try install again:
 
     ```bash
-    Python_EXECUTABLE=$(which python) uv pip install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
+    Python_EXECUTABLE=$(which python) pip3 install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
     ```
 
 2. Git clone the repo and manually hardcode python path:
@@ -91,7 +92,7 @@ If those didn't help, please try the below solutions:
     ```
 
     ```bash
-    uv pip install . --no-build-isolation --no-deps
+    pip3 install . --no-build-isolation --no-deps
     ```
 
 ## Optimization Guides

examples/arcee/README.md

@@ -17,7 +17,8 @@ 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
 
-uv pip install --no-build-isolation -e '.[flash-attn]'
+pip3 install packaging==26.0 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

examples/devstral/README.md

@@ -16,7 +16,8 @@ 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)
-uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 
 2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage

examples/gemma3n/README.md

@@ -10,16 +10,17 @@ 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)
-uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 
 2. In addition to Axolotl's requirements, Gemma-3n requires:
 
 ```bash
-uv pip install timm==1.0.17
+pip3 install timm==1.0.17
 
 # for loading audio data
-uv pip install librosa==0.11.0
+pip3 install librosa==0.11.0
 ```
 
 3. Download sample dataset files

examples/gpt-oss/README.md

@@ -14,7 +14,8 @@ 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)
-uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 
 2. Choose one of the following configs below for training the 20B model. (for 120B, see [below](#training-120b))
@@ -86,7 +87,7 @@ for more information about using a special vllm-openai docker image for inferenc
 Optionally, vLLM can be installed from nightly:
 
 ```bash
-uv pip install --no-build-isolation --pre -U vllm --extra-index-url https://wheels.vllm.ai/nightly
+pip install --no-build-isolation --pre -U vllm --extra-index-url https://wheels.vllm.ai/nightly
 ```
 and the vLLM server can be started with the following command (modify `--tensor-parallel-size 8` to match your environment):
 ```bash

examples/granite4/README.md

@@ -15,7 +15,8 @@ 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
 
-uv pip install --no-build-isolation -e '.[flash-attn]'
+pip3 install packaging==26.0 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

examples/hunyuan/README.md

@@ -13,7 +13,8 @@ Tencent released a family of opensource models called HunYuan with varying param
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
 
-uv pip install --no-build-isolation -e '.[flash-attn]'
+pip3 install packaging==26.0 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

examples/internvl3_5/README.md

@@ -11,7 +11,7 @@ This guide shows how to fine-tune it with Axolotl.
 2. Install `timm` for vision model support:
 
     ```bash
-    uv pip install timm==1.0.19
+    pip install timm==1.0.19
     ```
 
 3. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.

examples/magistral/README.md

@@ -14,7 +14,8 @@ 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)
-uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 
 2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage

examples/magistral/vision/README.md

@@ -12,7 +12,7 @@ Before starting, ensure you have:
 
 1. Install the required vision lib:
     ```bash
-    uv pip install 'mistral-common[opencv]==1.8.5'
+    pip install 'mistral-common[opencv]==1.8.5'
     ```
 
 2. Download the example dataset image:

examples/ministral3/README.md

@@ -23,7 +23,7 @@ Note: This is still experimental given it is based on transformers v5 RC.
     git checkout transformers-v5
 
     # Install packages for transformers v5
-    uv pip install -e .
+    pip install -e .
     ```
 
 4. Run the fine-tuning:

examples/ministral3/vision/README.md

@@ -12,7 +12,7 @@ Before starting, ensure you have:
 
 1. Install the required vision lib:
     ```bash
-    uv pip install 'mistral-common[opencv]==1.8.6'
+    pip install 'mistral-common[opencv]==1.8.6'
     ```
 
 2. Download the example dataset image:

examples/mistral-small/README.md

@@ -12,7 +12,7 @@ Before starting, ensure you have:
 
 1. Install the required vision lib:
     ```bash
-    uv pip install 'mistral-common[opencv]==1.8.5'
+    pip install 'mistral-common[opencv]==1.8.5'
     ```
 
 2. Download the example dataset image:

examples/mistral4/README.md

@@ -13,7 +13,7 @@ Thanks to the team at MistralAI for giving us early access to prepare for this r
 3. Install transformers from main
 
   ```bash
-  uv pip install git+https://github.com/huggingface/transformers.git
+  pip install git+https://github.com/huggingface/transformers.git
   ```
 
 4. Run one of the example configs:

examples/qwen3-next/README.md

@@ -12,7 +12,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
 
 3. Install FLA for improved performance
 ```bash
-uv pip uninstall causal-conv1d && uv pip install flash-linear-attention==0.4.1
+pip3 uninstall -y causal-conv1d && pip3 install flash-linear-attention==0.4.1
 ```
 
 4. Run the finetuning example:

examples/qwen3.5/README.md

@@ -10,7 +10,7 @@
 
 3. Install FLA for sample packing support with the Gated DeltaNet linear attention layers:
   ```bash
-  uv pip uninstall causal-conv1d && uv pip install flash-linear-attention==0.4.1
+  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.
 

examples/seed-oss/README.md

@@ -11,7 +11,8 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
     Here is an example of how to install from pip:
     ```bash
     # Ensure you have a compatible version of Pytorch installed
-    uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+    pip3 install packaging setuptools wheel ninja
+    pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 
     # Install Cut Cross Entropy
     python scripts/cutcrossentropy_install.py | sh

examples/smolvlm2/README.md

@@ -13,13 +13,14 @@ This guide shows how to fine-tune SmolVLM2 models with Axolotl.
     Here is an example of how to install from pip:
     ```bash
     # Ensure you have a compatible version of Pytorch installed
-    uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+    pip3 install packaging setuptools wheel ninja
+    pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
     ```
 
 2. Install an extra dependency:
 
     ```bash
-    uv pip install num2words==0.5.14
+    pip3 install num2words==0.5.14
     ```
 
 3.  Run the finetuning example:

examples/voxtral/README.md

@@ -12,15 +12,16 @@ 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)
-uv pip install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 
 2. Please install the below.
 
 ```bash
 # audio
-uv pip install librosa==0.11.0
+pip3 install librosa==0.11.0
-uv pip install 'mistral_common[audio]==1.8.3'
+pip3 install 'mistral_common[audio]==1.8.3'
 
 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
 python scripts/cutcrossentropy_install.py | sh

pyproject.toml

@@ -1,165 +1,15 @@
 [build-system]
-requires = ["setuptools>=64", "wheel", "setuptools_scm>=8"]
+requires = ["setuptools>=64", "wheel", "setuptools_scm>=8", "packaging==26.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "axolotl"
-dynamic = ["version"]
+dynamic = ["version", "dependencies", "optional-dependencies"]
 description = "LLM Trainer"
 readme = "README.md"
 requires-python = ">=3.10"
 # license = "Apache-2.0"
 
-dependencies = [
-    # Core ML stack
-    "torch>=2.6.0",
-    "packaging==26.0",
-    "huggingface_hub>=1.1.7",
-    "peft>=0.19.1,<0.20.0",
-    "tokenizers>=0.22.1",
-    "transformers==5.5.4",
-    "accelerate==1.13.0",
-    "datasets>=4.8.4,<4.9.0",
-    "trl==1.1.0",
-    "hf_xet==1.4.3",
-    "kernels==0.13.0",
-    "trackio>=0.16.1",
-    "typing-extensions>=4.15.0",
-    "optimum==1.16.2",
-    "hf_transfer",
-    "sentencepiece",
-    "gradio>=6.2.0,<7.0",
-    "modal==1.3.0.post1",
-    "pydantic>=2.10.6",
-    "addict",
-    "fire",
-    "PyYAML>=6.0",
-    "requests",
-    "wandb",
-    "einops",
-    "colorama",
-    "numba>=0.61.2",
-    "numpy>=2.2.6",
-
-    # Evaluation & metrics
-    "evaluate==0.4.1",
-    "scipy",
-    "nvidia-ml-py==12.560.30",
-    "art",
-    "tensorboard",
-    "python-dotenv==1.0.1",
-
-    # Remote filesystems
-    "s3fs>=2024.5.0",
-    "gcsfs>=2025.3.0",
-    "adlfs>=2024.5.0",
-    "ocifs==1.3.2",
-
-    "zstandard==0.22.0",
-    "fastcore",
-
-    # lm eval harness
-    "lm_eval==0.4.11",
-    "langdetect==1.0.9",
-    "immutabledict==4.2.0",
-    "antlr4-python3-runtime==4.13.2",
-
-    "schedulefree==1.4.1",
-    "openenv-core==0.1.0",
-
-    # Axolotl contribs
-    "axolotl-contribs-lgpl==0.0.7",
-    "axolotl-contribs-mit==0.0.6",
-
-    # Telemetry
-    "posthog==6.7.11",
-
-    "mistral-common==1.11.0",
-
-    # Platform-specific (Linux only)
-    "bitsandbytes==0.49.1 ; sys_platform != 'darwin'",
-    "triton>=3.4.0 ; sys_platform != 'darwin'",
-    "xformers>=0.0.23.post1 ; sys_platform != 'darwin'",
-    "liger-kernel==0.7.0 ; sys_platform != 'darwin'",
-    "torchao==0.17.0 ; sys_platform != 'darwin' and platform_machine != 'aarch64'",
-
-    # Architecture-specific
-    "fla-core==0.4.1 ; platform_machine != 'aarch64'",
-    "flash-linear-attention==0.4.1 ; platform_machine != 'aarch64'",
-]
-
-[project.optional-dependencies]
-flash-attn = ["flash-attn==2.8.3"]
-ring-flash-attn = [
-    "flash-attn==2.8.3",
-    "ring-flash-attn>=0.1.7",
-]
-deepspeed = [
-    "deepspeed>=0.18.6,<0.19.0",
-    "deepspeed-kernels",
-]
-mamba-ssm = [
-    "mamba-ssm==1.2.0.post1",
-    "causal_conv1d",
-]
-auto-gptq = [
-    "auto-gptq==0.5.1",
-]
-mlflow = [
-    "mlflow",
-]
-galore = [
-    "galore_torch",
-]
-apollo = [
-    "apollo-torch",
-]
-optimizers = [
-    "galore_torch",
-    "apollo-torch",
-    "lomo-optim==0.1.1",
-    "torch-optimi==0.2.1",
-    "came_pytorch==0.1.3",
-]
-ray = [
-    "ray[train]>=2.52.1",
-]
-vllm = [
-    "vllm>=0.15.0",
-]
-llmcompressor = [
-    "llmcompressor>=0.10.0",
-]
-fbgemm-gpu = ["fbgemm-gpu-genai>=1.3.0"]
-opentelemetry = [
-    "opentelemetry-api",
-    "opentelemetry-sdk",
-    "opentelemetry-exporter-prometheus",
-    "prometheus-client",
-]
-
-[dependency-groups]
-dev = [
-    "black",
-    "mypy",
-    "pre-commit",
-    "types-requests",
-    "quartodoc",
-    "jupyter",
-    "blobfile",
-    "tiktoken",
-]
-test = [
-    "codecov",
-    "codecov-cli",
-    "pytest",
-    "pytest-cov",
-    "pytest-retry",
-    "pytest-sugar",
-    "pytest-xdist",
-    "tbparse",
-]
-
 [project.scripts]
 axolotl = "axolotl.cli.main:main"
 
@@ -168,15 +18,18 @@ Homepage = "https://axolotl.ai/"
 Documentation = "https://docs.axolotl.ai/"
 Repository = "https://github.com/axolotl-ai-cloud/axolotl.git"
 
-[tool.setuptools]
+[tool.setuptools_scm]
-include-package-data = true
 
-[tool.setuptools.packages.find]
+[tool.setuptools]
-where = ["src"]
+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"
+
 [tool.ruff]
 line-length = 88
 target-version = "py310"
@@ -214,43 +67,5 @@ markers = [
     "slow: marks tests as slow",
 ]
 
-# UV specific configuration
-[tool.uv]
-prerelease = "allow"
-conflicts = [
-    [
-        { package = "axolotl" },
-        { extra = "vllm" },
-    ],
-    [
-        { package = "axolotl" },
-        { extra = "flash-attn" },
-    ],
-    [
-        { package = "axolotl" },
-        { extra = "ring-flash-attn" },
-    ],
-    [
-        { package = "axolotl" },
-        { extra = "mamba-ssm" },
-    ],
-    [
-        { package = "axolotl" },
-        { extra = "auto-gptq" },
-    ],
-    [
-        { package = "axolotl" },
-        { extra = "fbgemm-gpu" },
-    ],
-    [
-        { package = "axolotl" },
-        { extra = "llmcompressor" },
-    ],
-]
-
 [tool.uv.extra-build-dependencies]
-mamba-ssm = [{ requirement = "torch", match-runtime = true }]
+axolotl = ["huggingface_hub"]
-causal-conv1d = [{ requirement = "torch", match-runtime = true }]
-flash-attn = [{ requirement = "torch", match-runtime = true }]
-deepspeed = [{ requirement = "torch", match-runtime = true }]
-auto-gptq = [{ requirement = "torch", match-runtime = true }]

requirements-dev.txt

@@ -0,0 +1,8 @@
+black
+mypy
+pre-commit
+types-requests
+quartodoc
+jupyter
+blobfile
+tiktoken

requirements-tests.txt

@@ -0,0 +1,8 @@
+codecov
+codecov-cli
+pytest
+pytest-cov
+pytest-retry
+pytest-sugar
+pytest-xdist
+tbparse

requirements.txt

@@ -0,0 +1,78 @@
+--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
+mamba-ssm==1.2.0.post1
+xformers>=0.0.23.post1
+liger-kernel==0.7.0
+# END section
+
+packaging==26.0
+huggingface_hub>=1.1.7
+peft>=0.19.0,<0.20.0
+tokenizers>=0.22.1
+transformers==5.5.4
+accelerate==1.13.0
+datasets>=4.8.4,<4.9.0
+deepspeed>=0.18.6,<0.19.0
+trl==1.1.0
+hf_xet==1.4.3
+kernels==0.13.0
+
+fla-core==0.4.1
+flash-linear-attention==0.4.1
+
+trackio>=0.16.1
+typing-extensions>=4.15.0
+
+optimum==1.16.2
+hf_transfer
+sentencepiece
+gradio>=6.2.0,<7.0
+
+modal==1.3.0.post1
+pydantic>=2.10.6
+addict
+fire
+PyYAML>=6.0
+requests
+wandb
+einops
+colorama
+numba>=0.61.2
+numpy>=2.2.6
+
+# qlora things
+evaluate==0.4.1
+scipy
+nvidia-ml-py==12.560.30
+art
+tensorboard
+python-dotenv==1.0.1
+
+# remote filesystems
+s3fs>=2024.5.0
+gcsfs>=2025.3.0
+adlfs>=2024.5.0
+ocifs==1.3.2
+
+zstandard==0.22.0
+fastcore
+
+# lm eval harness
+lm_eval==0.4.11
+langdetect==1.0.9
+immutabledict==4.2.0
+antlr4-python3-runtime==4.13.2
+
+torchao==0.17.0
+openenv-core==0.1.0
+schedulefree==1.4.1
+
+axolotl-contribs-lgpl==0.0.7
+axolotl-contribs-mit==0.0.6
+# telemetry
+posthog==6.7.11
+
+mistral-common==1.11.0

scripts/build_scattermoe_lora_kernel.py

@@ -1,479 +0,0 @@
-#!/usr/bin/env python3
-"""Build a disposable Hugging Face Kernel Hub package for ScatterMoE LoRA.
-
-This script does not move or edit the in-tree Axolotl kernel sources. It copies
-``src/axolotl/integrations/kernels/libs/scattermoe_lora`` into an ignored
-build directory and emits a universal HF kernels project that can be pushed to
-the Hub.
-"""
-
-from __future__ import annotations
-
-import argparse
-import fnmatch
-import hashlib
-import json
-import os
-import shutil
-import subprocess
-import sys
-from importlib import metadata
-from pathlib import Path
-
-PACKAGE_NAME = "scattermoe_lora"
-BUILD_VARIANT = "torch-universal"
-DEFAULT_REPO_ID = "kernels-community/scattermoe-lora"
-HF_REPO_TYPE = "kernel"
-HF_KERNEL_URL_PREFIX = "https://hf.co/kernels"
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-DEFAULT_SOURCE_DIR = (
-    REPO_ROOT / "src" / "axolotl" / "integrations" / "kernels" / "libs" / PACKAGE_NAME
-)
-DEFAULT_OUTPUT_DIR = REPO_ROOT / "build" / "hf-kernels" / PACKAGE_NAME
-
-EXCLUDED_DIRS = {
-    "__pycache__",
-    ".mypy_cache",
-    ".pytest_cache",
-    ".ruff_cache",
-}
-EXCLUDED_FILE_PATTERNS = {
-    "*.pyc",
-    "*.pyo",
-    "*.so",
-    ".DS_Store",
-}
-
-TEXT_REPLACEMENTS = {
-    "from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant import": (
-        "from .selective_dequant import"
-    ),
-    "from axolotl.integrations.kernels.libs.scattermoe_lora.selective_dequant_kernel import": (
-        "from .selective_dequant_kernel import"
-    ),
-    "from axolotl.integrations.kernels.libs.scattermoe_lora.kernels.ops import": (
-        "from .ops import"
-    ),
-}
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description=(
-            "Copy Axolotl's ScatterMoE LoRA Triton kernels into a disposable "
-            "HF Kernel Hub universal package."
-        )
-    )
-    parser.add_argument(
-        "--source-dir",
-        type=Path,
-        default=DEFAULT_SOURCE_DIR,
-        help=f"ScatterMoE LoRA source package to copy. Default: {DEFAULT_SOURCE_DIR}",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=Path,
-        default=DEFAULT_OUTPUT_DIR,
-        help=f"Destination build/dist directory. Default: {DEFAULT_OUTPUT_DIR}",
-    )
-    parser.add_argument(
-        "--repo-id",
-        default=DEFAULT_REPO_ID,
-        help=f"HF Hub repo id to write into build.toml. Default: {DEFAULT_REPO_ID}",
-    )
-    parser.add_argument(
-        "--version",
-        type=int,
-        default=1,
-        help="Kernel major version written to build.toml and metadata.json.",
-    )
-    parser.add_argument(
-        "--force",
-        action="store_true",
-        help="Delete the output directory first if it already exists.",
-    )
-    parser.add_argument(
-        "--no-source-layout",
-        action="store_true",
-        help="Only write the shippable build/ tree, not torch-ext/ sources.",
-    )
-    parser.add_argument(
-        "--upload",
-        action="store_true",
-        help=(
-            "Upload the generated universal kernel package with huggingface_hub. "
-            "This bypasses kernel-builder and is intended for pure Python/Triton "
-            "universal kernels."
-        ),
-    )
-    parser.add_argument(
-        "--private",
-        action="store_true",
-        help="Create the HF Hub repo as private when used with --upload.",
-    )
-    parser.add_argument(
-        "--skip-version-branch",
-        action="store_true",
-        help="With --upload, only upload main and skip the v<version> branch.",
-    )
-    return parser.parse_args()
-
-
-def should_skip_file(path: Path) -> bool:
-    return any(
-        fnmatch.fnmatch(path.name, pattern) for pattern in EXCLUDED_FILE_PATTERNS
-    )
-
-
-def iter_source_files(source_dir: Path) -> list[Path]:
-    files: list[Path] = []
-    for root, dirs, filenames in os.walk(source_dir):
-        dirs[:] = sorted(d for d in dirs if d not in EXCLUDED_DIRS)
-        for filename in sorted(filenames):
-            path = Path(root) / filename
-            if not should_skip_file(path):
-                files.append(path)
-    return files
-
-
-def content_hash(source_dir: Path) -> str:
-    digest = hashlib.sha1()
-    for path in iter_source_files(source_dir):
-        rel = path.relative_to(source_dir).as_posix()
-        digest.update(rel.encode("utf-8"))
-        digest.update(b"\0")
-        digest.update(path.read_bytes())
-        digest.update(b"\0")
-    return digest.hexdigest()[:10]
-
-
-def git_revision() -> str:
-    try:
-        result = subprocess.run(
-            ["git", "rev-parse", "--short", "HEAD"],
-            cwd=REPO_ROOT,
-            check=True,
-            capture_output=True,
-            text=True,
-        )
-    except (OSError, subprocess.CalledProcessError):
-        return "unknown"
-    return result.stdout.strip() or "unknown"
-
-
-def transform_python_source(text: str, rel_path: Path, op_namespace: str) -> str:
-    for old, new in TEXT_REPLACEMENTS.items():
-        text = text.replace(old, new)
-
-    if rel_path.as_posix() == "gemma4_experts.py":
-        text = text.replace(
-            "    from axolotl.integrations.kernels.constants import resolve_experts_class",
-            (
-                "    raise RuntimeError(\n"
-                '        "patch_gemma4_scattermoe is only available from the in-tree Axolotl "\n'
-                '        "integration. Use register_scattermoe_experts() with the standalone "\n'
-                '        "HF kernel package."\n'
-                "    )"
-            ),
-        )
-
-    return text.replace("scattermoe::", f"{op_namespace}::")
-
-
-def copy_package(source_dir: Path, package_dir: Path, op_namespace: str) -> None:
-    for source in iter_source_files(source_dir):
-        rel_path = source.relative_to(source_dir)
-        destination = package_dir / rel_path
-        destination.parent.mkdir(parents=True, exist_ok=True)
-
-        if source.suffix == ".py":
-            text = source.read_text(encoding="utf-8")
-            text = transform_python_source(text, rel_path, op_namespace)
-            destination.write_text(text, encoding="utf-8")
-        else:
-            shutil.copy2(source, destination)
-
-    write_ops_module(package_dir / "_ops.py", op_namespace)
-
-
-def write_ops_module(path: Path, op_namespace: str) -> None:
-    path.write_text(
-        "\n".join(
-            [
-                "import torch",
-                "",
-                f"ops = torch.ops.{op_namespace}",
-                "",
-                "",
-                "def add_op_namespace_prefix(op_name: str) -> str:",
-                f'    return f"{op_namespace}::{{op_name}}"',
-                "",
-            ]
-        ),
-        encoding="utf-8",
-    )
-
-
-def write_build_toml(path: Path, repo_id: str, version: int) -> None:
-    lines = [
-        "[general]",
-        f'name = "{PACKAGE_NAME}"',
-        "universal = true",
-        f"version = {version}",
-        "",
-    ]
-    if repo_id:
-        lines.extend(
-            [
-                "[general.hub]",
-                f'repo-id = "{repo_id}"',
-                "",
-            ]
-        )
-    path.write_text("\n".join(lines), encoding="utf-8")
-
-
-def write_flake(path: Path) -> None:
-    path.write_text(
-        """{
-  description = "Flake for scattermoe_lora kernel";
-
-  inputs = {
-    builder.url = "github:huggingface/kernels";
-  };
-
-  outputs =
-    {
-      self,
-      builder,
-    }:
-    builder.lib.genKernelFlakeOutputs {
-      inherit self;
-      path = ./.;
-    };
-}
-""",
-        encoding="utf-8",
-    )
-
-
-def write_readme(path: Path, repo_id: str, source_hash: str, op_namespace: str) -> None:
-    repo_display = repo_id or "<your-org>/scattermoe-lora"
-    path.write_text(
-        f"""---
-library_name: kernels
-license: apache-2.0
-tags:
-- kernel
-- kernels
----
-
-# ScatterMoE LoRA
-
-Standalone Hugging Face Kernel Hub package for Axolotl's ScatterMoE LoRA Triton kernels.
-
-This package is generated from Axolotl's in-tree `scattermoe_lora` sources and is exported as a universal kernel because the implementation is Python/Triton rather than a precompiled C++/CUDA extension.
-
-```python
-from kernels import get_kernel
-
-scattermoe_lora = get_kernel("{repo_display}")
-```
-
-Export metadata:
-
-- source package: `src/axolotl/integrations/kernels/libs/scattermoe_lora`
-- source revision: `{git_revision()}`
-- source content hash: `{source_hash}`
-- torch custom op namespace: `{op_namespace}`
-
-The generated `build/torch-universal/{PACKAGE_NAME}` directory is the shippable Hub artifact. `torch-ext/{PACKAGE_NAME}` is included so `kernel-builder build-and-copy` can regenerate the universal build tree if desired.
-""",
-        encoding="utf-8",
-    )
-
-
-def write_metadata(path: Path, version: int) -> None:
-    path.write_text(
-        json.dumps({"version": version}, indent=2, sort_keys=True) + "\n",
-        encoding="utf-8",
-    )
-
-
-def prepare_output_dir(output_dir: Path, force: bool) -> None:
-    if output_dir.exists():
-        if not force:
-            raise FileExistsError(
-                f"{output_dir} already exists. Re-run with --force to replace it."
-            )
-        shutil.rmtree(output_dir)
-    output_dir.mkdir(parents=True)
-
-
-def build_package(args: argparse.Namespace) -> Path:
-    source_dir = args.source_dir.resolve()
-    output_dir = args.output_dir.resolve()
-
-    if not source_dir.is_dir():
-        raise FileNotFoundError(f"source package does not exist: {source_dir}")
-    if not (source_dir / "__init__.py").is_file():
-        raise FileNotFoundError(f"source package is missing __init__.py: {source_dir}")
-
-    source_hash = content_hash(source_dir)
-    op_namespace = f"_{PACKAGE_NAME}_{source_hash}"
-
-    prepare_output_dir(output_dir, args.force)
-
-    write_build_toml(output_dir / "build.toml", args.repo_id, args.version)
-    write_flake(output_dir / "flake.nix")
-    write_readme(output_dir / "README.md", args.repo_id, source_hash, op_namespace)
-
-    if not args.no_source_layout:
-        copy_package(source_dir, output_dir / "torch-ext" / PACKAGE_NAME, op_namespace)
-
-    build_package_dir = output_dir / "build" / BUILD_VARIANT / PACKAGE_NAME
-    copy_package(source_dir, build_package_dir, op_namespace)
-    write_metadata(build_package_dir.parent / "metadata.json", args.version)
-
-    return output_dir
-
-
-def upload_package(args: argparse.Namespace, output_dir: Path) -> None:
-    if not args.repo_id:
-        raise ValueError("--repo-id is required when using --upload")
-
-    try:
-        from huggingface_hub import HfApi, constants as hf_constants
-    except ImportError as exc:
-        raise RuntimeError(
-            "--upload requires huggingface_hub. Install it or run the upload "
-            "manually with the Hugging Face CLI."
-        ) from exc
-
-    try:
-        hub_version = metadata.version("huggingface_hub")
-    except metadata.PackageNotFoundError:
-        hub_version = "unknown"
-
-    accepted_repo_types = getattr(
-        hf_constants,
-        "REPO_TYPES_WITH_KERNEL",
-        getattr(hf_constants, "REPO_TYPES", ()),
-    )
-    if HF_REPO_TYPE not in accepted_repo_types:
-        raise RuntimeError(
-            "Your huggingface_hub installation does not support "
-            f"repo_type={HF_REPO_TYPE!r} (found huggingface_hub {hub_version}). "
-            f"Upgrade this interpreter with: {sys.executable} -m pip install --upgrade "
-            "'huggingface_hub>=1.10.0'"
-        )
-
-    # huggingface_hub 1.11.0 has partial kernel support: create_repo accepts
-    # "kernel", but upload_folder/create_commit still validate against the
-    # older REPO_TYPES list. Extend it in-process so those helpers use the
-    # /api/kernels/... endpoints until upstream broadens that check.
-    if HF_REPO_TYPE not in hf_constants.REPO_TYPES:
-        hf_constants.REPO_TYPES.append(HF_REPO_TYPE)
-
-    api = HfApi()
-    try:
-        repo_id = api.create_repo(
-            repo_id=args.repo_id,
-            repo_type=HF_REPO_TYPE,
-            private=args.private,
-            exist_ok=True,
-        ).repo_id
-    except ValueError as exc:
-        if "Invalid repo type" in str(exc):
-            raise RuntimeError(
-                "huggingface_hub rejected repo_type='kernel'. "
-                f"This usually means the command is running with an older Hub "
-                f"client than expected (found huggingface_hub {hub_version} at "
-                f"{sys.executable}). Upgrade with: {sys.executable} -m pip "
-                "install --upgrade 'huggingface_hub>=1.10.0'"
-            ) from exc
-        raise
-
-    delete_patterns = [
-        "build/**",
-        "torch-ext/**",
-        "build.toml",
-        "flake.nix",
-        "README.md",
-    ]
-
-    api.upload_folder(
-        repo_id=repo_id,
-        repo_type=HF_REPO_TYPE,
-        folder_path=output_dir,
-        revision="main",
-        delete_patterns=delete_patterns,
-        commit_message="Upload ScatterMoE LoRA universal kernel",
-    )
-    print(f"Uploaded main branch: {HF_KERNEL_URL_PREFIX}/{repo_id}")
-
-    if args.skip_version_branch:
-        return
-
-    version_branch = f"v{args.version}"
-    api.create_branch(
-        repo_id=repo_id,
-        repo_type=HF_REPO_TYPE,
-        branch=version_branch,
-        revision="main",
-        exist_ok=True,
-    )
-    api.upload_folder(
-        repo_id=repo_id,
-        repo_type=HF_REPO_TYPE,
-        folder_path=output_dir,
-        revision=version_branch,
-        delete_patterns=delete_patterns,
-        commit_message=f"Upload ScatterMoE LoRA universal kernel {version_branch}",
-    )
-    print(
-        f"Uploaded version branch: "
-        f"{HF_KERNEL_URL_PREFIX}/{repo_id}/tree/{version_branch}"
-    )
-
-
-def main() -> int:
-    args = parse_args()
-    try:
-        output_dir = build_package(args)
-        if args.upload:
-            upload_package(args, output_dir)
-    except Exception as exc:
-        print(f"error: {exc}", file=sys.stderr)
-        return 1
-
-    print(f"Wrote ScatterMoE LoRA HF kernel package to: {output_dir}")
-    print(f"Shippable artifact: {output_dir / 'build' / BUILD_VARIANT / PACKAGE_NAME}")
-    if args.upload:
-        print(f'Load it with: get_kernel("{args.repo_id}", version={args.version})')
-        print(f"Uploaded as Hugging Face repo_type={HF_REPO_TYPE!r}.")
-        return 0
-
-    print("Next step:")
-    print("  upload this universal Python/Triton kernel directly:")
-    print(
-        f"    python3 {Path(__file__).as_posix()} "
-        f"--repo-id {args.repo_id} --force --upload"
-    )
-    if shutil.which("kernel-builder") is None:
-        print("  optional: install kernel-builder for full Nix-based builds:")
-        print(
-            "    curl -fsSL "
-            "https://raw.githubusercontent.com/huggingface/kernels/main/install.sh "
-            "| bash"
-        )
-    else:
-        print("  optional: upload with kernel-builder:")
-        print(f"    cd {output_dir}")
-        print("    kernel-builder build-and-upload")
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

scripts/unsloth_install.py

@@ -0,0 +1,40 @@
+# noqa
+import sys
+
+try:
+    import torch
+except ImportError as error:
+    raise ImportError("Install torch via `pip install torch`") from error
+from packaging.version import Version as V
+
+use_uv = "--uv" in sys.argv[1:]
+
+v = V(torch.__version__)
+cuda = str(torch.version.cuda)
+try:
+    is_ampere = torch.cuda.get_device_capability()[0] >= 8
+except RuntimeError:
+    is_ampere = False
+if cuda != "12.1" and cuda != "11.8" and cuda != "12.4":
+    raise RuntimeError(f"CUDA = {cuda} not supported!")
+if v <= V("2.1.0"):
+    raise RuntimeError(f"Torch = {v} too old!")
+elif v <= V("2.1.1"):
+    x = "cu{}{}-torch211"
+elif v <= V("2.1.2"):
+    x = "cu{}{}-torch212"
+elif v < V("2.3.0"):
+    x = "cu{}{}-torch220"
+elif v < V("2.4.0"):
+    x = "cu{}{}-torch230"
+elif v < V("2.5.0"):
+    x = "cu{}{}-torch240"
+elif v < V("2.6.0"):
+    x = "cu{}{}-torch250"
+else:
+    raise RuntimeError(f"Torch = {v} too new!")
+x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
+uv_prefix = "uv " if use_uv else ""
+print(
+    f'{uv_prefix}pip install unsloth-zoo==2024.12.1 && {uv_prefix}pip install --no-deps "unsloth[{x}]==2024.12.4"'
+)

setup.py

@@ -0,0 +1,230 @@
+"""setup.py for axolotl"""
+
+import os
+import platform
+import re
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+
+from setuptools import find_packages, setup
+
+
+def parse_requirements(extras_require_map):
+    _install_requires = []
+    _dependency_links = []
+    with open("./requirements.txt", encoding="utf-8") as requirements_file:
+        lines = [r.strip() for r in requirements_file.readlines()]
+        for line in lines:
+            is_extras = "deepspeed" in line or "mamba-ssm" in line
+            if line.startswith("--extra-index-url"):
+                # Handle custom index URLs
+                _, url = line.split()
+                _dependency_links.append(url)
+            elif not is_extras and line and line[0] != "#":
+                # Handle standard packages
+                _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 = [
+                "bitsandbytes",
+                "triton",
+                "mamba-ssm",
+                "xformers",
+                "liger-kernel",
+            ]
+            _install_requires = [
+                req
+                for req in _install_requires
+                if re.split(r"[>=<]", req)[0].strip() not in skip_packages
+            ]
+            print(
+                _install_requires, [req in skip_packages for req in _install_requires]
+            )
+        else:
+            # detect the version of torch already installed
+            # and set it so dependencies don't clobber the torch version
+            try:
+                torch_version = version("torch")
+            except PackageNotFoundError:
+                torch_version = "2.8.0"  # default to torch 2.8.0
+            _install_requires.append(f"torch=={torch_version}")
+
+            version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version)
+            if version_match:
+                major, minor, patch = version_match.groups()
+                major, minor = int(major), int(minor)
+                patch = (
+                    int(patch) if patch is not None else 0
+                )  # Default patch to 0 if not present
+            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, 10):
+                extras_require_map.pop("fbgemm-gpu")
+                extras_require_map["fbgemm-gpu"] = [
+                    "fbgemm-gpu==1.5.0",
+                    "fbgemm-gpu-genai==1.5.0",
+                ]
+                if not install_xformers:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                extras_require_map["vllm"] = ["vllm>=0.19.1"]
+            elif (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",
+                ]
+                if not install_xformers:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                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")
+                    # 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")
+                    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")
+                # 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")
+                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")
+            else:
+                raise ValueError("axolotl requires torch>=2.4")
+
+    except PackageNotFoundError:
+        pass
+    return _install_requires, _dependency_links, extras_require_map
+
+
+def get_package_version():
+    with open(
+        Path(os.path.dirname(os.path.abspath(__file__))) / "VERSION",
+        "r",
+        encoding="utf-8",
+    ) as fin:
+        version_ = fin.read().strip()
+    return version_
+
+
+extras_require = {
+    "flash-attn": ["flash-attn==2.8.3"],
+    "ring-flash-attn": [
+        "flash-attn==2.8.3",
+        "ring-flash-attn>=0.1.7",
+    ],
+    "deepspeed": [
+        "deepspeed==0.18.2",
+        "deepspeed-kernels",
+    ],
+    "mamba-ssm": [
+        "mamba-ssm==1.2.0.post1",
+        "causal_conv1d",
+    ],
+    "auto-gptq": [
+        "auto-gptq==0.5.1",
+    ],
+    "mlflow": [
+        "mlflow",
+    ],
+    "galore": [
+        "galore_torch",
+    ],
+    "apollo": [
+        "apollo-torch",
+    ],
+    "optimizers": [
+        "galore_torch",
+        "apollo-torch",
+        "lomo-optim==0.1.1",
+        "torch-optimi==0.2.1",
+        "came_pytorch==0.1.3",
+    ],
+    "ray": [
+        "ray[train]>=2.52.1",
+    ],
+    "vllm": [
+        "vllm==0.10.0",
+    ],
+    "llmcompressor": [
+        "llmcompressor==0.5.1",
+    ],
+    "fbgemm-gpu": ["fbgemm-gpu-genai==1.3.0"],
+    "opentelemetry": [
+        "opentelemetry-api",
+        "opentelemetry-sdk",
+        "opentelemetry-exporter-prometheus",
+        "prometheus-client",
+    ],
+}
+install_requires, dependency_links, extras_require_build = parse_requirements(
+    extras_require
+)
+
+setup(
+    version=get_package_version(),
+    package_dir={"": "src"},
+    packages=find_packages("src"),
+    install_requires=install_requires,
+    dependency_links=dependency_links,
+    entry_points={
+        "console_scripts": [
+            "axolotl=axolotl.cli.main:main",
+        ],
+    },
+    extras_require=extras_require_build,
+)

src/axolotl/cli/utils/lora_merge.py

@@ -339,11 +339,7 @@ def _build_peft_layer_and_get_delta(
             )
         layer.lora_A[adapter_name].weight.data = lora_a
         layer.lora_B[adapter_name].weight.data = lora_b
-        delta = layer.get_delta_weight(adapter_name)
+        return layer.get_delta_weight(adapter_name)
-        # peft >=0.19.1 may return delta with transposed dims for 3D params
-        if delta.shape != base_tensor.shape and delta.ndim == 3:
-            delta = delta.transpose(1, 2).contiguous()
-        return delta
     elif (
         layer_type and "Conv" in layer_type or (layer_type is None and lora_a.ndim > 2)
     ):

src/axolotl/integrations/kernels/libs/scattermoe_lora/layers.py

@@ -60,14 +60,49 @@ def peft_lora_B_to_scattermoe(peft_B, num_experts, rank):
 
 
 def peft_lora_to_scattermoe(peft_A, peft_B, num_experts, rank):
-    """Convert peft LoRA weights to scattermoe layout.
+    """Convert peft LoRA weights to scattermoe layout (with A<->B swap).
 
-    peft >=0.19.1 assigns in/out features for 3D params such that
+    peft operates on the parameter in its native storage layout ``[E, dim1, dim2]``
-    A and B already align with scattermoe's convention (no A<->B swap).
+    where ``in_features=dim1, out_features=dim2``.  ScatterMoE transposes the
-    Only B needs rank-major → expert-major layout conversion.
+    parameter (``W = param.transpose(2, 1)``) giving ``[E, dim2, dim1]`` with
+    ``K=dim2, N=dim1``.  Because of this transposition, peft's A and B roles
+    are swapped relative to scattermoe's convention.
+
+    peft gives:
+        lora_A ``[r*E, dim1]``, lora_B ``[dim2, r*E]``
+
+    scattermoe needs:
+        lora_A ``[r*E, K=dim2]``, lora_B ``[N=dim1, r*E]``
+
+    This function swaps A<->B and converts B from rank-major to expert-major.
+    Uses vectorized tensor operations (no Python loop over experts).
+
+    Works for **both** gate_up_proj and down_proj since the transposition
+    issue is the same for any parameter.
     """
-    smoe_A = peft_A
+    peft_B_em = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
-    smoe_B = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
+
+    dim1 = peft_A.shape[1]  # peft in_features -> scattermoe N
+    dim2 = peft_B_em.shape[0]  # peft out_features -> scattermoe K
+
+    # smoe_A: per expert, transpose B_e [dim2, r] -> [r, dim2]
+    # [dim2, E*r] -> [dim2, E, r] -> [E, r, dim2] -> [E*r, dim2]
+    smoe_A = (
+        peft_B_em.reshape(dim2, num_experts, rank)
+        .permute(1, 2, 0)
+        .contiguous()
+        .reshape(rank * num_experts, dim2)
+    )
+
+    # smoe_B: per expert, transpose A_e [r, dim1] -> [dim1, r]
+    # [E*r, dim1] -> [E, r, dim1] -> [dim1, E, r] -> [dim1, E*r]
+    smoe_B = (
+        peft_A.reshape(num_experts, rank, dim1)
+        .permute(2, 0, 1)
+        .contiguous()
+        .reshape(dim1, num_experts * rank)
+    )
+
     return smoe_A, smoe_B
 
 

src/axolotl/loaders/patch_manager.py

@@ -162,6 +162,7 @@ class PatchManager:
     def apply_post_model_load_patches(self, model: PreTrainedModel):
         """Apply patches that require the model instance."""
         self._apply_llama_flash_attn_patches(model)
+        self._apply_unsloth_patches(model)
         self._apply_lora_kernel_patch(model)
         self._apply_scaling_softmax_patch(model)
 
@@ -395,16 +396,7 @@ class PatchManager:
                     patch_gemma4_fused_attn,
                 )
 
-                # Shared-KV side channel when activation checkpointing (PR #3611).
+                patch_gemma4_fused_attn()
-                fsdp_cfg = self.cfg.fsdp_config
-                needs_shared_kv_workaround = (not self.inference) and bool(
-                    self.cfg.gradient_checkpointing
-                    or self.cfg.activation_offloading
-                    or (fsdp_cfg is not None and fsdp_cfg.activation_checkpointing)
-                )
-                patch_gemma4_fused_attn(
-                    install_shared_kv_workaround=needs_shared_kv_workaround
-                )
 
     @staticmethod
     def _fix_nemotron_h_conversion_mapping():
@@ -682,10 +674,24 @@ class PatchManager:
             )
 
             patch_fa_llama_cross_entropy()
+        elif self.cfg.unsloth_cross_entropy_loss:
+            from axolotl.monkeypatch.unsloth_ import integrate_cross_entropy_loss_patch
+
+            integrate_cross_entropy_loss_patch(model_type="llama")
+
         if self.cfg.flash_attn_rms_norm and self.has_flash_attn:
             from axolotl.monkeypatch.llama_attn_hijack_flash import patch_llama_rms_norm
 
             patch_llama_rms_norm()
+        elif self.cfg.unsloth_rms_norm:
+            from axolotl.monkeypatch.unsloth_ import patch_unsloth_layernorm
+
+            patch_unsloth_layernorm()
+
+        if self.cfg.unsloth_lora_qkv or self.cfg.unsloth_lora_o:
+            from axolotl.monkeypatch.unsloth_ import patch_self_attn_lora
+
+            patch_self_attn_lora()
 
     def _patch_llama_flash_attention(self):
         """Apply Flash Attention patches for LLaMA models."""
@@ -752,6 +758,23 @@ class PatchManager:
                 LOG.info("Patching with SwiGLU...")
                 replace_llama_mlp_with_swiglu(model)
 
+    def _apply_unsloth_patches(self, model):
+        """Apply unsloth optimization patches."""
+        if self.cfg.unsloth_lora_mlp:
+            from axolotl.monkeypatch.unsloth_ import integrate_lora_mlp_patch
+
+            integrate_lora_mlp_patch(peft_model=model)
+
+        if self.cfg.unsloth_lora_qkv or self.cfg.unsloth_lora_o:
+            from axolotl.monkeypatch.unsloth_ import integrate_lora_patch
+
+            integrate_lora_patch(peft_model=model, cfg=self.cfg)
+
+        if self.cfg.unsloth_rope:
+            from axolotl.monkeypatch.unsloth_ import integrate_rope_embeddings
+
+            integrate_rope_embeddings()
+
     def _apply_lora_kernel_patch(self, model):
         """Apply LoRA kernel patches."""
         if (

src/axolotl/monkeypatch/models/gemma4/fused_attn.py

@@ -6,29 +6,15 @@ kernels, eliminating intermediate tensor allocations from rotate_half / apply_ro
 
 Usage:
     from axolotl.monkeypatch.models.gemma4.fused_attn import patch_gemma4_fused_attn
-    # Pass install_shared_kv_workaround=True when activation checkpointing is enabled.
+    patch_gemma4_fused_attn()
-    patch_gemma4_fused_attn(install_shared_kv_workaround=True)
 """
 
+import logging
 from typing import Callable
 
 import torch
 
-from axolotl.utils.logging import get_logger
+logger = logging.getLogger(__name__)
-
-logger = get_logger(__name__)
-
-# Module-level dict used as a side channel for shared KV states avoiding kwarg and TLS
-# to prevent memory leak on gradient checkpoint enabled training (PR #3611)
-_GEMMA4_SHARED_KV_STORE: dict = {"store": None}
-
-
-def _set_shared_kv_states(store):
-    _GEMMA4_SHARED_KV_STORE["store"] = store
-
-
-def _get_shared_kv_states():
-    return _GEMMA4_SHARED_KV_STORE["store"]
 
 
 def _make_fused_forward(original_forward):
@@ -44,7 +30,7 @@ def _make_fused_forward(original_forward):
         hidden_states: torch.Tensor,
         position_embeddings: torch.Tensor,
         attention_mask: torch.Tensor | None,
-        shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None = None,
+        shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]],
         past_key_values=None,
         **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -53,10 +39,6 @@ def _make_fused_forward(original_forward):
             eager_attention_forward,
         )
 
-        store = _get_shared_kv_states()
-        if store is not None:
-            shared_kv_states = store
-
         input_shape = hidden_states.shape[:-1]
         hidden_shape = (*input_shape, -1, self.head_dim)
         eps = self.config.rms_norm_eps
@@ -151,44 +133,15 @@ def _make_fused_forward(original_forward):
     return fused_forward
 
 
-def _patch_decoder_layer_call():
+def patch_gemma4_fused_attn():
-    """Strip `shared_kv_states` from decoder-layer kwargs and route via the
-    module-level side channel so the checkpoint partial cannot pin it (PR #3611).
     """
-    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextDecoderLayer
+    Monkeypatch Gemma4TextAttention.forward to use fused RMSNorm+RoPE kernels.
-
-    if getattr(Gemma4TextDecoderLayer, "_axolotl_shared_kv_patched", False):
-        return
-
-    original_call = Gemma4TextDecoderLayer.__call__
-
-    def patched_call(self, *args, **kwargs):
-        shared_kv = kwargs.pop("shared_kv_states", None)
-        # Overwrite unconditionally (including with None) so a previous step's
-        # dict cannot leak into a later call without shared_kv_states (PR #3611).
-        _set_shared_kv_states(shared_kv)
-        return original_call(self, *args, **kwargs)
-
-    Gemma4TextDecoderLayer.__call__ = patched_call
-    Gemma4TextDecoderLayer._axolotl_shared_kv_patched = True
-
-
-def patch_gemma4_fused_attn(install_shared_kv_workaround: bool = False):
-    """
-    Monkeypatch Gemma4TextAttention.forward to use fused RMSNorm+RoPE kernels,
-    and optionally route `shared_kv_states` via a module-level side channel to
-    avoid a VRAM leak under activation checkpointing (PR #3611).
     """
     from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
 
     original_forward = Gemma4TextAttention.forward
     Gemma4TextAttention.forward = _make_fused_forward(original_forward)
 
-    if install_shared_kv_workaround:
-        _patch_decoder_layer_call()
-
     logger.info(
         "Patched Gemma4TextAttention.forward with fused RMSNorm+RoPE Triton kernels"
     )
-    if install_shared_kv_workaround:
-        logger.info("Installed Gemma4 shared_kv_states side channel (PR #3611)")

src/axolotl/monkeypatch/unsloth_.py

@@ -0,0 +1,252 @@
+"""module for patching with unsloth optimizations"""
+
+import inspect
+import types
+
+import torch
+from peft import PeftModelForCausalLM
+from torch import nn
+from transformers.models.llama.modeling_llama import LlamaFlashAttention2
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+ORIGINAL_QKV_CODE = """
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+""".lstrip("\n")
+
+PATCHED_QKV_CODE = """
+    query_states, key_states, value_states = self.apply_qkv(self, hidden_states)
+""".lstrip("\n")
+
+ORIGINAL_O_CODE = """
+    attn_output = self.o_proj(attn_output)
+""".lstrip("\n")
+
+PATCHED_O_CODE = """
+    attn_output = self.apply_o(self, attn_output)
+""".lstrip("\n")
+
+
+def original_apply_qkv(self, hidden_states):
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+    return query_states, key_states, value_states
+
+
+def original_apply_o(self, hidden_states):
+    attn_output = self.o_proj(hidden_states)
+    return attn_output
+
+
+def get_self_attn_code() -> str:
+    forward = inspect.getsource(LlamaFlashAttention2.forward)
+    return forward
+
+
+def check_self_attn_is_patchable() -> bool:
+    qkv = get_self_attn_code()
+    qkv, _ = detab_code(qkv)
+    return ORIGINAL_QKV_CODE in qkv and ORIGINAL_O_CODE in qkv
+
+
+def integrate_cross_entropy_loss_patch(model_type: str = "llama") -> None:
+    from unsloth.kernels.cross_entropy_loss import fast_cross_entropy_loss
+
+    def UnslothForCausalLMLoss(
+        logits,
+        labels,
+        vocab_size: int,
+        num_items_in_batch: int = None,
+        ignore_index: int = -100,
+        **kwargs,
+    ):
+        # Upcast to float if we need to compute the loss to avoid potential precision issues
+        logits = logits.float()
+        # Shift so that tokens < n predict n
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        loss = fast_cross_entropy_loss(
+            logits=shift_logits, labels=shift_labels, n_items=num_items_in_batch
+        )
+        return loss
+
+    if model_type == "llama":
+        from transformers.loss import loss_utils
+
+        loss_utils.ForCausalLMLoss = UnslothForCausalLMLoss  # type: ignore[assignment]
+    else:
+        raise ValueError("Unsupported model type")
+
+
+self_attn_lora_patched = False
+
+
+def patch_self_attn_lora():
+    global self_attn_lora_patched
+    if self_attn_lora_patched:
+        # prevent patching multiple times
+        return
+    self_attn_forward = get_self_attn_code()
+    LlamaFlashAttention2._original_forward = self_attn_forward
+    self_attn_forward, _ = detab_code(self_attn_forward)
+    assert ORIGINAL_QKV_CODE in self_attn_forward, "Original qkv code not found"
+    assert ORIGINAL_O_CODE in self_attn_forward, "Original o code not found"
+
+    self_attn_forward = self_attn_forward.replace(ORIGINAL_QKV_CODE, PATCHED_QKV_CODE)
+    self_attn_forward = self_attn_forward.replace(ORIGINAL_O_CODE, PATCHED_O_CODE)
+    self_attn_forward = self_attn_forward.replace(
+        "def forward(",
+        "def unsloth_attn_forward(",
+        1,
+    )
+
+    # load imports necessary
+    import transformers.models.llama.modeling_llama
+
+    items_to_import = []
+    for item in dir(transformers.models.llama.modeling_llama):
+        if item in self_attn_forward:
+            items_to_import.append(item)
+
+    exec(
+        "from transformers.models.llama.modeling_llama import ("
+        + ", ".join(x for x in items_to_import)
+        + ")",
+        globals(),
+    )
+    exec(self_attn_forward, globals())
+    self_attn_lora_patched = True
+    LOG.info("patching unsloth attn lora")
+    LlamaFlashAttention2.forward = unsloth_attn_forward
+
+
+def integrate_rope_embeddings():
+    import transformers.models.llama.modeling_llama
+    from unsloth.kernels.rope_embedding import fast_rope_embedding
+
+    def apply_rotary_pos_emb(
+        q,
+        k,
+        cos,
+        sin,
+        position_ids=None,
+        unsqueeze_dim=1,
+    ):
+        return fast_rope_embedding(q, k, cos, sin)
+
+    LOG.info("patching unsloth RoPE embeddings")
+    transformers.models.llama.modeling_llama.apply_rotary_pos_emb = apply_rotary_pos_emb
+
+
+def integrate_lora_mlp_patch(peft_model: PeftModelForCausalLM):
+    if peft_model.base_model.config.model_type in ["llama", "mistral"]:
+        from unsloth.kernels import apply_lora_mlp_swiglu
+
+        apply_lora_mlp = apply_lora_mlp_swiglu
+    elif peft_model.base_model.config.model_type == "gemma":
+        from unsloth.kernels import apply_lora_mlp_geglu_approx
+
+        apply_lora_mlp = apply_lora_mlp_geglu_approx
+    else:
+        raise NotImplementedError(
+            f"Model type {peft_model.base_model.config.model_type} not supported"
+        )
+
+    for idx, layer in enumerate(peft_model.model.model.layers):
+        layer_modules = [
+            getattr(layer.mlp, linear_proj)
+            for linear_proj in ["gate_proj", "up_proj", "down_proj"]
+        ]
+        is_mlp_lora = all(hasattr(module, "lora_A") for module in layer_modules)
+        mlp_no_bias = all(
+            getattr(module, "base_layer", module).bias is None
+            for module in layer_modules
+        )
+        mlp_not_dora = all(
+            len(getattr(module, "lora_magnitude_vector", []) or []) == 0
+            for module in layer_modules
+        )
+
+        if is_mlp_lora and mlp_no_bias and mlp_not_dora:
+            layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp)
+        else:
+            LOG.warning(f"unable to apply unsloth lora mlp patch to layer {idx}")
+
+
+def integrate_lora_patch(peft_model: PeftModelForCausalLM, cfg):
+    from unsloth.kernels import apply_lora_o, apply_lora_qkv
+
+    for idx, layer in enumerate(peft_model.model.model.layers):
+        if cfg.unsloth_lora_qkv:
+            layer_modules = [
+                getattr(layer.self_attn, linear_proj)
+                for linear_proj in ["q_proj", "k_proj", "v_proj"]
+            ]
+            is_qkv_lora = all(hasattr(module, "lora_A") for module in layer_modules)
+            qkv_no_bias = all(
+                getattr(module, "base_layer", module).bias is None
+                for module in layer_modules
+            )
+            qkv_not_dora = all(
+                len(getattr(module, "lora_magnitude_vector", []) or []) == 0
+                for module in layer_modules
+            )
+
+            if is_qkv_lora and qkv_no_bias and qkv_not_dora:
+                layer.self_attn.apply_qkv = apply_lora_qkv
+            else:
+                layer.self_attn.apply_qkv = original_apply_qkv
+                LOG.warning(f"unable to apply unsloth lora qkv patch to layer {idx}")
+        if cfg.unsloth_lora_o:
+            layer_modules = [
+                getattr(layer.self_attn, linear_proj) for linear_proj in ["o_proj"]
+            ]
+            is_o_lora = all(hasattr(module, "lora_A") for module in layer_modules)
+            o_no_bias = all(
+                getattr(module, "base_layer", module).bias is None
+                for module in layer_modules
+            )
+            o_not_dora = all(
+                len(getattr(module, "lora_magnitude_vector", []) or []) == 0
+                for module in layer_modules
+            )
+
+            if is_o_lora and o_no_bias and o_not_dora:
+                layer.self_attn.apply_o = apply_lora_o
+            else:
+                layer.self_attn.apply_o = original_apply_o
+                LOG.warning(f"unable to apply unsloth lora o_proj patch to layer {idx}")
+
+
+def patch_unsloth_layernorm():
+    try:
+        import transformers.models.llama.modeling_llama
+        from unsloth.kernels.rms_layernorm import Fast_RMS_Layernorm
+
+        class LlamaRMSNorm(nn.Module):
+            """LlamaRMSNorm"""
+
+            def __init__(self, hidden_size, eps=1e-6):
+                """
+                LlamaRMSNorm is equivalent to T5LayerNorm
+                """
+                super().__init__()
+                self.weight = nn.Parameter(torch.ones(hidden_size))
+                self.variance_epsilon = eps
+
+            def forward(self, hidden_states):
+                return Fast_RMS_Layernorm.apply(
+                    hidden_states, self.weight, self.variance_epsilon, False
+                )
+
+        LOG.info("patching with unsloth.kernels.rms_layernorm")
+        transformers.models.llama.modeling_llama.LlamaRMSNorm = LlamaRMSNorm
+    except ImportError:
+        LOG.warning("missing unsloth library")

src/axolotl/utils/schemas/config.py

@@ -823,6 +823,13 @@ class AxolotlInputConfig(
         },
     )
 
+    unsloth_cross_entropy_loss: bool | None = None
+    unsloth_lora_mlp: bool | None = None
+    unsloth_lora_qkv: bool | None = None
+    unsloth_lora_o: bool | None = None
+    unsloth_rms_norm: bool | None = None
+    unsloth_rope: bool | None = None
+
     lora_mlp_kernel: bool | None = Field(
         default=None,
         json_schema_extra={
@@ -1462,6 +1469,21 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
             )
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_multigpu_unsloth(cls, data):
+        if (
+            data.get("unsloth_lora_mlp")
+            or data.get("unsloth_lora_qkv")
+            or data.get("unsloth_lora_o")
+        ):
+            capabilities = data.get("capabilities")
+            if capabilities and capabilities.get("n_gpu", 0) > 1:
+                raise ValueError(
+                    "unsloth_lora_mlp, unsloth_lora_qkv, and unsloth_lora_o are not compatible with multi-GPU training."
+                )
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_multigpu_lora_kernels(cls, data):
@@ -1515,7 +1537,8 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
             # RL trainers not tested so don't enable kernels by default
             return data
         if data.get("adapter") in ["lora", "qlora"]:
-            # Skip if already set or using 8-bit
+            # Skip if already set, using unsloth optimizations, or using 8-bit
+            unsloth_fields = ["unsloth_lora_mlp", "unsloth_lora_qkv", "unsloth_lora_o"]
             kernel_fields = [
                 "lora_mlp_kernel",
                 "lora_qkv_kernel",
@@ -1524,6 +1547,7 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
             ]
             if (
                 any(data.get(k) is not None for k in kernel_fields)
+                or any(data.get(k) for k in unsloth_fields)
                 or data.get("adapter") == "lora"
                 and data.get("load_in_8bit")
             ):

src/axolotl/utils/schemas/validation.py

@@ -52,26 +52,6 @@ class DatasetValidationMixin:
 
         return datasets
 
-    @model_validator(mode="before")
-    @classmethod
-    def check_deprecated_unsloth_fields(cls, data):
-        deprecated_fields = [
-            "unsloth_cross_entropy_loss",
-            "unsloth_lora_mlp",
-            "unsloth_lora_qkv",
-            "unsloth_lora_o",
-            "unsloth_rms_norm",
-            "unsloth_rope",
-        ]
-        found = [f for f in deprecated_fields if data.get(f)]
-        if found:
-            raise ValueError(
-                f"`{'`, `'.join(found)}` {'has' if len(found) == 1 else 'have'} been removed. "
-                "Please use `lora_mlp_kernel`, `lora_qkv_kernel`, `lora_o_kernel` instead. "
-                "See: https://docs.axolotl.ai/docs/lora_optims.html"
-            )
-        return data
-
     @model_validator(mode="before")
     @classmethod
     def check_dataset_or_pretraining_dataset(cls, data):
@@ -627,6 +607,36 @@ class LoRAValidationMixin:
             )
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_qlora_unsloth(cls, data):
+        if (
+            data.get("unsloth_lora_mlp")
+            or data.get("unsloth_lora_qkv")
+            or data.get("unsloth_lora_o")
+        ):
+            if data.get("adapter") == "lora" and data.get("load_in_8bit"):
+                raise ValueError(
+                    "unsloth_lora_mlp, unsloth_lora_qkv, and unsloth_lora_o are not compatible with 8-bit LoRA"
+                )
+        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_lora_axolotl_unsloth(cls, data):
+        is_lora_kernel = any(
+            data.get(k) for k in ["lora_mlp_kernel", "lora_qkv_kernel", "lora_o_kernel"]
+        )
+        is_unsloth_lora = any(
+            data.get(k)
+            for k in ["unsloth_lora_mlp", "unsloth_lora_qkv", "unsloth_lora_o"]
+        )
+        if is_lora_kernel and is_unsloth_lora:
+            raise ValueError(
+                "both lora_mlp_kernel and unsloth_lora_mlp cannot be true (similarly for lora_qkv_kernel, lora_o_kernel)"
+            )
+        return data
+
     @model_validator(mode="after")
     def check_fused_lora(self):
         if self.adapter in ["lora", "qlora"] and self.flash_attn_fuse_mlp:
@@ -850,6 +860,17 @@ class OptimizationValidationMixin:
 
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_xentropy_patch_conflicts(cls, data):
+        if data.get("flash_attn_cross_entropy") and data.get(
+            "unsloth_cross_entropy_loss"
+        ):
+            raise ValueError(
+                "flash_attn_cross_entropy and unsloth_cross_entropy_loss cannot be both enabled"
+            )
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_cross_entropy_conflicts(cls, data):

src/setuptools_axolotl_dynamic_dependencies.py

@@ -0,0 +1,102 @@
+"""
+dynamic requirements for axolotl
+"""
+
+import platform
+import re
+from importlib.metadata import PackageNotFoundError, version
+
+from setuptools.command.build_py import build_py as _build_py
+
+
+def parse_requirements():
+    _install_requires = []
+    _dependency_links = []
+    with open("./requirements.txt", encoding="utf-8") as requirements_file:
+        lines = [r.strip() for r in requirements_file.readlines()]
+        for line in lines:
+            is_extras = (
+                "flash-attn" in line
+                or "flash-attention" in line
+                or "deepspeed" in line
+                or "mamba-ssm" in line
+                or "lion-pytorch" in line
+            )
+            if line.startswith("--extra-index-url"):
+                # Handle custom index URLs
+                _, url = line.split()
+                _dependency_links.append(url)
+            elif not is_extras and line and line[0] != "#":
+                # Handle standard packages
+                _install_requires.append(line)
+
+    try:
+        xformers_version = [req for req in _install_requires if "xformers" in req][0]
+        torchao_version = [req for req in _install_requires if "torchao" in req][0]
+
+        if "Darwin" in platform.system():
+            # don't install xformers on MacOS
+            _install_requires.pop(_install_requires.index(xformers_version))
+        else:
+            # detect the version of torch already installed
+            # and set it so dependencies don't clobber the torch version
+            try:
+                torch_version = version("torch")
+            except PackageNotFoundError:
+                torch_version = "2.5.1"
+            _install_requires.append(f"torch=={torch_version}")
+
+            version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version)
+            if version_match:
+                major, minor, patch = version_match.groups()
+                major, minor = int(major), int(minor)
+                patch = (
+                    int(patch) if patch is not None else 0
+                )  # Default patch to 0 if not present
+            else:
+                raise ValueError("Invalid version format")
+
+            if (major, minor) >= (2, 5):
+                _install_requires.pop(_install_requires.index(xformers_version))
+                if patch == 0:
+                    _install_requires.append("xformers==0.0.28.post2")
+                else:
+                    _install_requires.append("xformers==0.0.28.post3")
+            elif (major, minor) >= (2, 4):
+                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")
+            elif (major, minor) >= (2, 3):
+                _install_requires.pop(_install_requires.index(torchao_version))
+                if patch == 0:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                    _install_requires.append("xformers>=0.0.26.post1")
+                else:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                    _install_requires.append("xformers>=0.0.27")
+            elif (major, minor) >= (2, 2):
+                _install_requires.pop(_install_requires.index(torchao_version))
+                _install_requires.pop(_install_requires.index(xformers_version))
+                _install_requires.append("xformers>=0.0.25.post1")
+            else:
+                _install_requires.pop(_install_requires.index(torchao_version))
+                _install_requires.pop(_install_requires.index(xformers_version))
+                _install_requires.append("xformers>=0.0.23.post1")
+
+    except PackageNotFoundError:
+        pass
+    return _install_requires, _dependency_links
+
+
+class BuildPyCommand(_build_py):
+    """
+    custom build_py command to parse dynamic requirements
+    """
+
+    def finalize_options(self):
+        super().finalize_options()
+        install_requires, _ = parse_requirements()
+        self.distribution.install_requires = install_requires

tests/conftest.py

@@ -325,10 +325,10 @@ def download_phi_4_reasoning_model_fixture():
 
 
 @pytest.fixture(scope="session", autouse=True)
-def download_phi_3_mini_model_fixture():
+def download_phi_3_medium_model_fixture():
     # download the tokenizer only
     snapshot_download_w_retry(
-        "microsoft/Phi-3-mini-4k-instruct",
+        "microsoft/Phi-3-medium-128k-instruct",
         repo_type="model",
         allow_patterns=["*token*", "config.json"],
     )

tests/e2e/integrations/test_scattermoe_lora_olmoe.py

@@ -54,8 +54,24 @@ except (ImportError, ModuleNotFoundError):
         )
 
     def peft_lora_to_scattermoe(peft_A, peft_B, num_experts, rank):
-        smoe_A = peft_A
+        peft_B_em = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
-        smoe_B = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
+        K_inter, N_hidden = peft_B.shape[0], peft_A.shape[1]
+        smoe_A = torch.zeros(
+            rank * num_experts,
+            K_inter,
+            device=peft_A.device,
+            dtype=peft_A.dtype,
+        )
+        smoe_B = torch.zeros(
+            N_hidden,
+            rank * num_experts,
+            device=peft_A.device,
+            dtype=peft_A.dtype,
+        )
+        for e in range(num_experts):
+            s = e * rank
+            smoe_A[s : s + rank, :] = peft_B_em[:, s : s + rank].T
+            smoe_B[:, s : s + rank] = peft_A[s : s + rank, :].T
         return smoe_A, smoe_B
 
     def _unwrap_experts_lora(experts_module):
@@ -306,14 +322,12 @@ class TestLoRABLayoutConversion:
         hidden, inter = 32, 16
         scaling = 2.0
 
-        # peft >=0.19.1 for down_proj [E, hidden, inter]:
+        peft_A = torch.randn(E * r, hidden)
-        # swaps in/out, lora_A [r*E, inter], lora_B [hidden, r*E]
+        peft_B = torch.randn(inter, E * r)
-        peft_A = torch.randn(E * r, inter)
-        peft_B = torch.randn(hidden, E * r)
 
-        A_r = peft_A.reshape(E, r, inter)
+        A_r = peft_A.reshape(E, r, hidden)
-        B_r = peft_B.reshape(hidden, r, E)
+        B_r = peft_B.reshape(inter, r, E)
-        delta_peft = torch.einsum("o r e, e r i -> e o i", B_r, A_r) * scaling
+        delta_peft = torch.einsum("o r e, e r i -> e i o", B_r, A_r) * scaling
 
         smoe_A, smoe_B = peft_lora_to_scattermoe(peft_A, peft_B, E, r)
         for e in range(E):
@@ -325,22 +339,31 @@ class TestLoRABLayoutConversion:
             )
 
     def test_gate_up_proj_conversion(self):
-        """Verify gate_up_proj LoRA conversion with non-square dims.
+        """Verify gate_up_proj LoRA conversion with non-square dims (Qwen3-like).
 
         gate_up_proj param: [E, 2*inter, hidden].
-        peft swaps in/out for 3D: lora_A [r*E, hidden], lora_B [2*inter, r*E].
+        peft: in_features=2*inter, out_features=hidden.
+        peft lora_A: [r*E, 2*inter], lora_B: [hidden, r*E].
+
+        scattermoe W = param.T = [E, hidden, 2*inter], K=hidden, N=2*inter.
         scattermoe needs: lora_A [r*E, K=hidden], lora_B [N=2*inter, r*E].
+
+        Uses non-square dims (hidden=32 != 2*inter=24) to catch A<->B swap bugs.
         """
         E, r = 4, 2
         hidden, inter = 32, 12  # 2*inter=24 != hidden=32
         scaling = 2.0
 
-        peft_A = torch.randn(E * r, hidden)  # [r*E, in=hidden]
+        # peft assigns: in_features=2*inter, out_features=hidden
-        peft_B = torch.randn(2 * inter, E * r)  # [out=2*inter, r*E]
+        peft_A = torch.randn(E * r, 2 * inter)  # [r*E, in_features=2*inter]
+        peft_B = torch.randn(hidden, E * r)  # [out_features=hidden, r*E]
 
-        A_r = peft_A.reshape(E, r, hidden)
+        # peft delta via einsum: "o r e, e r i -> e i o"
-        B_r = peft_B.reshape(2 * inter, r, E)
+        A_r = peft_A.reshape(E, r, 2 * inter)
-        delta_peft = torch.einsum("o r e, e r i -> e o i", B_r, A_r) * scaling
+        B_r = peft_B.reshape(hidden, r, E)
+        delta_peft = torch.einsum("o r e, e r i -> e i o", B_r, A_r) * scaling
+        # delta_peft[e] has shape [in_features, out_features] = [2*inter, hidden]
+        # = param[e] shape [2*inter, hidden]
 
         smoe_A, smoe_B = peft_gate_up_lora_to_scattermoe(peft_A, peft_B, E, r)
         # smoe_A should be [r*E, K=hidden], smoe_B should be [N=2*inter, r*E]
@@ -398,7 +421,8 @@ class TestPeftLoRAWeightExtraction:
             r,
         )
 
-        # gate_up_proj [E, 2*inter, hidden] — peft swaps in/out for 3D
+        # gate_up_proj [E, 2*inter, hidden]
+        # peft: in_features=hidden (last dim), out_features=2*inter (middle dim)
         assert trainable[
             "base_model.model.moe.experts.base_layer.lora_A.default.weight"
         ].shape == (E * r, config.hidden_size)
@@ -406,7 +430,8 @@ class TestPeftLoRAWeightExtraction:
             "base_model.model.moe.experts.base_layer.lora_B.default.weight"
         ].shape == (2 * config.intermediate_size, E * r)
 
-        # down_proj [E, hidden, inter] — peft swaps in/out for 3D
+        # down_proj [E, hidden, inter]
+        # peft: in_features=inter (last dim), out_features=hidden (middle dim)
         assert trainable[
             "base_model.model.moe.experts.lora_A.default.weight"
         ].shape == (E * r, config.intermediate_size)
@@ -463,26 +488,29 @@ class TestPeftLoRAWeightExtraction:
         assert gup_lora is not None, "gate_up_proj LoRA not detected"
         assert down_lora is not None, "down_proj LoRA not detected"
 
-        # gate_up_proj: K=hidden, N=2*inter
+        # Check shapes (after peft->scattermoe conversion with A<->B swap)
+        # gate_up_proj: peft A [E*r, hidden] / B [2*inter, E*r]
+        # After swap: smoe_A [E*r, 2*inter], smoe_B [hidden, E*r]
         E, r = config.num_experts, 4
         gup_A, gup_B, gup_s = gup_lora
-        assert gup_A.shape == (E * r, config.hidden_size), (
+        assert gup_A.shape == (E * r, 2 * config.intermediate_size), (
-            f"gate_up_proj smoe_A: expected [r*E, K=hidden]={(E * r, config.hidden_size)}, "
+            f"gate_up_proj smoe_A: expected [r*E, 2*inter]={(E * r, 2 * config.intermediate_size)}, "
             f"got {gup_A.shape}"
         )
-        assert gup_B.shape == (2 * config.intermediate_size, E * r), (
+        assert gup_B.shape == (config.hidden_size, E * r), (
-            f"gate_up_proj smoe_B: expected [N=2*inter, r*E]="
+            f"gate_up_proj smoe_B: expected [hidden, r*E]="
-            f"{(2 * config.intermediate_size, E * r)}, got {gup_B.shape}"
+            f"{(config.hidden_size, E * r)}, got {gup_B.shape}"
         )
 
-        # down_proj: K=inter, N=hidden
+        # down_proj: peft A [E*r, inter] / B [hidden, E*r]
+        # After swap: smoe_A [E*r, hidden], smoe_B [inter, E*r]
         down_A, down_B, down_s = down_lora
-        assert down_A.shape == (E * r, config.intermediate_size), (
+        assert down_A.shape == (E * r, config.hidden_size), (
-            f"down_proj smoe_A: expected [r*E, K=inter]={(E * r, config.intermediate_size)}, "
+            f"down_proj smoe_A: expected [r*E, hidden]={(E * r, config.hidden_size)}, "
             f"got {down_A.shape}"
         )
-        assert down_B.shape == (config.hidden_size, E * r), (
+        assert down_B.shape == (config.intermediate_size, E * r), (
-            f"down_proj smoe_B: expected [N=hidden, r*E]={(config.hidden_size, E * r)}, "
+            f"down_proj smoe_B: expected [inter, r*E]={(config.intermediate_size, E * r)}, "
             f"got {down_B.shape}"
         )
 

tests/e2e/patched/test_unsloth_integration.py

@@ -0,0 +1,21 @@
+"""Test module for checking whether the integration of Unsloth with Hugging Face Transformers is working as expected."""
+
+import unittest
+
+import pytest
+
+
+@pytest.mark.skip(
+    reason="Unsloth integration will be broken going into latest transformers"
+)
+class TestUnslothIntegration(unittest.TestCase):
+    """Unsloth monkeypatch integration tests."""
+
+    def test_is_self_attn_patchable(self):
+        from axolotl.monkeypatch.unsloth_ import check_self_attn_is_patchable
+
+        # ensures the current version of transformers has loss code that matches our patching code
+        self.assertTrue(
+            check_self_attn_is_patchable(),
+            "HF transformers self attention code has changed and isn't patchable",
+        )

tests/e2e/patched/test_unsloth_qlora.py

@@ -0,0 +1,184 @@
+"""
+e2e tests for unsloth qlora
+"""
+
+import pytest
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from ..utils import check_model_output_exists, check_tensorboard
+
+
+@pytest.mark.skip(
+    reason="Unsloth integration will be broken going into latest transformers"
+)
+class TestUnslothQLoRA:
+    """
+    Test class for Unsloth QLoRA Llama models
+    """
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_unsloth_llama_qlora_fa2(self, temp_dir, sample_packing):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "unsloth_lora_mlp": True,
+                "unsloth_lora_qkv": True,
+                "unsloth_lora_o": True,
+                "load_in_4bit": True,
+                "adapter": "qlora",
+                "lora_r": 16,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "save_steps": 10,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "use_tensorboard": True,
+                "bf16": "auto",
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss (%s) is too high"
+        )
+
+    def test_unsloth_llama_qlora_unpacked(self, temp_dir):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "unsloth_lora_mlp": True,
+                "unsloth_lora_qkv": True,
+                "unsloth_lora_o": True,
+                "sample_packing": False,
+                "load_in_4bit": True,
+                "adapter": "qlora",
+                "lora_r": 16,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "save_steps": 10,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "use_tensorboard": True,
+                "bf16": "auto",
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss (%s) is too high"
+        )
+
+    @pytest.mark.parametrize(
+        "sdp_attention",
+        [True, False],
+    )
+    def test_unsloth_llama_qlora_unpacked_no_fa2_fp16(self, temp_dir, sdp_attention):
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "sequence_len": 1024,
+                "unsloth_lora_mlp": True,
+                "unsloth_lora_qkv": True,
+                "unsloth_lora_o": True,
+                "sample_packing": False,
+                "load_in_4bit": True,
+                "adapter": "qlora",
+                "lora_r": 16,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0.05,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 5,
+                "save_steps": 10,
+                "micro_batch_size": 4,
+                "gradient_accumulation_steps": 2,
+                "sdp_attention": sdp_attention,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_8bit",
+                "lr_scheduler": "cosine",
+                "use_tensorboard": True,
+                "fp16": True,
+                "save_first_step": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+        check_tensorboard(
+            temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss (%s) is too high"
+        )

tests/monkeypatch/test_gemma4_fused_attn_patch.py

@@ -1,171 +0,0 @@
-"""Unit tests for the Gemma4 fused-attention shared_kv_states routing patch."""
-
-import pytest
-
-gemma4_modeling = pytest.importorskip("transformers.models.gemma4.modeling_gemma4")
-
-
-@pytest.fixture
-def clean_decoder_layer_patch_slate():
-    """Save and restore Gemma4TextDecoderLayer.__call__ and the sentinel."""
-    from axolotl.monkeypatch.models.gemma4 import fused_attn
-
-    cls = gemma4_modeling.Gemma4TextDecoderLayer
-    original_call = cls.__call__
-    had_sentinel = getattr(cls, "_axolotl_shared_kv_patched", False)
-
-    if had_sentinel:
-        del cls._axolotl_shared_kv_patched
-
-    try:
-        yield cls, fused_attn
-    finally:
-        cls.__call__ = original_call
-        if had_sentinel:
-            cls._axolotl_shared_kv_patched = True
-        elif hasattr(cls, "_axolotl_shared_kv_patched"):
-            del cls._axolotl_shared_kv_patched
-        fused_attn._set_shared_kv_states(None)
-
-
-class TestPatchedDecoderLayerCall:
-    def test_pops_shared_kv_states_and_populates_store(
-        self, clean_decoder_layer_patch_slate
-    ):
-        cls, fused_attn = clean_decoder_layer_patch_slate
-
-        captured = {}
-
-        def spy(self, *args, **kwargs):
-            captured["args"] = args
-            captured["kwargs"] = dict(kwargs)
-            return "spy_return"
-
-        cls.__call__ = spy
-        fused_attn._patch_decoder_layer_call()
-
-        assert getattr(cls, "_axolotl_shared_kv_patched", False) is True
-        assert cls.__call__ is not spy
-
-        shared_kv = {"layer_0": ("k", "v")}
-        result = cls.__call__(
-            object(),
-            "positional_arg",
-            shared_kv_states=shared_kv,
-            other_kwarg="keep_me",
-        )
-
-        assert result == "spy_return"
-        assert captured["args"] == ("positional_arg",)
-        assert "shared_kv_states" not in captured["kwargs"]
-        assert captured["kwargs"] == {"other_kwarg": "keep_me"}
-        assert fused_attn._get_shared_kv_states() is shared_kv
-
-    def test_clears_store_when_kwarg_absent(self, clean_decoder_layer_patch_slate):
-        """Regression for commit 251021e1: a prior step's dict must not leak
-        into a later call that omits `shared_kv_states`."""
-        cls, fused_attn = clean_decoder_layer_patch_slate
-
-        def spy(self, *args, **kwargs):
-            return None
-
-        cls.__call__ = spy
-        fused_attn._patch_decoder_layer_call()
-
-        stale = {"stale_step": True}
-        fused_attn._set_shared_kv_states(stale)
-        assert fused_attn._get_shared_kv_states() is stale
-
-        cls.__call__(object())
-
-        assert fused_attn._get_shared_kv_states() is None
-
-    def test_store_visible_across_threads(self):
-        """Regression for commit e3669b2c: the store must be readable from
-        threads other than the one that set it. `threading.local()` failed
-        this invariant, crashing with 'NoneType' object is not subscriptable'
-        on MoE Gemma4 variants when autograd worker threads ran backward
-        recompute under HF-Trainer gradient_checkpointing."""
-        import threading
-
-        from axolotl.monkeypatch.models.gemma4 import fused_attn
-
-        sentinel = {"layer_0": ("k", "v")}
-        try:
-            fused_attn._set_shared_kv_states(sentinel)
-
-            seen = {}
-
-            def worker():
-                seen["value"] = fused_attn._get_shared_kv_states()
-
-            t = threading.Thread(target=worker)
-            t.start()
-            t.join()
-
-            assert seen["value"] is sentinel
-        finally:
-            fused_attn._set_shared_kv_states(None)
-
-
-@pytest.fixture
-def clean_entry_point_patch_slate():
-    """Save and restore Gemma4TextAttention.forward and Gemma4TextDecoderLayer.__call__."""
-    from axolotl.monkeypatch.models.gemma4 import fused_attn
-
-    decoder_cls = gemma4_modeling.Gemma4TextDecoderLayer
-    attn_cls = gemma4_modeling.Gemma4TextAttention
-
-    original_call = decoder_cls.__call__
-    original_forward = attn_cls.forward
-    had_sentinel = getattr(decoder_cls, "_axolotl_shared_kv_patched", False)
-
-    if had_sentinel:
-        del decoder_cls._axolotl_shared_kv_patched
-
-    try:
-        yield decoder_cls, attn_cls, original_call, original_forward, fused_attn
-    finally:
-        decoder_cls.__call__ = original_call
-        attn_cls.forward = original_forward
-        if had_sentinel:
-            decoder_cls._axolotl_shared_kv_patched = True
-        elif hasattr(decoder_cls, "_axolotl_shared_kv_patched"):
-            del decoder_cls._axolotl_shared_kv_patched
-        fused_attn._set_shared_kv_states(None)
-
-
-class TestPatchGemma4FusedAttnEntryPoint:
-    def test_default_flag_swaps_only_attention_forward(
-        self, clean_entry_point_patch_slate
-    ):
-        (
-            decoder_cls,
-            attn_cls,
-            original_call,
-            original_forward,
-            fused_attn,
-        ) = clean_entry_point_patch_slate
-
-        fused_attn.patch_gemma4_fused_attn()
-
-        assert attn_cls.forward is not original_forward
-        assert decoder_cls.__call__ is original_call
-        assert not getattr(decoder_cls, "_axolotl_shared_kv_patched", False)
-
-    def test_workaround_flag_installs_decoder_layer_patch(
-        self, clean_entry_point_patch_slate
-    ):
-        (
-            decoder_cls,
-            attn_cls,
-            original_call,
-            original_forward,
-            fused_attn,
-        ) = clean_entry_point_patch_slate
-
-        fused_attn.patch_gemma4_fused_attn(install_shared_kv_workaround=True)
-
-        assert attn_cls.forward is not original_forward
-        assert decoder_cls.__call__ is not original_call
-        assert getattr(decoder_cls, "_axolotl_shared_kv_patched", False) is True

tests/prompt_strategies/test_dpo_chat_templates.py

@@ -111,7 +111,7 @@ def fixture_argilla_chat_dataset():
 @pytest.fixture(name="phi3_tokenizer")
 @enable_hf_offline
 def fixture_phi3_tokenizer():
-    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-medium-128k-instruct")
 
     return tokenizer
 
@@ -214,8 +214,8 @@ class TestAssistantDPOChatTemplatePhi3:
             + "<|user|>\ngoodbye<|end|>\n"
             + "<|assistant|>\n"
         )
-        assert result["chosen"] == "goodbye<|end|>\n<|endoftext|>"
+        assert result["chosen"] == "goodbye<|end|>"
-        assert result["rejected"] == "party on<|end|>\n<|endoftext|>"
+        assert result["rejected"] == "party on<|end|>"
 
 
 class TestAssistantDPOChatTemplateGemma:
@@ -290,8 +290,8 @@ class TestArgillaChatDPOChatTemplate:
         )
         result = transform_fn(argilla_chat_dataset[0], tokenizer=phi3_tokenizer)
         assert result["prompt"] == "<|user|>\nhello<|end|>\n" + "<|assistant|>\n"
-        assert result["chosen"] == "goodbye<|end|>\n<|endoftext|>"
+        assert result["chosen"] == "goodbye<|end|>"
-        assert result["rejected"] == "party on<|end|>\n<|endoftext|>"
+        assert result["rejected"] == "party on<|end|>"
 
 
 class TestDPOChatTemplateToolRole: