set version to v0.13.1 (#3363)

* upgrade transformers to 4.57.5

* explicitly set versions for fbgemm-gpu

* handle index url for cuda version

* explicitly set cu version for fbgemm deps, skip for 130

* cu suffix not needed on version if using whl subpath

.github/PULL_REQUEST_TEMPLATE.md

@@ -15,6 +15,11 @@
 <!--- 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,6 +21,8 @@ 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:
@@ -32,6 +34,7 @@ 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: ""
@@ -39,6 +42,7 @@ 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: ""
@@ -46,6 +50,7 @@ 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: "130"
             cuda_version: 13.0.0
             cudnn_version: ""
@@ -53,6 +58,7 @@ jobs:
             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: ""
@@ -79,7 +85,7 @@ jobs:
             axolotlai/axolotl-base
       - name: Login to Docker Hub
         uses: docker/login-action@v2
-        if: ${{ github.event_name != 'pull_request' && secrets.DOCKERHUB_USERNAME != '' && secrets.DOCKERHUB_TOKEN != '' }}
+        if: ${{ github.event_name != 'pull_request' && env.HAS_DOCKERHUB_CREDS == 'true' }}
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
@@ -90,7 +96,7 @@ jobs:
         with:
           context: .
           file: ./docker/${{ matrix.dockerfile }}
-          platforms: linux/amd64,linux/arm64
+          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 }}
@@ -105,6 +111,8 @@ 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:
@@ -116,6 +124,7 @@ 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: ""
@@ -123,6 +132,7 @@ 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: ""
@@ -130,6 +140,7 @@ 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: "130"
             cuda_version: 13.0.0
             cudnn_version: ""
@@ -137,6 +148,7 @@ jobs:
             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
@@ -148,6 +160,7 @@ 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 }}
@@ -158,6 +171,7 @@ 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,22 +20,26 @@ jobs:
             python_version: "3.11"
             pytorch: 2.8.0
             axolotl_extras:
-            is_latest: true
+            platforms: "linux/amd64"
           - 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: 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
@@ -61,7 +65,7 @@ jobs:
         uses: docker/build-push-action@v5
         with:
           context: .
-          platforms: linux/amd64,linux/arm64
+          platforms: ${{ matrix.platforms }}
           build-args: |
             BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
             CUDA=${{ matrix.cuda }}
@@ -88,22 +92,26 @@ jobs:
             python_version: "3.11"
             pytorch: 2.8.0
             axolotl_extras:
-            is_latest: true
+            platforms: "linux/amd64"
           - 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: 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
@@ -128,7 +136,7 @@ jobs:
         uses: docker/build-push-action@v5
         with:
           context: .
-          platforms: linux/amd64,linux/arm64
+          platforms: ${{ matrix.platforms }}
           build-args: |
             BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
             CUDA=${{ matrix.cuda }}
@@ -149,11 +157,11 @@ jobs:
           - cuda: 128
             cuda_version: 12.8.1
             python_version: "3.11"
-            pytorch: 2.8.0
+            pytorch: 2.9.1
             axolotl_extras:
-            is_latest:
+            is_latest: true
-          - cuda: 128
+          - cuda: 130
-            cuda_version: 12.8.1
+            cuda_version: 13.0.0
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:

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

@@ -47,7 +47,8 @@ jobs:
             cuda_version: 13.0.0
             python_version: "3.11"
             pytorch: 2.9.1
-            axolotl_extras: fbgemm-gpu
+            axolotl_extras:
+#            axolotl_extras: fbgemm-gpu
             num_gpus: 2
             nightly_build: "true"
     runs-on: [self-hosted, modal]

docker/Dockerfile

@@ -6,6 +6,7 @@ ARG AXOLOTL_EXTRAS=""
 ARG AXOLOTL_ARGS=""
 ARG CUDA="118"
 ARG PYTORCH_VERSION="2.1.2"
+ARG TARGETARCH
 
 ENV PYTORCH_VERSION=$PYTORCH_VERSION
 
@@ -20,13 +21,17 @@ 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
+# If AXOLOTL_EXTRAS is set, append it in brackets; don't install deepspeed with arm64
-RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
+RUN if [ "$TARGETARCH" = "arm64" ]; then \
-        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        BASE_EXTRAS="flash-attn,ring-flash-attn,optimizers,ray"; \
     else \
-        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
+        BASE_EXTRAS="deepspeed,flash-attn,ring-flash-attn,optimizers,ray"; \
     fi && \
-    python scripts/unsloth_install.py | sh && \
+    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/cutcrossentropy_install.py | sh && \
     pip install pytest && \
     pip cache purge

docker/Dockerfile-uv-base

@@ -2,6 +2,7 @@ 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
 
@@ -31,20 +32,35 @@ 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 [ "$CUDA" = "128" ]; then \
+            if [ "$TARGETARCH" = "amd64" ]; 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; \
+                if [ "$CUDA" = "128" ]; then \
-                uv pip install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
+                    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; \
-                rm 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; \
-            elif [ "$CUDA" = "130" ]; then \
+                    rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
-                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; \
+                elif [ "$CUDA" = "130" ]; then \
-                uv pip install --no-cache-dir flash_attn-2.8.3+cu130torch2.9-cp311-cp311-linux_x86_64.whl; \
+                    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; \
-                rm 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 \
             fi \
             ;; \
     esac

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

