config fix -___-

linting v2
simplifying check
2025-03-18 11:35:20 +00:00 · 2025-03-18 11:33:54 +00:00 · 2025-03-18 11:26:53 +00:00 · 2025-03-18 11:23:48 +00:00 · 2025-03-18 11:23:23 +00:00 · 2025-03-18 11:22:00 +00:00
175 changed files with 1261 additions and 1887 deletions
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -40,12 +40,6 @@ jobs:
            python_version: "3.11"
            pytorch: 2.6.0
            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
-          - cuda: "128"
-            cuda_version: 12.8.1
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: nightly
-            torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
    steps:
      - name: Checkout
        uses: actions/checkout@v4
@@ -67,7 +61,7 @@ jobs:
        uses: docker/build-push-action@v4
        with:
          context: .
-          file: ${{ matrix.pytorch == 'nightly' && './docker/Dockerfile-base-nightly' || './docker/Dockerfile-base' }}
+          file: ./docker/Dockerfile-base
          push: ${{ github.event_name != 'pull_request' }}
          tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
          labels: ${{ steps.metadata.outputs.labels }}
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -20,12 +20,9 @@ jobs:
          uses: actions/setup-python@v5
          with:
            python-version: '3.11'
-        - name: Install dependencies
+        - name: install dependencies
          run: |
-            python3 -m pip install jupyter quartodoc
-            python3 -m pip install -e .
-        - name: Build autodoc
-          run: quartodoc build
+            python3 -m pip install jupyter
        - name: Publish to GitHub Pages (and render)
          uses: quarto-dev/quarto-actions/publish@v2
          with:
--- a/.github/workflows/precommit-autoupdate.yml
+++ b/.github/workflows/precommit-autoupdate.yml
@@ -1,49 +0,0 @@
-name: Pre-commit auto-update
-
-on:
-  schedule:
-    - cron: '0 0 * * 0'  # Run weekly
-  workflow_dispatch:  # Manual kickoff
-
-jobs:
-  auto-update:
-    runs-on: ubuntu-latest
-    permissions:
-      contents: write
-      pull-requests: write
-    steps:
-      - uses: actions/checkout@v4
-
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.11'
-
-      - name: Update pre-commit hooks
-        id: update
-        run: |
-          pip install pre-commit
-          pre-commit autoupdate
-          if [[ -n $(git status --porcelain) ]]; then
-            echo "changes=true" >> $GITHUB_OUTPUT
-            git diff .pre-commit-config.yaml > pre-commit-update.diff
-          fi
-
-      - name: Create Pull Request
-        if: steps.update.outputs.changes == 'true'
-        uses: peter-evans/create-pull-request@v6
-        with:
-          token: ${{ secrets.GITHUB_TOKEN }}
-          branch: update/pre-commit-hooks
-          delete-branch: true
-          title: "chore: update pre-commit hooks"
-          commit-message: "chore: update pre-commit hooks"
-          body: |
-            Automated PR to update pre-commit hooks to their latest versions.
-
-            <details>
-            <summary>Changes:</summary>
-
-            ```diff
-            ${{ steps.update.outputs.diff }}
-            ```
-            </details>
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -40,7 +40,7 @@ jobs:

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

--- a/.github/workflows/tests-nightly.yml
+++ b/.github/workflows/tests-nightly.yml
@@ -42,7 +42,7 @@ jobs:
      - name: upgrade pip
        run: |
          pip3 install --upgrade pip
-          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 wheel
+          pip3 install --upgrade packaging setuptools wheel

      - name: Install PyTorch
        run: |
@@ -59,7 +59,7 @@ jobs:
      - name: Install dependencies
        run: |
          pip3 install --upgrade pip
-          pip3 install --upgrade packaging==23.2
+          pip3 install --upgrade packaging
          pip3 install --no-build-isolation -U -e .
          python scripts/unsloth_install.py | sh
          python scripts/cutcrossentropy_install.py | sh
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -74,7 +74,7 @@ jobs:
      - name: upgrade pip
        run: |
          pip3 install --upgrade pip
-          pip3 install --upgrade packaging==23.2 setuptools==75.8.0 wheel
+          pip3 install --upgrade packaging setuptools wheel

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

      - name: Install PyTorch
        run: |
--- a/.gitignore
+++ b/.gitignore
@@ -181,10 +181,6 @@ prepared-datasets/
 submit.sh
 *.out*

-# Quartodoc generated files
-objects.json
-site_libs/
-
 typings/
 out/

--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -3,7 +3,7 @@ default_language_version:

 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v4.4.0
    hooks:
    -   id: check-yaml
    -   id: end-of-file-fixer
@@ -11,23 +11,23 @@ repos:
    -   id: no-commit-to-branch
        args: ['--branch', 'main']
 -   repo: https://github.com/psf/black
-    rev: 25.1.0
+    rev: 23.3.0
    hooks:
    -   id: black
 -   repo: https://github.com/pycqa/isort
-    rev: 6.0.1
+    rev: 5.12.0
    hooks:
      - id: isort
 -   repo: https://github.com/PyCQA/flake8
-    rev: 7.1.2
+    rev: 6.1.0
    hooks:
    - id: flake8
-   repo: https://github.com/pylint-dev/pylint
-    rev: v3.3.6
+-   repo: https://github.com/PyCQA/pylint
+    rev: v3.3.0
    hooks:
    - id: pylint
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.15.0
+    rev: v1.3.0
    hooks:
    - id: mypy
      additional_dependencies:
@@ -36,7 +36,7 @@ repos:
            'pydantic>=2.5.3',
        ]
 -   repo: https://github.com/PyCQA/bandit
