add validation for DFT

.github/PULL_REQUEST_TEMPLATE.md

@@ -15,11 +15,6 @@
 <!--- Include details of your testing environment, tests ran to see how -->
 <!--- your change affects other areas of the code, etc. -->
 
-## AI Usage Disclaimer
-
-<!--- Was AI (e.g., ChatGPT, Claude, Copilot) used to generate or assist with this PR? -->
-<!--- Please indicate: No / Yes (specify which tool and to what extent) -->
-
 ## Screenshots (if appropriate)
 
 ## Types of changes

.github/workflows/base.yml

@@ -21,8 +21,6 @@ jobs:
     timeout-minutes: 480
     # this job needs to be run on self-hosted GPU runners...
     runs-on: ubuntu-latest-m
-    env:
-      HAS_DOCKERHUB_CREDS: ${{ secrets.DOCKERHUB_USERNAME != '' && secrets.DOCKERHUB_TOKEN != '' }}
     strategy:
       fail-fast: false
       matrix:
@@ -34,7 +32,6 @@ jobs:
             pytorch: 2.8.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
             dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64"
           - cuda: "128"
             cuda_version: 12.8.1
             cudnn_version: ""
@@ -42,7 +39,6 @@ jobs:
             pytorch: 2.9.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
             dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
           - cuda: "128"
             cuda_version: 12.8.1
             cudnn_version: ""
@@ -50,15 +46,6 @@ jobs:
             pytorch: 2.9.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
             dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "129"
-            cuda_version: 12.9.1
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.9.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
           - cuda: "130"
             cuda_version: 13.0.0
             cudnn_version: ""
@@ -66,15 +53,6 @@ jobs:
             pytorch: 2.9.1
             torch_cuda_arch_list: "9.0+PTX"
             dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "130"
-            cuda_version: 13.0.0
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.9.1
-            torch_cuda_arch_list: "9.0+PTX"
-            dockerfile: "Dockerfile-base"
-            platforms: "linux/amd64,linux/arm64"
 #          - cuda: "128"
 #            cuda_version: 12.8.1
 #            cudnn_version: ""
@@ -101,7 +79,7 @@ jobs:
             axolotlai/axolotl-base
       - name: Login to Docker Hub
         uses: docker/login-action@v2
-        if: ${{ github.event_name != 'pull_request' && env.HAS_DOCKERHUB_CREDS == 'true' }}
+        if: ${{ github.event_name != 'pull_request' && secrets.DOCKERHUB_USERNAME != '' && secrets.DOCKERHUB_TOKEN != '' }}
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
@@ -112,7 +90,7 @@ jobs:
         with:
           context: .
           file: ./docker/${{ matrix.dockerfile }}
-          platforms: ${{ matrix.platforms }}
+          platforms: linux/amd64,linux/arm64
           push: ${{ github.event_name != 'pull_request' }}
           tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}
@@ -127,8 +105,6 @@ jobs:
     if: ${{ github.repository_owner == 'axolotl-ai-cloud' && (github.event_name != 'pull_request' || !github.event.pull_request.draft) }}
     timeout-minutes: 480
     runs-on: ubuntu-latest-m
-    env:
-      HAS_DOCKERHUB_CREDS: ${{ secrets.DOCKERHUB_USERNAME != '' && secrets.DOCKERHUB_TOKEN != '' }}
     strategy:
       fail-fast: false
       matrix:
@@ -140,7 +116,6 @@ jobs:
             pytorch: 2.8.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
             dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64"
           - cuda: "128"
             cuda_version: 12.8.1
             cudnn_version: ""
@@ -148,7 +123,6 @@ jobs:
             pytorch: 2.9.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
             dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
           - cuda: "128"
             cuda_version: 12.8.1
             cudnn_version: ""
@@ -156,15 +130,6 @@ jobs:
             pytorch: 2.9.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
             dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "129"
-            cuda_version: 12.9.1
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.9.1
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
           - cuda: "130"
             cuda_version: 13.0.0
             cudnn_version: ""
@@ -172,15 +137,6 @@ jobs:
             pytorch: 2.9.1
             torch_cuda_arch_list: "9.0+PTX"
             dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: "130"
-            cuda_version: 13.0.0
-            cudnn_version: ""
-            python_version: "3.12"
-            pytorch: 2.9.1
-            torch_cuda_arch_list: "9.0+PTX"
-            dockerfile: "Dockerfile-uv-base"
-            platforms: "linux/amd64,linux/arm64"
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -192,7 +148,6 @@ jobs:
             axolotlai/axolotl-base-uv
       - name: Login to Docker Hub
         uses: docker/login-action@v2
-        if: ${{ github.event_name != 'pull_request' && env.HAS_DOCKERHUB_CREDS == 'true' }}
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
@@ -203,7 +158,6 @@ jobs:
         with:
           context: .
           file: ./docker/${{ matrix.dockerfile }}
-          platforms: ${{ matrix.platforms }}
           push: ${{ github.event_name != 'pull_request' }}
           tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
           labels: ${{ steps.metadata.outputs.labels }}

.github/workflows/main.yml

@@ -20,32 +20,22 @@ jobs:
             python_version: "3.11"
             pytorch: 2.8.0
             axolotl_extras:
-            platforms: "linux/amd64"
+            is_latest: true
           - cuda: 128
             cuda_version: 12.8.1
             python_version: "3.11"
             pytorch: 2.9.0
             axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
           - cuda: 128
             cuda_version: 12.8.1
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
-            is_latest: true
-          - cuda: 129
-            cuda_version: 12.9.1
-            python_version: "3.12"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
           - cuda: 130
             cuda_version: 13.0.0
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -71,7 +61,7 @@ jobs:
         uses: docker/build-push-action@v5
         with:
           context: .
-          platforms: ${{ matrix.platforms }}
+          platforms: linux/amd64,linux/arm64
           build-args: |
             BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
             CUDA=${{ matrix.cuda }}
@@ -98,32 +88,22 @@ jobs:
             python_version: "3.11"
             pytorch: 2.8.0
             axolotl_extras:
-            platforms: "linux/amd64"
+            is_latest: true
           - cuda: 128
             cuda_version: 12.8.1
             python_version: "3.11"
             pytorch: 2.9.0
             axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
           - cuda: 128
             cuda_version: 12.8.1
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:
-            is_latest: true
-            platforms: "linux/amd64,linux/arm64"
-          - cuda: 129
-            cuda_version: 12.9.1
-            python_version: "3.12"
-            pytorch: 2.9.1
-            axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
           - cuda: 130
             cuda_version: 13.0.0
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:
-            platforms: "linux/amd64,linux/arm64"
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -148,7 +128,7 @@ jobs:
         uses: docker/build-push-action@v5
         with:
           context: .
-          platforms: ${{ matrix.platforms }}
+          platforms: linux/amd64,linux/arm64
           build-args: |
             BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
             CUDA=${{ matrix.cuda }}
@@ -169,11 +149,11 @@ jobs:
           - cuda: 128
             cuda_version: 12.8.1
             python_version: "3.11"
-            pytorch: 2.9.1
+            pytorch: 2.8.0
             axolotl_extras:
-            is_latest: true
-          - cuda: 130
-            cuda_version: 13.0.0
+            is_latest:
+          - cuda: 128
+            cuda_version: 12.8.1
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:

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

@@ -35,26 +35,21 @@ jobs:
             pytorch: 2.8.0
             axolotl_extras: fbgemm-gpu
             num_gpus: 2
+            nightly_build: "true"
           - cuda: 128
             cuda_version: 12.8.1
             python_version: "3.11"
             pytorch: 2.9.1
-            axolotl_extras: "fbgemm-gpu"
+            axolotl_extras: fbgemm-gpu
             num_gpus: 2
-          - cuda: 129
-            cuda_version: 12.9.1
-            python_version: "3.12"
-            pytorch: 2.9.1
-            axolotl_extras: "fbgemm-gpu"
-            num_gpus: 2
-            dockerfile: "Dockerfile-uv.jinja"
+            nightly_build: "true"
           - 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
+            nightly_build: "true"
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
     steps:
@@ -76,8 +71,8 @@ jobs:
           echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
           echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
-          echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         run: |
           modal run -m cicd.multigpu

.github/workflows/pypi.yml

@@ -40,7 +40,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          pip3 install wheel packaging==26.0
+          pip3 install wheel packaging==23.2
           pip3 install --no-build-isolation -e .
           pip3 install -r requirements-dev.txt -r requirements-tests.txt
 
@@ -48,9 +48,9 @@ jobs:
         id: tag
         run: echo ::set-output name=TAG_NAME::$(echo $GITHUB_REF | cut -d / -f 3)
 
-      - name: Update version in VERSION file
+      - name: Update version in setup.py
         run: |
-          echo "${{ steps.tag.outputs.TAG_NAME }}" | sed 's/^v//' > VERSION
+          sed -i -E 's/version="([0-9.]+)",/version="${{ steps.tag.outputs.TAG_NAME }}",/g' setup.py
 
       - name: Build a source dist
         run: |

.github/workflows/tests-nightly.yml

@@ -48,7 +48,7 @@ jobs:
       - name: upgrade pip
         run: |
           pip3 install --upgrade pip
-          pip3 install --upgrade packaging==26.0 setuptools==75.8.0 wheel
+          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 wheel
 
       - name: Install PyTorch
         run: |

.github/workflows/tests.yml

@@ -54,13 +54,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ["3.11", "3.12"]
+        python_version: ["3.11"]
         pytorch_version: ["2.8.0", "2.9.0", "2.9.1"]
-        exclude:
-          - python_version: "3.12"
-            pytorch_version: "2.8.0"
-          - python_version: "3.12"
-            pytorch_version: "2.9.0"
     timeout-minutes: 20
 
     steps:
@@ -87,7 +82,7 @@ jobs:
       - name: upgrade pip
         run: |
           pip3 install --upgrade pip
-          pip3 install --upgrade packaging==26.0 setuptools==75.8.0 wheel
+          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 wheel
 
       - name: Install PyTorch
         run: |
@@ -115,10 +110,10 @@ jobs:
 
       - name: Pre-Download dataset fixture
         run: |
-          hf download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
+          huggingface-cli download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
 
       - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan
 
       - name: Run tests
         run: |
@@ -132,7 +127,7 @@ jobs:
           pytest -v --durations=10 tests/cli/ --cov=axolotl --cov-append --cov-report=xml
 
       - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5
@@ -149,13 +144,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ["3.11", "3.12"]
+        python_version: ["3.11"]
         pytorch_version: ["2.8.0", "2.9.0", "2.9.1"]
-        exclude:
-          - python_version: "3.12"
-            pytorch_version: "2.8.0"
-          - python_version: "3.12"
-            pytorch_version: "2.9.0"
     timeout-minutes: 20
 
     steps:
@@ -182,7 +172,7 @@ jobs:
       - name: upgrade pip
         run: |
           pip3 install --upgrade pip
-          pip3 install --upgrade packaging==26.0 setuptools==75.8.0 setuptools_scm build wheel psutil
+          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 setuptools_scm build wheel psutil
 
       - name: Install PyTorch
         run: |
@@ -210,7 +200,7 @@ jobs:
           axolotl --help
 
       - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan
 
       - name: Run tests
         run: |
@@ -219,10 +209,10 @@ jobs:
           pytest -v --durations=10 tests/cli/
 
       - name: Show HF cache
-        run: hf cache ls
+        run: hf cache scan
 
   gate-skip-e2e:
-    needs: [pre-commit]
+    needs: [pre-commit, pytest, pytest-sdist]
     runs-on: ubuntu-latest
     outputs:
       skip: ${{ steps.compute.outputs.skip }}
@@ -258,16 +248,16 @@ jobs:
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
-    needs: [pre-commit, pytest]
+    needs: [pre-commit, pytest, pytest-sdist, gate-skip-e2e]
 
     strategy:
       fail-fast: false
       matrix:
         include:
-          - cuda: 129
-            cuda_version: 12.9.1
-            python_version: "3.12"
-            pytorch: 2.9.1
+          - cuda: 128
+            cuda_version: 12.8.1
+            python_version: "3.11"
+            pytorch: 2.8.0
             num_gpus: 1
             axolotl_extras:
             dockerfile: "Dockerfile-uv.jinja"
@@ -369,9 +359,9 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - cuda: 129
-            cuda_version: 12.9.1
-            python_version: "3.12"
+          - cuda: 128
+            cuda_version: 12.8.1
+            python_version: "3.11"
             pytorch: 2.9.1
             num_gpus: 1
             axolotl_extras:

.runpod/README.md

@@ -123,7 +123,7 @@ datasets:
 | --------------------------------- | -------------------------- | ----------------------------------- |
 | `dataset_prepared_path`           | `"data/last_run_prepared"` | Path for prepared dataset           |
 | `push_dataset_to_hub`             | `""`                       | Push dataset to HF hub              |
-| `dataset_num_proc`                | `4`                        | Number of preprocessing processes   |
+| `dataset_processes`               | `4`                        | Number of preprocessing processes   |
 | `dataset_keep_in_memory`          | `false`                    | Keep dataset in memory              |
 | `shuffle_merged_datasets`         | `true`                     | Shuffle merged datasets             |
 | `shuffle_before_merging_datasets` | `false`                    | Shuffle each dataset before merging |

.runpod/src/config/config.yaml

@@ -39,6 +39,7 @@
 #     type: # linear | dynamic
 #     factor: # float
 
+
 # # Whether you are training a 4-bit GPTQ quantized model
 # gptq: true
 # gptq_groupsize: 128 # group size
@@ -106,7 +107,7 @@
 # push_dataset_to_hub: # repo path
 # # The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()`
 # # if not set.
-# dataset_num_proc: # defaults to os.cpu_count() if not set
+# dataset_processes: # defaults to os.cpu_count() if not set
 # # push checkpoints to hub
 # hub_model_id: # repo path to push finetuned model
 # # how to push checkpoints to hub
@@ -223,6 +224,9 @@
 # eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
 # eval_table_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128
 
+# # Save model as safetensors (require safetensors package)
+# save_safetensors:
+
 # # Whether to mask out or include the human's prompt from the training labels
 # train_on_inputs: false
 # # Group similarly sized data to minimize padding.
@@ -348,6 +352,8 @@
 # # Allow overwrite yml config using from cli
 # strict:
 
+
+
 base_model: ${BASE_MODEL}
 base_model_ignore_patterns: ${BASE_MODEL_IGNORE_PATTERNS}
 base_model_config: ${BASE_MODEL_CONFIG}
@@ -406,7 +412,7 @@ chat_template_jinja: ${CHAT_TEMPLATE_JINJA}
 default_system_message: ${DEFAULT_SYSTEM_MESSAGE}
 dataset_prepared_path: ${DATASET_PREPARED_PATH}
 push_dataset_to_hub: ${PUSH_DATASET_TO_HUB}
-dataset_num_proc: ${DATASET_NUM_PROC}
+dataset_processes: ${DATASET_PROCESSES}
 dataset_keep_in_memory: ${DATASET_KEEP_IN_MEMORY}
 hub_model_id: ${HUB_MODEL_ID}
 hub_strategy: ${HUB_STRATEGY}
@@ -506,6 +512,7 @@ profiler_steps: ${PROFILER_STEPS}
 loss_watchdog_threshold: ${LOSS_WATCHDOG_THRESHOLD}
 loss_watchdog_patience: ${LOSS_WATCHDOG_PATIENCE}
 
+save_safetensors: ${SAVE_SAFETENSORS}
 train_on_inputs: ${TRAIN_ON_INPUTS}
 group_by_length: ${GROUP_BY_LENGTH}
 gradient_checkpointing: ${GRADIENT_CHECKPOINTING}

README.md