@@ -1,6 +1,7 @@
 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
@@ -29,7 +30,7 @@ output_dir: ./outputs/out
 adapter: qlora
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0.05
+lora_dropout: 0
 lora_target_linear: true
 
 sequence_len: 2048

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

@@ -1,6 +1,7 @@
 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
@@ -29,7 +30,7 @@ output_dir: ./outputs/out
 adapter: qlora
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0.05
+lora_dropout: 0
 lora_target_linear: true
 
 sequence_len: 2048

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

@@ -2,6 +2,7 @@ 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
 
@@ -32,8 +33,8 @@ sample_packing: true
 
 lora_r: 32
 lora_alpha: 16
-lora_dropout: 0.05
+lora_dropout: 0
-lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+lora_target_linear: true
 
 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.05
+lora_dropout: 0
 lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
 
 wandb_project:

requirements.txt

@@ -11,11 +11,11 @@ liger-kernel==0.6.4
 packaging==23.2
 
 huggingface_hub>=0.36.0
-peft>=0.18.0
+peft>=0.18.1
 tokenizers>=0.22.1
-transformers==4.57.1
+transformers==4.57.6
 accelerate==1.12.0
-datasets==4.4.2
+datasets==4.5.0
 deepspeed>=0.18.3
 trl==0.25.1
 hf_xet==1.2.0

setup.py

@@ -26,6 +26,7 @@ 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,44 +63,63 @@ 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-genai==1.4.1"]
+                extras_require_map["fbgemm-gpu"] = [
+                    "fbgemm-gpu==1.4.0",
+                    "fbgemm-gpu-genai==1.4.2",
+                ]
                 extras_require_map["vllm"] = ["vllm==0.11.1"]
+                if not install_xformers:
+                    _install_requires.pop(_install_requires.index(xformers_version))
             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:
-                    _install_requires.append("xformers==0.0.30")
+                    if install_xformers:
+                        _install_requires.append("xformers==0.0.30")
                     # vllm 0.9.x is incompatible with latest transformers
                     extras_require_map.pop("vllm")
                 else:
-                    _install_requires.append("xformers==0.0.31")
+                    if install_xformers:
+                        _install_requires.append("xformers==0.0.31")
                     extras_require_map["vllm"] = ["vllm==0.10.1"]
             elif (major, minor) >= (2, 6):
                 _install_requires.pop(_install_requires.index(xformers_version))
-                _install_requires.append("xformers==0.0.29.post3")
+                if install_xformers:
+                    _install_requires.append("xformers==0.0.29.post3")
                 # since we only support 2.6.0+cu126
                 _dependency_links.append("https://download.pytorch.org/whl/cu126")
                 extras_require_map.pop("vllm")
             elif (major, minor) >= (2, 5):
                 _install_requires.pop(_install_requires.index(xformers_version))
-                if patch == 0:
+                if install_xformers:
-                    _install_requires.append("xformers==0.0.28.post2")
+                    if patch == 0:
-                else:
+                        _install_requires.append("xformers==0.0.28.post2")
-                    _install_requires.append("xformers>=0.0.28.post3")
+                    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 patch == 0:
+                if install_xformers:
-                    _install_requires.pop(_install_requires.index(xformers_version))
+                    if patch == 0:
-                    _install_requires.append("xformers>=0.0.27")
+                        _install_requires.pop(_install_requires.index(xformers_version))
-                else:
+                        _install_requires.append("xformers>=0.0.27")
-                    _install_requires.pop(_install_requires.index(xformers_version))
+                    else:
-                    _install_requires.append("xformers==0.0.28.post1")
+                        _install_requires.pop(_install_requires.index(xformers_version))
+                        _install_requires.append("xformers==0.0.28.post1")
             else:
                 raise ValueError("axolotl requires torch>=2.4")
 

src/axolotl/__init__.py

@@ -4,4 +4,4 @@ import pkgutil
 
 __path__ = pkgutil.extend_path(__path__, __name__)  # Make this a namespace package
 
-__version__ = "0.13.0.dev"
+__version__ = "0.13.1"

src/axolotl/loaders/utils.py

@@ -5,6 +5,7 @@ from typing import Type
 
 import addict
 import torch
+import transformers
 from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
 
 from axolotl.utils.dict import DictDefault
@@ -153,6 +154,9 @@ 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.
 
@@ -174,8 +178,13 @@ 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 = AutoConfig.from_pretrained(
+        model_config = config_cls.from_pretrained(
             model_config_name,
             trust_remote_code=trust_remote_code,
             **config_kwargs,

src/axolotl/utils/schemas/model.py

@@ -25,7 +25,12 @@ class ModelInputConfig(BaseModel):
             "description": "If the base_model repo on hf hub doesn't include configuration .json files, You can set that here, or leave this empty to default to base_model"
         },
     )
-    cls_model_config: str | None = None
+    cls_model_config: str | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "transformers config class (e.g., 'LlamaConfig', 'MistralConfig'). Defaults to AutoConfig."
+        },
+    )
     tokenizer_config: str | None = Field(
         default=None,
         json_schema_extra={