-    rev: 1.8.3
+    rev: 1.7.5
    hooks:
    -   id: bandit
        args: [
--- a/README.md
+++ b/README.md
@@ -55,7 +55,7 @@ Features:
 ### Installation

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

 # Download example axolotl configs, deepspeed configs
@@ -97,7 +97,6 @@ That's it! Check out our [Getting Started Guide](https://axolotl-ai-cloud.github
 - [Multi-GPU Training](https://axolotl-ai-cloud.github.io/axolotl/docs/multi-gpu.html)
 - [Multi-Node Training](https://axolotl-ai-cloud.github.io/axolotl/docs/multi-node.html)
 - [Multipacking](https://axolotl-ai-cloud.github.io/axolotl/docs/multipack.html)
- [API Reference](https://axolotl-ai-cloud.github.io/axolotl/docs/api/) - Auto-generated code documentation
 - [FAQ](https://axolotl-ai-cloud.github.io/axolotl/docs/faq.html) - Frequently asked questions

 ## 🤝 Getting Help
--- a/_quarto.yml
+++ b/_quarto.yml
@@ -1,178 +1,6 @@
 project:
  type: website

-quartodoc:
-  dir: docs/api
-  package: axolotl
-  title: API Reference
-  parser: google
-
-  sections:
-    - title: Core
-      desc: Core functionality for training
-      contents:
-        - train
-        - evaluate
-        - datasets
-        - convert
-        - prompt_tokenizers
-        - logging_config
-        - core.trainer_builder
-        - core.training_args
-        - core.chat.messages
-        - core.chat.format.chatml
-        - core.chat.format.llama3x
-        - core.chat.format.shared
-        - core.datasets.chat
-        - core.datasets.transforms.chat_builder
-    - title: CLI
-      desc: Command-line interface
-      contents:
-        - cli.main
-        - cli.train
-        - cli.evaluate
-        - cli.args
-        - cli.checks
-        - cli.config
-        - cli.inference
-        - cli.merge_lora
-        - cli.merge_sharded_fsdp_weights
-        - cli.preprocess
-        - cli.sweeps
-        - cli.utils
-        - cli.cloud.base
-        - cli.cloud.modal_
-    - title: Trainers
-      desc: Training implementations
-      contents:
-        - core.trainers.base
-        - core.trainers.trl
-        - core.trainers.dpo.trainer
-        - core.trainers.grpo.trainer
-    - title: Prompt Strategies
-      desc: Prompt formatting strategies
-      contents:
-        - prompt_strategies.base
-        - prompt_strategies.chat_template
-        - prompt_strategies.alpaca_chat
-        - prompt_strategies.alpaca_instruct
-        - prompt_strategies.alpaca_w_system
-        - prompt_strategies.user_defined
-        - prompt_strategies.llama2_chat
-        - prompt_strategies.completion
-        - prompt_strategies.input_output
-        - prompt_strategies.stepwise_supervised
-        - prompt_strategies.metharme
-        - prompt_strategies.orcamini
-        - prompt_strategies.pygmalion
-        - prompt_strategies.messages.chat
-        - prompt_strategies.dpo.chat_template
-        - prompt_strategies.dpo.llama3
-        - prompt_strategies.dpo.chatml
-        - prompt_strategies.dpo.zephyr
-        - prompt_strategies.dpo.user_defined
-        - prompt_strategies.dpo.passthrough
-        - prompt_strategies.kto.llama3
-        - prompt_strategies.kto.chatml
-        - prompt_strategies.kto.user_defined
-        - prompt_strategies.orpo.chat_template
-        - prompt_strategies.bradley_terry.llama3
-    - title: Kernels
-      desc: Low-level performance optimizations
-      contents:
-        - kernels.lora
-        - kernels.geglu
-        - kernels.swiglu
-        - kernels.quantize
-        - kernels.utils
-    - title: MonkeyPatches
-      desc: Runtime patches for model optimizations
-      contents:
-        - monkeypatch.llama_attn_hijack_flash
-        - monkeypatch.llama_attn_hijack_xformers
-        - monkeypatch.mistral_attn_hijack_flash
-        - monkeypatch.multipack
-        - monkeypatch.relora
-        - monkeypatch.llama_expand_mask
-        - monkeypatch.lora_kernels
-        - monkeypatch.utils
-        - monkeypatch.btlm_attn_hijack_flash
-        - monkeypatch.llama_patch_multipack
-        - monkeypatch.stablelm_attn_hijack_flash
-        - monkeypatch.trainer_fsdp_optim
-        - monkeypatch.transformers_fa_utils
-        - monkeypatch.unsloth_
-        - monkeypatch.attention.mllama
-        - monkeypatch.data.batch_dataset_fetcher
-        - monkeypatch.mixtral
-    - title: Utils
-      desc: Utility functions
-      contents:
-        - utils.models
-        - utils.tokenization
-        - utils.chat_templates
-        - utils.lora
-        - utils.lora_embeddings
-        - utils.model_shard_quant
-        - utils.bench
-        - utils.freeze
-        - utils.trainer
-        - utils.schedulers
-        - utils.distributed
-        - utils.dict
-        - utils.optimizers.adopt
-        - utils.data.pretraining
-        - utils.data.sft
-        - utils.gradient_checkpointing.unsloth
-    - title: Schemas
-      desc: Pydantic data models for Axolotl config
-      contents:
-        - utils.schemas.config
-        - utils.schemas.model
-        - utils.schemas.training
-        - utils.schemas.datasets
-        - utils.schemas.peft
-        - utils.schemas.trl
-        - utils.schemas.integrations
-        - utils.schemas.enums
-        - utils.schemas.utils
-    - title: Integrations
-      desc: Third-party integrations and extensions
-      contents:
-        - integrations.base
-        - integrations.cut_cross_entropy.args
-        - integrations.grokfast.optimizer
-        - integrations.kd.trainer
-        - integrations.liger.args
-        - integrations.lm_eval.args
-        - integrations.spectrum.args
-    - title: Common
-      desc: Common utilities and shared functionality
-      contents:
-        - common.architectures
-        - common.const
-        - common.datasets
-    - title: Models
-      desc: Custom model implementations
-      contents:
-        - models.mamba.modeling_mamba
-    - title: Data Processing
-      desc: Data processing utilities
-      contents:
-        - utils.collators.core
-        - utils.collators.batching
-        - utils.collators.mamba
-        - utils.collators.mm_chat
-        - utils.samplers.multipack
-    - title: Callbacks
-      desc: Training callbacks
-      contents:
-        - utils.callbacks.perplexity
-        - utils.callbacks.profiler
-        - utils.callbacks.lisa
-        - utils.callbacks.mlflow_
-        - utils.callbacks.comet_
-
 website:
  title: "Axolotl"
  description: "We make fine-tuning accessible, scalable, and fun"
@@ -207,8 +35,6 @@ website:
            - docs/inference.qmd
            - docs/cli.qmd
            - docs/config.qmd
-            - text: "API Reference"
-              href: docs/api

        - section: "Dataset Formats"
          contents: docs/dataset-formats/*
@@ -254,22 +80,3 @@ format:
    theme: darkly
    css: styles.css
    toc: true
-    # Enable better handling of line breaks in markdown
-    preserve-tabs: true
-    html-math-method: mathjax
-    # Improved markdown processing options
-    md-extensions:
-      - markdown_it
-      - def_list
-      - attr_list
-      - fenced_divs
-      - tables
-      - html_admonition
-      - lineblocks
-      - fancy_lists
-    # Control whitespace handling
-    whitespace: preserve
-    # Process newlines in paragraphs
-    wrap: preserve
-    # Better line break handling
-    preserve-linebreaks: true
--- a/cicd/Dockerfile.jinja
+++ b/cicd/Dockerfile.jinja
@@ -31,7 +31,6 @@ 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==23.2 setuptools==75.8.0
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
        pip install --no-build-isolation -e .[deepspeed,flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
    else \
--- a/cicd/multigpu.py
+++ b/cicd/multigpu.py
@@ -1,7 +1,6 @@
 """
-modal application to run axolotl gpu tests in Modal
-"""
-
+ modal application to run axolotl gpu tests in Modal
+ """
 # pylint: disable=duplicate-code

 import os
--- a/cicd/tests.py
+++ b/cicd/tests.py
@@ -1,5 +1,4 @@
 """Modal app to run axolotl GPU tests"""
-
 # pylint: disable=duplicate-code

 import os
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -28,7 +28,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"

 WORKDIR /workspace

-RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==23.2 setuptools==75.8.0 wheel && \
+RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
    python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA && \
    python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \
    python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main"
--- a/docker/Dockerfile-base-nightly
+++ b/docker/Dockerfile-base-nightly
@@ -1,39 +0,0 @@
-ARG CUDA_VERSION="12.8.1"
-ARG CUDNN_VERSION="8"
-ARG UBUNTU_VERSION="22.04"
-ARG MAX_JOBS=4
-
-FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
-
-ENV PATH="/root/miniconda3/bin:${PATH}"
-
-ARG PYTHON_VERSION="3.11"
-ARG PYTORCH_VERSION="nightly"
-ARG CUDA="128"
-ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
-
-ENV PYTHON_VERSION=$PYTHON_VERSION
-ENV TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST
-
-RUN apt-get update \
-    && apt-get install -y wget git build-essential ninja-build git-lfs libaio-dev pkg-config && rm -rf /var/lib/apt/lists/* \
-    && wget \
-    https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
-    && mkdir /root/.conda \
-    && bash Miniconda3-latest-Linux-x86_64.sh -b \
-    && rm -f Miniconda3-latest-Linux-x86_64.sh \
-    && conda create -n "py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
-
-ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
-
-WORKDIR /workspace
-
-RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
-    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"
-
-RUN git lfs install --skip-repo && \
-    pip3 install awscli && \
-    # The base image ships with `pydantic==1.8.2` which is not working
-    pip3 install -U --no-cache-dir pydantic==1.10.10
--- a/docs/.gitignore
+++ b/docs/.gitignore
@@ -1,4 +1,2 @@
 /.quarto/
 _site/
-/api/*.qmd
-/api/*.html
--- a/docs/cli.qmd
+++ b/docs/cli.qmd
@@ -1,5 +1,5 @@
 ---
-title: "Command Line Interface (CLI)"
+title: "CLI Reference"
 format:
  html:
    toc: true
--- a/docs/config.qmd
+++ b/docs/config.qmd
@@ -85,12 +85,6 @@ gpu_memory_limit: 20GiB
 # Do the LoRA/PEFT loading on CPU -- this is required if the base model is so large it takes up most or all of the available GPU VRAM, e.g. during a model and LoRA merge
 lora_on_cpu: true

-# List[str]. Add plugins to extend the pipeline.
-# See `src/axolotl/integrations` for the available plugins or doc below for more details.
-# https://axolotl-ai-cloud.github.io/axolotl/docs/custom_integrations.html
-plugins:
-  # - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
 # A list of one or more datasets to finetune the model with
 datasets:
  # HuggingFace dataset repo | s3://,gs:// path | "json" for local dataset, make sure to fill data_files
--- a/docs/custom_integrations.qmd
+++ b/docs/custom_integrations.qmd
@@ -55,47 +55,3 @@ sections = [
 for section_name, folder_name in sections:
    print(print_section(section_name, folder_name))
 ```
-
-## Adding a new integration
-
-Plugins can be used to customize the behavior of the training pipeline through [hooks](https://en.wikipedia.org/wiki/Hooking). See [`axolotl.integrations.BasePlugin`](https://github.com/axolotl-ai-cloud/axolotl/blob/main/src/axolotl/integrations/base.py) for the possible hooks.
-
-To add a new integration, please follow these steps:
-
-1. Create a new folder in the `src/axolotl/integrations` directory.
-2. Add any relevant files (`LICENSE`, `README.md`, `ACKNOWLEDGEMENTS.md`, etc.) to the new folder.
-3. Add `__init__.py` and `args.py` files to the new folder.
-  - `__init__.py` should import the integration and hook into the appropriate functions.
-  - `args.py` should define the arguments for the integration.
-4. (If applicable) Add CPU tests under `tests/integrations` or GPU tests under `tests/e2e/integrations`.
-
-::: {.callout-tip}
-
-See [src/axolotl/integrations/cut_cross_entropy](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/cut_cross_entropy) for a minimal integration example.
-
-:::
-
-::: {.callout-warning}
-
-If you could not load your integration, please ensure you are pip installing in editable mode.
-
-```bash
-pip install -e .
-```
-
-and correctly spelled the integration name in the config file.
-
-```yaml
-plugins:
-  - axolotl.integrations.your_integration_name.YourIntegrationPlugin
-```
-
-:::
-
-::: {.callout-note}
-
-It is not necessary to place your integration in the `integrations` folder. It can be in any location, so long as it's installed in a package in your python env.
-
-See this repo for an example: [https://github.com/axolotl-ai-cloud/diff-transformer](https://github.com/axolotl-ai-cloud/diff-transformer)
-
-:::
--- a/docs/dataset_preprocessing.qmd
+++ b/docs/dataset_preprocessing.qmd
@@ -6,7 +6,7 @@ description: How datasets are processed
 ## Overview

 Dataset pre-processing is the step where Axolotl takes each dataset you've configured alongside
-the [dataset format](dataset-formats) and prompt strategies to:
+the [dataset format](docs/dataset-formats) and prompt strategies to:

 - parse the dataset based on the *dataset format*
 - transform the dataset to how you would interact with the model based on the *prompt strategy*
--- a/docs/installation.qmd
+++ b/docs/installation.qmd
@@ -79,7 +79,6 @@ For providers supporting Docker:
  - [Latitude.sh](https://latitude.sh/blueprint/989e0e79-3bf6-41ea-a46b-1f246e309d5c)
  - [JarvisLabs.ai](https://jarvislabs.ai/templates/axolotl)
  - [RunPod](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
-  - [Novita](https://novita.ai/gpus-console?templateId=311)

 ### Google Colab {#sec-colab}

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

 [project]
@@ -8,7 +8,6 @@ dynamic = ["version", "dependencies", "optional-dependencies"]
 description = "LLM Trainer"
 readme = "README.md"
 requires-python = ">=3.10"
-# license = "Apache-2.0"

 [project.scripts]
 axolotl = "axolotl.cli.main:main"
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -2,5 +2,3 @@ pre-commit
 black
 mypy
 types-requests
-quartodoc
-jupyter
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,7 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/

 # START section of dependencies that don't install on Darwin/MacOS
-bitsandbytes==0.45.3
+bitsandbytes==0.45.2
 triton>=3.0.0
 mamba-ssm==1.2.0.post1
 flash-attn==2.7.4.post1
@@ -12,12 +12,12 @@ liger-kernel==0.5.3

 packaging==23.2

-peft==0.15.0
+peft==0.14.0
 transformers==4.49.0
-tokenizers>=0.21.1
-accelerate==1.5.2
-datasets==3.4.1
-deepspeed==0.16.4
+tokenizers>=0.21.0
+accelerate==1.3.0
+datasets==3.2.0
+deepspeed==0.16.1
 trl==0.15.1

 optimum==1.16.2
--- a/scripts/chat_datasets.py
+++ b/scripts/chat_datasets.py
@@ -1,7 +1,6 @@
 """
 helper script to parse chat datasets into a usable yaml
 """
-
 import click
 import yaml
 from datasets import load_dataset
--- a/scripts/cutcrossentropy_install.py
+++ b/scripts/cutcrossentropy_install.py
@@ -1,5 +1,4 @@
 """Script to output the correct installation command for cut-cross-entropy."""
-
 import importlib.util
 import sys

--- a/setup.py
+++ b/setup.py
@@ -128,7 +128,7 @@ setup(
            "flash-attn==2.7.4.post1",
        ],
        "deepspeed": [
-            "deepspeed==0.16.4",
+            "deepspeed==0.16.1",
            "deepspeed-kernels",
        ],
        "mamba-ssm": [
--- a/src/axolotl/cli/cloud/init.py
+++ b/src/axolotl/cli/cloud/init.py
@@ -1,7 +1,6 @@
 """
 launch axolotl in supported cloud platforms
 """
-
 from pathlib import Path
 from typing import Union

--- a/src/axolotl/cli/cloud/base.py
+++ b/src/axolotl/cli/cloud/base.py
@@ -1,7 +1,6 @@
 """
 base class for cloud platforms from cli
 """
-
 from abc import ABC, abstractmethod


--- a/src/axolotl/cli/cloud/modal_.py
+++ b/src/axolotl/cli/cloud/modal_.py
@@ -1,7 +1,6 @@
 """
 Modal Cloud support from CLI
 """
-
 import copy
 import json
 import os
--- a/src/axolotl/cli/main.py
+++ b/src/axolotl/cli/main.py
@@ -1,5 +1,4 @@
 """Click CLI definitions for various axolotl commands."""
-
 # pylint: disable=redefined-outer-name

 import logging
@@ -25,7 +24,7 @@ from axolotl.cli.utils import (
 )
 from axolotl.integrations.lm_eval.cli import lm_eval
 from axolotl.utils import set_pytorch_cuda_alloc_conf
-from axolotl.utils.schemas.config import AxolotlInputConfig
+from axolotl.utils.config.models.input.v0_4_1 import AxolotlInputConfig


@click.group()
--- a/src/axolotl/cli/utils.py
+++ b/src/axolotl/cli/utils.py
@@ -5,6 +5,7 @@ import dataclasses
 import hashlib
 import json
 import logging
+import typing
 from functools import wraps
 from pathlib import Path
 from types import NoneType
@@ -23,7 +24,7 @@ configure_logging()
 LOG = logging.getLogger(__name__)


-def strip_optional_type(field_type: type | str | None):
+def strip_optional_type(field_type: type | typing._SpecialForm | None):
    """
    Extracts the non-`None` type from an `Optional` / `Union` type.

--- a/src/axolotl/convert.py
+++ b/src/axolotl/convert.py
@@ -1,5 +1,6 @@
 """Module containing File Reader, File Writer, Json Parser, and Jsonl Serializer classes"""

+
 import json
 import sys

--- a/src/axolotl/core/chat/format/chatml.py
+++ b/src/axolotl/core/chat/format/chatml.py
@@ -1,7 +1,6 @@
 """
 ChatML transformation functions for MessageContents
 """
-
 from typing import Optional

 from ..messages import MessageContents, Messages
--- a/src/axolotl/core/chat/format/llama3x.py
+++ b/src/axolotl/core/chat/format/llama3x.py
@@ -1,7 +1,6 @@
 """
 Llama 3.x chat formatting functions for MessageContents
 """
-
 from typing import Optional

 from ..messages import MessageContents, Messages
--- a/src/axolotl/core/chat/format/shared.py
+++ b/src/axolotl/core/chat/format/shared.py
@@ -1,7 +1,6 @@
 """
 shared functions for format transforms
 """
-
 from axolotl.core.chat.messages import MessageContents, Messages


--- a/src/axolotl/core/chat/messages.py
+++ b/src/axolotl/core/chat/messages.py
@@ -1,7 +1,6 @@
 """
 internal message representations of chat messages
 """
-
 import json
 from enum import Enum
 from typing import Any, Callable, List, Optional, Union
--- a/src/axolotl/core/datasets/chat.py
+++ b/src/axolotl/core/datasets/chat.py
@@ -1,7 +1,6 @@
 """
 chat dataset module
 """
-
 import os
 from typing import Callable, Optional, Union

--- a/src/axolotl/core/datasets/transforms/chat_builder.py
+++ b/src/axolotl/core/datasets/transforms/chat_builder.py
@@ -1,7 +1,6 @@
 """
 This module contains a function that builds a transform that takes a row from the dataset and converts it to a Chat.
 """
-
 from typing import Any, Mapping, Union


--- a/src/axolotl/core/trainer_builder.py
+++ b/src/axolotl/core/trainer_builder.py
@@ -13,7 +13,9 @@
 # limitations under the License.

 # pylint: disable=too-many-lines
-"""Builder for the training args and trainer"""
+"""
+Builder for the training args and trainer
+"""

 import abc
 import importlib
@@ -83,8 +85,8 @@ from axolotl.utils.collators import (
    V2BatchSamplerDataCollatorForSeq2Seq,
 )
 from axolotl.utils.collators.mm_chat import MultiModalChatDataCollator
+from axolotl.utils.config.models.input.v0_4_1 import CustomSupportedOptimizers
 from axolotl.utils.models import ensure_dtype
-from axolotl.utils.schemas.enums import CustomSupportedOptimizers

 try:
    import torch._dynamo  # pylint: disable=ungrouped-imports
@@ -330,9 +332,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
        training_arguments_kwargs = {}

        if self.cfg.include_tokens_per_second is not None:
-            training_arguments_kwargs["include_tokens_per_second"] = (
-                self.cfg.include_tokens_per_second
-            )
+            training_arguments_kwargs[
+                "include_tokens_per_second"
+            ] = self.cfg.include_tokens_per_second

        if self.cfg.bf16 == "full":
            training_arguments_kwargs["bf16_full_eval"] = True
@@ -349,13 +351,13 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
            training_arguments_kwargs["seed"] = self.cfg.seed

        if self.cfg.gradient_checkpointing:
-            training_arguments_kwargs["gradient_checkpointing"] = (
-                self.cfg.gradient_checkpointing
-            )
+            training_arguments_kwargs[
+                "gradient_checkpointing"
+            ] = self.cfg.gradient_checkpointing
            if self.cfg.gradient_checkpointing_kwargs is not None:
-                training_arguments_kwargs["gradient_checkpointing_kwargs"] = (
-                    self.cfg.gradient_checkpointing_kwargs
-                )
+                training_arguments_kwargs[
+                    "gradient_checkpointing_kwargs"
+                ] = self.cfg.gradient_checkpointing_kwargs
        if self.cfg.fsdp:
            training_arguments_kwargs["fsdp"] = self.cfg.fsdp
            if self.cfg.fsdp_config:
@@ -371,9 +373,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
            training_arguments_kwargs["deepspeed"] = self.cfg.deepspeed

        if self.cfg.lr_quadratic_warmup is not None:
-            training_arguments_kwargs["lr_quadratic_warmup"] = (
-                self.cfg.lr_quadratic_warmup
-            )
+            training_arguments_kwargs[
+                "lr_quadratic_warmup"
+            ] = self.cfg.lr_quadratic_warmup

        if self.cfg.adam_beta1:
            training_arguments_kwargs["adam_beta1"] = self.cfg.adam_beta1
@@ -397,28 +399,28 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
            training_arguments_kwargs["save_safetensors"] = self.cfg.save_safetensors

        if self.cfg.dataloader_pin_memory is not None:
-            training_arguments_kwargs["dataloader_pin_memory"] = (
-                self.cfg.dataloader_pin_memory
-            )
+            training_arguments_kwargs[
+                "dataloader_pin_memory"
+            ] = self.cfg.dataloader_pin_memory
        if self.cfg.dataloader_num_workers is not None:
-            training_arguments_kwargs["dataloader_num_workers"] = (
-                self.cfg.dataloader_num_workers
-            )
+            training_arguments_kwargs[
+                "dataloader_num_workers"
+            ] = self.cfg.dataloader_num_workers
        if self.cfg.dataloader_prefetch_factor is not None:
-            training_arguments_kwargs["dataloader_prefetch_factor"] = (
-                self.cfg.dataloader_prefetch_factor
-            )
+            training_arguments_kwargs[
+                "dataloader_prefetch_factor"
+            ] = self.cfg.dataloader_prefetch_factor
        if self.cfg.dataloader_drop_last is not None:
-            training_arguments_kwargs["dataloader_drop_last"] = (
-                self.cfg.dataloader_drop_last
-            )
+            training_arguments_kwargs[
+                "dataloader_drop_last"
+            ] = self.cfg.dataloader_drop_last
        elif self.cfg.sample_packing and self.cfg.eval_sample_packing is False:
            training_arguments_kwargs["dataloader_drop_last"] = True

        if self.cfg.remove_unused_columns is not None:
-            training_arguments_kwargs["remove_unused_columns"] = (
-                self.cfg.remove_unused_columns
-            )
+            training_arguments_kwargs[
+                "remove_unused_columns"
+            ] = self.cfg.remove_unused_columns

        if not self.cfg.test_datasets and self.cfg.val_set_size == 0:
            # no eval set, so don't eval
@@ -450,9 +452,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
        if self.cfg.do_causal_lm_eval:
            training_arguments_kwargs["do_causal_lm_eval"] = self.cfg.do_causal_lm_eval
        if self.cfg.metric_for_best_model:
-            training_arguments_kwargs["metric_for_best_model"] = (
-                self.cfg.metric_for_best_model
-            )
+            training_arguments_kwargs[
+                "metric_for_best_model"
+            ] = self.cfg.metric_for_best_model
        if self.cfg.greater_is_better:
            training_arguments_kwargs["greater_is_better"] = self.cfg.greater_is_better

@@ -465,13 +467,13 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                )
                training_arguments_kwargs["torch_compile"] = self.cfg.torch_compile
                if self.cfg.torch_compile_backend:
-                    training_arguments_kwargs["torch_compile_backend"] = (
-                        self.cfg.torch_compile_backend
-                    )
+                    training_arguments_kwargs[
+                        "torch_compile_backend"
+                    ] = self.cfg.torch_compile_backend
                if self.cfg.torch_compile_mode:
-                    training_arguments_kwargs["torch_compile_mode"] = (
-                        self.cfg.torch_compile_mode
-                    )
+                    training_arguments_kwargs[
+                        "torch_compile_mode"
+                    ] = self.cfg.torch_compile_mode

        # DDP Config
        if self.cfg.ddp_timeout:
@@ -480,32 +482,32 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
        if self.cfg.ddp_bucket_cap_mb:
            training_arguments_kwargs["ddp_bucket_cap_mb"] = self.cfg.ddp_bucket_cap_mb
        if self.cfg.ddp_broadcast_buffers is not None:
-            training_arguments_kwargs["ddp_broadcast_buffers"] = (
-                self.cfg.ddp_broadcast_buffers
-            )
+            training_arguments_kwargs[
+                "ddp_broadcast_buffers"
+            ] = self.cfg.ddp_broadcast_buffers

        # these are all the "standard" kwargs that are def used
        training_arguments_kwargs["max_steps"] = (
            total_num_steps if self.cfg.max_steps else -1
        )
        training_arguments_kwargs["max_seq_length"] = self.cfg.sequence_len
-        training_arguments_kwargs["per_device_train_batch_size"] = (
-            self.cfg.micro_batch_size
-        )
+        training_arguments_kwargs[
+            "per_device_train_batch_size"
+        ] = self.cfg.micro_batch_size
        if self.cfg.eval_batch_size:
-            training_arguments_kwargs["per_device_eval_batch_size"] = (
-                self.cfg.eval_batch_size
-            )
+            training_arguments_kwargs[
+                "per_device_eval_batch_size"
+            ] = self.cfg.eval_batch_size
        if self.cfg.auto_find_batch_size is not None:
-            training_arguments_kwargs["auto_find_batch_size"] = (
-                self.cfg.auto_find_batch_size
-            )
-        training_arguments_kwargs["gradient_accumulation_steps"] = (
-            self.cfg.gradient_accumulation_steps
-        )
-        training_arguments_kwargs["eval_accumulation_steps"] = (
-            self.cfg.gradient_accumulation_steps
-        )
+            training_arguments_kwargs[
+                "auto_find_batch_size"
+            ] = self.cfg.auto_find_batch_size
+        training_arguments_kwargs[
+            "gradient_accumulation_steps"
+        ] = self.cfg.gradient_accumulation_steps
+        training_arguments_kwargs[
+            "eval_accumulation_steps"
+        ] = self.cfg.gradient_accumulation_steps
        training_arguments_kwargs["num_train_epochs"] = self.cfg.num_epochs
        training_arguments_kwargs["learning_rate"] = self.cfg.learning_rate
        training_arguments_kwargs["output_dir"] = self.cfg.output_dir
@@ -552,9 +554,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):

        if self.cfg.lr_scheduler in ["one_cycle", "rex", "log_sweep"]:
            training_arguments_kwargs["lr_scheduler_type"] = "cosine"
-            training_arguments_kwargs["alternate_lr_scheduler_type"] = (
-                self.cfg.lr_scheduler
-            )
+            training_arguments_kwargs[
+                "alternate_lr_scheduler_type"
+            ] = self.cfg.lr_scheduler
        else:
            training_arguments_kwargs["lr_scheduler_type"] = (
                self.cfg.lr_scheduler if self.cfg.lr_scheduler else "cosine"
@@ -563,9 +565,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
            self.cfg.lr_scheduler_kwargs if self.cfg.lr_scheduler_kwargs else {}
        )
        training_arguments_kwargs["cosine_min_lr_ratio"] = self.cfg.cosine_min_lr_ratio
-        training_arguments_kwargs["cosine_constant_lr_ratio"] = (
-            self.cfg.cosine_constant_lr_ratio
-        )
+        training_arguments_kwargs[
+            "cosine_constant_lr_ratio"
+        ] = self.cfg.cosine_constant_lr_ratio
        training_arguments_kwargs["weight_decay"] = (
            self.cfg.weight_decay if self.cfg.weight_decay is not None else 0.0
        )
@@ -578,40 +580,40 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
            self.cfg.eval_sample_packing
        )
        if self.cfg.sample_packing_bin_size is not None:
-            training_arguments_kwargs["sample_packing_bin_size"] = (
-                self.cfg.sample_packing_bin_size
-            )
+            training_arguments_kwargs[
+                "sample_packing_bin_size"
+            ] = self.cfg.sample_packing_bin_size
        if self.cfg.sample_packing_group_size is not None:
-            training_arguments_kwargs["sample_packing_group_size"] = (
-                self.cfg.sample_packing_group_size
-            )
+            training_arguments_kwargs[
+                "sample_packing_group_size"
+            ] = self.cfg.sample_packing_group_size
        if self.cfg.sample_packing_eff_est:
-            training_arguments_kwargs["sample_packing_efficiency"] = (
-                self.cfg.sample_packing_eff_est
-            )
+            training_arguments_kwargs[
+                "sample_packing_efficiency"
+            ] = self.cfg.sample_packing_eff_est

        if self.cfg.relora_steps:
            training_arguments_kwargs["relora_steps"] = self.cfg.relora_steps
-            training_arguments_kwargs["relora_warmup_steps"] = (
-                self.cfg.relora_warmup_steps
-            )
+            training_arguments_kwargs[
+                "relora_warmup_steps"
+            ] = self.cfg.relora_warmup_steps
            if self.cfg.relora_anneal_steps:
-                training_arguments_kwargs["relora_anneal_steps"] = (
-                    self.cfg.relora_anneal_steps
-                )
+                training_arguments_kwargs[
+                    "relora_anneal_steps"
+                ] = self.cfg.relora_anneal_steps
            if self.cfg.relora_prune_ratio:
-                training_arguments_kwargs["relora_prune_ratio"] = (
-                    self.cfg.relora_prune_ratio
-                )
+                training_arguments_kwargs[
+                    "relora_prune_ratio"
+                ] = self.cfg.relora_prune_ratio

        if self.cfg.lisa_step_interval and self.cfg.lisa_n_layers:
            training_arguments_kwargs["lisa_n_layers"] = self.cfg.lisa_n_layers
-            training_arguments_kwargs["lisa_step_interval"] = (
-                self.cfg.lisa_step_interval
-            )
-            training_arguments_kwargs["lisa_layers_attribute"] = (
-                self.cfg.lisa_layers_attribute
-            )
+            training_arguments_kwargs[
+                "lisa_step_interval"
+            ] = self.cfg.lisa_step_interval
+            training_arguments_kwargs[
+                "lisa_layers_attribute"
+            ] = self.cfg.lisa_layers_attribute

        training_arguments_kwargs = self.hook_pre_create_training_args(
            training_arguments_kwargs
@@ -625,9 +627,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
            )

        if self.cfg.neftune_noise_alpha is not None:
-            training_arguments_kwargs["neftune_noise_alpha"] = (
-                self.cfg.neftune_noise_alpha
-            )
+            training_arguments_kwargs[
+                "neftune_noise_alpha"
+            ] = self.cfg.neftune_noise_alpha

        trainer_kwargs = {}

@@ -729,23 +731,23 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                importlib.import_module("torchdistx")

        if self.cfg.optim_target_modules:
-            training_arguments_kwargs["optim_target_modules"] = (
-                self.cfg.optim_target_modules
-            )
+            training_arguments_kwargs[
+                "optim_target_modules"
+            ] = self.cfg.optim_target_modules

        training_arguments_kwargs["embedding_lr"] = self.cfg.embedding_lr
        training_arguments_kwargs["embedding_lr_scale"] = self.cfg.embedding_lr_scale

        training_arguments_kwargs["loraplus_lr_ratio"] = self.cfg.loraplus_lr_ratio
-        training_arguments_kwargs["loraplus_lr_embedding"] = (
-            self.cfg.loraplus_lr_embedding
-        )
+        training_arguments_kwargs[
+            "loraplus_lr_embedding"
+        ] = self.cfg.loraplus_lr_embedding
        training_arguments_kwargs["lr_groups"] = self.cfg.lr_groups

        if self.cfg.accelerator_config:
-            training_arguments_kwargs["accelerator_config"] = (
-                self.cfg.accelerator_config
-            )
+            training_arguments_kwargs[
+                "accelerator_config"
+            ] = self.cfg.accelerator_config

        if self.cfg.kd_ce_alpha is not None:
            training_arguments_kwargs["kd_ce_alpha"] = self.cfg.kd_ce_alpha
@@ -754,13 +756,13 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
        if self.cfg.kd_temperature is not None:
            training_arguments_kwargs["kd_temperature"] = self.cfg.kd_temperature
        if self.cfg.kd_zscore_base_temp is not None:
-            training_arguments_kwargs["kd_zscore_base_temp"] = (
-                self.cfg.kd_zscore_base_temp
-            )
+            training_arguments_kwargs[
+                "kd_zscore_base_temp"
+            ] = self.cfg.kd_zscore_base_temp
        if self.cfg.kd_top_k_before_softmax is not None:
-            training_arguments_kwargs["kd_top_k_before_softmax"] = (
-                self.cfg.kd_top_k_before_softmax
-            )
+            training_arguments_kwargs[
+                "kd_top_k_before_softmax"
+            ] = self.cfg.kd_top_k_before_softmax

        if self.cfg.reward_model:
            training_args_cls = AxolotlRewardConfig
@@ -970,32 +972,32 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
            self.cfg.lr_scheduler_kwargs if self.cfg.lr_scheduler_kwargs else {}
        )
        if self.cfg.remove_unused_columns is not None:
-            training_args_kwargs["remove_unused_columns"] = (
-                self.cfg.remove_unused_columns
-            )
+            training_args_kwargs[
+                "remove_unused_columns"
+            ] = self.cfg.remove_unused_columns
        else:
            training_args_kwargs["remove_unused_columns"] = False

        if self.cfg.dataloader_pin_memory is not None:
-            training_args_kwargs["dataloader_pin_memory"] = (
-                self.cfg.dataloader_pin_memory
-            )
+            training_args_kwargs[
+                "dataloader_pin_memory"
+            ] = self.cfg.dataloader_pin_memory
        if self.cfg.dataloader_num_workers is not None:
-            training_args_kwargs["dataloader_num_workers"] = (
-                self.cfg.dataloader_num_workers
-            )
+            training_args_kwargs[
+                "dataloader_num_workers"
+            ] = self.cfg.dataloader_num_workers
        if self.cfg.dataloader_prefetch_factor is not None:
-            training_args_kwargs["dataloader_prefetch_factor"] = (
-                self.cfg.dataloader_prefetch_factor
-            )
+            training_args_kwargs[
+                "dataloader_prefetch_factor"
+            ] = self.cfg.dataloader_prefetch_factor
        if self.cfg.gradient_checkpointing:
-            training_args_kwargs["gradient_checkpointing"] = (
-                self.cfg.gradient_checkpointing
-            )
+            training_args_kwargs[
+                "gradient_checkpointing"
+            ] = self.cfg.gradient_checkpointing
            if self.cfg.gradient_checkpointing_kwargs is not None:
-                training_args_kwargs["gradient_checkpointing_kwargs"] = (
-                    self.cfg.gradient_checkpointing_kwargs
-                )
+                training_args_kwargs[
+                    "gradient_checkpointing_kwargs"
+                ] = self.cfg.gradient_checkpointing_kwargs
            else:
                training_args_kwargs["gradient_checkpointing_kwargs"] = {
                    "use_reentrant": False
@@ -1069,9 +1071,9 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
            if self.cfg.dpo_use_weighting is not None:
                training_args_kwargs["use_weighting"] = self.cfg.dpo_use_weighting
            if self.cfg.dpo_use_logits_to_keep is not None:
-                training_args_kwargs["use_logits_to_keep"] = (
-                    self.cfg.dpo_use_logits_to_keep
-                )
+                training_args_kwargs[
+                    "use_logits_to_keep"
+                ] = self.cfg.dpo_use_logits_to_keep

        for blocklist_key in blocklist_args_kwargs:
            if blocklist_key in training_args_kwargs:
@@ -1106,9 +1108,9 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
        if self.cfg.adapter and self.peft_config:
            dpo_trainer_kwargs["peft_config"] = self.peft_config
        if self.cfg.precompute_ref_log_probs is not None:
-            dpo_trainer_kwargs["precompute_ref_log_probs"] = (
-                self.cfg.precompute_ref_log_probs
-            )
+            dpo_trainer_kwargs[
+                "precompute_ref_log_probs"
+            ] = self.cfg.precompute_ref_log_probs
        if self.cfg.rl == "grpo":
            trainer_cls = GRPOStrategy.get_trainer_class()
            trainer_cls_args = [self.model]
--- a/src/axolotl/core/trainers/base.py
+++ b/src/axolotl/core/trainers/base.py
@@ -462,9 +462,9 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, Trainer):
                "pin_memory": self.args.dataloader_pin_memory,
            }
            if self.args.dataloader_prefetch_factor:
-                dataloader_params["prefetch_factor"] = (
-                    self.args.dataloader_prefetch_factor
-                )
+                dataloader_params[
+                    "prefetch_factor"
+                ] = self.args.dataloader_prefetch_factor

            sampler = self._get_train_sampler()
            if isinstance(sampler, BatchSampler):
@@ -509,9 +509,9 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, Trainer):
                "pin_memory": self.args.dataloader_pin_memory,
            }
            if self.args.dataloader_prefetch_factor:
-                dataloader_params["prefetch_factor"] = (
-                    self.args.dataloader_prefetch_factor
-                )
+                dataloader_params[
+                    "prefetch_factor"
+                ] = self.args.dataloader_prefetch_factor

            if isinstance(eval_sampler, BatchSampler):
                dataloader_params["batch_sampler"] = eval_sampler
--- a/src/axolotl/core/trainers/dpo/init.py
+++ b/src/axolotl/core/trainers/dpo/init.py
@@ -1,7 +1,6 @@
 """
 DPO Specific Strategy for training
 """
-
 from axolotl.core.trainers.dpo.trainer import AxolotlDPOTrainer


--- a/src/axolotl/core/trainers/dpo/args.py
+++ b/src/axolotl/core/trainers/dpo/args.py
@@ -1,7 +1,6 @@
 """
 Axolotl specific DPO args
 """
-
 from dataclasses import dataclass

 from trl import DPOConfig
--- a/src/axolotl/core/trainers/dpo/trainer.py
+++ b/src/axolotl/core/trainers/dpo/trainer.py
@@ -1,7 +1,6 @@
 """
 DPO trainer for axolotl
 """
-
 import gc
 from functools import wraps
 from typing import Any, Dict, Union
--- a/src/axolotl/core/trainers/grpo/init.py
+++ b/src/axolotl/core/trainers/grpo/init.py
@@ -9,7 +9,7 @@ import logging
 from trl.trainer.grpo_trainer import RewardFunc

 from axolotl.core.trainers.grpo.trainer import AxolotlGRPOTrainer
-from axolotl.utils.schemas.trl import TRLConfig
+from axolotl.utils.config.models.input.v0_4_1.trl import TRLConfig

 LOG = logging.getLogger("axolotl")

@@ -45,9 +45,9 @@ class GRPOStrategy:
            )

            if trl.vllm_gpu_memory_utilization:
-                grpo_args_kwargs["vllm_gpu_memory_utilization"] = (
-                    trl.vllm_gpu_memory_utilization
-                )
+                grpo_args_kwargs[
+                    "vllm_gpu_memory_utilization"
+                ] = trl.vllm_gpu_memory_utilization

            if trl.vllm_max_model_len:
                grpo_args_kwargs["vllm_max_model_len"] = trl.vllm_max_model_len
@@ -86,9 +86,9 @@ class GRPOStrategy:
    def set_trainer_kwargs(cls, cfg):
        trainer_kwargs = {}
        if cfg.trl and cfg.trl.reward_processing_classes:
-            trainer_kwargs["reward_processing_classes"] = (
-                cfg.trl.reward_processing_classes
-            )
+            trainer_kwargs[
+                "reward_processing_classes"
+            ] = cfg.trl.reward_processing_classes
        return trainer_kwargs

    @classmethod
--- a/src/axolotl/core/trainers/grpo/args.py
+++ b/src/axolotl/core/trainers/grpo/args.py
@@ -1,7 +1,6 @@
 """
 Axolotl Specific Training Args
 """
-
 from dataclasses import dataclass

 from trl import GRPOConfig
--- a/src/axolotl/core/trainers/grpo/trainer.py
+++ b/src/axolotl/core/trainers/grpo/trainer.py
@@ -1,7 +1,6 @@
 """
 Axolotl GRPO trainer
 """
-
 from accelerate.utils import is_peft_model
 from accelerate.utils.other import is_compiled_module
 from transformers import PreTrainedModel
--- a/src/axolotl/core/trainers/trl.py
+++ b/src/axolotl/core/trainers/trl.py
@@ -1,7 +1,6 @@
 """
 module for TRL PPO training
 """
-
 import torch
 from tqdm import tqdm
 from trl import PPOTrainer
--- a/src/axolotl/core/training_args.py
+++ b/src/axolotl/core/training_args.py
@@ -1,7 +1,6 @@
 """
 extra axolotl specific training args
 """
-
 from dataclasses import dataclass, field
 from typing import Optional

--- a/src/axolotl/evaluate.py
+++ b/src/axolotl/evaluate.py
@@ -8,8 +8,6 @@ from typing import Dict, Optional

 import torch
 from accelerate.logging import get_logger
-from datasets import Dataset
-from transformers.trainer import Trainer

 from axolotl.logging_config import configure_logging
 from axolotl.train import TrainDatasetMeta
@@ -27,18 +25,18 @@ LOG = get_logger("axolotl.evaluate")


 def evaluate_dataset(
-    trainer: Trainer, dataset: Dataset, dataset_type: str, flash_optimum: bool = False
+    trainer, dataset, dataset_type: str, flash_optimum: bool = False
 ) -> Optional[Dict[str, float]]:
-    """Helper function to evaluate a single dataset.
+    """Helper function to evaluate a single dataset safely.

    Args:
-        trainer: The trainer instance.
-        dataset: Dataset to evaluate.
-        dataset_type: Type of dataset ('train' or 'eval').
-        flash_optimum: Whether to use flash optimum.
+        trainer: The trainer instance
+        dataset: Dataset to evaluate
+        dataset_type: Type of dataset ('train' or 'eval')
+        flash_optimum: Whether to use flash optimum

    Returns:
-        Dictionary of metrics or None if dataset is None.
+        Dictionary of metrics or None if dataset is None
    """
    if dataset is None:
        return None
@@ -65,14 +63,17 @@ def evaluate_dataset(

 def evaluate(*, cfg: DictDefault, dataset_meta: TrainDatasetMeta) -> Dict[str, float]:
    """
-    Evaluate a model on training and validation datasets.
+    Evaluate a model on training and validation datasets

    Args:
        cfg: Dictionary mapping `axolotl` config keys to values.
        dataset_meta: Dataset metadata containing training and evaluation datasets.

    Returns:
-        Dictionary mapping metric names to their values.
+        Tuple containing:
+        - The model (either PeftModel or PreTrainedModel)
+        - The tokenizer
+        - Dictionary of evaluation metrics
    """
    # pylint: disable=duplicate-code
    # Enable expandable segments for cuda allocation to improve VRAM usage
--- a/src/axolotl/integrations/config.py
+++ b/src/axolotl/integrations/config.py
@@ -11,17 +11,19 @@
 # the License.

 """
-Module to handle merging the plugins' input arguments with the base configurations.
+module to handle merging the plugins' input arguments with the base configurations.

-This was moved here to prevent circular imports.
+this was moved here to prevent circular imports
 """

 from typing import Any, Dict, List

-from axolotl.utils.schemas.config import (
+from axolotl.utils.config.models.input.v0_4_1 import (
    AxolotlConfigWCapabilities as AxolotlConfigWCapabilitiesBase,
 )
-from axolotl.utils.schemas.config import AxolotlInputConfig as AxolotlInputConfigBase
+from axolotl.utils.config.models.input.v0_4_1 import (
+    AxolotlInputConfig as AxolotlInputConfigBase,
+)


 def merge_input_args():
--- a/src/axolotl/integrations/grokfast/init.py
+++ b/src/axolotl/integrations/grokfast/init.py
@@ -1,7 +1,6 @@
 """
 Grokfast plugin for Axolotl
 """
-
 import logging

 from transformers.trainer_callback import TrainerCallback
--- a/src/axolotl/integrations/grokfast/args.py
+++ b/src/axolotl/integrations/grokfast/args.py
@@ -1,7 +1,6 @@
 """
 config args for grokfast plugin
 """
-
 from typing import Optional

 from pydantic import BaseModel
--- a/src/axolotl/integrations/kd/args.py
+++ b/src/axolotl/integrations/kd/args.py
@@ -26,12 +26,12 @@ class KDArgs(BaseModel):
    """

    kd_trainer: Optional[bool] = None  # whether to use KD trainer
-    kd_ce_alpha: Optional[float] = (
-        None  # loss coefficient for cross-entropy loss during KD
-    )
+    kd_ce_alpha: Optional[
+        float
+    ] = None  # loss coefficient for cross-entropy loss during KD
    kd_alpha: Optional[float] = None  # loss coefficient for KD loss
    kd_temperature: Optional[float] = None  # temperature for sampling during KD
    kd_zscore_base_temp: Optional[float] = None  # base temperature for zscore scaling
-    kd_top_k_before_softmax: Optional[bool] = (
-        None  # whether to sample top k before softmax during KD
-    )
+    kd_top_k_before_softmax: Optional[
+        bool
+    ] = None  # whether to sample top k before softmax during KD
--- a/src/axolotl/integrations/liger/init.py
+++ b/src/axolotl/integrations/liger/init.py
@@ -55,9 +55,9 @@ class LigerPlugin(BasePlugin):
            if "cross_entropy" in liger_fn_sig.parameters:
                kwargs["cross_entropy"] = cfg.liger_cross_entropy
            if "fused_linear_cross_entropy" in liger_fn_sig.parameters:
-                kwargs["fused_linear_cross_entropy"] = (
-                    cfg.liger_fused_linear_cross_entropy
-                )
+                kwargs[
+                    "fused_linear_cross_entropy"
+                ] = cfg.liger_fused_linear_cross_entropy
            if "rms_norm" in liger_fn_sig.parameters:
                kwargs["rms_norm"] = cfg.liger_rms_norm
            if "layer_norm" in liger_fn_sig.parameters:
--- a/src/axolotl/integrations/liger/models/deepseekv2.py
+++ b/src/axolotl/integrations/liger/models/deepseekv2.py
@@ -1,7 +1,6 @@
 """
 DeepseekV2 model with LigerFusedLinearCrossEntropyLoss
 """
-
 # pylint: disable=duplicate-code

 from typing import List, Optional, Tuple, Union
--- a/src/axolotl/integrations/liger/models/jamba.py
+++ b/src/axolotl/integrations/liger/models/jamba.py
@@ -1,7 +1,6 @@
 """
 Jamba model with LigerFusedLinearCrossEntropyLoss
 """
-
 # pylint: disable=duplicate-code

 from typing import Optional, Tuple, Union
--- a/src/axolotl/integrations/lm_eval/init.py
+++ b/src/axolotl/integrations/lm_eval/init.py
@@ -1,7 +1,6 @@
 """
 Module for the Plugin for LM Eval Harness
 """
-
 import subprocess  # nosec

 from axolotl.integrations.base import BasePlugin
--- a/src/axolotl/integrations/lm_eval/args.py
+++ b/src/axolotl/integrations/lm_eval/args.py
@@ -1,7 +1,6 @@
 """
 Module for handling lm eval harness input arguments.
 """
-
 from typing import List, Optional

 from pydantic import BaseModel
--- a/src/axolotl/integrations/lm_eval/cli.py
+++ b/src/axolotl/integrations/lm_eval/cli.py
@@ -1,7 +1,6 @@
 """
 axolotl CLI for running lm_eval tasks
 """
-
 import subprocess  # nosec
 from collections import defaultdict
 from datetime import datetime
--- a/src/axolotl/kernels/geglu.py
+++ b/src/axolotl/kernels/geglu.py
@@ -5,7 +5,6 @@ See "GLU Variants Improve Transformer" (https://arxiv.org/abs/2002.05202).

 Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
 """
-
 # pylint: disable=invalid-name,unnecessary-lambda-assignment,duplicate-code

 import torch
--- a/src/axolotl/kernels/lora.py
+++ b/src/axolotl/kernels/lora.py
@@ -6,7 +6,6 @@ See "LoRA: Low-Rank Adaptation of Large Language Models"

 Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
 """
-
 # pylint: disable=invalid-name

 from typing import Callable
--- a/src/axolotl/kernels/quantize.py
+++ b/src/axolotl/kernels/quantize.py
@@ -1,5 +1,4 @@
 """Dequantization utilities for `bitsandbytes` integration."""
-
 # pylint: disable=invalid-name,global-statement

 import ctypes
--- a/src/axolotl/kernels/swiglu.py
+++ b/src/axolotl/kernels/swiglu.py
@@ -5,7 +5,6 @@ See "GLU Variants Improve Transformer" (https://arxiv.org/abs/2002.05202).

 Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
 """
-
 import torch
 import triton
 import triton.language as tl
--- a/src/axolotl/models/mamba/configuration_mamba.py
+++ b/src/axolotl/models/mamba/configuration_mamba.py
@@ -1,7 +1,6 @@
 """
 HF Transformers MambaConfig
 """
-
 from transformers import PretrainedConfig


--- a/src/axolotl/monkeypatch/attention/mllama.py
+++ b/src/axolotl/monkeypatch/attention/mllama.py
@@ -1,7 +1,6 @@
 """
 Monkeypatch for Vision Llama for FA2 support
 """
-
 # pylint: disable=duplicate-code

 from typing import Optional, Tuple
@@ -221,10 +220,10 @@ def patch_mllama():
        True
    )
    MLLAMA_TEXT_ATTENTION_CLASSES["flash_attention_2"] = MllamaTextSelfFlashAttention2
-    MLLAMA_TEXT_CROSS_ATTENTION_CLASSES["flash_attention_2"] = (
-        MllamaTextCrossFlashAttention2
-    )
+    MLLAMA_TEXT_CROSS_ATTENTION_CLASSES[
+        "flash_attention_2"
+    ] = MllamaTextCrossFlashAttention2
    # fallback to SDPA
-    MLLAMA_VISION_ATTENTION_CLASSES["flash_attention_2"] = (
-        MLLAMA_VISION_ATTENTION_CLASSES["sdpa"]
-    )
+    MLLAMA_VISION_ATTENTION_CLASSES[
+        "flash_attention_2"
+    ] = MLLAMA_VISION_ATTENTION_CLASSES["sdpa"]
--- a/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
+++ b/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
@@ -1,5 +1,4 @@
 """monkey patches for the dataset fetcher to handle batches of packed indexes"""
-
 # pylint: disable=protected-access

 import torch
--- a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -12,9 +12,7 @@ import transformers
 from einops import rearrange
 from flash_attn.bert_padding import pad_input, unpad_input
 from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.models.llama.modeling_llama import (
-    LlamaAttention,
-)
+from transformers.models.llama.modeling_llama import LlamaAttention
 from transformers.models.llama.modeling_llama import (
    LlamaDecoderLayer as OriginalLlamaDecoderLayer,
 )
@@ -492,11 +490,9 @@ def flashattn_forward(
            # We have disabled _prepare_decoder_attention_mask in LlamaModel
            # the attention_mask should be the same as the key_padding_mask
            key_padding_mask=attention_mask,
-            query_padding_mask=(
-                attention_mask[:, -query_states.size(1) :]
-                if attention_mask is not None
-                else None
-            ),
+            query_padding_mask=attention_mask[:, -query_states.size(1) :]
+            if attention_mask is not None
+            else None,
        )
        output_unpad = flash_attn_varlen_qkvpacked_func(
            qkv_unpad,
@@ -535,11 +531,9 @@ def flashattn_forward(
                value_states,
                kvpacked=True,
                key_padding_mask=attention_mask,
-                query_padding_mask=(
-                    attention_mask[:, -query_states.size(1) :]
-                    if attention_mask is not None
-                    else None
-                ),
+                query_padding_mask=attention_mask[:, -query_states.size(1) :]
+                if attention_mask is not None
+                else None,
            )
            if q_unpad.dtype != kv_unpad.dtype:
                kv_unpad = kv_unpad.to(q_unpad.dtype)
--- a/src/axolotl/monkeypatch/llama_expand_mask.py
+++ b/src/axolotl/monkeypatch/llama_expand_mask.py
@@ -1,7 +1,6 @@
 """
 expands the binary attention mask per 3.2.2 of https://arxiv.org/pdf/2107.02027.pdf
 """
-
 from typing import Optional

 import torch
--- a/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
@@ -1,5 +1,4 @@
 """Flash attention monkey patch for mistral model"""
-
 # pylint: disable=duplicate-code

 import logging
@@ -22,10 +21,7 @@ from transformers.models.mistral.modeling_mistral import (
 from transformers.models.mistral.modeling_mistral import (
    MistralDecoderLayer as OriginalMistralDecoderLayer,
 )
-from transformers.models.mistral.modeling_mistral import (
-    apply_rotary_pos_emb,
-    repeat_kv,
-)
+from transformers.models.mistral.modeling_mistral import apply_rotary_pos_emb, repeat_kv

 from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids

@@ -247,11 +243,9 @@ def flashattn_forward(
            # We have disabled _prepare_decoder_attention_mask in LlamaModel
            # the attention_mask should be the same as the key_padding_mask
            key_padding_mask=attention_mask,
-            query_padding_mask=(
-                attention_mask[:, -query_states.size(1) :]
-                if attention_mask is not None
-                else None
-            ),
+            query_padding_mask=attention_mask[:, -query_states.size(1) :]
+            if attention_mask is not None
+            else None,
        )
        output_unpad = flash_attn_varlen_qkvpacked_func(
            qkv_unpad,
@@ -292,11 +286,9 @@ def flashattn_forward(
                value_states,
                kvpacked=True,
                key_padding_mask=attention_mask,
-                query_padding_mask=(
-                    attention_mask[:, -query_states.size(1) :]
-                    if attention_mask is not None
-                    else None
-                ),
+                query_padding_mask=attention_mask[:, -query_states.size(1) :]
+                if attention_mask is not None
+                else None,
            )
            if q_unpad.dtype != kv_unpad.dtype:
                kv_unpad = kv_unpad.to(q_unpad.dtype)
--- a/src/axolotl/monkeypatch/mixtral/init.py
+++ b/src/axolotl/monkeypatch/mixtral/init.py
@@ -1,7 +1,6 @@
 """
 Patches to support multipack for mixtral
 """
-
 import torch


--- a/src/axolotl/monkeypatch/relora.py
+++ b/src/axolotl/monkeypatch/relora.py
@@ -1,5 +1,4 @@
 """Implements the ReLoRA training procedure from https://arxiv.org/abs/2307.05695, minus the initial full fine-tune."""
-
 import glob
 import json
 import logging
@@ -412,10 +411,7 @@ def merge_and_save(
        if shard_path.endswith(".safetensors"):
            in_tensors = st.load_file(str(Path(model_src) / shard_path))
        else:
-            in_tensors = torch.load(
-                Path(model_src) / shard_path,
-                weights_only=True,  # to prevent arbitrary code execution
-            )
+            in_tensors = torch.load(Path(model_src) / shard_path)
            if "state_dict" in in_tensors:
                in_tensors = in_tensors["state_dict"]

--- a/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
@@ -17,7 +17,7 @@
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
 # pylint: disable=duplicate-code
-"""PyTorch StableLM Epoch model."""
+""" PyTorch StableLM Epoch model. """
 import importlib
 import math
 from typing import Optional, Tuple, Union
--- a/src/axolotl/monkeypatch/trainer_fsdp_optim.py
+++ b/src/axolotl/monkeypatch/trainer_fsdp_optim.py
@@ -1,7 +1,6 @@
 """
 fix for FSDP optimizer save in trainer w 4.47.0
 """
-
 import inspect
 import logging

--- a/src/axolotl/monkeypatch/utils.py
+++ b/src/axolotl/monkeypatch/utils.py
@@ -1,7 +1,6 @@
 """
 Shared utils for the monkeypatches
 """
-
 import re
 from typing import Optional, Tuple

--- a/src/axolotl/monkeypatch/xformers_/init.py
+++ b/src/axolotl/monkeypatch/xformers_/init.py
@@ -1,7 +1,6 @@
 """
 Fused MLP layer for incrementally improved training efficiency
 """
-
 import torch
 from transformers.models.llama.modeling_llama import LlamaMLP
 from xformers.ops import SwiGLU
--- a/src/axolotl/prompt_strategies/alpaca_w_system.py
+++ b/src/axolotl/prompt_strategies/alpaca_w_system.py
@@ -1,7 +1,6 @@
 """
 Prompt strategies loader for alpaca instruction datasets with system prompts
 """
-
 from typing import Generator, Tuple, Union

 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
--- a/src/axolotl/prompt_strategies/chat_template.py
+++ b/src/axolotl/prompt_strategies/chat_template.py
@@ -13,7 +13,7 @@ from axolotl.prompt_strategies.jinja_template_analyzer import JinjaTemplateAnaly
 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
 from axolotl.prompters import IGNORE_TOKEN_ID, Prompter
 from axolotl.utils.chat_templates import get_chat_template_from_config
-from axolotl.utils.schemas.datasets import DatasetConfig
+from axolotl.utils.config.models.input.v0_4_1 import DatasetConfig

 # Configure the logger
 LOG = logging.getLogger("axolotl")
--- a/src/axolotl/prompt_strategies/completion.py
+++ b/src/axolotl/prompt_strategies/completion.py
@@ -1,7 +1,6 @@
 """
 Basic completion text
 """
-
 from collections import defaultdict
 from typing import Any, Dict, Generator, Optional, Tuple

--- a/src/axolotl/prompt_strategies/context_qa.py
+++ b/src/axolotl/prompt_strategies/context_qa.py
@@ -1,5 +1,4 @@
 """Module containing the classes for Context QA Prompt Tokenization Strategies"""
-
 from typing import Tuple

 from axolotl.prompt_tokenizers import InstructionPromptTokenizingStrategy
--- a/src/axolotl/prompt_strategies/dpo/init.py
+++ b/src/axolotl/prompt_strategies/dpo/init.py
@@ -1,7 +1,6 @@
 """
 module for DPO style dataset transform strategies
 """
-
 from functools import partial

 from ..base import load as load_base
--- a/src/axolotl/prompt_strategies/dpo/chat_template.py
+++ b/src/axolotl/prompt_strategies/dpo/chat_template.py
@@ -3,7 +3,7 @@ DPO prompt strategies for using tokenizer chat templates.
 """

 from axolotl.utils.chat_templates import extract_chat_template_args, get_chat_template
-from axolotl.utils.schemas.utils import handle_legacy_message_fields_logic
+from axolotl.utils.config.models.input.v0_4_1 import handle_legacy_message_fields_logic


 def default(
--- a/src/axolotl/prompt_strategies/dpo/chatml.py
+++ b/src/axolotl/prompt_strategies/dpo/chatml.py
@@ -33,9 +33,9 @@ def default(
                f"<|im_start|>user\n{sample[prompt_key]}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample[prompt_key]}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample[prompt_key]}<|im_end|>\n<|im_start|>assistant\n"
        sample["chosen"] = f"{sample[chosen_key]}<|im_end|>"
        sample["rejected"] = f"{sample[rejected_key]}<|im_end|>"
        return sample
@@ -52,9 +52,9 @@ def argilla_chat(
    """

    def transform_fn(sample):
-        sample["prompt"] = (
-            f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
-        )
+        sample[
+            "prompt"
+        ] = f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
        sample["chosen"] = f"{sample['chosen'][1]['content']}<|im_end|>"
        sample["rejected"] = f"{sample['rejected'][1]['content']}<|im_end|>"
        return sample
@@ -78,9 +78,9 @@ def icr(
                f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
        sample["chosen"] = f"{sample['chosen']}<|im_end|>"
        sample["rejected"] = f"{sample['rejected']}<|im_end|>"
        return sample
@@ -100,9 +100,9 @@ def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
        sample["chosen"] = f"{sample['chosen']}<|im_end|>"
        sample["rejected"] = f"{sample['rejected']}<|im_end|>"
        return sample
@@ -120,9 +120,9 @@ def prompt_pairs(
                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
        sample["chosen"] = f"{sample['chosen']}<|im_end|>"
        sample["rejected"] = f"{sample['rejected']}<|im_end|>"
        return sample
@@ -142,9 +142,9 @@ def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
        sample["chosen"] = f"{sample['chosen'][1]['content']}<|im_end|>"
        sample["rejected"] = f"{sample['rejected'][1]['content']}<|im_end|>"
        return sample
--- a/src/axolotl/prompt_strategies/dpo/llama3.py
+++ b/src/axolotl/prompt_strategies/dpo/llama3.py
@@ -34,9 +34,9 @@ def default(
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample[prompt_key]}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample[prompt_key]}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample[prompt_key]}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["chosen"] = f"{sample[chosen_key]}<|eot_id|>"
        sample["rejected"] = f"{sample[rejected_key]}<|eot_id|>"
        return sample
@@ -53,9 +53,9 @@ def argilla_chat(
    """

    def transform_fn(sample):
-        sample["prompt"] = (
-            f"<|start_header_id|>user<|end_header_id|>\n\n{sample['chosen'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-        )
+        sample[
+            "prompt"
+        ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['chosen'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["chosen"] = f"{sample['chosen'][1]['content']}<|eot_id|>"
        sample["rejected"] = f"{sample['rejected'][1]['content']}<|eot_id|>"
        return sample
@@ -79,9 +79,9 @@ def icr(
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
        sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
        return sample
@@ -101,9 +101,9 @@ def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
        sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
        return sample
@@ -121,9 +121,9 @@ def prompt_pairs(
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
        sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
        return sample
@@ -143,9 +143,9 @@ def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["chosen"] = f"{sample['chosen'][1]['content']}<|eot_id|>"
        sample["rejected"] = f"{sample['rejected'][1]['content']}<|eot_id|>"
        return sample
--- a/src/axolotl/prompt_strategies/input_output.py
+++ b/src/axolotl/prompt_strategies/input_output.py
@@ -1,5 +1,4 @@
 """Module for plain input/output prompt pairs"""
-
 from typing import Generator, Tuple

 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
--- a/src/axolotl/prompt_strategies/jinja_template_analyzer.py
+++ b/src/axolotl/prompt_strategies/jinja_template_analyzer.py
@@ -1,5 +1,4 @@
 """Module for inspect jinja templates for the variables they use"""
-
 from typing import Dict, Optional, Set, TypedDict, Union

 from jinja2 import Environment, meta, nodes
--- a/src/axolotl/prompt_strategies/kto/chatml.py
+++ b/src/axolotl/prompt_strategies/kto/chatml.py
@@ -1,7 +1,6 @@
 """
 KTO strategies for chatml
 """
-
 # pylint: disable=duplicate-code


@@ -16,9 +15,9 @@ def argilla(
                f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
        sample["completion"] = f"{sample['completion']}<|im_end|>"
        return sample

@@ -34,9 +33,9 @@ def argilla_chat(
    """

    def transform_fn(sample):
-        sample["prompt"] = (
-            f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
-        )
+        sample[
+            "prompt"
+        ] = f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
        sample["completion"] = f"{sample['completion'][1]['content']}<|im_end|>"
        return sample

@@ -56,9 +55,9 @@ def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
        sample["completion"] = f"{sample['completion']}<|im_end|>"
        return sample

@@ -75,9 +74,9 @@ def prompt_pairs(
                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
        sample["completion"] = f"{sample['completion']}<|im_end|>"
        return sample

@@ -97,9 +96,9 @@ def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
        sample["completion"] = f"{sample['completion']}<|im_end|>"
        return sample

--- a/src/axolotl/prompt_strategies/kto/llama3.py
+++ b/src/axolotl/prompt_strategies/kto/llama3.py
@@ -1,7 +1,6 @@
 """
 KTO strategies for llama-3 chat template
 """
-
 # pylint: disable=duplicate-code


@@ -16,9 +15,9 @@ def argilla(
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["completion"] = f"{sample['completion']}<|eot_id|>"
        return sample

@@ -34,9 +33,9 @@ def argilla_chat(
    """

    def transform_fn(sample):
-        sample["prompt"] = (
-            f"<|start_header_id|>user<|end_header_id|>\n\n{sample['completion'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-        )
+        sample[
+            "prompt"
+        ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['completion'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["completion"] = f"{sample['completion'][1]['content']}<|eot_id|>"
        return sample

@@ -56,9 +55,9 @@ def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["completion"] = f"{sample['completion']}<|eot_id|>"
        return sample

@@ -75,9 +74,9 @@ def prompt_pairs(
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["completion"] = f"{sample['completion']}<|eot_id|>"
        return sample

@@ -97,9 +96,9 @@ def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-arg
                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
            )
        else:
-            sample["prompt"] = (
-                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            )
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        sample["completion"] = f"{sample['completion']}<|eot_id|>"
        return sample

--- a/src/axolotl/prompt_strategies/kto/user_defined.py
+++ b/src/axolotl/prompt_strategies/kto/user_defined.py
@@ -1,7 +1,6 @@
 """
 User-defined KTO strategies
 """
-
 # pylint: disable=duplicate-code


--- a/src/axolotl/prompt_strategies/messages/chat.py
+++ b/src/axolotl/prompt_strategies/messages/chat.py
@@ -1,7 +1,6 @@
 """
 Chat dataset wrapping strategy for new internal messages representations
 """
-
 from typing import Any, Callable, Dict, Optional

 from axolotl.core.datasets.chat import TokenizedChatDataset
--- a/src/axolotl/prompt_strategies/orcamini.py
+++ b/src/axolotl/prompt_strategies/orcamini.py
@@ -9,7 +9,6 @@ this one specifies the system prompt with "### System:".

 Not suited/tested for multiple-turn conversations without further adjustments.
 """
-
 from typing import Generator, Union

 from axolotl.prompt_strategies.alpaca_w_system import OpenOrcaPromptTokenizingStrategy
--- a/src/axolotl/prompt_strategies/orpo/chat_template.py
+++ b/src/axolotl/prompt_strategies/orpo/chat_template.py
@@ -1,5 +1,4 @@
 """chatml prompt tokenization strategy for ORPO"""
-
 from typing import Any, Dict, Generator, List, Optional, Tuple

 from pydantic import BaseModel
--- a/src/axolotl/prompt_strategies/pretrain.py
+++ b/src/axolotl/prompt_strategies/pretrain.py
@@ -1,5 +1,4 @@
 """pretraining prompt strategies"""
-
 from typing import Generator

 from transformers import BatchEncoding
--- a/src/axolotl/train.py
+++ b/src/axolotl/train.py
@@ -406,7 +406,9 @@ def handle_untrained_tokens_fix(
        )


-def setup_model_and_trainer(cfg: DictDefault, dataset_meta: TrainDatasetMeta) -> tuple[
+def setup_model_and_trainer(
+    cfg: DictDefault, dataset_meta: TrainDatasetMeta
+) -> tuple[
    HFRLTrainerBuilder | HFCausalTrainerBuilder,
    PeftModel | PreTrainedModel,
    PreTrainedTokenizer,
--- a/src/axolotl/utils/init.py
+++ b/src/axolotl/utils/init.py
@@ -40,6 +40,6 @@ def set_pytorch_cuda_alloc_conf():
    torch_major, torch_minor = int(torch_version[0]), int(torch_version[1])
    if torch_major == 2 and torch_minor >= 2:
        if os.getenv("PYTORCH_CUDA_ALLOC_CONF") is None:
-            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = (
-                "expandable_segments:True,roundup_power2_divisions:16"
-            )
+            os.environ[
+                "PYTORCH_CUDA_ALLOC_CONF"
+            ] = "expandable_segments:True,roundup_power2_divisions:16"
--- a/src/axolotl/utils/bench.py
+++ b/src/axolotl/utils/bench.py
@@ -1,5 +1,4 @@
 """Benchmarking and measurement utilities"""
-
 import functools

 import torch
--- a/src/axolotl/utils/callbacks/init.py
+++ b/src/axolotl/utils/callbacks/init.py
@@ -33,6 +33,7 @@ from trl.models import unwrap_model_for_generation
 from axolotl.utils import is_comet_available, is_mlflow_available
 from axolotl.utils.bench import log_gpu_memory_usage
 from axolotl.utils.callbacks.perplexity import Perplexity
+from axolotl.utils.config.models.input.v0_4_1 import AxolotlInputConfig
 from axolotl.utils.distributed import (
    barrier,
    broadcast_dict,
@@ -42,7 +43,6 @@ from axolotl.utils.distributed import (
    is_main_process,
    zero_first,
 )
-from axolotl.utils.schemas.config import AxolotlInputConfig

 if TYPE_CHECKING:
    from axolotl.core.trainer_builder import AxolotlTrainingArguments
@@ -343,9 +343,9 @@ def bench_eval_callback_factory(trainer, tokenizer):
                    bench_refs.extend(combined_bench_names[bench_name]["refs"])
                    bench_preds.extend(combined_bench_names[bench_name]["preds"])
                    if not pd.isna(bench_score):
-                        results[f"{bench_split}_bench_accuracy_{bench_name}"] = (
-                            bench_score
-                        )
+                        results[
+                            f"{bench_split}_bench_accuracy_{bench_name}"
+                        ] = bench_score
                        bench_scores.append(bench_score)
                    else:
                        results[f"{bench_split}_bench_accuracy_{bench_name}"] = 0.0
--- a/src/axolotl/utils/callbacks/mlflow_.py
+++ b/src/axolotl/utils/callbacks/mlflow_.py
@@ -1,5 +1,4 @@
 """MLFlow module for trainer callbacks"""
-
 import logging
 from shutil import copyfile
 from tempfile import NamedTemporaryFile
--- a/src/axolotl/utils/callbacks/perplexity.py
+++ b/src/axolotl/utils/callbacks/perplexity.py
@@ -1,5 +1,4 @@
 """callback to calculate perplexity as an evaluation metric."""
-
 from typing import Dict, List, Optional

 import torch
--- a/src/axolotl/utils/callbacks/profiler.py
+++ b/src/axolotl/utils/callbacks/profiler.py
@@ -1,7 +1,6 @@
 """
 HF Trainer callback for creating pytorch profiling snapshots
 """
-
 from pathlib import Path
 from pickle import dump  # nosec B403

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Salman Mohammadi	0c36a6fea6	config fix -___-	2025-03-18 11:35:20 +00:00
Salman Mohammadi	64aca3c23c	linting v2	2025-03-18 11:33:54 +00:00
Salman Mohammadi	22abfd6170	simplifying check	2025-03-18 11:26:53 +00:00
Salman Mohammadi	0658c458b7	Merge branch 'fix_kto' of github.com:axolotl-ai-cloud/axolotl into fix_kto	2025-03-18 11:23:48 +00:00
Salman Mohammadi	690908cf2f	linting	2025-03-18 11:23:23 +00:00
salman	b9378e9b39	Merge branch 'main' into fix_kto	2025-03-18 11:22:00 +00:00
Salman Mohammadi	57b0ad1467	adding adapter check	2025-03-18 11:21:42 +00:00
Salman Mohammadi	ec4ead6e3e	adding error	2025-03-18 11:20:34 +00:00
Salman Mohammadi	a319ac7d3e	removing artifacts	2025-03-17 20:00:09 +00:00
Salman Mohammadi	09d3f2cffa	WIP	2025-03-17 19:59:19 +00:00