@@ -88,7 +88,7 @@ Features:
 #### Using pip
 
 ```bash
-pip3 install -U packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install -U packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]
 
 # Download example axolotl configs, deepspeed configs

VERSION

@@ -1 +0,0 @@
-0.15.0.dev0

_quarto.yml

@@ -251,6 +251,7 @@ website:
                 - docs/models/olmo3.qmd
                 - docs/models/trinity.qmd
                 - docs/models/arcee.qmd
+                - docs/models/mistral.qmd
                 - section: "Ministral3"
                   contents:
                     - docs/models/ministral3.qmd
@@ -265,7 +266,6 @@ website:
                 - docs/models/mistral-small.qmd
                 - docs/models/voxtral.qmd
                 - docs/models/devstral.qmd
-                - docs/models/mistral.qmd
                 - docs/models/llama-4.qmd
                 - docs/models/llama-2.qmd
                 - docs/models/qwen3-next.qmd
@@ -320,7 +320,6 @@ website:
             - docs/multipack.qmd
             - docs/mixed_precision.qmd
             - docs/optimizers.qmd
-            - docs/attention.qmd
 
         - section: "Advanced Features"
           contents:

cicd/Dockerfile-uv.jinja

@@ -31,7 +31,7 @@ RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
         sed -i 's#^datasets.*#datasets @ git+https://github.com/huggingface/datasets.git@main#' requirements.txt; \
     fi
 
-RUN uv pip install packaging==26.0 setuptools==75.8.0
+RUN uv pip install packaging==23.2 setuptools==75.8.0
 RUN uv pip install torchvision
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
         uv pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \

cicd/Dockerfile.jinja

@@ -32,7 +32,7 @@ RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
         sed -i 's#^datasets.*#datasets @ git+https://github.com/huggingface/datasets.git@main#' requirements.txt; \
     fi
 
-RUN pip install packaging==26.0 setuptools==75.8.0 psutil
+RUN pip install packaging==23.2 setuptools==75.8.0 psutil
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
         pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \

cicd/multigpu.py

@@ -17,8 +17,7 @@ template_loader = jinja2.FileSystemLoader(searchpath=cicd_path)
 template_env = jinja2.Environment(
     loader=template_loader, autoescape=select_autoescape()
 )
-dockerfile = os.environ.get("E2E_DOCKERFILE", "Dockerfile.jinja")
-df_template = template_env.get_template(dockerfile)
+df_template = template_env.get_template("Dockerfile.jinja")
 
 df_args = {
     "AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
@@ -28,11 +27,8 @@ df_args = {
     "CUDA": os.environ.get("CUDA", "126"),
     "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
     "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
-    "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
     "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
     "HF_HOME": "/workspace/data/huggingface-cache/hub",
-    "PYTHONUNBUFFERED": os.environ.get("PYTHONUNBUFFERED", "1"),
-    "DEEPSPEED_LOG_LEVEL": os.environ.get("DEEPSPEED_LOG_LEVEL", "WARNING"),
 }
 
 dockerfile_contents = df_template.render(**df_args)

cicd/multigpu.sh

@@ -2,7 +2,7 @@
 set -e
 
 # Only run two tests at a time to avoid OOM on GPU (with coverage collection)
-pytest -v --durations=10 -n2 --maxfail=3 \
+pytest -v --durations=10 -n2 --maxfail=4 \
   --ignore=/workspace/axolotl/tests/e2e/multigpu/solo/ \
   --ignore=/workspace/axolotl/tests/e2e/multigpu/patched/ \
   /workspace/axolotl/tests/e2e/multigpu/ \

docker/Dockerfile

@@ -6,7 +6,6 @@ ARG AXOLOTL_EXTRAS=""
 ARG AXOLOTL_ARGS=""
 ARG CUDA="118"
 ARG PYTORCH_VERSION="2.1.2"
-ARG TARGETARCH
 
 ENV PYTORCH_VERSION=$PYTORCH_VERSION
 
@@ -21,17 +20,13 @@ RUN git clone --depth=1 https://github.com/axolotl-ai-cloud/axolotl.git
 
 WORKDIR /workspace/axolotl
 
-# If AXOLOTL_EXTRAS is set, append it in brackets; don't install deepspeed with arm64
-RUN if [ "$TARGETARCH" = "arm64" ]; then \
-        BASE_EXTRAS="flash-attn,ring-flash-attn,optimizers,ray"; \
+# If AXOLOTL_EXTRAS is set, append it in brackets
+RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \
-        BASE_EXTRAS="deepspeed,flash-attn,ring-flash-attn,optimizers,ray"; \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
     fi && \
-    if [ "$AXOLOTL_EXTRAS" != "" ]; then \
-        pip install --no-build-isolation -e .[$BASE_EXTRAS,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
-    else \
-        pip install --no-build-isolation -e .[$BASE_EXTRAS] $AXOLOTL_ARGS; \
-    fi && \    python scripts/unsloth_install.py | sh && \
+    python scripts/unsloth_install.py | sh && \
     python scripts/cutcrossentropy_install.py | sh && \
     pip install pytest && \
     pip cache purge

docker/Dockerfile-base

@@ -43,7 +43,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
 
 WORKDIR /workspace
 
-RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==26.0 setuptools==75.8.0 wheel psutil && \
+RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==23.2 setuptools==75.8.0 wheel psutil && \
     python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} torchvision --extra-index-url https://download.pytorch.org/whl/cu$CUDA && \
     python3 -m pip cache purge
 

docker/Dockerfile-base-nightly

@@ -30,7 +30,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
 
 WORKDIR /workspace
 
-RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==26.0 setuptools==75.8.0 wheel && \
+RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==23.2 setuptools==75.8.0 wheel && \
     python3 -m pip install --no-cache-dir -U torch --extra-index-url https://download.pytorch.org/whl/nightly/cu$CUDA && \
     python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \
     python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" && \

docker/Dockerfile-uv-base

@@ -2,7 +2,6 @@ ARG CUDA_VERSION="12.6.3"
 ARG CUDNN_VERSION=""
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
-ARG TARGETARCH
 
 FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
 
@@ -32,35 +31,20 @@ ENV PATH="/workspace/axolotl-venv/bin:${PATH}"
 
 RUN uv pip install packaging setuptools wheel psutil \
     && uv pip install torch==${PYTORCH_VERSION} torchvision \
+    && uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" \
+    && uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" \
     && uv pip install awscli pydantic
 
-RUN if [ "$TARGETARCH" = "amd64" ]; then \
-        uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main"; \
-        uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main"; \
-    fi
-
 RUN case "$PYTORCH_VERSION" in \
         2.9.[0-9]*) \
-            if [ "$TARGETARCH" = "amd64" ]; then \
-                if [ "$CUDA" = "128" ]; then \
-                    wget -nv https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
-                    uv pip install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
-                    rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
-                elif [ "$CUDA" = "130" ]; then \
-                    wget -nv https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_x86_64.whl; \
-                    uv pip install --no-cache-dir flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_x86_64.whl; \
-                    rm flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_x86_64.whl; \
-                fi \
-            elif [ "$TARGETARCH" = "arm64" ]; then \
-                if [ "$CUDA" = "128" ]; then \
-                    wget -nv https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.6.4/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_aarch64.whl; \
-                    uv pip install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_aarch64.whl; \
-                    rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_aarch64.whl; \
-                elif [ "$CUDA" = "130" ]; then \
-                    wget -nv https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.6.4/flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_aarch64.whl; \
-                    uv pip install --no-cache-dir flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_aarch64.whl; \
-                    rm flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_aarch64.whl; \
-                fi \
+            if [ "$CUDA" = "128" ]; then \
+                wget -nv https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+                uv pip install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+                rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+            elif [ "$CUDA" = "130" ]; then \
+                wget -nv https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_x86_64.whl; \
+                uv pip install --no-cache-dir flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_x86_64.whl; \
+                rm flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_x86_64.whl; \
             fi \
             ;; \
     esac

docs/amd_hpc.qmd

@@ -86,7 +86,7 @@ export HF_DATASETS_OFFLINE=1
 Download a base model using the Hugging Face CLI:
 
 ```bash
-hf download meta-llama/Meta-Llama-3.1-8B --local-dir ~/hfdata/llama3.1-8B
+huggingface-cli download meta-llama/Meta-Llama-3.1-8B --local-dir ~/hfdata/llama3.1-8B
 ```
 
 ### 10. Create Axolotl Configuration

docs/attention.qmd

@@ -1,140 +0,0 @@
----
-title: Attention
-description: Supported attention modules in Axolotl
----
-
-## SDP Attention
-
-This is the default built-in attention in PyTorch.
-
-```yaml
-sdp_attention: true
-```
-
-For more details: [PyTorch docs](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)
-
-## Flash Attention 2
-
-Uses efficient kernels to compute attention.
-
-```yaml
-flash_attention: true
-```
-
-For more details: [Flash Attention](https://github.com/Dao-AILab/flash-attention/)
-
-### Nvidia
-
-Requirements: Ampere, Ada, or Hopper GPUs
-
-Note: For Turing GPUs or lower, please use other attention methods.
-
-```bash
-pip install flash-attn --no-build-isolation
-```
-
-::: {.callout-tip}
-
-If you get `undefined symbol` while training, ensure you installed PyTorch prior to Axolotl. Alternatively, try reinstall or downgrade a version.
-
-:::
-
-#### Flash Attention 3
-
-Requirements: Hopper only and CUDA 12.8 (recommended)
-
-```bash
-git clone https://github.com/Dao-AILab/flash-attention.git
-cd flash-attention/hopper
-
-python setup.py install
-```
-
-### AMD
-
-Requirements: ROCm 6.0 and above.
-
-See [Flash Attention AMD docs](https://github.com/Dao-AILab/flash-attention/tree/main?tab=readme-ov-file#amd-rocm-support).
-
-## Flex Attention
-
-A flexible PyTorch API for attention used in combination with `torch.compile`.
-
-```yaml
-flex_attention: true
-
-# recommended
-torch_compile: true
-```
-
-::: {.callout-note}
-
-We recommend using latest stable version of PyTorch for best performance.
-
-:::
-
-For more details: [PyTorch docs](https://pytorch.org/blog/flexattention/)
-
-## SageAttention
-
-Attention kernels with QK Int8 and PV FP16 accumulator.
-
-```yaml
-sage_attention: true
-```
-
-Requirements: Ampere, Ada, or Hopper GPUs
-
-```bash
-pip install sageattention==2.2.0 --no-build-isolation
-```
-
-::: {.callout-warning}
-
-Only LoRA/QLoRA recommended at the moment. We found loss drop to 0 for full finetuning. See [GitHub Issue](https://github.com/thu-ml/SageAttention/issues/198).
-
-:::
-
-For more details: [Sage Attention](https://github.com/thu-ml/SageAttention)
-
-::: {.callout-note}
-
-We do not support SageAttention 3 at the moment. If you are interested on adding this or improving SageAttention implementation, please make an Issue.
-
-:::
-
-
-## xFormers
-
-```yaml
-xformers_attention: true
-```
-
-::: {.callout-tip}
-
-We recommend using with Turing GPUs or below (such as on Colab).
-
-:::
-
-For more details: [xFormers](https://github.com/facebookresearch/xformers)
-
-## Shifted Sparse Attention
-
-::: {.callout-warning}
-
-We plan to deprecate this! If you use this feature, we recommend switching to methods above.
-
-:::
-
-Requirements: LLaMA model architecture
-
-```yaml
-flash_attention: true
-s2_attention: true
-```
-
-::: {.callout-tip}
-
-No sample packing support!
-
-:::

docs/cli.qmd

@@ -210,8 +210,6 @@ axolotl lm-eval config.yml
 Configuration options:
 
 ```yaml
-lm_eval_model: # model to evaluate (local or hf path)
-
 # List of tasks to evaluate
 lm_eval_tasks:
   - arc_challenge
@@ -220,7 +218,7 @@ lm_eval_batch_size: # Batch size for evaluation
 output_dir: # Directory to save evaluation results
 ```
 
-See [LM Eval Harness integration docs](https://docs.axolotl.ai/docs/custom_integrations.html#language-model-evaluation-harness-lm-eval) for full configuration details.
+See [LM Eval Harness](https://github.com/EleutherAI/lm-evaluation-harness) for more details.
 
 ### delinearize-llama4
 

docs/installation.qmd

@@ -165,7 +165,7 @@ We recommend using WSL2 (Windows Subsystem for Linux) or Docker.
    ```
 4. (Optional) Login to Hugging Face:
    ```{.bash}
-   hf auth login
+   huggingface-cli login
    ```
 
 ## Troubleshooting {#sec-troubleshooting}

docs/lora_optims.qmd

@@ -89,10 +89,6 @@ lora_o_kernel: true
 Currently, LoRA kernels are not supported for RLHF training, only SFT.
 :::
 
-::: {.callout-warning}
-LoRA kernels do not support remote modeling code.
-:::
-
 ## Requirements
 
 - One or more NVIDIA or AMD GPUs (in order to use the Triton kernels)

docs/multimodal.qmd

@@ -19,7 +19,6 @@ format:
 - [Gemma-3n](#sec-gemma-3n)
 - [Qwen2-VL](#sec-qwen2-vl)
 - [Qwen2.5-VL](#sec-qwen25-vl)
-- [GLM-4.6V](#sec-glm-4-6v)
 - [SmolVLM2](#sec-smolvlm2)
 - [LFM2-VL](#sec-lfm2-vl)
 - [Intern-VL](#sec-intern-vl)
@@ -184,18 +183,6 @@ base_model: Qwen/Qwen3-VL-4B-Instruct
 chat_template: qwen2_vl  # same as qwen2-vl
 ```
 
-### GLM-4.6V {#sec-glm-4-6v}
-
-Both GLM-4.6V (106B MoE) and GLM-4.6V-Flash (9B) are supported.
-
-```yaml
-# GLM-4.6V (106B MoE version)
-base_model: zai-org/GLM-4.6V
-
-# OR GLM-4.6V-Flash (9B version)
-base_model: zai-org/GLM-4.6V-Flash
-```
-
 ### SmolVLM2 {#sec-smolvlm2}
 
 ::: {.callout-tip}

docs/rlhf.qmd

@@ -17,7 +17,6 @@ feedback. Various methods include, but not limited to:
 - [Kahneman-Tversky Optimization (KTO)](#kto)
 - [Odds Ratio Preference Optimization (ORPO)](#orpo)
 - [Group Relative Policy Optimization (GRPO)](#grpo)
-- [Group Reward-Decoupled Policy Optimization (GDPO)](#gdpo)
 
 
 ## RLHF using Axolotl
@@ -721,102 +720,6 @@ trl:
 
 For more information, see [GRPO docs](https://huggingface.co/docs/trl/v0.17.0/en/grpo_trainer#loss-types).
 
-### GDPO
-
-GDPO (Group Reward-Decoupled Policy Optimization) extends GRPO for multi-reward training. It addresses the **reward advantage collapse** problem by normalizing each reward function independently before combining them.
-
-::: {.callout-tip}
-Use GDPO when training with multiple reward functions. For single reward, GRPO and GDPO produce equivalent results.
-:::
-
-Paper: [https://arxiv.org/pdf/2501.05242](https://arxiv.org/pdf/2501.05242)
-
-GDPO uses TRL's native `multi_objective_aggregation` parameter under the hood. When you set `rl: gdpo`, axolotl automatically configures TRL to use `normalize_then_sum` aggregation.
-
-```yaml
-base_model: Qwen/Qwen2.5-1.5B-Instruct
-
-vllm:
-    host: 0.0.0.0
-    port: 8000
-    tensor_parallel_size: 2
-    gpu_memory_utilization: 0.85
-
-rl: gdpo
-
-trl:
-    beta: 0.001
-    max_completion_length: 256
-    use_vllm: true
-    num_generations: 4
-    reward_funcs:
-        - rewards.format_reward
-        - rewards.correctness_reward
-    reward_weights: [1.0, 2.0]
-
-datasets:
-    - path: openai/gsm8k
-      name: main
-      type: rewards.oai_gsm8k_transform
-```
-
-You can also use GRPO with explicit aggregation control:
-
-```yaml
-rl: grpo
-trl:
-    multi_objective_aggregation: normalize_then_sum  # GDPO behavior
-    # or: sum_then_normalize  # Default GRPO behavior
-```
-
-#### GDPO vs GRPO
-
-| Aspect | GRPO | GDPO |
-|--------|------|------|
-| **Aggregation** | `sum_then_normalize` | `normalize_then_sum` |
-| **Multi-reward** | May collapse advantages | Preserves reward signals |
-| **Single reward** | Standard behavior | Equivalent to GRPO |
-
-#### Why GDPO?
-
-When using multiple rewards with GRPO, different reward combinations can produce identical advantages:
-
-```
-# Example: format + correctness rewards
-[format=0, correct=3] → sum=3
-[format=1, correct=2] → sum=3  ← GRPO sees these as equal!
-[format=2, correct=1] → sum=3
-[format=3, correct=0] → sum=3
-```
-
-GDPO normalizes each reward independently, preserving their relative differences.
-
-#### Reward Functions
-
-GDPO uses the same reward function format as GRPO:
-
-```python
-# rewards.py
-def format_reward(completions, **kwargs) -> list[float]:
-    return [1.0 if len(c) > 10 else 0.0 for c in completions]
-
-def correctness_reward(completions, answers, **kwargs) -> list[float]:
-    rewards = []
-    for completion, answer in zip(completions, answers):
-        # Your scoring logic here
-        rewards.append(score)
-    return rewards
-```
-
-#### Sequence Parallelism
-
-GDPO supports sequence parallelism for long-context training:
-
-```yaml
-rl: gdpo
-context_parallel_size: 2
-```
-
 ### SimPO
 
 SimPO uses [CPOTrainer](https://huggingface.co/docs/trl/main/en/cpo_trainer) but with alternative loss function.

examples/apertus/README.md

@@ -15,7 +15,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
 
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'
 
 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy

examples/arcee/README.md

@@ -17,7 +17,7 @@ Thanks to the team at Arcee.ai for using Axolotl in supervised fine-tuning the A
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
 
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'
 
 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy

examples/colab-notebooks/colab-axolotl-example.ipynb

@@ -40,7 +40,7 @@
     "%%capture\n",
     "# This step can take ~5-10 minutes to install dependencies\n",
     "!pip install --no-build-isolation axolotl[flash-attn]>=0.9.1\n",
-    "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0d4ce4b\""
+    "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@318b7e2\""
    ]
   },
   {

examples/devstral/README.md

@@ -16,7 +16,7 @@ Thanks to the team at MistralAI for giving us early access to prepare for this r
 
 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 

examples/eaft/eaft-example.yml

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

examples/gemma3/gemma-3-1b-qlora.yml

@@ -1,7 +1,6 @@
 base_model: google/gemma-3-1b-it
 
 model_type: Gemma3ForCausalLM
-cls_model_config: Gemma3TextConfig
 
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
@@ -30,7 +29,7 @@ output_dir: ./outputs/out
 adapter: qlora
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0
+lora_dropout: 0.05
 lora_target_linear: true
 
 sequence_len: 2048

examples/gemma3/gemma-3-270m-qlora.yml

@@ -1,7 +1,6 @@
 base_model: google/gemma-3-270m-it
 
 model_type: Gemma3ForCausalLM
-cls_model_config: Gemma3TextConfig
 
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
@@ -30,7 +29,7 @@ output_dir: ./outputs/out
 adapter: qlora
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0
+lora_dropout: 0.05
 lora_target_linear: true
 
 sequence_len: 2048

examples/gemma3/gemma-3-4b-qlora.yml

@@ -2,7 +2,6 @@ base_model: google/gemma-3-4b-it
 
 # Need to set else transformers tries to load vision too
 model_type: Gemma3ForCausalLM
-cls_model_config: Gemma3TextConfig
 
 load_in_4bit: true
 
@@ -33,8 +32,8 @@ sample_packing: true
 
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0
-lora_target_linear: true
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
 
 wandb_project:
 wandb_entity:

examples/gemma3/gemma-3-4b-vision-qlora.yml

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

examples/gemma3n/README.md

@@ -10,7 +10,7 @@ Gemma-3n is a family of multimodal models from Google found on [HuggingFace](htt
 
 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 

examples/glm46v/README.md

@@ -1,44 +0,0 @@
-# Finetune GLM-4.6V with Axolotl
-
-GLM-4.6V is a family of vision-language models from ZhipuAI found on [HuggingFace](https://huggingface.co/zai-org/GLM-4.6V). This guide shows how to fine-tune it with Axolotl for vision-language tasks.
-
-
-
-## Getting started
-
-1. Install Axolotl from source following the [installation guide](https://docs.axolotl.ai/docs/installation.html#sec-edge-build).
-
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage.
-
-
-3. Run the fine-tuning:
-
-    glm-4-6v-flash(9B)
-    ```bash
-    axolotl train examples/glm46v/glm-4-6v-flash-qlora.yaml
-    ```
-
-Let us know how it goes. Happy finetuning! 🚀
-
-## Tips
-
-- Vision datasets should follow the format described in the [multimodal docs](https://docs.axolotl.ai/docs/multimodal.html#dataset-format)
-- You can run a **full finetuning** by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
-- Read more on how to load your own dataset in the [dataset loading docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-
-## Supported Models
-
-- **GLM-4.6V**: Full vision-language model (`zai-org/GLM-4.6V`)
-- **GLM-4.6V-Flash**: Faster variant (`zai-org/GLM-4.6V-Flash`)
-
-## Optimization Guides
-
-Please check the [Optimizations doc](https://docs.axolotl.ai/docs/optimizations.html).
-
-## Related Resources
-
-- [ZhipuAI GLM-4.6V](https://huggingface.co/zai-org/GLM-4.6V)
-- [Axolotl Docs](https://docs.axolotl.ai)
-- [Axolotl Website](https://axolotl.ai)
-- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
-- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)

examples/glm46v/glm-4-6v-flash-ddp.yaml

@@ -1,53 +0,0 @@
-base_model: zai-org/GLM-4.6V-Flash
-trust_remote_code: true
-
-processor_type: AutoProcessor
-load_in_4bit: true
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-ddp_find_unused_parameters: true
-
-output_dir: ./outputs/glm-4-6v-flash-qlora
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-sequence_len: 2048
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-logging_steps: 1
-sdp_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 0
-saves_per_epoch: 1
-weight_decay: 0.0

examples/glm46v/glm-4-6v-flash-qlora.yaml

@@ -1,50 +0,0 @@
-base_model: zai-org/GLM-4.6V-Flash
-trust_remote_code: true
-
-processor_type: AutoProcessor
-load_in_4bit: true
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-output_dir: ./outputs/glm-4-6v-flash-qlora
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-sequence_len: 2048
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-logging_steps: 1
-sdp_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 0
-saves_per_epoch: 1
-weight_decay: 0.0

examples/gpt-oss/README.md

@@ -14,7 +14,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
 
 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 

examples/granite4/README.md

@@ -15,7 +15,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
 
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'
 
 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy

examples/hunyuan/README.md

@@ -13,7 +13,7 @@ Tencent released a family of opensource models called HunYuan with varying param
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
 
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'
 
 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy

examples/jamba/qlora_fsdp_large.yaml

@@ -19,6 +19,7 @@ datasets:
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
 output_dir: jamba-large-fsdp-qlora-ft
+save_safetensors: true
 adapter: qlora
 sequence_len: 2048
 sample_packing: true

examples/llama-3/qlora-1b-gdpo.yaml

@@ -1,68 +0,0 @@
-base_model: meta-llama/Llama-3.2-1B-Instruct
-
-chat_template: llama3
-
-rl: gdpo
-
-trl:
-  beta: 0.001
-  max_completion_length: 128
-  num_generations: 2
-  temperature: 0.7
-  top_p: 0.95
-
-  use_vllm: false
-
-
-  multi_objective_aggregation: normalize_then_sum
-
-  reward_funcs:
-    - rwd.format_reward
-    - rwd.correctness_reward
-  reward_weights: [1.0, 2.0]
-
-  log_completions: true
-  num_completions_to_print: 3
-  scale_rewards: true
-
-datasets:
-  - path: openai/gsm8k
-    name: main
-    split: train[:1000]
-    type: rwd.gsm8k_transform
-
-val_set_size: 0.0
-output_dir: ./outputs/llama3-gdpo-out
-
-sequence_len: 512
-sample_packing: false
-pad_to_sequence_len: false
-
-gradient_accumulation_steps: 8
-micro_batch_size: 1
-num_epochs: 1
-max_steps: 100
-
-optimizer: adamw_torch_fused
-lr_scheduler: cosine
-learning_rate: 5e-5
-weight_decay: 0.01
-warmup_steps: 10
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-
-flash_attention: true
-logging_steps: 1
-save_steps: 50
-save_safetensors: true
-
-special_tokens:
-  pad_token: "<|end_of_text|>"
-
-
-seed: 42

examples/llama-3/qlora-fsdp-405b.yaml

@@ -12,6 +12,7 @@ datasets:
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
 output_dir: ./outputs/out/qlora-llama3_1-405b
+save_safetensors: true
 
 adapter: qlora
 

examples/magistral/README.md

@@ -14,7 +14,7 @@ Thanks to the team at MistralAI for giving us early access to prepare for these
 
 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.7.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 

examples/mamba/config.yml

@@ -47,5 +47,6 @@ saves_per_epoch: 1
 weight_decay: 0.0
 special_tokens:
 tokens:
+save_safetensors: False
 
 # save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/qwen3-next/README.md

@@ -15,7 +15,7 @@ This guide shows how to fine-tune it with Axolotl with multi-turn conversations
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
 
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'
 
 # Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy

examples/voxtral/README.md

@@ -12,7 +12,7 @@ Thanks to the team at MistralAI for giving us early access to prepare for this r
 
 ```bash
 # Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-pip3 install packaging==26.0 setuptools==75.8.0 wheel ninja
+pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=64", "wheel", "setuptools_scm>=8", "packaging==26.0"]
+requires = ["setuptools>=64", "wheel", "setuptools_scm>=8", "packaging==23.2"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -24,9 +24,6 @@ Repository = "https://github.com/axolotl-ai-cloud/axolotl.git"
 py-modules = ["setuptools_axolotl_dynamic_dependencies"]
 include-package-data = true
 
-[tool.setuptools.dynamic]
-version = { file = "VERSION" }
-
 [tool.setuptools.cmdclass]
 build_py = "setuptools_axolotl_dynamic_dependencies.BuildPyCommand"
 
@@ -60,6 +57,3 @@ indent-style = "space"
 skip-magic-trailing-comma = false
 line-ending = "auto"
 docstring-code-format = false
-
-[tool.uv.extra-build-dependencies]
-axolotl = ["huggingface_hub"]

requirements.txt

@@ -8,18 +8,18 @@ xformers>=0.0.23.post1
 liger-kernel==0.6.4
 # END section
 
-packaging==26.0
-huggingface_hub>=1.1.7
-peft>=0.18.1
+packaging==23.2
+
+huggingface_hub>=0.36.0
+peft>=0.18.0
 tokenizers>=0.22.1
-transformers==5.0.0
+transformers==4.57.1
 accelerate==1.12.0
-datasets==4.5.0
+datasets==4.4.2
 deepspeed>=0.18.3
-trl==0.27.1
+trl==0.25.1
 hf_xet==1.2.0
 kernels==0.11.5
-
 trackio>=0.13.0
 typing-extensions>=4.15.0
 
@@ -72,4 +72,4 @@ axolotl-contribs-mit==0.0.6
 # telemetry
 posthog==6.7.11
 
-mistral-common==1.8.8
+mistral-common==1.8.6

scripts/cutcrossentropy_install.py

@@ -29,5 +29,5 @@ UV_PREFIX = "uv " if USE_UV else ""
 
 print(
     UNINSTALL_PREFIX
-    + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0d4ce4b"'
+    + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@318b7e2"'
 )

setup.py

@@ -1,5 +1,6 @@
 """setup.py for axolotl"""
 
+import ast
 import os
 import platform
 import re
@@ -25,7 +26,6 @@ def parse_requirements(extras_require_map):
                 _install_requires.append(line)
     try:
         xformers_version = [req for req in _install_requires if "xformers" in req][0]
-        install_xformers = platform.machine() != "aarch64"
         if "Darwin" in platform.system():
             # skip packages not compatible with OSX
             skip_packages = [
@@ -62,68 +62,44 @@ def parse_requirements(extras_require_map):
             else:
                 raise ValueError("Invalid version format")
 
-            torch_parts = torch_version.split("+")
-            if len(torch_parts) == 2:
-                torch_cuda_version = torch_parts[1]
-                _dependency_links.append(
-                    f"https://download.pytorch.org/whl/{torch_cuda_version}"
-                )
-
             if (major, minor) >= (2, 9):
                 extras_require_map.pop("fbgemm-gpu")
-                extras_require_map["fbgemm-gpu"] = [
-                    "fbgemm-gpu==1.4.0",
-                    "fbgemm-gpu-genai==1.4.2",
-                ]
+                extras_require_map["fbgemm-gpu"] = ["fbgemm-gpu-genai==1.4.1"]
                 extras_require_map["vllm"] = ["vllm==0.11.1"]
-                if not install_xformers:
-                    _install_requires.pop(_install_requires.index(xformers_version))
-                extras_require_map["vllm"] = ["vllm==0.13.0"]
-                if patch == 0:
-                    extras_require_map["vllm"] = ["vllm==0.13.0"]
-                else:
-                    extras_require_map["vllm"] = ["vllm==0.14.0"]
             elif (major, minor) >= (2, 8):
                 extras_require_map.pop("fbgemm-gpu")
                 extras_require_map["fbgemm-gpu"] = ["fbgemm-gpu-genai==1.3.0"]
                 extras_require_map["vllm"] = ["vllm==0.11.0"]
-                if not install_xformers:
-                    _install_requires.pop(_install_requires.index(xformers_version))
             elif (major, minor) >= (2, 7):
                 _install_requires.pop(_install_requires.index(xformers_version))
                 if patch == 0:
-                    if install_xformers:
-                        _install_requires.append("xformers==0.0.30")
+                    _install_requires.append("xformers==0.0.30")
                     # vllm 0.9.x is incompatible with latest transformers
                     extras_require_map.pop("vllm")
                 else:
-                    if install_xformers:
-                        _install_requires.append("xformers==0.0.31")
+                    _install_requires.append("xformers==0.0.31")
                     extras_require_map["vllm"] = ["vllm==0.10.1"]
             elif (major, minor) >= (2, 6):
                 _install_requires.pop(_install_requires.index(xformers_version))
-                if install_xformers:
-                    _install_requires.append("xformers==0.0.29.post3")
+                _install_requires.append("xformers==0.0.29.post3")
                 # since we only support 2.6.0+cu126
                 _dependency_links.append("https://download.pytorch.org/whl/cu126")
                 extras_require_map.pop("vllm")
             elif (major, minor) >= (2, 5):
                 _install_requires.pop(_install_requires.index(xformers_version))
-                if install_xformers:
-                    if patch == 0:
-                        _install_requires.append("xformers==0.0.28.post2")
-                    else:
-                        _install_requires.append("xformers>=0.0.28.post3")
+                if patch == 0:
+                    _install_requires.append("xformers==0.0.28.post2")
+                else:
+                    _install_requires.append("xformers>=0.0.28.post3")
                 extras_require_map.pop("vllm")
             elif (major, minor) >= (2, 4):
                 extras_require_map.pop("vllm")
-                if install_xformers:
-                    if patch == 0:
-                        _install_requires.pop(_install_requires.index(xformers_version))
-                        _install_requires.append("xformers>=0.0.27")
-                    else:
-                        _install_requires.pop(_install_requires.index(xformers_version))
-                        _install_requires.append("xformers==0.0.28.post1")
+                if patch == 0:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                    _install_requires.append("xformers>=0.0.27")
+                else:
+                    _install_requires.pop(_install_requires.index(xformers_version))
+                    _install_requires.append("xformers==0.0.28.post1")
             else:
                 raise ValueError("axolotl requires torch>=2.4")
 
@@ -134,11 +110,15 @@ def parse_requirements(extras_require_map):
 
 def get_package_version():
     with open(
-        Path(os.path.dirname(os.path.abspath(__file__))) / "VERSION",
+        Path(os.path.dirname(os.path.abspath(__file__)))
+        / "src"
+        / "axolotl"
+        / "__init__.py",
         "r",
         encoding="utf-8",
     ) as fin:
-        version_ = fin.read().strip()
+        version_match = re.search(r"^__version__\s*=\s*(.*)$", fin.read(), re.MULTILINE)
+    version_ = ast.literal_eval(version_match.group(1))
     return version_
 
 

src/axolotl/__init__.py

@@ -1,11 +1,7 @@
 """Axolotl - Train and fine-tune large language models"""
 
 import pkgutil
-from importlib.metadata import PackageNotFoundError, version
 
 __path__ = pkgutil.extend_path(__path__, __name__)  # Make this a namespace package
 
-try:
-    __version__ = version("axolotl")
-except PackageNotFoundError:
-    __version__ = "unknown"
+__version__ = "0.13.0.dev"

src/axolotl/cli/__init__.py

@@ -5,6 +5,6 @@ import os
 from axolotl.logging_config import configure_logging
 
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
 
 configure_logging()

src/axolotl/cli/checks.py

@@ -44,7 +44,7 @@ def check_user_token() -> bool:
         return bool(user_info)
     except LocalTokenNotFoundError:
         LOG.warning(
-            "Error verifying HuggingFace token. Remember to log in using `hf auth login` and get your access token from https://huggingface.co/settings/tokens if you want to use gated models or datasets."
+            "Error verifying HuggingFace token. Remember to log in using `huggingface-cli login` and get your access token from https://huggingface.co/settings/tokens if you want to use gated models or datasets."
         )
         return False
     except HTTPError:

src/axolotl/cli/merge_lora.py

@@ -24,6 +24,7 @@ def do_merge_lora(*, cfg: DictDefault) -> None:
         cfg: Dictionary mapping `axolotl` config keys to values.
     """
     model, tokenizer, processor = load_model_and_tokenizer(cfg=cfg)
+    safe_serialization = cfg.save_safetensors is True
 
     LOG.info("Running merge of LoRA with base model...")
     model = model.merge_and_unload(progressbar=True)
@@ -41,6 +42,7 @@ def do_merge_lora(*, cfg: DictDefault) -> None:
         LOG.info(f"Saving merged model to: {str(Path(cfg.output_dir) / 'merged')}...")
         model.save_pretrained(
             str(Path(cfg.output_dir) / "merged"),
+            safe_serialization=safe_serialization,
             progressbar=True,
         )
         tokenizer.save_pretrained(

src/axolotl/cli/merge_sharded_fsdp_weights.py

@@ -14,6 +14,8 @@ from accelerate import PartialState
 from accelerate.utils import (
     SAFE_WEIGHTS_INDEX_NAME,
     SAFE_WEIGHTS_NAME,
+    WEIGHTS_INDEX_NAME,
+    WEIGHTS_NAME,
     is_torch_version,
 )
 from huggingface_hub import split_torch_state_dict_into_shards
@@ -38,15 +40,17 @@ class BFloat16CastPlanner(_EmptyStateDictLoadPlanner):
 def _distributed_checkpoint_to_merged_weights(
     checkpoint_dir: Union[str, Path],
     save_path: str,
+    safe_serialization: bool = False,
     max_shard_size: str = "5GB",
 ) -> Path:
     """
     Passthrough to `torch.distributed.checkpoint.format_utils.dcp_to_torch_save`. Will
-    save under `save_path` as `model.safetensors`.
+    save under `save_path` as either `model.safetensors` or `pytorch_model.bin`.
 
     Args:
         checkpoint_dir: Directory where distributed checkpoint is saved.
         save_path: Path to save model to.
+        safe_serialization: Whether to save in safetensors format.
         max_shard_size: Max size of model shards to save.
 
     Returns:
@@ -72,7 +76,11 @@ def _distributed_checkpoint_to_merged_weights(
         if isinstance(value, torch.Tensor) and value.dtype != torch.bfloat16:
             state_dict[key] = value.to(torch.bfloat16)
 
-    filename_pattern = SAFE_WEIGHTS_NAME.replace(".safetensors", "{suffix}.safetensors")
+    weights_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
+
+    filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(
+        ".safetensors", "{suffix}.safetensors"
+    )
     state_dict_split = split_torch_state_dict_into_shards(
         state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size
     )
@@ -90,12 +98,19 @@ def _distributed_checkpoint_to_merged_weights(
 
     for shard_file, tensors in filename_to_tensors:
         shard = {tensor: state_dict[tensor] for tensor in tensors}
-        safe_save_file(
-            shard, os.path.join(save_path_, shard_file), metadata={"format": "pt"}
-        )
+
+        if safe_serialization:
+            safe_save_file(
+                shard, os.path.join(save_path_, shard_file), metadata={"format": "pt"}
+            )
+        else:
+            torch.save(shard, os.path.join(save_path_, shard_file))
 
     if index is not None:
-        save_index_file = os.path.join(save_path_, SAFE_WEIGHTS_INDEX_NAME)
+        save_index_file = (
+            SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME
+        )
+        save_index_file = os.path.join(save_path_, save_index_file)
         # Save the index as well
         with open(save_index_file, "w", encoding="utf-8") as fout:
             content = json.dumps(index, indent=2, sort_keys=True) + "\n"
@@ -108,11 +123,13 @@ def _distributed_checkpoint_to_merged_weights(
 def merge_fsdp_weights(
     checkpoint_dir: str,
     output_path: str,
+    safe_serialization: bool = False,
     remove_checkpoint_dir: bool = False,
 ):
     """
     Merge the weights from sharded FSDP model checkpoints into a single combined checkpoint. Should be used if
-    `SHARDED_STATE_DICT` was used for the model. Weights will be saved to `{output_path}/model.safetensors`.
+    `SHARDED_STATE_DICT` was used for the model. Weights will be saved to `{output_path}/model.safetensors` if
+    `safe_serialization` else `pytorch_model.bin`.
 
     Note: this is a CPU-bound process.
 
@@ -121,6 +138,8 @@ def merge_fsdp_weights(
             The directory containing the FSDP checkpoints (can be either the model or optimizer).
         output_path (`str`):
             The path to save the merged checkpoint.
+        safe_serialization (`bool`, *optional*, defaults to `True`):
+            Whether to save the merged weights with safetensors (recommended).
         remove_checkpoint_dir (`bool`, *optional*, defaults to `False`):
             Whether to remove the checkpoint directory after merging.
 
@@ -158,7 +177,7 @@ def merge_fsdp_weights(
     if state.is_main_process:
         LOG.info(f"Merging FSDP weights from {checkpoint_dir_}")
         save_path = _distributed_checkpoint_to_merged_weights(
-            checkpoint_dir_, output_path
+            checkpoint_dir_, output_path, safe_serialization
         )
         LOG.info(f"Successfully merged FSDP weights and saved to {save_path}")
         if remove_checkpoint_dir:
@@ -191,6 +210,7 @@ def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
     merge_fsdp_weights(
         checkpoint_dir=str(fsdp_dir),
         output_path=output_path,
+        safe_serialization=True,
     )
     state = PartialState()
     state.wait_for_everyone()

src/axolotl/cli/quantize.py

@@ -102,10 +102,12 @@ def do_quantize(
     LOG.info(f"Saving quantized model to: {str(Path(output_dir) / 'quantized')}.")
     model.save_pretrained(
         str(Path(output_dir) / "quantized"),
+        safe_serialization=False,
         progressbar=True,
     )
     tokenizer.save_pretrained(
         str(Path(output_dir) / "quantized"),
+        safe_serialization=False,
         progressbar=True,
         save_jinja_files=cfg.tokenizer_save_jinja_files,
     )
@@ -119,7 +121,7 @@ def do_quantize(
             hub_model_id.rstrip("-")
             + f"-{quantization_config_to_str[type(quantization_config)]}"
         )
-        model.push_to_hub(hub_model_id)
+        model.push_to_hub(hub_model_id, safe_serialization=False)
         tokenizer.push_to_hub(hub_model_id)
         if processor:
             processor.push_to_hub(hub_model_id)

src/axolotl/core/builders/base.py

@@ -216,7 +216,7 @@ class TrainerBuilderBase(abc.ABC):
     def _configure_warmup_and_logging(
         self, total_num_steps: int, training_args_kwargs: dict
     ):
-        warmup_steps: int | float = 0
+        warmup_steps = 0
         warmup_ratio = 0.0
         if self.cfg.warmup_steps is not None:
             warmup_steps = self.cfg.warmup_steps
@@ -230,10 +230,6 @@ class TrainerBuilderBase(abc.ABC):
         else:
             warmup_ratio = 0.03
 
-        # transformers v5
-        if warmup_ratio > 0.0 and warmup_steps == 0:
-            warmup_steps = warmup_ratio
-
         if warmup_steps == 1:
             warmup_steps = 2
 
@@ -246,6 +242,7 @@ class TrainerBuilderBase(abc.ABC):
                 else max(min(int(0.005 * total_num_steps), 10), 1)
             )
 
+        training_args_kwargs["warmup_ratio"] = warmup_ratio
         training_args_kwargs["warmup_steps"] = warmup_steps
 
     def _configure_precision_settings(self, training_args_kwargs: dict):
@@ -409,9 +406,6 @@ class TrainerBuilderBase(abc.ABC):
             if self.cfg.hub_strategy:
                 training_args_kwargs["hub_strategy"] = self.cfg.hub_strategy
 
-            if self.cfg.hub_revision:
-                training_args_kwargs["hub_revision"] = self.cfg.hub_revision
-
     def _configure_save_and_eval_strategy(self, training_args_kwargs: dict):
         # save_strategy and save_steps
         if self.cfg.save_steps:
@@ -536,7 +530,9 @@ class TrainerBuilderBase(abc.ABC):
             "loraplus_lr_ratio",
             "loraplus_lr_embedding",
             "output_dir",
+            "save_safetensors",
             "save_only_model",
+            "include_tokens_per_second",
             "weight_decay",
             "seed",
             "dion_momentum",
@@ -549,7 +545,6 @@ class TrainerBuilderBase(abc.ABC):
 
         arg_map = {
             "dion_learning_rate": "dion_lr",
-            "include_num_input_tokens_seen": "include_tokens_per_second",
         }
         for kwarg, cfg_arg in arg_map.items():
             if hasattr(self.cfg, cfg_arg) and getattr(self.cfg, cfg_arg) is not None:

src/axolotl/core/builders/causal.py

@@ -373,18 +373,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
             # https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
             data_collator_kwargs["pad_to_multiple_of"] = multiple
 
-        if self.cfg.use_eaft:
-            from functools import partial
-
-            from axolotl.monkeypatch.loss.eaft import eaft_loss
-
-            configured_eaft_loss = partial(
-                eaft_loss,
-                alpha=self.cfg.eaft_alpha if self.cfg.eaft_alpha is not None else 1.0,
-                k=self.cfg.eaft_k if self.cfg.eaft_k is not None else 20,
-            )
-            trainer_kwargs["compute_loss_func"] = configured_eaft_loss
-
         trainer_cls = self._get_trainer_cls()
 
         trainer_kwargs, trainer_cls = self.hook_pre_create_trainer(
@@ -449,9 +437,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                 or self.cfg.micro_batch_size > 1
             ):
                 return DataCollatorForSeq2Seq(self.tokenizer, **kwargs)
-            if not (self.cfg.sample_packing and self.cfg.pretrain_multipack_attn) or (
-                self.cfg.micro_batch_size == 1 and is_eval is False
-            ):
+            if not (self.cfg.sample_packing and self.cfg.pretrain_multipack_attn):
                 return None
 
         if self.cfg.model_config_type == "mamba":

src/axolotl/core/builders/rl.py

@@ -52,11 +52,12 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
         trainer_cls = None
         trainer_cls_args = [self.model]
 
-        if self.cfg.rl in {RLType.GRPO, RLType.GDPO}:
+        if self.cfg.rl is RLType.GRPO:
             trainer_cls = GRPOStrategy.get_trainer_class(
                 sequence_parallel=self.cfg.context_parallel_size > 1
             )
             trainer_cls_args.extend(GRPOStrategy.set_trainer_args(self.cfg))
+
             trainer_kwargs.update(GRPOStrategy.set_trainer_kwargs(self.cfg))
 
         elif self.cfg.rl in [RLType.DPO, RLType.IPO]:
@@ -146,8 +147,6 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
 
         elif self.cfg.rl is RLType.KTO:
             training_args_cls = AxolotlKTOConfig
-            # KTOConfig in TRL >= 0.27.0 no longer accepts max_prompt_length
-            blocklist_args_kwargs = ["max_prompt_length"]
 
             training_args_kwargs["desirable_weight"] = (
                 self.cfg.kto_desirable_weight or 1.0
@@ -156,14 +155,10 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
                 self.cfg.kto_undesirable_weight or 1.0
             )
 
-        elif self.cfg.rl in {RLType.GRPO, RLType.GDPO}:
+        elif self.cfg.rl is RLType.GRPO:
             training_args_cls = GRPOStrategy.get_training_args_class()
             training_args_kwargs.update(GRPOStrategy.set_training_args_kwargs(self.cfg))
             blocklist_args_kwargs = GRPOStrategy.get_blocklist_args_kwargs()
-            if self.cfg.rl is RLType.GDPO:
-                training_args_kwargs.setdefault(
-                    "multi_objective_aggregation", "normalize_then_sum"
-                )
 
         elif self.cfg.rl in [RLType.DPO, RLType.IPO]:
             training_args_cls = AxolotlDPOConfig

src/axolotl/core/trainers/base.py

@@ -25,7 +25,7 @@ from torch.utils.data import (
 from transformers import PreTrainedModel, Trainer
 from transformers.trainer import TRAINING_ARGS_NAME
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, has_length, seed_worker
-from transformers.utils import SAFE_WEIGHTS_NAME, is_peft_available
+from transformers.utils import SAFE_WEIGHTS_NAME, WEIGHTS_NAME, is_peft_available
 from trl.trainer.utils import pad_to_length
 from typing_extensions import override
 
@@ -719,13 +719,6 @@ class AxolotlTrainer(
         output_dir = output_dir if output_dir is not None else self.args.output_dir
         os.makedirs(output_dir, exist_ok=True)
         LOG.info(f"Saving model checkpoint to {output_dir}")
-        if state_dict is None:
-            state_dict = self.accelerator.get_state_dict(self.model)
-        if state_dict is not None:
-            state_dict = {
-                k: v.clone() if isinstance(v, torch.Tensor) else v
-                for k, v in state_dict.items()
-            }
         supported_classes = (
             (PreTrainedModel,)
             if not is_peft_available()
@@ -745,38 +738,43 @@ class AxolotlTrainer(
                 ).save_pretrained(
                     output_dir,
                     state_dict=state_dict,
+                    safe_serialization=self.args.save_safetensors,
                 )
             else:
                 LOG.info(
                     "Trainer.model is not a `PreTrainedModel`, only saving its state dict."
                 )
-                safetensors.torch.save_file(
-                    state_dict,
-                    os.path.join(output_dir, SAFE_WEIGHTS_NAME),
-                    metadata={"format": "pt"},
-                )
+                if self.args.save_safetensors:
+                    safetensors.torch.save_file(
+                        state_dict,
+                        os.path.join(output_dir, SAFE_WEIGHTS_NAME),
+                        metadata={"format": "pt"},
+                    )
+                else:
+                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
         else:
             self.model.save_pretrained(
                 output_dir,
                 state_dict=state_dict,
+                safe_serialization=self.args.save_safetensors,
                 is_main_process=self.accelerator.is_main_process,
             )
 
-        if self.processing_class is not None:
-            self.processing_class.save_pretrained(output_dir)
-        elif (
-            self.data_collator is not None
-            and hasattr(self.data_collator, "tokenizer")
-            and self.data_collator.tokenizer is not None
-        ):
-            LOG.info(
-                "Saving Trainer.data_collator.tokenizer by default as Trainer.processing_class is `None`"
-            )
-            save_jinja_files = True
-            if self.axolotl_cfg:
-                save_jinja_files = self.axolotl_cfg.tokenizer_save_jinja_files
-            self.data_collator.tokenizer.save_pretrained(
-                output_dir, save_jinja_files=save_jinja_files
-            )
-        # Good practice: save your training arguments together with the trained model
-        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+            if self.processing_class is not None:
+                self.processing_class.save_pretrained(output_dir)
+            elif (
+                self.data_collator is not None
+                and hasattr(self.data_collator, "tokenizer")
+                and self.data_collator.tokenizer is not None
+            ):
+                LOG.info(
+                    "Saving Trainer.data_collator.tokenizer by default as Trainer.processing_class is `None`"
+                )
+                save_jinja_files = True
+                if self.axolotl_cfg:
+                    save_jinja_files = self.axolotl_cfg.tokenizer_save_jinja_files
+                self.data_collator.tokenizer.save_pretrained(
+                    output_dir, save_jinja_files=save_jinja_files
+                )
+            # Good practice: save your training arguments together with the trained model
+            torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

src/axolotl/core/trainers/grpo/__init__.py

@@ -126,10 +126,8 @@ class GRPOStrategy:
         if trl.use_liger_loss is not None:
             grpo_args_kwargs["use_liger_loss"] = trl.use_liger_loss
 
-        if trl.multi_objective_aggregation is not None:
-            grpo_args_kwargs["multi_objective_aggregation"] = (
-                trl.multi_objective_aggregation
-            )
+        if trl.rollout_func:
+            grpo_args_kwargs["rollout_func"] = cls.get_rollout_func(trl.rollout_func)
 
         return grpo_args_kwargs
 
@@ -151,8 +149,6 @@ class GRPOStrategy:
             trainer_kwargs["reward_processing_classes"] = (
                 cfg.trl.reward_processing_classes
             )
-        if cfg.trl and cfg.trl.rollout_func:
-            trainer_kwargs["rollout_func"] = cls.get_rollout_func(cfg.trl.rollout_func)
 
         return trainer_kwargs
 
@@ -163,12 +159,7 @@ class GRPOStrategy:
 
     @classmethod
     def get_blocklist_args_kwargs(cls) -> list[str]:
-        return [
-            "dataset_num_proc",
-            "max_length",
-            "include_tokens_per_second",
-            "max_prompt_length",
-        ]
+        return ["dataset_num_proc", "max_length", "include_tokens_per_second"]
 
     @classmethod
     def get_reward_func(cls, reward_func_fqn: str) -> RewardFunc:

src/axolotl/core/trainers/trl.py

@@ -1,10 +1,12 @@
 """Module for TRL RL trainers"""
 
-from trl import RewardTrainer
-from trl.experimental.cpo import CPOTrainer
-from trl.experimental.kto import KTOTrainer
-from trl.experimental.orpo import ORPOTrainer
-from trl.experimental.prm import PRMTrainer
+from trl import (
+    CPOTrainer,
+    KTOTrainer,
+    ORPOTrainer,
+    PRMTrainer,
+    RewardTrainer,
+)
 
 from axolotl.core.trainers.mixins import DistributedParallelMixin, RngLoaderMixin
 from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin

src/axolotl/core/training_args.py

@@ -8,11 +8,7 @@ from dataclasses import dataclass, field
 from typing import Optional, Type
 
 from transformers import TrainingArguments
-from trl import RewardConfig
-from trl.experimental.cpo import CPOConfig
-from trl.experimental.kto import KTOConfig
-from trl.experimental.orpo import ORPOConfig
-from trl.experimental.prm import PRMConfig
+from trl import CPOConfig, KTOConfig, ORPOConfig, PRMConfig, RewardConfig
 
 from axolotl.integrations.config import merge_training_args
 

src/axolotl/integrations/cut_cross_entropy/README.md

@@ -19,7 +19,7 @@ python scripts/cutcrossentropy_install.py | sh
 
 - If you are installing from pip
 ```bash
-pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0d4ce4b"
+pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@318b7e2"
 ```
 
 ## Usage
@@ -36,7 +36,6 @@ plugins:
 - cohere
 - cohere2
 - deepseek_v3
-- exaone4
 - gemma
 - gemma2
 - gemma3
@@ -46,16 +45,13 @@ plugins:
 - glm
 - glm4
 - glm4_moe
-- glm4_moe_lite
-- glm46v
 - glm4v
 - glm4v_moe
-- glm_image
 - gpt_oss
 - granite
 - granitemoe
-- granitemoehybrid
 - granitemoeshared
+- granitemoehybrid
 - hunyuan_v1_dense
 - hunyuan_v1_moe
 - internvl
@@ -80,17 +76,16 @@ plugins:
 - phi3
 - phi4_multimodal
 - qwen2
-- qwen2_moe
 - qwen2_vl
+- qwen2_moe
 - qwen2_5_vl
 - qwen3
 - qwen3_moe
-- qwen3_next
 - qwen3_vl
 - qwen3_vl_moe
-- seed_oss
+- qwen3_next
 - smollm3
-- step3p5
+- seed_oss
 - voxtral
 
 ## Citation

src/axolotl/integrations/cut_cross_entropy/__init__.py

@@ -35,7 +35,7 @@ LOG = get_logger(__name__)
 
 _CCE_INSTALL_MESSAGE = (
     "Please install Axolotl's fork of cut_cross_entropy with transformers support using "
-    '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@0d4ce4b"`'
+    '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@318b7e2"`'
 )
 
 
@@ -104,7 +104,7 @@ class CutCrossEntropyPlugin(BasePlugin):
 
     def patch_llama_like(
         self,
-        model_type_to_patch: str,
+        model_type: str,
     ) -> None:
         """
         Generic patch for model architectures with causal lm similar to llama
@@ -112,10 +112,7 @@ class CutCrossEntropyPlugin(BasePlugin):
         from cut_cross_entropy.transformers.patch import PATCH_FNS
 
         def patch_generic(
-            maybe_model,
-            patch_options,
-            remote_model_id: str | None,
-            model_type: str,
+            maybe_model, patch_options, model_type: str, remote_model_id: str | None
         ):
             import cut_cross_entropy.transformers.llama
             from cut_cross_entropy.transformers.llama import cce_forward
@@ -139,13 +136,11 @@ class CutCrossEntropyPlugin(BasePlugin):
                     f"Error: {str(e)}"
                 ) from e
 
-        if model_type_to_patch not in PATCH_FNS:
+        if model_type not in PATCH_FNS:
             LOG.warning_once(
-                "Setting up generic cce patch for model type: %s", model_type_to_patch
+                "Setting up generic cce patch for model type: %s", model_type
             )
             LOG.warning_once(
-                f"Generic Cut Cross Entropy + {model_type_to_patch} support is experimental and may not work as expected."
-            )
-            PATCH_FNS[model_type_to_patch] = partial(
-                patch_generic, model_type=model_type_to_patch
+                f"Generic Cut Cross Entropy + {model_type} support is experimental and may not work as expected."
             )
+            PATCH_FNS[model_type] = partial(patch_generic, model_type=model_type)

src/axolotl/integrations/kernels/README.md

@@ -1,44 +0,0 @@
-# Kernels Integration
-
-MoE (Mixture of Experts) kernels speed up training for MoE layers and reduce VRAM costs. In transformers v5, `batched_mm` and `grouped_mm` were integrated as built-in options via the `experts_implementation` config kwarg:
-
-```python
-class ExpertsInterface(GeneralInterface):
-    _global_mapping = {
-        "batched_mm": batched_mm_experts_forward,
-        "grouped_mm": grouped_mm_experts_forward,
-    }
-```
-
-In our custom integration, we add support for **ScatterMoE**, which is even more efficient and faster than `grouped_mm`.
-
-## Usage
-
-Add the following to your axolotl YAML config:
-
-```yaml
-plugins:
-  - axolotl.integrations.kernels.KernelsPlugin
-
-use_kernels: true
-use_scattermoe: true
-```
-
-**Important:** Setting `experts_implementation` is incompatible with `use_scattermoe`.
-
-## How It Works
-
-The `KernelsPlugin` runs before model loading and:
-
-1. Registers the ScatterMoE kernel from the [`axolotl-ai-co/scattermoe`](https://huggingface.co/axolotl-ai-co/scattermoe) Hub repo.
-2. Patches the model's `SparseMoeBlock` forward method with the optimized ScatterMoE implementation.
-
-This works for any MoE model in transformers that uses a `SparseMoeBlock` class (Mixtral, Qwen2-MoE, OLMoE, etc.).
-
-## Limitations
-
-ScatterMoE uses a softmax -> topk routing, so results may be different for some model arch as baseline (GPT-OSS, GLM_MOE_DSA).
-
-## Note on MegaBlocks
-
-We tested [MegaBlocks](https://huggingface.co/kernels-community/megablocks) but were unable to ensure numerical accuracy, so we did not integrate it. It was also incompatible with many newer model architectures in transformers.

src/axolotl/integrations/kernels/__init__.py

@@ -1,7 +0,0 @@
-from .args import KernelsArgs
-from .plugin import KernelsPlugin
-
-__all__ = [
-    "KernelsArgs",
-    "KernelsPlugin",
-]

src/axolotl/integrations/kernels/args.py

@@ -1,35 +0,0 @@
-from pydantic import BaseModel, model_validator
-
-from axolotl.utils.logging import get_logger
-
-LOG = get_logger(__name__)
-
-
-class KernelsArgs(BaseModel):
-    use_scattermoe: bool | None = True
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_use_kernels(cls, data):
-        if data.get("use_kernels") is not True:
-            LOG.warning(
-                "`use_kernels` must be set to True to use this. Automatically setting it to True."
-            )
-            data["use_kernels"] = True
-
-        return data
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_experts_implementation(cls, data):
-        experts_implementation = data.get("experts_implementation")
-        if experts_implementation is None:
-            # transformers may default to batched_mm when unset
-            data["experts_implementation"] = "eager"
-        elif experts_implementation != "eager":
-            LOG.warning(
-                "`experts_implementation` must be set to 'eager' to use this. Automatically setting it to 'eager'."
-            )
-            data["experts_implementation"] = "eager"
-
-        return data

src/axolotl/integrations/kernels/plugin.py

@@ -1,61 +0,0 @@
-from kernels import (
-    LayerRepository,
-    Mode,
-    register_kernel_mapping,
-    replace_kernel_forward_from_hub,
-)
-
-from axolotl.integrations.base import BasePlugin
-from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
-
-
-class KernelsPlugin(BasePlugin):
-    def get_input_args(self):
-        return "axolotl.integrations.kernels.KernelsArgs"
-
-    def pre_model_load(self, cfg):
-        if cfg.use_scattermoe:
-            self._register_kernels()
-            self._kernelize_model(cfg.model_config_type)
-
-    def _register_kernels(self):
-        register_kernel_mapping(
-            {
-                "HFScatterMoEParallelExperts": {
-                    "cuda": {
-                        Mode.TRAINING: LayerRepository(
-                            repo_id="axolotl-ai-co/scattermoe",
-                            layer_name="HFScatterMoEGatedMLP",
-                        ),
-                        Mode.INFERENCE: LayerRepository(
-                            repo_id="axolotl-ai-co/scattermoe",
-                            layer_name="HFScatterMoEGatedMLP",
-                        ),
-                    },
-                }
-            }
-        )
-
-    def _kernelize_model(self, model_type: str):
-        if model_type == "olmoe":
-            from transformers.models.olmoe.modeling_olmoe import OlmoeSparseMoeBlock
-
-            replace_kernel_forward_from_hub(
-                OlmoeSparseMoeBlock, "HFScatterMoEParallelExperts"
-            )
-        else:
-            try:
-                model_moe_cls = get_model_moe_block(model_type)
-                replace_kernel_forward_from_hub(
-                    model_moe_cls, "HFScatterMoEParallelExperts"
-                )
-            except Exception as err:
-                raise ValueError(f"Unsupported model type: {model_type}") from err
-
-
-def get_model_moe_block(model_type: str):
-    module_path = f"transformers.models.{model_type}.modeling_{model_type}"
-    model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
-    module = __import__(module_path, fromlist=[f"{model_cls_prefix}SparseMoeBlock"])
-    model_cls = getattr(module, f"{model_cls_prefix}SparseMoeBlock")
-    return model_cls

src/axolotl/integrations/llm_compressor/utils.py

@@ -12,6 +12,7 @@ def save_compressed_model(
     model: PreTrainedModel,
     output_dir: Union[str, bytes],
     trainer: Trainer,
+    safe_serialization: bool = False,
     save_compressed: bool = False,
 ) -> None:
     """
@@ -21,6 +22,7 @@ def save_compressed_model(
         model (PreTrainedModel): The model to be saved.
         output_dir (str or bytes): Path where the model files will be written.
         trainer (Trainer): Hugging Face Trainer for process synchronization.
+        safe_serialization (bool): Use safe serialization if True.
         save_compressed (bool): Write compressed tensors if True.
     """
     trainer.accelerator.wait_for_everyone()
@@ -32,6 +34,7 @@ def save_compressed_model(
     modify_save_pretrained(model)
     model.save_pretrained(
         output_dir,
+        safe_serialization=safe_serialization,
         save_compressed=save_compressed,
         skip_sparsity_compression_stats=not save_compressed,
     )

src/axolotl/integrations/lm_eval/README.md

@@ -6,12 +6,6 @@ See https://github.com/EleutherAI/lm-evaluation-harness
 
 ## Usage
 
-There are two ways to use the LM Eval integration:
-
-### 1. Post-Training Evaluation
-
-When training with the plugin enabled, evaluation runs automatically after training completes:
-
 ```yaml
 plugins:
   - axolotl.integrations.lm_eval.LMEvalPlugin
@@ -22,50 +16,9 @@ lm_eval_tasks:
   - arc_easy
 
 lm_eval_batch_size: # Batch size for evaluation
-
-# Directory to save evaluation results.
-# The final model is loaded from this directory
-# unless specified otherwise (see below)
-output_dir:
+output_dir: # Directory to save evaluation results
 ```
 
-Run training as usual:
-```bash
-axolotl train config.yml
-```
-
-### 2. Standalone CLI Evaluation
-
-Evaluate any model directly without training:
-
-```yaml
-lm_eval_model: meta-llama/Llama-2-7b-hf
-
-plugins:
-  - axolotl.integrations.lm_eval.LMEvalPlugin
-
-lm_eval_tasks:
-  - gsm8k
-  - hellaswag
-  - arc_easy
-
-lm_eval_batch_size: 8
-output_dir: ./outputs
-```
-
-Run evaluation:
-```bash
-axolotl lm-eval config.yml
-```
-
-## Model Selection Priority
-
-The model to evaluate is selected in the following priority order:
-
-1. **`lm_eval_model`** - Explicit model path or HuggingFace repo (highest priority)
-2. **`hub_model_id`** - Trained model pushed to HuggingFace Hub
-3. **`output_dir`** - Local checkpoint directory containing trained model weights
-
 ## Citation
 
 ```bib

src/axolotl/integrations/lm_eval/__init__.py

@@ -5,7 +5,7 @@ Module for the Plugin for LM Eval Harness
 import subprocess  # nosec
 
 from axolotl.integrations.base import BasePlugin
-from axolotl.integrations.lm_eval.cli import build_lm_eval_command, get_model_path
+from axolotl.integrations.lm_eval.cli import build_lm_eval_command
 
 from .args import LMEvalArgs as LMEvalArgs
 
@@ -29,7 +29,7 @@ class LMEvalPlugin(BasePlugin):
                 wandb_project=cfg.wandb_project,
                 wandb_entity=cfg.wandb_entity,
                 wandb_name=cfg.wandb_name,
-                model=get_model_path(cfg),
+                model=cfg.lm_eval_model or cfg.hub_model_id,
             ):
                 subprocess.run(  # nosec
                     lm_eval_args,

src/axolotl/integrations/lm_eval/cli.py

@@ -13,21 +13,6 @@ import yaml
 from axolotl.utils.dict import DictDefault
 
 
-def get_model_path(cfg: DictDefault) -> str | None:
-    """
-    Determine which model path to use for evaluation.
-
-    Priority order (highest to lowest):
-    1. lm_eval_model - Explicit model path override
-    2. hub_model_id - Model pushed to HuggingFace Hub
-    3. None - Falls back to output_dir in build_lm_eval_command
-
-    Returns:
-        Model path string or None to use output_dir fallback
-    """
-    return cfg.lm_eval_model or cfg.hub_model_id or None
-
-
 def build_lm_eval_command(
     tasks: list[str],
     bfloat16=True,
@@ -123,7 +108,7 @@ def lm_eval(config: str, cloud: Optional[str] = None):
             wandb_project=cfg.wandb_project,
             wandb_entity=cfg.wandb_entity,
             wandb_name=cfg.wandb_name,
-            model=get_model_path(cfg),
+            model=cfg.lm_eval_model or cfg.hub_model_id,
             revision=cfg.revision,
             apply_chat_template=cfg.apply_chat_template,
             fewshot_as_multiturn=cfg.fewshot_as_multiturn,

src/axolotl/kernels/lora.py

@@ -15,7 +15,7 @@ from torch import nn
 from torch.distributed.tensor import DTensor
 
 from .geglu import geglu_backward, geglu_forward
-from .quantize import dequantize_weight
+from .quantize import dequantize
 from .swiglu import swiglu_backward, swiglu_forward
 from .utils import torch_amp_custom_bwd, torch_amp_custom_fwd
 
@@ -46,12 +46,6 @@ def get_lora_parameters(
     W = base_layer.weight
     b = base_layer.bias
 
-    # Unwrap DTensor if FSDP2 left the weight wrapped -- DTensor does not proxy
-    # attribute access to the underlying tensor subclass, so torchao methods like
-    # .dequantize() or .get_original_weight() would not be visible.
-    if isinstance(W, DTensor):
-        W = W.full_tensor()
-
     if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
         quant_state = getattr(W, "quant_state", None)
         return W, b, quant_state, None, None, None
@@ -92,7 +86,6 @@ def matmul_lora(
     B: torch.Tensor | None,
     s: float | None,
     out: torch.Tensor | None = None,
-    transpose: bool = True,
 ) -> torch.Tensor:
     """
     Efficient fused matmul + LoRA computation.
@@ -105,15 +98,12 @@ def matmul_lora(
         B: LoRA B matrix [out_features, rank]
         s: LoRA scaling factor
         out: Optional output tensor for inplace operations
-        transpose: If True (default), transpose W before matmul (forward path).
-            Set to False for backward paths where W is already in the correct layout.
 
     Returns:
         Result of X @ W + X @ A @ B
     """
     dtype = X.dtype
-    is_quantized = W_quant is not None or type(W) is not torch.Tensor
-    W = dequantize_weight(W, W_quant, transpose=transpose)
+    W = dequantize(W.t(), W_quant)
 
     reshape = False
     if X.dim() == 3:
@@ -122,7 +112,7 @@ def matmul_lora(
         reshape = True
 
     out = torch.matmul(X, W, out=out)
-    if is_quantized:
+    if W_quant is not None:
         del W
 
     if A is not None:
@@ -302,16 +292,15 @@ class LoRA_MLP(torch.autograd.Function):
         up = up.view(-1, up.shape[-1])
         dtype = X.dtype
 
-        # Down projection (backward: no transpose needed, W is already [out, in])
+        # Down projection
         grad_down = matmul_lora(
             grad_output,
-            down_weight,
+            down_weight.t(),
             None,
             down_quant,
             down_B,
             down_A,
             down_scale,
-            transpose=False,
         )
 
         # Activation backward
@@ -343,7 +332,7 @@ class LoRA_MLP(torch.autograd.Function):
 
         if dX is not None:
             # Up projection gradients
-            up_weight = dequantize_weight(up_weight, up_quant, transpose=True)
+            up_weight = dequantize(up_weight.t(), up_quant)
             if ctx.inplace:
                 dX = torch.matmul(grad_up, up_weight.t(), out=X)
             else:
@@ -355,7 +344,7 @@ class LoRA_MLP(torch.autograd.Function):
                 dX += grad_up @ up_B.to(dtype).t() @ (up_scale * up_A.to(dtype).t())
 
             # Gate projection gradients
-            gate_weight = dequantize_weight(gate_weight, gate_quant)
+            gate_weight = dequantize(gate_weight, gate_quant)
             dX += grad_gate @ gate_weight
             del gate_weight
 
@@ -642,7 +631,7 @@ class LoRA_QKV(torch.autograd.Function):
         out_buffer = X if ctx.inplace else None
 
         # Q path
-        q_weight_t = dequantize_weight(q_weight, q_quant)
+        q_weight_t = dequantize(q_weight, q_quant)
         grad_X = torch.mm(q_grad, q_weight_t, out=out_buffer)
         del q_weight
         del q_weight_t
@@ -650,7 +639,7 @@ class LoRA_QKV(torch.autograd.Function):
             grad_X.addmm_(q_grad, torch.mm(B_q_scaled, A_q_scaled))
 
         # K path
-        k_weight_t = dequantize_weight(k_weight, k_quant)
+        k_weight_t = dequantize(k_weight, k_quant)
         grad_X.addmm_(k_grad, k_weight_t)
         del k_weight
         del k_weight_t
@@ -658,7 +647,7 @@ class LoRA_QKV(torch.autograd.Function):
             grad_X.addmm_(k_grad, torch.mm(B_k_scaled, A_k_scaled))
 
         # V path
-        v_weight_t = dequantize_weight(v_weight, v_quant)
+        v_weight_t = dequantize(v_weight, v_quant)
         grad_X.addmm_(v_grad, v_weight_t)
         del v_weight
         del v_weight_t
@@ -821,7 +810,7 @@ class LoRA_O(torch.autograd.Function):
         d_B = s * A @ dY_X
 
         # Get derivative for dX
-        W = dequantize_weight(W, W_quant, transpose=True)
+        W = dequantize(W.t(), W_quant)
         dX = dY @ W.t()
         del W
 

src/axolotl/kernels/quantize.py

@@ -146,43 +146,3 @@ def dequantize(
     # Handle transposed data
     is_transposed: bool = W.shape[0] == 1
     return out.t() if is_transposed else out
-
-
-def dequantize_weight(
-    W: torch.Tensor,
-    quant_state: QuantState | list | None = None,
-    transpose: bool = False,
-) -> torch.Tensor:
-    """Unified dequantization for both torchao and bnb quantized weights.
-
-    For torchao tensor subclasses (AffineQuantizedTensor, NF4Tensor), dequantizes
-    using the appropriate instance method. For bnb Params4bit, delegates to the
-    optimized CUDA kernel in ``dequantize``.
-
-    Args:
-        W: Quantized weight tensor ``[out_features, in_features]``.
-        quant_state: bnb ``QuantState`` (None for torchao / unquantized).
-        transpose: If True, return ``[in_features, out_features]``.
-
-    Returns:
-        Dequantized float tensor, optionally transposed.
-    """
-    # torchao path: tensor subclass with embedded quantization state
-    if quant_state is None and type(W) is not torch.Tensor:
-        result = None
-        # NF4Tensor (check first — NF4Tensor.dequantize is a static method)
-        if hasattr(W, "get_original_weight"):
-            result = W.get_original_weight()
-        else:
-            # AffineQuantizedTensor (INT4, etc.)
-            try:
-                result = W.dequantize()
-            except (TypeError, RuntimeError):
-                pass
-        if result is not None:
-            return result.t() if transpose else result
-
-    # bnb path: transpose input before the CUDA kernel (existing convention)
-    if transpose:
-        return dequantize(W.t(), quant_state)
-    return dequantize(W, quant_state)

src/axolotl/loaders/adapter.py

@@ -23,7 +23,6 @@ from axolotl.loaders.utils import get_linear_embedding_layers
 from axolotl.telemetry.errors import send_errors
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.logging import get_logger
-from axolotl.utils.schemas.enums import TorchAOQuantDType
 
 LOG = get_logger(__name__)
 
@@ -135,13 +134,11 @@ def load_lora(
 
     rank = int(os.environ.get("LOCAL_RANK", 0))
 
-    is_torchao = cfg.peft and cfg.peft.backend == "torchao"
     if (
         cfg.fsdp_config
         and cfg.adapter
         and cfg.fsdp_config.cpu_ram_efficient_loading
         and rank != 0
-        and not is_torchao
     ):
         setup_quantized_meta_for_peft(model)
 
@@ -149,15 +146,6 @@ def load_lora(
     if cfg.peft_autocast_adapter_dtype is not None:
         model_kwargs["autocast_adapter_dtype"] = cfg.peft_autocast_adapter_dtype
 
-    # Patch PEFT's torchao dispatch before any model creation/loading.
-    # Must happen before both get_peft_model and PeftModel.from_pretrained,
-    # as both trigger LoRA layer dispatch that would fail for INT4/NF4 weights.
-    # INT8 is natively supported by PEFT's TorchaoLoraLinear, so skip the patch.
-    if is_torchao and cfg.peft.weight_dtype != TorchAOQuantDType.int8:
-        from axolotl.monkeypatch.peft.utils import patch_peft_torchao_dispatch
-
-        patch_peft_torchao_dispatch()
-
     if cfg.lora_model_dir:
         LOG.debug("Loading pretrained PEFT - LoRA")
         if cfg.lora_on_cpu:
@@ -184,7 +172,6 @@ def load_lora(
         and cfg.adapter
         and cfg.fsdp_config.cpu_ram_efficient_loading
         and rank != 0
-        and not is_torchao
     ):
         setup_quantized_peft_meta_for_training(model)
 

src/axolotl/loaders/model.py

@@ -26,6 +26,7 @@ from torch.distributed import DeviceMesh
 from transformers import (
     AutoModelForCausalLM,
     AutoModelForImageTextToText,
+    AutoModelForVision2Seq,
     AwqConfig,
     BitsAndBytesConfig,
     GPTQConfig,
@@ -158,15 +159,6 @@ class ModelLoader:
         """Property that determines if FSDP with QLoRA is enabled."""
         return self.is_fsdp_enabled and self.cfg.adapter == "qlora"
 
-    @property
-    def is_torchao_qlora(self):
-        """Property that determines if torchao backend is used for QLoRA."""
-        return (
-            self.cfg.adapter == "qlora"
-            and self.cfg.peft
-            and self.cfg.peft.backend == "torchao"
-        )
-
     @send_errors
     def load(self) -> tuple[PreTrainedModel | PeftModelForCausalLM, PeftConfig | None]:
         """Load and prepare the model with all configurations and patches.
@@ -234,7 +226,6 @@ class ModelLoader:
         ):
             self.model = self.model.merge_and_unload()
 
-        self._configure_experts_implementation()
         self._apply_activation_checkpointing()
         self._resize_token_embeddings()
         self._adjust_model_config()
@@ -242,10 +233,6 @@ class ModelLoader:
         self._configure_qat()
         log_gpu_memory_usage(LOG, "Memory usage after model load", 0)
 
-    def _configure_experts_implementation(self):
-        if self.cfg.experts_implementation is not None:
-            self.model.set_experts_implementation(self.cfg.experts_implementation)
-
     def _apply_activation_checkpointing(self):
         if self.cfg.activation_offloading is True:
             from axolotl.core.trainers.mixins.activation_checkpointing import (
@@ -347,12 +334,7 @@ class ModelLoader:
             # LlamaRMSNorm layers are in fp32 after kbit_training or full finetune, so
             # we need to convert them back to fp16/bf16 for flash-attn compatibility.
             (
-                (
-                    needs_fa2_dtype
-                    or self.cfg.flash_attention
-                    or self.cfg.flex_attention
-                    or self.cfg.sage_attention
-                )
+                (needs_fa2_dtype or self.cfg.flash_attention or self.cfg.flex_attention)
                 and not self.is_qlora_and_fsdp_enabled
             )
             or (
@@ -452,7 +434,7 @@ class ModelLoader:
         """
         if self.cfg.is_multimodal:
             self.auto_model_loader = MULTIMODAL_AUTO_MODEL_MAPPING.get(
-                self.model_config.model_type, AutoModelForImageTextToText
+                self.model_config.model_type, AutoModelForVision2Seq
             )
             if isinstance(self.auto_model_loader, str):
                 self.auto_model_loader = AutoModelForImageTextToText
@@ -494,15 +476,13 @@ class ModelLoader:
             max_memory = None
 
         self.model_kwargs["torch_dtype"] = self.cfg.torch_dtype
-        self.model_kwargs["dtype"] = self.cfg.torch_dtype
 
         is_ds_zero3 = is_deepspeed_zero3_enabled()
 
         # FSDP requires control over device placement, so don't set device_map when FSDP is enabled
         if self.is_fsdp_enabled:
-            # For QLoRA + FSDP with bnb, we still need to set device_map for proper initialization
-            # torchao tensors work natively with FSDP2, no device_map override needed
-            if self.is_qlora_and_fsdp_enabled and not self.is_torchao_qlora:
+            # For QLoRA + FSDP, we still need to set device_map to "auto" for proper initialization
+            if self.is_qlora_and_fsdp_enabled:
                 self.model_kwargs["device_map"] = {
                     "": int(os.environ.get("LOCAL_RANK", 0))
                 }
@@ -571,44 +551,6 @@ class ModelLoader:
                 self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
                     **self.model_config.quantization_config
                 )
-        elif (
-            self.cfg.adapter == "qlora"
-            and self.cfg.peft
-            and self.cfg.peft.backend == "torchao"
-            and not self.cfg.merge_lora
-        ):
-            from transformers import TorchAoConfig
-
-            from axolotl.utils.schemas.enums import TorchAOQuantDType
-
-            weight_dtype = self.cfg.peft.weight_dtype
-            if weight_dtype == TorchAOQuantDType.int4:
-                group_size = self.cfg.peft.group_size or 128
-                self.model_kwargs["quantization_config"] = TorchAoConfig(
-                    quant_type="int4_weight_only",
-                    group_size=group_size,
-                )
-            elif weight_dtype == TorchAOQuantDType.int8:
-                group_size = self.cfg.peft.group_size or 128
-                self.model_kwargs["quantization_config"] = TorchAoConfig(
-                    quant_type="int8_weight_only",
-                    group_size=group_size,
-                )
-            elif weight_dtype == TorchAOQuantDType.nf4:
-                from torchao.dtypes._nf4tensor_api import NF4WeightOnlyConfig
-
-                block_size = self.cfg.peft.group_size or 64
-                self.model_kwargs["quantization_config"] = TorchAoConfig(
-                    quant_type=NF4WeightOnlyConfig(
-                        block_size=block_size,
-                        scaler_block_size=256,
-                    ),
-                )
-            else:
-                raise ValueError(
-                    f"Unsupported torchao weight_dtype for QLoRA: {weight_dtype}. "
-                    "Supported: int4, int8, nf4"
-                )
         elif self.cfg.adapter == "qlora" and self.cfg.load_in_4bit:
             bnb_config = {
                 "load_in_4bit": True,
@@ -665,10 +607,6 @@ class ModelLoader:
         elif self.cfg.sdp_attention:
             self.model_kwargs["attn_implementation"] = "sdpa"
             self.model_config._attn_implementation = "sdpa"
-        elif self.cfg.sage_attention:
-            # sets FA2 attention to re-use same internal handling like masking
-            self.model_kwargs["attn_implementation"] = "flash_attention_2"
-            self.model_config._attn_implementation = "flash_attention_2"
         elif self.cfg.eager_attention:
             self.model_kwargs["attn_implementation"] = "eager"
             self.model_config._attn_implementation = "eager"
@@ -732,7 +670,7 @@ class ModelLoader:
         Uses the selected loader when provided; otherwise falls back to the auto loader.
         """
         loader = model_loader_class or self.auto_model_loader
-        if loader in [AutoModelForCausalLM, AutoModelForImageTextToText]:
+        if loader in [AutoModelForCausalLM, AutoModelForVision2Seq]:
             model = loader.from_config(
                 config=self.model_config,
                 trust_remote_code=self.cfg.trust_remote_code or False,
@@ -850,7 +788,6 @@ class ModelLoader:
                 # Use auto model loader (handles gptq and default cases)
                 model_loader_class = self.auto_model_loader
 
-            self.model_kwargs["dtype"] = self.model_kwargs["torch_dtype"]
             if self.cfg.reinit_weights:
                 self.model = self._load_model_from_config(model_loader_class)
             else:
@@ -908,10 +845,6 @@ class ModelLoader:
             # Make sure everything is in the same dtype
             skip_prepare_model_for_kbit_training = True
 
-        # torchao quantized models don't use Params4bit and don't need kbit preparation
-        if self.is_torchao_qlora:
-            skip_prepare_model_for_kbit_training = True
-
         if (
             not skip_prepare_model_for_kbit_training
             and self.cfg.adapter in ["lora", "qlora"]

src/axolotl/loaders/patch_manager.py

@@ -96,7 +96,6 @@ class PatchManager:
         # self._apply_flex_attention_patches()
         self._apply_flash_attention_patches()
         self._apply_chunked_cross_entropy_patch()
-        self._apply_sageattn_patches()
         self._apply_fsdp_patches()
         self._apply_adapter_patches()
         self._apply_model_specific_patches()
@@ -154,9 +153,12 @@ class PatchManager:
             from axolotl.monkeypatch.loss.chunked import patch_chunked_ce_loss_fn
 
             if self.cfg.chunked_cross_entropy_num_chunks:
-                patch_chunked_ce_loss_fn(self.cfg.chunked_cross_entropy_num_chunks)
+                patch_chunked_ce_loss_fn(
+                    self.cfg.chunked_cross_entropy_num_chunks,
+                    use_dft=self.cfg.use_dynamic_finetuning,
+                )
             else:
-                patch_chunked_ce_loss_fn()
+                patch_chunked_ce_loss_fn(use_dft=self.cfg.use_dynamic_finetuning)
 
     def _apply_fsdp_patches(self):
         """Apply patches for FSDP configurations."""
@@ -202,13 +204,6 @@ class PatchManager:
             flex_attn_compile_kwargs = self.cfg.flex_attn_compile_kwargs or {}
             patch_flex_wrapper(**flex_attn_compile_kwargs)
 
-    def _apply_sageattn_patches(self):
-        """Apply patches for SageAttention."""
-        if self.cfg.sage_attention:
-            from axolotl.monkeypatch.attention.sage_attn import patch_sageattn
-
-            patch_sageattn()
-
     def _apply_model_specific_patches(self):
         """Apply patches specific to model architectures."""
         if (
@@ -228,6 +223,13 @@ class PatchManager:
 
             patch_qwen3_next_modeling_packing()
 
+        if self.cfg.model_config_type == "mistral3" and self.cfg.processor_type:
+            from axolotl.monkeypatch.models.mistral3.mistral_common_tokenizer import (
+                apply_mistral_tokenizer_image_patch,
+            )
+
+            apply_mistral_tokenizer_image_patch()
+
         if self.cfg.model_config_type == "kimi_linear":
             from axolotl.monkeypatch.models.kimi_linear.patch_kimi_linear import (
                 patch_kimi_model,
@@ -348,12 +350,10 @@ class PatchManager:
 
     def _apply_fsdp2_bnb_patches(self):
         """Apply FSDP2 BNB patches."""
-        is_torchao = self.cfg.peft and self.cfg.peft.backend == "torchao"
         if (
             self.cfg.fsdp_config
             and str(self.cfg.fsdp_version) == "2"
             and self.cfg.adapter == "qlora"
-            and not is_torchao
         ):
             from axolotl.monkeypatch.fsdp2_qlora import (
                 apply_init_sharded_param_patch,

src/axolotl/loaders/processor.py

@@ -31,7 +31,7 @@ def load_processor(cfg: DictDefault, tokenizer: PreTrainedTokenizerBase):
 
             from axolotl.utils.mistral import HFMistralTokenizer
 
-            tokenization_mistral_common.MistralCommonBackend = HFMistralTokenizer
+            tokenization_mistral_common.MistralCommonTokenizer = HFMistralTokenizer
 
         _patch_mistralcommontokenizer()
 

src/axolotl/loaders/utils.py

@@ -5,7 +5,6 @@ from typing import Type
 
 import addict
 import torch
-import transformers
 from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
 
 from axolotl.utils.dict import DictDefault
@@ -154,9 +153,6 @@ def load_model_config(cfg: DictDefault) -> PretrainedConfig | addict.Dict:
     This function determines the appropriate model config source, loads it, applies any
     necessary overrides, and validates it for compatibility with the `axolotl` config.
 
-    If `cfg.cls_model_config` is set, a custom config class from transformers will be
-    used instead of `AutoConfig` (e.g., 'LlamaConfig', 'MistralConfig').
-
     Args:
         cfg: Dictionary mapping `axolotl` config keys to values.
 
@@ -178,13 +174,8 @@ def load_model_config(cfg: DictDefault) -> PretrainedConfig | addict.Dict:
     if cfg.num_labels:
         # num_labels is used to initialize classifier models
         config_kwargs["num_labels"] = cfg.num_labels
-
-    config_cls = AutoConfig
-    if cfg.cls_model_config:
-        config_cls = getattr(transformers, cfg.cls_model_config)
-
     try:
-        model_config = config_cls.from_pretrained(
+        model_config = AutoConfig.from_pretrained(
             model_config_name,
             trust_remote_code=trust_remote_code,
             **config_kwargs,

src/axolotl/models/mamba/modeling_mamba.py

@@ -111,6 +111,7 @@ class MambaLMHeadModel(nn.Module, GenerationMixin):
         self,
         save_directory: Union[str, os.PathLike],
         state_dict: Optional[dict] = None,
+        safe_serialization: Optional[bool] = None,
     ):
         if state_dict is None:
             state_dict = self.state_dict()

src/axolotl/monkeypatch/attention/sage_attn.py

@@ -1,211 +0,0 @@
-"""
-Monkeypatch for SageAttention for use with transformers.
-
-https://github.com/thu-ml/SageAttention/
-"""
-
-import torch
-from transformers.integrations.sdpa_attention import repeat_kv
-
-from axolotl.utils.logging import get_logger
-
-LOG = get_logger(__name__)
-
-sageattn = None  # pylint: disable=invalid-name
-sageattn_varlen = None  # pylint: disable=invalid-name
-
-
-def _is_sageattn_available():
-    """Determine if SageAttention is available"""
-    try:
-        import sageattention  # noqa: F401 # pylint: disable=unused-import
-
-        return True
-    except ImportError:
-        return False
-
-
-if _is_sageattn_available():
-    # import sageattn here if available
-    from sageattention import sageattn, sageattn_varlen
-
-
-def _check_sageattn_imported():
-    """Check if SageAttention is imported. Raises an ImportError if not."""
-    if sageattn is None:
-        raise ImportError(
-            "SageAttention is not installed. Please install it from source: "
-            "`pip install git+https://github.com/thu-ml/SageAttention.git@1718ddc06dbc694bcf3c6b49ac28c1921aa2d8bd`"
-        )
-
-
-def sage_attention_forward(
-    module: torch.nn.Module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: torch.Tensor | None = None,
-    dropout: float = 0.0,
-    scaling: float | None = None,
-    is_causal: bool | None = None,
-    **kwargs,
-) -> tuple[torch.Tensor, None]:
-    """
-    Forward pass for SageAttention compatible with transformers attention interfaces.
-
-    https://github.com/thu-ml/SageAttention/
-    """
-
-    _check_sageattn_imported()
-
-    if kwargs.get("output_attentions", False) or kwargs.get("head_mask") is not None:
-        raise NotImplementedError(
-            "SageAttention does not support `output_attentions=True` or `head_mask`."
-        )
-
-    # The base sageattn API does not support dropout.
-    if dropout > 0.0:
-        raise NotImplementedError("SageAttention does not support dropout.")
-
-    # Handle Grouped-Query Attention (GQA) and Multi-Query Attention (MQA)
-    if hasattr(module, "num_key_value_groups"):
-        key = repeat_kv(key, module.num_key_value_groups)
-        value = repeat_kv(value, module.num_key_value_groups)
-
-    # Calculate is_causal following transformers
-    assert is_causal is not False, "is_causal must be True or None"
-    is_causal = True
-
-    position_ids = kwargs.get("position_ids", None)
-    query_length = query.shape[2]
-
-    cu_seqlens_q = kwargs.get("cu_seqlens_q", None)
-    cu_seqlens_k = kwargs.get("cu_seqlens_k", None)
-    max_length_q = kwargs.get("max_length_q", None)
-    max_length_k = kwargs.get("max_length_k", None)
-
-    # Sample packing uses position_ids, so we check for it first
-    if position_ids is not None and (
-        max_length_q is not None
-        or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all())
-    ):
-        # transpose inputs to NHD layout for use with FA2 utils
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
-
-        batch_size = query.size(0)
-
-        from transformers.modeling_flash_attention_utils import (
-            prepare_fa2_from_position_ids,
-        )
-
-        if cu_seqlens_q is None or cu_seqlens_k is None:
-            query, key, value, indices_q, cu_seq_lens, max_seq_lens = (
-                prepare_fa2_from_position_ids(query, key, value, position_ids)
-            )
-
-            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
-            max_length_q, max_length_k = max_seq_lens
-
-        else:
-            query = query.reshape(-1, query.size(-2), query.size(-1))
-            key = key.reshape(-1, key.size(-2), key.size(-1))
-            value = value.reshape(-1, value.size(-2), value.size(-1))
-
-        attn_output_unpad = sageattn_varlen(
-            q=query,
-            k=key,
-            v=value,
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            max_seqlen_q=max_length_q,
-            max_seqlen_k=max_length_k,
-            is_causal=is_causal,
-            sm_scale=scaling,
-            smooth_k=False,  # reduces loss 0 / nan grad norms
-            tensor_layout="NHD",
-        )
-
-        attn_output = attn_output_unpad.view(
-            batch_size, -1, attn_output_unpad.size(-2), attn_output_unpad.size(-1)
-        )
-
-    elif attention_mask is not None:
-        # NOTE: When used without `pad_to_sequence_len`, the loss becomes unstable after a few steps.
-
-        assert attention_mask.ndim == 2, "Attention mask must be 2D"
-
-        from transformers.modeling_flash_attention_utils import (
-            _upad_input,
-        )
-
-        # transpose inputs to NHD layout for use with FA2 utils
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
-
-        batch_size = query.shape[0]
-
-        query, key, value, indices_q, cu_seq_lens, max_seq_lens = _upad_input(
-            query, key, value, attention_mask, query_length
-        )
-        cu_seqlens_q, cu_seqlens_k = cu_seq_lens
-        max_seqlen_q, max_seqlen_k = max_seq_lens
-
-        attn_output_unpad = sageattn_varlen(
-            q=query,
-            k=key,
-            v=value,
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            max_seqlen_q=max_seqlen_q,
-            max_seqlen_k=max_seqlen_k,
-            is_causal=is_causal,
-            sm_scale=scaling,
-            tensor_layout="NHD",
-        )
-
-        from flash_attn.bert_padding import pad_input
-
-        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
-    else:
-        # Use standard sageattn
-        # The input layout for transformers models is (batch_size, num_heads, seq_len, head_dim),
-        # which corresponds to SageAttention's "HND" layout.
-        attn_output = sageattn(
-            q=query,
-            k=key,
-            v=value,
-            tensor_layout="HND",
-            is_causal=is_causal,
-            sm_scale=scaling,
-        )
-
-        # SageAttention with "HND" returns (batch, heads, seq_len, head_dim)
-        # Transformers expects (batch, seq_len, heads, head_dim) for the output
-        # So we need to transpose dimensions 1 and 2
-        attn_output = attn_output.transpose(1, 2).contiguous()
-
-    return attn_output, None
-
-
-def patch_sageattn():
-    """Patch SageAttention for use with transformers."""
-
-    _check_sageattn_imported()
-
-    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
-
-    # Replace flash attention with sage attention
-    ALL_ATTENTION_FUNCTIONS.register("flash_attention_2", sage_attention_forward)
-
-    # Note: New method after transformers refactor to use ALL_MASK_ATTENTION_FUNCTIONS
-    # Register sage_attention with the global attention interface
-    # ALL_ATTENTION_FUNCTIONS.register("sage_attention", sage_attention_forward)
-
-    # from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS, flash_attention_mask
-
-    # ALL_MASK_ATTENTION_FUNCTIONS.register("sage_attention", flash_attention_mask)
-
-    LOG.info("SageAttention patched successfully")

src/axolotl/monkeypatch/lora_kernels.py

@@ -169,8 +169,7 @@ def get_attention_cls_from_config(cfg: DictDefault) -> Type[nn.Module]:
         return attention_cls
     except (ImportError, AttributeError) as e:
         raise ValueError(
-            f"Axolotl could not import attention class for model_type: {model_type}. "
-            "Please raise an Issue and turn off lora kernels to continue training. "
+            f"Could not import attention class for model_type: {model_type}. "
             f"Error: {str(e)}"
         ) from e
 

src/axolotl/monkeypatch/loss/chunked.py

@@ -16,10 +16,16 @@ class CEWithChunkedOutputLoss(torch.nn.Module):
     For more details, please refer to: https://github.com/pytorch/torchtune/pull/1390
     """
 
-    def __init__(self, num_output_chunks: int = 8, ignore_index: int = -100):
+    def __init__(
+        self,
+        num_output_chunks: int = 8,
+        ignore_index: int = -100,
+        use_dft: bool = False,
+    ):
         super().__init__()
         self.num_output_chunks = num_output_chunks
         self.ignore_index = ignore_index
+        self.use_dft = use_dft
 
     def compute_cross_entropy(
         self,
@@ -30,10 +36,30 @@ class CEWithChunkedOutputLoss(torch.nn.Module):
         """
         Upcast logits to fp32 and compute cross entropy loss.
         """
-        return F.cross_entropy(
-            logits.float(), labels, ignore_index=self.ignore_index, reduction="sum"
+        ce_loss = F.cross_entropy(
+            logits.float(), labels, ignore_index=self.ignore_index, reduction="none"
         )
 
+        if self.use_dft:
+            # Compute probabilities and gather the ones corresponding to labels
+            with torch.no_grad():  # Stop gradient
+                probs = torch.softmax(logits.float(), dim=-1)
+                # Create mask for valid tokens (not ignore_index)
+                valid_mask = labels != self.ignore_index
+                # Gather probabilities for the correct tokens
+                label_probs = probs.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
+                # Apply mask to only scale valid tokens
+                label_probs = label_probs * valid_mask
+                # Avoid multiplication by 0 for ignored tokens
+                label_probs = torch.where(
+                    valid_mask, label_probs, torch.ones_like(label_probs)
+                )
+
+            # Scale the loss by the probability (DFT)
+            ce_loss = ce_loss * label_probs
+
+        return ce_loss.sum()
+
     def forward(
         self, logits: List[torch.Tensor], labels: torch.Tensor, reduction="sum"
     ) -> torch.Tensor:
@@ -71,16 +97,20 @@ class CEWithChunkedOutputLoss(torch.nn.Module):
         return total_loss / total_elements
 
 
-def _build_chunked_ce_loss_fn(num_output_chunks: int = 8, ignore_index: int = -100):
-    loss_fn_ce = CEWithChunkedOutputLoss(num_output_chunks, ignore_index)
+def _build_chunked_ce_loss_fn(
+    num_output_chunks: int = 8, ignore_index: int = -100, use_dft: bool = False
+):
+    loss_fn_ce = CEWithChunkedOutputLoss(num_output_chunks, ignore_index, use_dft)
     loss_fn_ce.compute_cross_entropy = torch.compile(
         loss_fn_ce.compute_cross_entropy, backend="inductor"
     )
     return loss_fn_ce
 
 
-def get_causal_lm_loss(num_output_chunks: int = 8, ignore_index: int = -100):
-    loss_fn_ce = _build_chunked_ce_loss_fn(num_output_chunks, ignore_index)
+def get_causal_lm_loss(
+    num_output_chunks: int = 8, ignore_index: int = -100, use_dft: bool = False
+):
+    loss_fn_ce = _build_chunked_ce_loss_fn(num_output_chunks, ignore_index, use_dft)
 
     def chunked_fix_cross_entropy(
         source,
@@ -124,10 +154,14 @@ def get_causal_lm_loss(num_output_chunks: int = 8, ignore_index: int = -100):
     return for_causal_lm_chunked_loss
 
 
-def patch_chunked_ce_loss_fn(num_output_chunks: int = 8, ignore_index: int = -100):
+def patch_chunked_ce_loss_fn(
+    num_output_chunks: int = 8, ignore_index: int = -100, use_dft: bool = False
+):
     import transformers.loss.loss_utils
 
-    for_causal_lm_chunked_loss = get_causal_lm_loss(num_output_chunks, ignore_index)
+    for_causal_lm_chunked_loss = get_causal_lm_loss(
+        num_output_chunks, ignore_index, use_dft
+    )
     transformers.loss.loss_utils.ForCausalLMLoss = for_causal_lm_chunked_loss
     transformers.loss.loss_utils.LOSS_MAPPING["ForCausalLM"] = (
         for_causal_lm_chunked_loss

src/axolotl/monkeypatch/loss/eaft.py

@@ -1,51 +0,0 @@
-"""
-eaft (entropy-aware focal training) loss implementation
-weights examples by entropy approximation from top-k logits
-
-Reference: https://github.com/ymxyll/LlamaFactory-EAFT/blob/e2ce19e8efcc226450ee8f2b81dfe4e69f1f945d/src/llamafactory/train/trainer_utils.py
-"""
-
-import torch
-import torch.nn.functional as F
-
-
-def eaft_loss(outputs, labels, num_items_in_batch=None, alpha=1.0, k=20):
-    """
-    compute eaft loss with entropy weighting
-
-    args:
-        outputs: model outputs containing logits
-        labels: target labels for computing loss
-        num_items_in_batch: for sample packing support
-        alpha: exponent for entropy weighting (default 1.0)
-        k: number of top logits for entropy approximation (default 20)
-    """
-    logits = outputs.logits
-
-    shift_logits = logits[..., :-1, :].contiguous()
-    shift_labels = labels[..., 1:].contiguous()
-
-    vocab_size = shift_logits.size(-1)
-    shift_logits_view = shift_logits.view(-1, vocab_size)
-    shift_labels_view = shift_labels.view(-1)
-
-    mask = shift_labels_view != -100
-
-    with torch.no_grad():
-        top_k_logits, _ = torch.topk(
-            shift_logits_view[mask].float(), k=min(k, vocab_size), dim=-1
-        )
-        top_k_probs = F.softmax(top_k_logits, dim=-1)
-        entropy = -(top_k_probs * torch.log(top_k_probs + 1e-10)).sum(dim=-1)
-        weights = torch.pow(entropy, alpha)
-
-    loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
-    per_token_loss = loss_fct(shift_logits_view[mask], shift_labels_view[mask])
-    weighted_loss = per_token_loss * weights
-
-    if num_items_in_batch is not None:
-        loss = weighted_loss.sum() / num_items_in_batch
-    else:
-        loss = weighted_loss.mean()
-
-    return loss

src/axolotl/monkeypatch/models/mistral3/mistral_common_tokenizer.py

@@ -1,5 +1,5 @@
 """
-Monkeypatch to fix inefficient tensor conversion in MistralCommonBackend.apply_chat_template
+Monkeypatch to fix inefficient tensor conversion in MistralCommonTokenizer.apply_chat_template
 """
 
 import importlib
@@ -12,11 +12,11 @@ LOG = get_logger(__name__)
 
 
 def apply_mistral_tokenizer_image_patch():
-    """Apply patch to MistralCommonBackend.apply_chat_template to fix image tensor conversion."""
-    from transformers.tokenization_mistral_common import MistralCommonBackend
+    """Apply patch to MistralCommonTokenizer.apply_chat_template to fix image tensor conversion."""
+    from transformers.tokenization_mistral_common import MistralCommonTokenizer
 
     # Get original source
-    original_source = inspect.getsource(MistralCommonBackend.apply_chat_template)
+    original_source = inspect.getsource(MistralCommonTokenizer.apply_chat_template)
     original_source, _ = detab_code(original_source)
 
     # Define the replacement
@@ -41,7 +41,7 @@ def apply_mistral_tokenizer_image_patch():
         )
 
         # Load necessary imports from the module
-        module_name = MistralCommonBackend.__module__
+        module_name = MistralCommonTokenizer.__module__
         module = importlib.import_module(module_name)
 
         # Detect what needs to be imported
@@ -79,7 +79,7 @@ def apply_mistral_tokenizer_image_patch():
         exec(patched_source, globals())  # nosec B102
 
         # Replace the method
-        MistralCommonBackend.apply_chat_template = patched_apply_chat_template
-        LOG.info("Successfully applied MistralCommonBackend tensor conversion patch")
+        MistralCommonTokenizer.apply_chat_template = patched_apply_chat_template
+        LOG.info("Successfully applied MistralCommonTokenizer tensor conversion patch")
     else:
-        LOG.warning("Could not find target code for MistralCommonBackend patching")
+        LOG.warning("Could not find target code for MistralCommonTokenizer patching")

src/axolotl/monkeypatch/peft/utils.py

@@ -78,30 +78,3 @@ def patch_peft_prep_code():
     axolotl.loaders.model.prepare_model_for_kbit_training = (
         fixed_prepare_model_for_kbit_training
     )
-
-
-def patch_peft_torchao_dispatch():
-    """Skip PEFT's TorchaoLoraLinear for non-INT8 torchao weights.
-
-    PEFT's dispatch_torchao() matches AffineQuantizedTensor but then errors in
-    _check_dtype_supported() because it only allows INT8. Our LoRA kernels handle
-    dequantization explicitly, so we bypass PEFT's torchao dispatch entirely and
-    let it fall back to standard Linear LoRA layers.
-    """
-    try:
-        from peft.tuners.lora import torchao as peft_torchao
-    except ImportError:
-        LOG.warning("Could not import peft.tuners.lora.torchao for patching")
-        return
-
-    if getattr(peft_torchao, "_axolotl_patched", False):
-        return
-
-    def patched_dispatch(target, adapter_name, lora_config, **kwargs):
-        # Return None so PEFT falls back to standard Linear LoRA layers.
-        # Our LoRA kernels handle torchao dequantization explicitly.
-        return None
-
-    peft_torchao.dispatch_torchao = patched_dispatch
-    peft_torchao._axolotl_patched = True
-    LOG.info("Patched PEFT dispatch_torchao to skip TorchaoLoraLinear")

src/axolotl/monkeypatch/relora.py

@@ -155,6 +155,7 @@ class ReLoRACallback(TrainerCallback):
                     f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
                     "adapter",
                 ),
+                safe_serialization=True,
             )
             with torch.no_grad():
                 merge_and_save(
@@ -213,7 +214,7 @@ class ReLoRACallback(TrainerCallback):
 
                     self.last_full_model = checkpoint_folder
             else:
-                model.model.save_pretrained(checkpoint_folder)
+                model.model.save_pretrained(checkpoint_folder, safe_serialization=True)
 
         return control
 

src/axolotl/monkeypatch/transformers/trainer_context_parallel.py

@@ -52,15 +52,9 @@ def patch_prepare_context_parallel_inputs() -> None:
         if item in patched_source:
             items_to_import.append(item)
 
-    # Use a separate namespace to capture the exec'd function
-    namespace = {}
-    exec(f"from {module_name} import ({', '.join(items_to_import)})", namespace)
-    exec(patched_source, namespace)
+    exec(f"from {module_name} import ({', '.join(items_to_import)})", globals())
+    exec(patched_source, globals())
 
-    # Explicitly get the function from the namespace
-    axolotl_prepare_context_parallel_inputs = namespace[
-        "axolotl_prepare_context_parallel_inputs"
-    ]
     Trainer._original_prepare_context_parallel_inputs = (
         Trainer._prepare_context_parallel_inputs
     )

src/axolotl/processing_strategies.py

@@ -14,6 +14,7 @@ from transformers.models.voxtral import VoxtralProcessor
 
 from axolotl.utils.dict import remove_none_values
 from axolotl.utils.logging import get_logger
+from axolotl.utils.mistral.mistral3_processor import Mistral3Processor
 
 LOG = get_logger(__name__)
 
@@ -429,7 +430,7 @@ class Mistral3ProcessingStrategy(ProcessingStrategy):
 
     def __init__(
         self,
-        processor,
+        processor: Mistral3Processor,
         chat_template: Optional[str] = None,
         image_size: int | tuple[int, int] | None = None,
         image_resize_algorithm: Resampling | None = None,
@@ -485,58 +486,6 @@ class InternVLProcessingStrategy(ProcessingStrategy):
         return labels
 
 
-class Glm4vProcessingStrategy(ProcessingStrategy):
-    """Processing Strategy class for GLM4V and GLM4V-MoE vision models."""
-
-    def __init__(
-        self,
-        processor: ProcessorMixin,
-        chat_template: Optional[str] = None,
-        image_size: int | tuple[int, int] | None = None,
-        image_resize_algorithm: Resampling | None = None,
-    ):
-        super().__init__(processor, chat_template, image_size, image_resize_algorithm)
-
-        self.tokenizer = getattr(processor, "tokenizer", processor)
-
-        self.image_token = "<|image|>"  # nosec
-        self.begin_image_token = "<|begin_of_image|>"  # nosec
-        self.end_image_token = "<|end_of_image|>"  # nosec
-        self.video_token = "<|video|>"  # nosec
-        self.begin_video_token = "<|begin_of_video|>"  # nosec
-        self.end_video_token = "<|end_of_video|>"  # nosec
-
-        self.image_token_id = self.tokenizer.convert_tokens_to_ids(self.image_token)
-        self.begin_image_token_id = self.tokenizer.convert_tokens_to_ids(
-            self.begin_image_token
-        )
-        self.end_image_token_id = self.tokenizer.convert_tokens_to_ids(
-            self.end_image_token
-        )
-        self.video_token_id = self.tokenizer.convert_tokens_to_ids(self.video_token)
-        self.begin_video_token_id = self.tokenizer.convert_tokens_to_ids(
-            self.begin_video_token
-        )
-        self.end_video_token_id = self.tokenizer.convert_tokens_to_ids(
-            self.end_video_token
-        )
-
-    def process_labels(self, input_ids):
-        labels = input_ids.clone()
-
-        labels[labels == self.tokenizer.pad_token_id] = -100
-
-        labels[labels == self.image_token_id] = -100
-        labels[labels == self.begin_image_token_id] = -100
-        labels[labels == self.end_image_token_id] = -100
-
-        labels[labels == self.video_token_id] = -100
-        labels[labels == self.begin_video_token_id] = -100
-        labels[labels == self.end_video_token_id] = -100
-
-        return labels
-
-
 def get_processing_strategy(
     processor: ProcessorMixin,
     chat_template,
@@ -544,8 +493,6 @@ def get_processing_strategy(
     image_size: int | tuple[int, int] | None = None,
     image_resize_algorithm: Resampling | None = None,
 ):
-    from axolotl.utils.mistral.mistral3_processor import Mistral3Processor
-
     processing_kwargs = {
         "processor": processor,
         "chat_template": chat_template,
@@ -553,10 +500,10 @@ def get_processing_strategy(
         "image_resize_algorithm": image_resize_algorithm,
     }
 
-    if chat_template_type in [None, "tokenizer_default"]:
-        tokenizer = getattr(processor, "tokenizer", processor)
-        if hasattr(tokenizer, "chat_template"):
-            processing_kwargs["chat_template"] = tokenizer.chat_template
+    if chat_template_type in [None, "tokenizer_default"] and hasattr(
+        processor.tokenizer, "chat_template"
+    ):
+        processing_kwargs["chat_template"] = processor.tokenizer.chat_template
 
     if chat_template_type == "qwen2_vl":
         return Qwen2VLProcessingStrategy(
@@ -585,15 +532,6 @@ def get_processing_strategy(
         return Mistral3ProcessingStrategy(
             **processing_kwargs,
         )
-    try:
-        from transformers.models.glm46v.processing_glm46v import Glm46VProcessor
-
-        if isinstance(processor, Glm46VProcessor):
-            return Glm4vProcessingStrategy(
-                **processing_kwargs,
-            )
-    except ImportError:
-        pass
 
     if isinstance(processor, InternVLProcessor):
         return InternVLProcessingStrategy(

src/axolotl/prompt_strategies/chat_template.py

@@ -150,8 +150,6 @@ class ChatTemplatePrompter(Prompter):
 
         return self.tokenizer.apply_chat_template(
             conversation,
-            tokenize=True,
-            return_dict=False,
             **chat_template_kwargs,
         )
 

src/axolotl/telemetry/callbacks.py

@@ -153,27 +153,13 @@ class TelemetryCallback(TrainerCallback):
                 self.last_report_step = step
 
     def _extract_last_metrics(self, state: TrainerState) -> dict:
-        """Extract last loss, learning_rate, grad_norm, and token metrics from log history."""
+        """Extract last loss, learning_rate, and grad_norm from log history."""
         if not state.log_history:
-            return {
-                "loss": 0,
-                "ppl": 0,
-                "learning_rate": 0,
-                "grad_norm": 0,
-                "tokens/total": 0,
-                "tokens/trainable": 0,
-                "tokens/train_per_sec_per_gpu": 0,
-            }
+            return {"loss": 0, "learning_rate": 0, "grad_norm": 0}
 
         last_log = state.log_history[-1]
         return {
             "loss": last_log.get("loss", 0),
-            "ppl": last_log.get("ppl", 0),
             "learning_rate": last_log.get("learning_rate", 0),
             "grad_norm": last_log.get("grad_norm", 0),
-            "tokens/total": last_log.get("tokens/total", 0),
-            "tokens/trainable": last_log.get("tokens/trainable", 0),
-            "tokens/train_per_sec_per_gpu": last_log.get(
-                "tokens/train_per_sec_per_gpu", 0
-            ),
         }

src/axolotl/telemetry/errors.py

@@ -155,10 +155,6 @@ def send_errors(func: Callable) -> Callable:
                     },
                 )
 
-                LOG.error(
-                    f"Error captured in telemetry. Run ID: {telemetry_manager.run_id}"
-                )
-
             raise
 
     return wrapper

src/axolotl/telemetry/manager.py

@@ -5,6 +5,7 @@ import importlib
 import logging
 import os
 import platform
+import time
 import uuid
 from pathlib import Path
 from typing import Any
@@ -19,6 +20,21 @@ LOG = logging.getLogger(__name__)
 POSTHOG_HOST = "https://app.posthog.com"
 POSTHOG_WRITE_KEY = "phc_1kUR0o04oJKKTTeSsIz2Mfm5mpiVsQEf2WOlzljMD7y"
 
+OPT_OUT_WARNING_SLEEP_SECONDS = 10
+OPT_OUT_WARNING = (
+    "\nTelemetry is now enabled by default to help improve Axolotl. "
+    "If you'd like to disable it, set AXOLOTL_DO_NOT_TRACK=1 in your environment.\n\n"
+    "Telemetry data helps us understand:\n"
+    "- Which features are most used\n"
+    "- What hardware configurations to prioritize\n"
+    "- Where users encounter errors\n\n"
+    "Personally identifiable information (PII) is not collected.\n\n"
+    "To remove this warning, explicitly set AXOLOTL_DO_NOT_TRACK=0 (enable telemetry) "
+    "or AXOLOTL_DO_NOT_TRACK=1 (disable telemetry).\n\n"
+    "For details, see: https://docs.axolotl.ai/docs/telemetry.html\n\n"
+    f"Sleeping for {OPT_OUT_WARNING_SLEEP_SECONDS}s..."
+)
+
 WHITELIST_PATH = str(Path(__file__).parent / "whitelist.yaml")
 
 # NOTE: Need to keep these up to date with any config schema changes
@@ -30,8 +46,8 @@ FIELDS_TO_REDACT = {
     "resume_from_checkpoint",
     "hub_model_id",
 }
-PREFIXES_TO_REDACT = {"wandb_", "comet_", "mlflow_", "gradio_", "trackio_", "swanlab_"}
-PATH_INDICATORS = {"path", "dir", "data_files"}
+PREFIXES_TO_REDACT = {"wandb_", "comet_", "mlflow_", "gradio_"}
+PATH_INDICATORS = {"path", "dir"}
 
 # pylint: disable=duplicate-code
 RELEVANT_PACKAGES = {
@@ -167,6 +183,11 @@ class TelemetryManager:
             "false",
             "true",
         ):
+            # Print opt-out info message for main process only
+            if is_main_process():
+                LOG.warning(OPT_OUT_WARNING)
+            time.sleep(OPT_OUT_WARNING_SLEEP_SECONDS)
+
             return True
 
         # Only rank 0 will send telemetry

src/axolotl/telemetry/whitelist.yaml

@@ -31,10 +31,3 @@ organizations:
   - "mistral-community"
   - "llava-hf"
   - "ByteDance-Seed"
-  - "ACE-Step"
-  - "openbmb"
-  - "MiniMaxAI"
-  - "stepfun-ai"
-  - "internlm"
-  - "katanemo"
-  - "XiaomiMiMo"

src/axolotl/train.py

@@ -135,13 +135,16 @@ def setup_reference_model(
     return model_ref
 
 
-def setup_signal_handler(cfg: DictDefault, model: PreTrainedModel):
+def setup_signal_handler(
+    cfg: DictDefault, model: PreTrainedModel, safe_serialization: bool
+):
     """
     Set up signal handler for graceful termination.
 
     Args:
         cfg: Dictionary mapping `axolotl` config keys to values.
         model: The model to save on termination
+        safe_serialization: Whether to use safe serialization when saving
     """
     # ray workers don't have access to this signal
     if cfg.local_rank == 0 and not cfg.use_ray:
@@ -149,7 +152,9 @@ def setup_signal_handler(cfg: DictDefault, model: PreTrainedModel):
         def terminate_handler(_, __, model_weakref):
             if model_weakref() is not None:
                 _model = model_weakref()
-                _model.save_pretrained(cfg.output_dir)
+                _model.save_pretrained(
+                    cfg.output_dir, safe_serialization=safe_serialization
+                )
 
             cleanup_distributed()
             sys.exit(0)
@@ -214,6 +219,7 @@ def save_trained_model(
     cfg: DictDefault,
     trainer: Any,
     model: PreTrainedModel,
+    safe_serialization: bool,
 ):
     """
     Save the trained model according to configuration and training setup.
@@ -222,6 +228,7 @@ def save_trained_model(
         cfg: Dictionary mapping `axolotl` config keys to values.
         trainer: The trainer object.
         model: The trained model to save.
+        safe_serialization: Whether to use safe serialization.
     """
     LOG.info(f"Training completed! Saving trained model to {cfg.output_dir}.")
 
@@ -276,6 +283,7 @@ def save_trained_model(
                 merge_fsdp_weights(
                     checkpoint_dir=str(fsdp_dir),
                     output_path=merged_path,
+                    safe_serialization=True,
                 )
                 trainer.accelerator.wait_for_everyone()
                 if trainer.accelerator.is_main_process:
@@ -322,9 +330,11 @@ def save_trained_model(
                 pass
     elif cfg.local_rank == 0:
         if cfg.rl and cfg.adapter and not cfg.rl_adapter_ref_model:
-            trainer.model.save_pretrained(cfg.output_dir)
+            trainer.model.save_pretrained(
+                cfg.output_dir, safe_serialization=safe_serialization
+            )
 
-        model.save_pretrained(cfg.output_dir)
+        model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
 
     if hasattr(cfg, "llmcompressor") and cfg.llmcompressor:
         # TODO: add integration support so this can be implemented completely within the plugin
@@ -334,6 +344,7 @@ def save_trained_model(
             model=model,
             output_dir=cfg.output_dir,
             trainer=trainer,
+            safe_serialization=safe_serialization,
             save_compressed=cfg.llmcompressor.save_compressed,
         )
 
@@ -438,6 +449,7 @@ def handle_untrained_tokens_fix(
     model: PreTrainedModel,
     tokenizer: PreTrainedTokenizer,
     train_dataset: Dataset,
+    safe_serialization: bool,
 ):
     """
     Apply fixes for untrained tokens if configured.
@@ -447,6 +459,7 @@ def handle_untrained_tokens_fix(
         model: The model to apply fixes to.
         tokenizer: The tokenizer for token identification.
         train_dataset: The training dataset to use.
+        safe_serialization: Whether to use safe serialization when saving.
     """
     if not cfg.fix_untrained_tokens:
         return
@@ -470,7 +483,9 @@ def handle_untrained_tokens_fix(
     fix_untrained_tokens(model, tokenizer, train_dataset, **fix_kwargs)
 
     if cfg.local_rank == 0:
-        model.save_pretrained(str(Path(cfg.output_dir)))
+        model.save_pretrained(
+            str(Path(cfg.output_dir)), safe_serialization=safe_serialization
+        )
 
 
 def setup_model_and_trainer(
@@ -567,12 +582,15 @@ def train(
     ) = setup_model_and_trainer(cfg, dataset_meta)
 
     # Handle untrained tokens if configured
+    safe_serialization = cfg.save_safetensors is True
     train_dataset = dataset_meta.train_dataset
-    handle_untrained_tokens_fix(cfg, model, tokenizer, train_dataset)
+    handle_untrained_tokens_fix(
+        cfg, model, tokenizer, train_dataset, safe_serialization
+    )
 
     # Additional setup
     save_initial_configs(cfg, tokenizer, model, peft_config, processor)
-    setup_signal_handler(cfg, model)
+    setup_signal_handler(cfg, model, safe_serialization)
     setup_model_card(cfg)
 
     # Execute the training
@@ -584,7 +602,7 @@ def train(
         torch.cuda.empty_cache()
 
     # Save the trained model and cleanup
-    save_trained_model(cfg, trainer, model)
+    save_trained_model(cfg, trainer, model, safe_serialization)
     tokenizer.save_pretrained(
         str(Path(cfg.output_dir)), save_jinja_files=cfg.tokenizer_save_jinja_files
     )