remove ref_model when peft model is passed into grpo trainer

2025-02-20 21:53:20 -05:00
226 changed files with 2192 additions and 4633 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/main.yml
+++ b/.github/workflows/main.yml
@@ -88,11 +88,6 @@ jobs:
            pytorch: 2.5.1
            axolotl_extras:
            is_latest: true
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.6.0
-            axolotl_extras:
    runs-on: axolotl-gpu-runner
    steps:
      - name: Checkout
--- a/.github/workflows/multi-gpu-e2e.yml
+++ b/.github/workflows/multi-gpu-e2e.yml
@@ -4,10 +4,6 @@ on:
  pull_request:
    paths:
      - 'tests/e2e/multigpu/*.py'
-      - 'requirements.txt'
-      - 'setup.py'
-      - 'pyproject.toml'
-      - '.github/workflows/multi-gpu-e2e.yml'
  workflow_dispatch:
  schedule:
    - cron: '0 0 * * 1,4'  # Runs at 00:00 UTC every monday & thursday
--- a/.github/workflows/nightlies.yml
+++ b/.github/workflows/nightlies.yml
@@ -80,11 +80,6 @@ jobs:
            python_version: "3.11"
            pytorch: 2.5.1
            axolotl_extras:
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.6.0
-            axolotl_extras:
    runs-on: axolotl-gpu-runner
    steps:
      - name: Checkout
--- 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
@@ -19,6 +19,9 @@
    <br/>
    <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/tests-nightly.yml/badge.svg" alt="tests-nightly">
    <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/multi-gpu-e2e.yml/badge.svg" alt="multigpu-semi-weekly tests">
+    <a href="https://www.phorm.ai/query?projectId=e315ba4a-4e14-421f-ab05-38a1f9076f25">
+    <img alt="phorm.ai" src="https://img.shields.io/badge/Phorm-Ask_AI-%23F2777A.svg?&logo=data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iNSIgaGVpZ2h0PSI0IiBmaWxsPSJub25lIiB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciPgogIDxwYXRoIGQ9Ik00LjQzIDEuODgyYTEuNDQgMS40NCAwIDAgMS0uMDk4LjQyNmMtLjA1LjEyMy0uMTE1LjIzLS4xOTIuMzIyLS4wNzUuMDktLjE2LjE2NS0uMjU1LjIyNmExLjM1MyAxLjM1MyAwIDAgMS0uNTk1LjIxMmMtLjA5OS4wMTItLjE5Mi4wMTQtLjI3OS4wMDZsLTEuNTkzLS4xNHYtLjQwNmgxLjY1OGMuMDkuMDAxLjE3LS4xNjkuMjQ2LS4xOTFhLjYwMy42MDMgMCAwIDAgLjItLjEwNi41MjkuNTI5IDAgMCAwIC4xMzgtLjE3LjY1NC42NTQgMCAwIDAgLjA2NS0uMjRsLjAyOC0uMzJhLjkzLjkzIDAgMCAwLS4wMzYtLjI0OS41NjcuNTY3IDAgMCAwLS4xMDMtLjIuNTAyLjUwMiAwIDAgMC0uMTY4LS4xMzguNjA4LjYwOCAwIDAgMC0uMjQtLjA2N0wyLjQzNy43MjkgMS42MjUuNjcxYS4zMjIuMzIyIDAgMCAwLS4yMzIuMDU4LjM3NS4zNzUgMCAwIDAtLjExNi4yMzJsLS4xMTYgMS40NS0uMDU4LjY5Ny0uMDU4Ljc1NEwuNzA1IDRsLS4zNTctLjA3OUwuNjAyLjkwNkMuNjE3LjcyNi42NjMuNTc0LjczOS40NTRhLjk1OC45NTggMCAwIDEgLjI3NC0uMjg1Ljk3MS45NzEgMCAwIDEgLjMzNy0uMTRjLjExOS0uMDI2LjIyNy0uMDM0LjMyNS0uMDI2TDMuMjMyLjE2Yy4xNTkuMDE0LjMzNi4wMy40NTkuMDgyYTEuMTczIDEuMTczIDAgMCAxIC41NDUuNDQ3Yy4wNi4wOTQuMTA5LjE5Mi4xNDQuMjkzYTEuMzkyIDEuMzkyIDAgMCAxIC4wNzguNThsLS4wMjkuMzJaIiBmaWxsPSIjRjI3NzdBIi8+CiAgPHBhdGggZD0iTTQuMDgyIDIuMDA3YTEuNDU1IDEuNDU1IDAgMCAxLS4wOTguNDI3Yy0uMDUuMTI0LS4xMTQuMjMyLS4xOTIuMzI0YTEuMTMgMS4xMyAwIDAgMS0uMjU0LjIyNyAxLjM1MyAxLjM1MyAwIDAgMS0uNTk1LjIxNGMtLjEuMDEyLS4xOTMuMDE0LS4yOC4wMDZsLTEuNTYtLjEwOC4wMzQtLjQwNi4wMy0uMzQ4IDEuNTU5LjE1NGMuMDkgMCAuMTczLS4wMS4yNDgtLjAzM2EuNjAzLjYwMyAwIDAgMCAuMi0uMTA2LjUzMi41MzIgMCAwIDAgLjEzOS0uMTcyLjY2LjY2IDAgMCAwIC4wNjQtLjI0MWwuMDI5LS4zMjFhLjk0Ljk0IDAgMCAwLS4wMzYtLjI1LjU3LjU3IDAgMCAwLS4xMDMtLjIwMi41MDIuNTAyIDAgMCAwLS4xNjgtLjEzOC42MDUuNjA1IDAgMCAwLS4yNC0uMDY3TDEuMjczLjgyN2MtLjA5NC0uMDA4LS4xNjguMDEtLjIyMS4wNTUtLjA1My4wNDUtLjA4NC4xMTQtLjA5Mi4yMDZMLjcwNSA0IDAgMy45MzhsLjI1NS0yLjkxMUExLjAxIDEuMDEgMCAwIDEgLjM5My41NzIuOTYyLjk2MiAwIDAgMSAuNjY2LjI4NmEuOTcuOTcgMCAwIDEgLjMzOC0uMTRDMS4xMjIuMTIgMS4yMy4xMSAxLjMyOC4xMTlsMS41OTMuMTRjLjE2LjAxNC4zLjA0Ny40MjMuMWExLjE3IDEuMTcgMCAwIDEgLjU0NS40NDhjLjA2MS4wOTUuMTA5LjE5My4xNDQuMjk1YTEuNDA2IDEuNDA2IDAgMCAxIC4wNzcuNTgzbC0uMDI4LjMyMloiIGZpbGw9IndoaXRlIi8+CiAgPHBhdGggZD0iTTQuMDgyIDIuMDA3YTEuNDU1IDEuNDU1IDAgMCAxLS4wOTguNDI3Yy0uMDUuMTI0LS4xMTQuMjMyLS4xOTIuMzI0YTEuMTMgMS4xMyAwIDAgMS0uMjU0LjIyNyAxLjM1MyAxLjM1MyAwIDAgMS0uNTk1LjIxNGMtLjEuMDEyLS4xOTMuMDE0LS4yOC4wMDZsLTEuNTYtLjEwOC4wMzQtLjQwNi4wMy0uMzQ4IDEuNTU5LjE1NGMuMDkgMCAuMTczLS4wMS4yNDgtLjAzM2EuNjAzLjYwMyAwIDAgMCAuMi0uMTA2LjUzMi41MzIgMCAwIDAgLjEzOS0uMTcyLjY2LjY2IDAgMCAwIC4wNjQtLjI0MWwuMDI5LS4zMjFhLjk0Ljk0IDAgMCAwLS4wMzYtLjI1LjU3LjU3IDAgMCAwLS4xMDMtLjIwMi41MDIuNTAyIDAgMCAwLS4xNjgtLjEzOC42MDUuNjA1IDAgMCAwLS4yNC0uMDY3TDEuMjczLjgyN2MtLjA5NC0uMDA4LS4xNjguMDEtLjIyMS4wNTUtLjA1My4wNDUtLjA4NC4xMTQtLjA5Mi4yMDZMLjcwNSA0IDAgMy45MzhsLjI1NS0yLjkxMUExLjAxIDEuMDEgMCAwIDEgLjM5My41NzIuOTYyLjk2MiAwIDAgMSAuNjY2LjI4NmEuOTcuOTcgMCAwIDEgLjMzOC0uMTRDMS4xMjIuMTIgMS4yMy4xMSAxLjMyOC4xMTlsMS41OTMuMTRjLjE2LjAxNC4zLjA0Ny40MjMuMWExLjE3IDEuMTcgMCAwIDEgLjU0NS40NDhjLjA2MS4wOTUuMTA5LjE5My4xNDQuMjk1YTEuNDA2IDEuNDA2IDAgMCAxIC4wNzcuNTgzbC0uMDI4LjMyMloiIGZpbGw9IndoaXRlIi8+Cjwvc3ZnPgo=">
+  </a>
 </p>

 Axolotl is a tool designed to streamline post-training for various AI models.
@@ -47,15 +50,13 @@ Features:
 ## 🚀 Quick Start

 **Requirements**:
-
 - NVIDIA GPU (Ampere or newer for `bf16` and Flash Attention) or AMD GPU
 - Python 3.11
 - PyTorch ≥2.4.1

 ### Installation

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

 # Download example axolotl configs, deepspeed configs
@@ -67,7 +68,7 @@ Other installation approaches are described [here](https://axolotl-ai-cloud.gith

 ### Your First Fine-tune

-```bash
+```shell
 # Fetch axolotl examples
 axolotl fetch examples

@@ -97,7 +98,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,186 +1,12 @@
 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"
+  description: "Fine-tuning"
  favicon: favicon.jpg
-
  navbar:
-    logo: image/axolotl_logo_digital_white.svg
-    title: false
+    title: Axolotl
    background: dark
    pinned: false
    collapse: false
@@ -199,77 +25,33 @@ website:
      contents:
        - text: Home
          href: index.qmd
-
-        - section: "Getting Started"
+        - section: "How-To Guides"
          contents:
+          # TODO Edit folder structure after we have more docs.
            - docs/getting-started.qmd
            - docs/installation.qmd
+            - docs/debugging.qmd
            - docs/inference.qmd
-            - docs/cli.qmd
-            - docs/config.qmd
-            - text: "API Reference"
-              href: docs/api
-
-        - section: "Dataset Formats"
-          contents: docs/dataset-formats/*
-
-        - section: "Deployments"
-          contents:
-            - docs/docker.qmd
+            - docs/multipack.qmd
+            - docs/fsdp_qlora.qmd
+            - docs/input_output.qmd
+            - docs/rlhf.qmd
+            - docs/nccl.qmd
+            - docs/mac.qmd
            - docs/multi-gpu.qmd
            - docs/multi-node.qmd
-            - docs/ray-integration.qmd
-            - docs/amd_hpc.qmd
-            - docs/mac.qmd
-
-        - section: "How To Guides"
-          contents:
-            - docs/multimodal.qmd
-            - docs/rlhf.qmd
-            - docs/reward_modelling.qmd
-            - docs/lr_groups.qmd
-            - docs/lora_optims.qmd
-
-        - section: "Core Concepts"
-          contents:
-            - docs/batch_vs_grad.qmd
-            - docs/dataset_preprocessing.qmd
-            - docs/multipack.qmd
-
-        - section: "Advanced Features"
-          contents:
-            - docs/fsdp_qlora.qmd
            - docs/unsloth.qmd
-            - docs/torchao.qmd
-            - docs/custom_integrations.qmd
-
-        - section: "Troubleshooting"
+            - docs/amd_hpc.qmd
+            - docs/ray-integration.qmd
+        - section: "Dataset Formats"
+          contents: docs/dataset-formats/*
+        - section: "Reference"
          contents:
-            - docs/faq.qmd
-            - docs/debugging.qmd
-            - docs/nccl.qmd
+            - docs/config.qmd
+        - docs/faq.qmd

 format:
  html:
-    theme: darkly
+    theme: materia
    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
@@ -38,11 +37,15 @@ temp_dir = tempfile.mkdtemp()
 with open(pathlib.Path(temp_dir) / "Dockerfile", "w", encoding="utf-8") as f:
    f.write(dockerfile_contents)

-cicd_image = Image.from_dockerfile(
-    pathlib.Path(temp_dir) / "Dockerfile",
-    force_build=True,
-    gpu="A10G",
-).env(df_args)
+cicd_image = (
+    Image.from_dockerfile(
+        pathlib.Path(temp_dir) / "Dockerfile",
+        force_build=True,
+        gpu="A10G",
+    )
+    .env(df_args)
+    .pip_install("fastapi==0.110.0", "pydantic==2.6.3")
+)

 app = App("Axolotl CI/CD", secrets=[])

--- 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/docker/Dockerfile-cloud
+++ b/docker/Dockerfile-cloud
@@ -14,7 +14,7 @@ COPY scripts/motd /etc/motd

 RUN pip install jupyterlab notebook ipywidgets && \
    jupyter lab clean
-RUN apt install --yes --no-install-recommends openssh-server tmux iproute2 nvtop && \
+RUN apt install --yes --no-install-recommends openssh-server tmux && \
    mkdir -p ~/.ssh && \
    chmod 700 ~/.ssh && \
    printf "\n[[ -z \"\$TMUX\"  ]] && { tmux attach-session -t ssh_tmux || tmux new-session -s ssh_tmux; exit; }\n" >> ~/.bashrc && \
--- a/docs/.gitignore
+++ b/docs/.gitignore
@@ -1,4 +1,2 @@
 /.quarto/
 _site/
-/api/*.qmd
-/api/*.html
--- a/docs/amd_hpc.qmd
+++ b/docs/amd_hpc.qmd
@@ -1,5 +1,5 @@
 ---
-title: AMD GPUs on HPC Systems
+title: Training with AMD GPUs on HPC Systems
 description: A comprehensive guide for using Axolotl on distributed systems with AMD GPUs
 ---

--- a/docs/cli.qmd
+++ b/docs/cli.qmd
@@ -1,19 +1,28 @@
---
-title: "Command Line Interface (CLI)"
-format:
-  html:
-    toc: true
-    toc-expand: 2
-    toc-depth: 3
-execute:
-  enabled: false
---
+# Axolotl CLI Documentation

 The Axolotl CLI provides a streamlined interface for training and fine-tuning large language models. This guide covers
 the CLI commands, their usage, and common examples.

+### Table of Contents

-## Basic Commands
+- Basic Commands
+- Command Reference
+  - fetch
+  - preprocess
+  - train
+  - inference
+  - merge-lora
+  - merge-sharded-fsdp-weights
+  - evaluate
+  - lm-eval
+- Legacy CLI Usage
+- Remote Compute with Modal Cloud
+  - Cloud Configuration
+  - Running on Modal Cloud
+  - Cloud Configuration Options
+
+
+### Basic Commands

 All Axolotl commands follow this general structure:

@@ -23,9 +32,9 @@ axolotl <command> [config.yml] [options]

 The config file can be local or a URL to a raw YAML file.

-## Command Reference
+### Command Reference

-### fetch
+#### fetch

 Downloads example configurations and deepspeed configs to your local machine.

@@ -40,7 +49,7 @@ axolotl fetch deepspeed_configs
 axolotl fetch examples --dest path/to/folder
 ```

-### preprocess
+#### preprocess

 Preprocesses and tokenizes your dataset before training. This is recommended for large datasets.

@@ -65,7 +74,7 @@ dataset_prepared_path: Local folder for saving preprocessed data
 push_dataset_to_hub: HuggingFace repo to push preprocessed data (optional)
 ```

-### train
+#### train

 Trains or fine-tunes a model using the configuration specified in your YAML file.

@@ -86,38 +95,7 @@ axolotl train config.yml --no-accelerate
 axolotl train config.yml --resume-from-checkpoint path/to/checkpoint
 ```

-It is possible to run sweeps over multiple hyperparameters by passing in a sweeps config.
-
-```bash
-# Basic training with sweeps
-axolotl train config.yml --sweep path/to/sweep.yaml
-```
-
-Example sweep config:
-```yaml
-_:
-  # This section is for dependent variables we need to fix
-  - load_in_8bit: false
-    load_in_4bit: false
-    adapter: lora
-  - load_in_8bit: true
-    load_in_4bit: false
-    adapter: lora
-
-# These are independent variables
-learning_rate: [0.0003, 0.0006]
-lora_r:
-  - 16
-  - 32
-lora_alpha:
-  - 16
-  - 32
-  - 64
-```
-
-
-
-### inference
+#### inference

 Runs inference using your trained model in either CLI or Gradio interface mode.

@@ -137,7 +115,7 @@ cat prompt.txt | axolotl inference config.yml \
    --base-model="./completed-model"
 ```

-### merge-lora
+#### merge-lora

 Merges trained LoRA adapters into the base model.

@@ -159,7 +137,7 @@ gpu_memory_limit: Limit GPU memory usage
 lora_on_cpu: Load LoRA weights on CPU
 ```

-### merge-sharded-fsdp-weights
+#### merge-sharded-fsdp-weights

 Merges sharded FSDP model checkpoints into a single combined checkpoint.

@@ -168,7 +146,7 @@ Merges sharded FSDP model checkpoints into a single combined checkpoint.
 axolotl merge-sharded-fsdp-weights config.yml
 ```

-### evaluate
+#### evaluate

 Evaluates a model's performance using metrics specified in the config.

@@ -177,27 +155,27 @@ Evaluates a model's performance using metrics specified in the config.
 axolotl evaluate config.yml
 ```

-### lm-eval
+#### lm-eval

 Runs LM Evaluation Harness on your model.

 ```bash
 # Basic evaluation
 axolotl lm-eval config.yml
+
+# Evaluate specific tasks
+axolotl lm-eval config.yml --tasks arc_challenge,hellaswag
 ```

 Configuration options:

 ```yaml
-# List of tasks to evaluate
-lm_eval_tasks:
-  - arc_challenge
-  - hellaswag
-lm_eval_batch_size: # Batch size for evaluation
-output_dir: # Directory to save evaluation results
+lm_eval_tasks: List of tasks to evaluate
+lm_eval_batch_size: Batch size for evaluation
+output_dir: Directory to save evaluation results
 ```

-## Legacy CLI Usage
+### Legacy CLI Usage

 While the new Click-based CLI is preferred, Axolotl still supports the legacy module-based CLI:

@@ -217,18 +195,12 @@ accelerate launch -m axolotl.cli.inference config.yml \
    --lora_model_dir="./outputs/lora-out" --gradio
 ```

-::: {.callout-important}
-When overriding CLI parameters in the legacy CLI, use same notation as in yaml file (e.g., `--lora_model_dir`).
-
-**Note:** This differs from the new Click-based CLI, which uses dash notation (e.g., `--lora-model-dir`). Keep this in mind if you're referencing newer documentation or switching between CLI versions.
-:::
-
-## Remote Compute with Modal Cloud
+### Remote Compute with Modal Cloud

 Axolotl supports running training and inference workloads on Modal cloud infrastructure. This is configured using a
 cloud YAML file alongside your regular Axolotl config.

-### Cloud Configuration
+#### Cloud Configuration

 Create a cloud config YAML with your Modal settings:

@@ -243,17 +215,13 @@ branch: main    # Git branch to use (optional)
 volumes:        # Persistent storage volumes
  - name: axolotl-cache
    mount: /workspace/cache
-  - name: axolotl-data
-    mount: /workspace/data
-  - name: axolotl-artifacts
-    mount: /workspace/artifacts

 env:            # Environment variables
  - WANDB_API_KEY
  - HF_TOKEN
 ```

-### Running on Modal Cloud
+#### Running on Modal Cloud

 Commands that support the --cloud flag:

@@ -271,18 +239,18 @@ axolotl train config.yml --cloud cloud_config.yml --no-accelerate
 axolotl lm-eval config.yml --cloud cloud_config.yml
 ```

-### Cloud Configuration Options
+#### Cloud Configuration Options

 ```yaml
-provider: # compute provider, currently only `modal` is supported
-gpu: # GPU type to use
-gpu_count: # Number of GPUs (default: 1)
-memory: # RAM in GB (default: 128)
-timeout: # Maximum runtime in seconds
-timeout_preprocess: # Preprocessing timeout
-branch: # Git branch to use
-docker_tag: # Custom Docker image tag
-volumes: # List of persistent storage volumes
-env: # Environment variables to pass
-secrets: # Secrets to inject
+provider: compute provider, currently only `modal` is supported
+gpu: GPU type to use
+gpu_count: Number of GPUs (default: 1)
+memory: RAM in GB (default: 128)
+timeout: Maximum runtime in seconds
+timeout_preprocess: Preprocessing timeout
+branch: Git branch to use
+docker_tag: Custom Docker image tag
+volumes: List of persistent storage volumes
+env: Environment variables to pass
+secrets: Secrets to inject
 ```
--- a/docs/config.qmd
+++ b/docs/config.qmd
@@ -1,5 +1,5 @@
 ---
-title: Config Reference
+title: Config options
 description: A complete list of all configuration options.
 ---

@@ -30,8 +30,6 @@ tokenizer_legacy:
 # Resize the model embeddings when new tokens are added to multiples of 32
 # This is reported to improve training speed on some models
 resize_token_embeddings_to_32x:
-# Optional[bool] Whether to shrink the embeddings to len(tokenizer). By default, we won't shrink.
-shrink_embeddings:

 # (Internal use only)
 # Used to identify which the model is based on
@@ -85,12 +83,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
@@ -162,6 +154,8 @@ datasets:
      content: value
      # ...

+    message_property_mappings:
+
    # Optional[Dict[str, List]]. Roles mapping in the messages. The default is:
    roles:
      user: ["human", "user"]
@@ -169,16 +163,10 @@ datasets:
      system: ["system"]
      tool: ["tool"]

-    # Optional[bool]. Whether to drop the system turn from the dataset. Only works with chat_template.
-    # This does not drop the default system message from chat_template if it exists. If you wish to,
-    # we recommend using a custom jinja template with the default system message removed or
-    # adding a system turn with empty content.
-    drop_system_message:
-
    # IMPORTANT: The following fields determine which parts of the conversation to train on.
    # Priority order: message_field_training > message_field_training_detail > train_on_inputs or role in roles_to_train
    # See examples at `docs/dataset-formats/conversation.qmd`
-    # Note: If the below 4 fields are set to empty, defaults to training only on the last message.
+    # Note: If the below 4 fields are empty, defaults to training only on the last message.

    # Optional[List[str]]. Roles to train on. The tokens from these roles will be considered for the loss.
    roles_to_train: ["assistant"]  # default
@@ -186,7 +174,6 @@ datasets:
    # - all: train on all EOS tokens
    # - turn (default): train on the EOS token at the end of each trainable turn
    # - last: train on the last EOS token in the conversation
-    # TIP: Please make sure that your `tokenizer.eos_token` is same as EOS/EOT token in template. Otherwise, set `eos_token` under `special_tokens`.
    train_on_eos: last
    # The key in the message turn that indicates via boolean whether tokens of a turn should be considered for training. Useful to selectively train on certain turns besides the `roles_to_train`.
    message_field_training: training
@@ -213,46 +200,10 @@ test_datasets:
    data_files:
      - /workspace/data/eval.jsonl

-# use RL training: 'dpo', 'ipo', 'kto', 'simpo', 'orpo', 'grpo'
+# use RL training: 'dpo', 'ipo', 'kto'
 rl:
-rl_beta:  # Optional[float]. The beta parameter for the RL training.
-
-# dpo
-dpo_use_weighting:  # Optional[bool]. Whether to perform weighting.
-rpo_alpha: # Optional[float]. Weighting of NLL term in loss from RPO paper.
-
-# orpo
-orpo_alpha: 0.1  # Parameter controlling the relative ratio loss weight in the ORPO loss. Passed to `beta` in `ORPOConfig` due to trl mapping.
-
-# kto
-kto_desirable_weight: # Optional[float]. Factor for desirable loss term in KTO loss.
-kto_undesirable_weight: # Optional[float]. Factor for undesirable loss term in KTO loss.
-
-# simpo
-cpo_alpha: 1.0  # Weight of the BC regularizer
-simpo_gamma: 0.5  # Target reward margin for the SimPO loss
-
-# grpo
-trl:
-  use_vllm: # Optional[bool]. Whether to use VLLM for RL training.
-  vllm_device: # Optional[str]. Device to use for VLLM.
-  vllm_gpu_memory_utilization: # Optional[float]. GPU memory utilization for VLLM.
-  vllm_max_model_len: # Optional[int]. Maximum length of the model for VLLM.
-  vllm_dtype: # Optional[str]. Data type for VLLM.
-
-  beta: # Optional[float]. Beta parameter for the RL training. Same as `rl_beta`. Use
-  max_completion_length: # Optional[int]. Maximum length of the completion for RL training.
-
-  reward_funcs: # Optional[list[str]]. List of reward functions to load. Paths must be importable from current dir.
-  reward_weights: # Optional[list[float]]. List of reward weights for the reward functions.
-
-  num_generations: # Optional[int]. Number of generations to sample.
-  log_completions: # Optional[bool]. Whether to log completions.
-
-  sync_ref_model: # Optional[bool]. Whether to sync the reference model.
-  ref_model_mixup_alpha: # Optional[float]. Mixup alpha for the reference model.
-  ref_model_sync_steps: # Optional[int]. Sync steps for the reference model.
-
+# whether to perform weighting if doing DPO training. Boolean.
+dpo_use_weighting:

 # reward modelling: `True` or `False`
 reward_model:
@@ -270,13 +221,13 @@ process_reward_model:
 chat_template: tokenizer_default
 # custom jinja template for chat template. This will be only used if chat_template is set to `jinja` or `null` (in which case chat_template is automatically set to `jinja`). Default is null.
 chat_template_jinja: null
-# Changes the default system message. Currently only supports chatml.
-default_system_message: You are a helpful assistant. Please give a long and detailed answer.
+# Changes the default system message
+default_system_message: You are a helpful assistant. Please give a long and detailed answer. # Currently only supports chatml.
 # Axolotl attempts to save the dataset as an arrow after packing the data together so
 # subsequent training attempts load faster, relative path
 dataset_prepared_path: data/last_run_prepared
 # Push prepared dataset to hub
-push_dataset_to_hub: # Optional[str] repo_org/repo_name
+push_dataset_to_hub: # repo path
 # The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()`
 # if not set.
 dataset_processes: # defaults to os.cpu_count() if not set
@@ -456,10 +407,7 @@ save_total_limit: # Checkpoints saved at a time
 max_steps:

 # bool of whether to include tokens trainer per second in the training metrics. This iterates over the entire dataset once, so it takes some time.
-include_tokens_per_second: # Optional[bool]
-
-# whether to find batch size that fits in memory. Passed to underlying transformers Trainer
-auto_find_batch_size: # Optional[bool]
+include_tokens_per_second:

 eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
 eval_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128
@@ -493,7 +441,7 @@ gradient_checkpointing: false
 early_stopping_patience: 3

 # Specify a scheduler and kwargs to use with the optimizer
-lr_scheduler: # 'one_cycle' | 'rex' | 'log_sweep' | empty for cosine
+lr_scheduler: # 'one_cycle' | 'log_sweep' | empty for cosine
 lr_scheduler_kwargs:
 cosine_min_lr_ratio: # decay lr to some percentage of the peak lr, e.g. cosine_min_lr_ratio=0.1 for 10% of peak lr
 cosine_constant_lr_ratio: # freeze lr at some percentage of the step, e.g. cosine_constant_lr_ratio=0.8 means start cosine_min_lr at 80% of training step (https://arxiv.org/pdf/2308.04014.pdf)
@@ -576,8 +524,6 @@ flash_attn_fuse_mlp: # Whether to fuse part of the MLP into a single operation
 sdp_attention:
 # Shifted-sparse attention (only llama) - https://arxiv.org/pdf/2309.12307.pdf
 s2_attention:
-# Optional[bool]. Whether to use low_cpu_mem_usage
-low_cpu_mem_usage:
 # Resume from a specific checkpoint dir
 resume_from_checkpoint:
 # If resume_from_checkpoint isn't set and you simply want it to start where it left off.
@@ -598,13 +544,6 @@ special_tokens:
 # Add extra tokens.
 tokens:

-# Mapping token_id to new_token_string to override reserved added_tokens in the tokenizer.
-# Only works for tokens that are not part of the base vocab (aka are added_tokens).
-# Can be checked if they exist in tokenizer.json added_tokens.
-added_tokens_overrides:  # Dict[int, str]
-#  128041: "<|im_start|>"
-#  128042: "<|im_end|>"
-
 # FSDP
 fsdp:
 fsdp_config:
--- a/docs/custom_integrations.qmd
+++ b/docs/custom_integrations.qmd
@@ -1,101 +0,0 @@
---
-title: Custom Integrations
-toc: true
-toc-depth: 3
---
-
-```{python}
-#| echo: false
-
-import re
-
-def process_readme(integration_name):
-    try:
-        path = f'../src/axolotl/integrations/{integration_name}/README.md'
-        with open(path, 'r') as f:
-            txt = f.read()
-            # Remove h1 headings
-            txt = re.sub(r'^# .*\n?', '', txt, flags=re.MULTILINE)
-            # Convert h2 to h3
-            txt = re.sub(r'^## ', '### ', txt, flags=re.MULTILINE)
-            return txt
-    except FileNotFoundError:
-        return None
-
-def print_section(name, folder_name):
-    output = f"\n## {name}\n"
-    content = process_readme(folder_name)
-    if content:
-        output += content
-    output += f"\nPlease see reference [here](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/{folder_name})\n"
-    return output
-```
-
-```{python}
-#| output: asis
-#| echo: false
-
-# Introduction text
-print("""
-Axolotl adds custom features through `integrations`. They are located within the `src/axolotl/integrations` directory.
-
-To enable them, please check the respective documentations.
-""")
-
-# Sections
-sections = [
-    ("Cut Cross Entropy", "cut_cross_entropy"),
-    ("Grokfast", "grokfast"),
-    ("Knowledge Distillation (KD)", "kd"),
-    ("Liger Kernels", "liger"),
-    ("Language Model Evaluation Harness (LM Eval)", "lm_eval"),
-    ("Spectrum", "spectrum")
-]
-
-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-formats/conversation.qmd
+++ b/docs/dataset-formats/conversation.qmd
@@ -6,9 +6,7 @@ order: 3

 ## sharegpt

-::: {.callout-important}
-ShareGPT is deprecated!. Please see [chat_template](#chat_template) section below.
-:::
+IMPORTANT: ShareGPT is deprecated!. Please see [chat_template](#chat_template) section below.

 ## pygmalion

@@ -74,10 +72,6 @@ datasets:
    train_on_eos:
 ```

-::: {.callout-tip}
-If you receive an error like "`chat_template` choice is `tokenizer_default` but tokenizer's `chat_template` is null.", it means the tokenizer does not have a default `chat_template`. Follow the examples below instead to set a custom `chat_template`.
-:::
-
 2. Using the `gemma` chat template to override the tokenizer_config.json's chat template on OpenAI messages format, training on all assistant messages.

 ```yaml
@@ -108,10 +102,6 @@ datasets:
    type: chat_template
 ```

-::: {.callout-important}
-Please make sure that your `tokenizer.eos_token` is same as EOS/EOT token in template. Otherwise, set `eos_token` under `special_tokens`.
-:::
-
 5. (Advanced) Using fine-grained control over tokens and turns to train in a conversation

 For a data sample that looks like:
@@ -159,6 +149,4 @@ datasets:
    message_field_training_detail: train_detail
 ```

-::: {.callout-tip}
-It is not necessary to set both `message_field_training` and `message_field_training_detail` at once.
-:::
+Tip: It is not necessary to use both `message_field_training` and `message_field_training_detail` at a time.
--- a/docs/dataset-formats/index.qmd
+++ b/docs/dataset-formats/index.qmd
@@ -13,7 +13,7 @@ As there are a lot of available options in Axolotl, this guide aims to provide a

 Axolotl supports 3 kinds of training methods: pre-training, supervised fine-tuning, and preference-based post-training (e.g. DPO, ORPO, PRMs). Each method has their own dataset format which are described below.

-## Pre-training
+## [Pre-training](pretraining.qmd)

 When aiming to train on large corpora of text datasets, pre-training is your go-to choice. Due to the size of these datasets, downloading the entire-datasets before beginning training would be prohibitively time-consuming. Axolotl supports [streaming](https://huggingface.co/docs/datasets/en/stream) to only load batches into memory at a time.

@@ -96,10 +96,6 @@ One step is equal to `sequence_len * micro_batch_size * gradient_accumulation_st

 It is recommended to leave this off if downloading from Hugging Face hub as it would download the entire dataset which can be very large.

-### Reference
-
-Please see docs [here](pretraining.qmd).
-
 ## Supervised fine-tuning (SFT)

 Supervised fine-tuning is the process of training models to respond to an instruction or chat input.
@@ -124,12 +120,11 @@ If you went through the flow chart and did not find one that matches, it is reco
 You can mix and match within each approach or across approaches to train a model on a variety of datasets.
 :::

-### Pre-Tokenized Dataset
+### [Pre-Tokenized Dataset](tokenized.qmd)

 We suggest this approach when you want to bring your own tokenized dataset.

 Axolotl expects the dataset to have three keys:
-
 - `input_ids`: from tokenizing formatted prompt
 - `attention_mask`: for masking padding. If you don't add padding, it would be equal to `len(input_ids) * [1]`
 - `labels`: this is the same as `input_ids`, however, if you want to mask certain tokens, you would set those indices to `-100`.
@@ -150,9 +145,7 @@ datasets:
 `type: ` is empty!
 :::

-Reference: [Pre-Tokenized Dataset Documentation](tokenized.qmd).
-
-### Template Free Dataset
+### [Template Free Dataset](template_free.qmd)

 We reccomend this approach when you want granular control over the prompt formatting, special tokens, and masking, whilst letting Axolotl handle the tokenization. This is very useful if your dataset has unique prompts that differ across samples and where one single general template wouldn't suffice.

@@ -189,9 +182,7 @@ datasets:
    type: input_output
 ```

-Reference: [Template Free Documentation](template_free.qmd).
-
-### Conversation Dataset
+### [Conversation Dataset](conversation.qmd)

 `conversation` messages are a list of messages which usually contain a `role` and `content` key.

@@ -267,7 +258,7 @@ Newer conversation datasets usually follow the OpenAI format.

 Axolotl supports both as well as allowing customization of any kind of key.

-#### Chat Template Usage
+#### [Chat Template Usage](conversation.qmd#chat_template)

 To properly use this method, it is important to identify three things:

@@ -349,19 +340,9 @@ datasets:
      narrator: ["narrator"]
 ```

-::: {.callout-tip}
-As chat_templates may use hardcoded EOS/EOT tokens that are different from the tokenizer's EOS, it is highly recommended to set them. For example, `ChatML` uses `<|im_end|>` to end turns.
+#### Applying `chat_template`

-```yaml
-special_tokens:
-  eos_token: <|im_end|>
-```
-
-:::
-
-##### Applying `chat_template`
-
-Once all the above steps are completed, you could combine all these configs together to form a bespoke configuration for your custom dataset.
+Once all the above steps are completed, you could combine all these configs together to form a bespoke configuration for your  custom dataset. The final step would be to correctly set the EOS token in your config:

 ```yaml
 datasets:
@@ -410,17 +391,7 @@ If this config were to be applied to the sample dataset above, the output would

 The first number refers to the label, the second refers to the `token_id`. For example, `-100` labels appear on non-assistant portions, meaning that they are masked during. For assistant portions, the label is the same as the `token_id`.

-::: {.callout-note}
-
-If during `preprocess`, there are a lot of warnings of `Could not find content __ boundary`, please check the FAQ section for [chat_templates](../faq.qmd#chat-templates).
-
-:::
-
-#### Reference
-
-Please see docs [here](conversation.qmd).
-
-### Instruction Dataset
+### [Instruction Dataset](inst_tune.qmd)

 Instruction datasets are used to train instruction-following models and comprise a prompt, containing an instruction, and a single response. In contrast to chat datasets which may be multi-turn, instruct datasets are typically single-turn.

@@ -452,9 +423,6 @@ datasets:

 Axolotl supports many kinds of instruction dataset. All of them can be found here (https://axolotl-ai-cloud.github.io/axolotl/docs/dataset-formats/inst_tune.html) with their respective type and sample row format.

-
-Reference: [Instruction Dataset Documentation](inst_tune.qmd).
-
 #### Custom Instruct Prompt Format

 Due to the myriad possibilities of instruction formats, Axolotl allows customizing your own instruction format without having to dive into the code directly.
@@ -485,8 +453,6 @@ datasets:

 The config sets that the `field_instruction` is actually named `input`, and the `field_input` is empty as we don't have an `input` in this sample. Generally, `instruction` can be thought as the question to the model, and `input` as the additional information with `output` being the response. It is not necessary to have an `input` nor `system`. In the end, the most important part is to understand what format you want it to look like and how you can customize this to your use case.

-Reference: [Custom Instruct Prompt Format Documentation](inst_tune.qmd#how-to-add-custom-prompt-format).
-
 ## Reinforcement Learning from Human Feedback (RLHF)

-As there are multiple RLHF methods with their own dataset requirements. Please see [RLHF documentation](../rlhf.qmd) for more detail.
+As there are multiple RLHF methods with their own dataset requirements. Please see [RLHF datasets](../rlhf.qmd) documentation for more detail.
--- a/docs/dataset-formats/pretraining.qmd
+++ b/docs/dataset-formats/pretraining.qmd
@@ -27,6 +27,7 @@ pretraining_dataset:
    type: pretrain
    trust_remote_code:
    skip: # number of rows of data to skip over from the beginning
+...
 ```

 :::
--- a/docs/dataset-formats/template_free.qmd
+++ b/docs/dataset-formats/template_free.qmd
@@ -1,239 +1,7 @@
 ---
 title: Template-Free
 description: Construct prompts without a template.
-toc: true
-toc-depth: 3
 order: 4
 ---

-## Background {#sec-background}
-
-### Masking Inputs {#masking-inputs}
-
-One of the most popular features of
-[axolotl](https://github.com/axolotl-ai-cloud/axolotl) is
-setting the following configuration value:
-
-
-```yaml
-train_on_inputs: false
-```
-
-If you declare a [dataset formats](https://github.com/axolotl-ai-cloud/axolotl?tab=readme-ov-file#dataset)
-such as `alpaca` or `chatml`, axolotl knows what is an input
-(i.e. human) vs. an output (i.e. the assistant) and masks the input
-labels so that your model can focus on predicting the outputs only.
-
-### You may not want prompt templates {#sec-you-may-not-want-prompt-templates}
-
-However, there are many situations where you don't want to use one of
-these formats or templates. This is because they can:
-
-   Add unnecessary boilerplate to your prompts.
-   Create artifacts like special delimiters `<|im_start|>` that can
-    quickly become footguns if you don't include them correctly at
-    inference time.
-   Enforce a *chat* interface when you do not want one. Sometimes you
-    just want to fine-tune a model to a very specific task and do NOT
-    want multi-turn conversations, roles, etc.
-   Limit you to only certain roles that the template allows.
-
-### The `input_output` format {#sec-the-inputoutput-format}
-
-You can construct your prompts without a template by using the
-`input_output` format, by setting `type: input_output` in your
-configuration file like this:
-
-**config.yml**
-
-```yaml
-train_on_inputs: false # Mask segments of your data
-datasets:
-  - path: output.jsonl
-    type: input_output  # use template free prompt construction
-```
-
-Unlike `type: completion`, which is also template-free,
-`type: input_output` allows you to mask segments of your text. More
-details on how this works are described below.
-
-## Usage {#sec-usage}
-
-This is how you can use the `input_output` format:
-
-### 1. Prepare Data {#sec-1-prepare-data}
-
-To use the `input_output` format, collect your data in the following
-format into a jsonl file (below is the first row from the file
-`output`.jsonl` pretty printed):
-
-```bash
-$ head -n1 output.jsonl | python -m json.tool
-```
-
-:::{.cell-output .cell-output-stdout}
-    {
-        "segments": [
-            {
-                "label": true,
-                "text": "<s>Hello\n"
-            },
-            {
-                "label": true,
-                "text": "hi there!. "
-            },
-            {
-                "label": false,
-                "text": "goodbye "
-            },
-            {
-                "label": true,
-                "text": "farewell</s>"
-            }
-        ]
-    }
-:::
-
-Set `label:false` when you want to mask a segment of text so that the
-model isn't trained on it. Some things to keep in mind:
-
-> [!IMPORTANT]
-> 1.  **EOS, BOS, spaces, newlines etc. are entirely up to you. Axolotl
-    concatenates all the segments as-is.** The tokenizer doesn't add
-    anything additional. Notice how I added spaces, newlines, `<s>`
-    (BOS), and `</s>` (EOS) myself.
-> 2.  Make sure you check the materialized output to validate that the
-    prompt is getting assembled how you like.
-
-### 2. Use `type: input_output` {#sec-2-use-type-inputoutput}
-
-Let's materialize data with our `output.jsonl` file by setting
-`type: input_output` in our axolotl config:
-
-```yaml
-# training_config.yaml
-base_model: mistralai/Mistral-7B-v0.1
-data_seed: 49
-seed: 49
-
-datasets:
-  - path: output.jsonl
-    type: input_output
-val_set_size: 0.1
-
-sequence_len: 896
-sample_packing: false
-
-micro_batch_size: 2
-gradient_accumulation_steps: 3
-eval_batch_size: 2
-num_epochs: 1
-learning_rate: 0.0002
-
-train_on_inputs: false
-special_tokens:
-  bos_token: "<s>"
-  eos_token: "</s>"
-  unk_token: "<unk>"
-```
-
-You can use the following command to materialize your data. The
-`--debug` flag will print the tokens, along with the labels so you can
-verify that the correct items are being ignored:
-
-```bash
-axolotl preprocess training_config.yaml --debug
-
-...
-[2024-03-05 23:36:46,969] [INFO] [axolotl.check_example_labels:35] [PID:607731] [RANK:0] <s>(1, 1) Hello(22557, 22557)
-(13, 13) hi(12014, 12014) there(736, 736) !(28808, 28808) .(28723, 28723) (28705, 28705) good(-100, 1179) bye(-100, 17664) (-100, 28705) fare(19111, 19111) well(5458, 5458) </s>(2, 2)
-
-```
-
-The format is `decoded_token`(`label`, `token_id`), for example,
-`<s>(1, 1)` means that the token is `<s>`, the label is `1` and the
-token_id is `1`. When the label is `-100` then that token is ignored for
-training.
-
-### 3. Check the prompts {#sec-3-check-the-prompts}
-
-Here is another way to check the materialized output:
-
-```python
-from transformers import AutoTokenizer
-from datasets import load_from_disk
-import yaml
-
-directory = !ls last_run_prepared/
-with open('training_config.yaml', 'r') as f:
-    cfg = yaml.safe_load(f)
-model_id = cfg['base_model']
-tok = AutoTokenizer.from_pretrained(model_id)
-ds = load_from_disk(f'last_run_prepared/{directory[0]}/')
-```
-
-```python
->>> row = ds[0]
->>> print(tok.decode(row['input_ids']))
-<s> Hello
-    hi there!.  goodbye  farewell</s>
-```
-
-We can check that the right tokens are ignored by comparing the labels
-to each token:
-
-```python
-import pandas as pd
-pd.DataFrame([{'token': tok.decode(i), 'label': l, 'id':i} for i,l in
-              zip(row['input_ids'], row['labels'])])
-```
-
-| token | label | id    |
-|-------|-------|-------|
-| 0     | \<s\> | 1     |
-| 1     | Hello | 22557 |
-| 2     | \\n   | 13    |
-| 3     | hi    | 12014 |
-| 4     | there | 736   |
-| 5     | !     | 28808 |
-| 6     | .     | 28723 |
-| 7     |       | 28705 |
-| 8     | good  | -100  |
-| 9     | bye   | -100  |
-| 10    |       | -100  |
-| 11    | fare  | 19111 |
-| 12    | well  | 5458  |
-| 13    | \</s\>| 2     |
-
-
-
-If we look at the input data, the above table seems correct! (The jsonl
-version is repeated below for reference):
-
-
-```bash
-$ head -n1 output.jsonl | python -m json.tool
-```
-
-:::{.cell-output .cell-output-stdout}
-    {
-        "segments": [
-            {
-                "label": true,
-                "text": "<s>Hello\n"
-            },
-            {
-                "label": true,
-                "text": "hi there!. "
-            },
-            {
-                "label": false,
-                "text": "goodbye "
-            },
-            {
-                "label": true,
-                "text": "farewell</s>"
-            }
-        ]
-    }
-:::
+See [these docs](../input_output.qmd).
--- a/docs/dataset_preprocessing.qmd
+++ b/docs/dataset_preprocessing.qmd
@@ -3,11 +3,8 @@ title: Dataset Preprocessing
 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)[../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*
 - tokenize the dataset based on the configured model & tokenizer
@@ -15,12 +12,10 @@ the [dataset format](dataset-formats) and prompt strategies to:

 The processing of the datasets can happen one of two ways:

-1. Before kicking off training by calling `axolotl preprocess config.yaml --debug`
+1. Before kicking off training by calling `python -m axolotl.cli.preprocess /path/to/your.yaml --debug`
 2. When training is started

-### What are the benefits of pre-processing?
-
-When training interactively or for sweeps
+What are the benefits of pre-processing? When training interactively or for sweeps
 (e.g. you are restarting the trainer often), processing the datasets can oftentimes be frustratingly
 slow. Pre-processing will cache the tokenized/formatted datasets according to a hash of dependent
 training parameters so that it will intelligently pull from its cache when possible.
@@ -33,12 +28,8 @@ default path of `./last_run_prepared/`, but will ignore anything already cached
 setting `dataset_prepared_path: ./last_run_prepared`, the trainer will use whatever pre-processed
 data is in the cache.

-### What are the edge cases?
-
-Let's say you are writing a custom prompt strategy or using a user-defined
+What are the edge cases? Let's say you are writing a custom prompt strategy or using a user-defined
 prompt template. Because the trainer cannot readily detect these changes, we cannot change the
-calculated hash value for the pre-processed dataset.
-
-If you have `dataset_prepared_path: ...` set
+calculated hash value for the pre-processed dataset. If you have `dataset_prepared_path: ...` set
 and change your prompt templating logic, it may not pick up the changes you made and you will be
 training over the old prompt.
--- a/docs/debugging.qmd
+++ b/docs/debugging.qmd
@@ -31,13 +31,11 @@ While debugging it's helpful to simplify your test scenario as much as possible.
    - Set `CUDA_VISIBLE_DEVICES` to a single GPU, ex: `export CUDA_VISIBLE_DEVICES=0`.
    - Set `dataset_processes: 1` in your axolotl config or run the training command with `--dataset_processes=1`.
 2. **Use a small dataset**: Construct or use a small dataset from HF Hub. When using a small dataset, you will often have to make sure `sample_packing: False` and `eval_sample_packing: False` to avoid errors.  If you are in a pinch and don't have time to construct a small dataset but want to use from the HF Hub, you can shard the data (this will still tokenize the entire dataset, but will only use a fraction of the data for training.  For example, to shard the dataset into 20 pieces, add the following to your axolotl config):
-
    ```yaml
-    datasets:
+    dataset:
        ...
        shards: 20
    ```
-
 3. **Use a small model**: A good example of a small model is [TinyLlama/TinyLlama-1.1B-Chat-v1.0](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0).
 4. **Minimize iteration time**: Make sure the training loop finishes as fast as possible, with these settings.
    - `micro_batch_size: 1`
@@ -87,7 +85,7 @@ The easiest way to get started is to modify the [.vscode/launch.json](../.vscode

 For example, to mimic the command `cd devtools && CUDA_VISIBLE_DEVICES=0 accelerate launch -m axolotl.cli.train dev_chat_template.yml`, you would use the below configuration[^1].  Note that we add additional flags that override the axolotl config and incorporate the tips above (see the comments). We also set the working directory to `devtools` and set the `env` variable `HF_HOME` to a temporary folder that is later partially deleted.  This is because we want to delete the HF dataset cache before each run in order to ensure that the data preprocessing code is run from scratch.

-```json
+```jsonc
 // .vscode/launch.json
 {
    "version": "0.2.0",
@@ -134,7 +132,7 @@ For example, to mimic the command `cd devtools && CUDA_VISIBLE_DEVICES=0 acceler

 Below is the [./vscode/tasks.json](../.vscode/tasks.json) file that defines the `cleanup-for-dataprep` task.  This task is run before each debugging session when you use the above configuration.  Note how there are two tasks that delete the two folders mentioned above.  The third task `cleanup-for-dataprep` is a composite task that combines the two tasks.  A composite task is necessary because VSCode does not allow you to specify multiple tasks in the `preLaunchTask` argument of the `launch.json` file.

-```json
+```jsonc
 // .vscode/tasks.json
 // this file is used by launch.json
 {
--- a/docs/docker.qmd
+++ b/docs/docker.qmd
@@ -1,140 +0,0 @@
---
-title: "Docker"
-format:
-  html:
-    toc: true
-    toc-depth: 4
---
-
-This section describes the different Docker images that are released by AxolotlAI at [Docker Hub](https://hub.docker.com/u/axolotlai).
-
-## Base
-
-The base image is the most minimal image that can install Axolotl. It is based on the `nvidia/cuda` image. It includes python, torch, git, git-lfs, awscli, pydantic, and more.
-
-#### Image
-
-```
-axolotlai/axolotl-base
-```
-
-Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl-base)
-
-#### Tags format
-
-```bash
-main-base-py{python_version}-cu{cuda_version}-{pytorch_version}
-```
-
-Tags examples:
-
- `main-base-py3.11-cu124-2.6.0`
- `main-base-py3.11-cu124-2.5.1`
- `main-base-py3.11-cu124-2.4.1`
-
-## Main
-
-The main image is the image that is used to run Axolotl. It is based on the `axolotlai/axolotl-base` image and includes the Axolotl codebase, dependencies, and more.
-
-#### Image
-
-```
-axolotlai/axolotl
-```
-
-Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl)
-
-#### Tags format {#sec-main-tags}
-
-```bash
-# on push to main
-main-py{python_version}-cu{cuda_version}-{pytorch_version}
-
-# latest main (currently torch 2.5.1, python 3.11, cuda 12.4)
-main-latest
-
-# nightly build
-{branch}-{date_in_YYYYMMDD}-py{python_version}-cu{cuda_version}-{pytorch_version}
-
-# tagged release
-{version}
-```
-
-:::{.callout-tip}
-
-There may be some extra tags appended to the image, like `-vllm` which installs those packages.
-
-:::
-
-Tags examples:
-
- `main-py3.11-cu124-2.6.0`
- `main-py3.11-cu124-2.5.1`
- `main-py3.11-cu124-2.4.1`
- `main-latest`
- `main-20250303-py3.11-cu124-2.6.0`
- `main-20250303-py3.11-cu124-2.5.1`
- `main-20250303-py3.11-cu124-2.4.1`
- `0.7.1`
-
-## Cloud
-
-The cloud image is the image that is used to run Axolotl in the cloud. It is based on the `axolotlai/axolotl` image and sets ENV variables like HuggingFace cache directories for volume mounts, tmux, and more for different cloud providers.
-
-:::{.callout-tip}
-
-Jupyter lab is run by default. Set `JUPYTER_DISABLE=1` in the environment variables to disable it.
-
-:::
-
-#### Image
-
-```
-axolotlai/axolotl-cloud
-```
-
-Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl-cloud)
-
-#### Tags format
-
-This uses the same tags as the [`main` image](#sec-main-tags).
-
-#### Environment variables
-
- `JUPYTER_DISABLE`: Disable Jupyter lab.
- `JUPYTER_PASSWORD`: Set a password for the Jupyter lab.
- `PUBLIC_KEY`: Add a public key for the SSH service.
- `SSH_KEY`: Add a private key for the SSH service.
-
-#### Volume mounts
-
-:::{.callout-tip}
-
-We recommend mounting volumes to `/workspace/data` for data persistence. `/workspace/axolotl` contains the source code and is ephemeral.
-
-:::
-
- `/workspace/data/axolotl-artifacts`: Directory to store Axolotl artifacts.
- `/workspace/data/huggingface-cache`: Directory to store HuggingFace cache.
-
-## Cloud-no-tmux
-
-This is the same as the [`cloud` image](#sec-cloud) but without tmux.
-
-#### Image
-
-```
-axolotlai/axolotl-cloud-term
-```
-
-Link: [Docker Hub](https://hub.docker.com/r/axolotlai/axolotl-cloud-term)
-
-:::{.callout-note}
-
-The naming may be a bit confusing as it has `-term` appended to the end.
-
-:::
-
-#### Tags format
-
-This uses the same tags as the [`cloud` image](#sec-cloud-tags).
--- a/docs/faq.qmd
+++ b/docs/faq.qmd
@@ -3,7 +3,6 @@ title: FAQ
 description: Frequently asked questions
 ---

-### General

 **Q: The trainer stopped and hasn't progressed in several minutes.**

@@ -19,50 +18,12 @@ description: Frequently asked questions

 **Q: AttributeError: 'DummyOptim' object has no attribute 'step'**

-**Q: ModuleNotFoundError: No module named 'mpi4py' using single GPU with deepspeed**
-
-> A: You may be using deepspeed with single gpu. Please remove the `deepspeed:` section in the yaml file or `--deepspeed` CLI flag.
+> A: You may be using deepspeed with single gpu. Please don't set `deepspeed:` in yaml or cli.

 **Q: The codes is stuck on saving preprocessed datasets.**

 > A: This is usually an issue with the GPU. This can be resolved through setting the os environment variable `CUDA_VISIBLE_DEVICES=0`. If you are on runpod, this is usually a pod issue. Starting a new pod should take care of it.

-**Q: Received mismatch error on merge adapters / loading adapters between torch.Size of checkpoint and model.**
-
-> A: This is likely due to vocab size mismatch. By default, Axolotl expands the model's embeddings if the tokenizer has more tokens than the model. Please use the `axolotl merge-lora` command to merge the adapters instead of using your own scripts.
-
-> On the other hand, if the model has more tokens than the tokenizer, Axolotl does not shrink the model's embeddings unless `shrink_embeddings: true` is set in the config.
-
-**Q: How to call Axolotl via custom python scripts?**
-
-> A: Yes, since Axolotl is just Python, please see `src/axolotl/cli/main.py` on how each command is called.
-
-### Chat templates
-
 **Q: `jinja2.exceptions.UndefinedError: 'dict object' has no attribute 'content' / 'role' / ____`**

 > A: This means that the property mapping for the stated attribute does not exist when building `chat_template` prompt. For example, if `no attribute 'content'`, please check you have added the correct mapping for `content` under `message_property_mappings`.
-
-**Q: `Empty template generated for turn ___`**
-
-> A: The `content` is empty for that turn.
-
-**Q: `Could not find content start/end boundary for turn __`**
-
-> A: The specific turn's start/end could not be detected. Please ensure you have set the `eos_token` following your `chat_template`. Otherwise, this could be a `chat_template` which doesn't use proper boundaries for each turn (like system). On the rare occurrence, make sure your content is not `[[dummy_message]]`. Please let us know about this.
-
-**Q: `Content end boundary is before start boundary for turn ___`**
-
-> A: This is an edge case which should not occur. Please create an Issue if this happens.
-
-**Q: `Content end boundary is the same as start boundary for turn ___. This is likely an empty turn.`**
-
-> A: This is likely an empty turn.
-
-**Q: The EOS/EOT token is incorrectly being masked or not being masked.**
-
-> A: This is because of the mismatch between `tokenizer.eos_token` and EOS/EOT token in template. Please make sure to set `eos_token` under `special_tokens` to the same EOS/EOT token as in template.
-
-**Q: "`chat_template` choice is `tokenizer_default` but tokenizer's `chat_template` is null. Please add a `chat_template` in tokenizer config"**
-
-> A: This is because the tokenizer does not have a chat template. Please add a chat template in the tokenizer config. See [chat_template](dataset-formats/conversation.qmd#chat-template) for more details.
--- a/docs/getting-started.qmd
+++ b/docs/getting-started.qmd
@@ -1,5 +1,5 @@
 ---
-title: "Quickstart"
+title: "Getting Started with Axolotl"
 format:
  html:
    toc: true
@@ -17,12 +17,12 @@ Let's start by fine-tuning a small language model using LoRA. This example uses
 Assuming `axolotl` is installed (if not, see our [Installation Guide](installation.qmd))

 1. Download example configs:
-```bash
+```shell
 axolotl fetch examples
 ```

 2. Run the training:
-```bash
+```shell
 axolotl train examples/llama-3/lora-1b.yml
 ```

@@ -36,9 +36,7 @@ The YAML configuration file controls everything about your training. Here's what

 ```yaml
 base_model: NousResearch/Llama-3.2-1B
-
-load_in_8bit: true
-adapter: lora
+# hub_model_id: username/custom_model_name

 datasets:
  - path: teknium/GPT4-LLM-Cleaned
@@ -46,15 +44,11 @@ datasets:
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.1
 output_dir: ./outputs/lora-out
+
+adapter: lora
+lora_model_dir:
 ```

-::: {.callout-tip}
-`load_in_8bit: true` and `adapter: lora` enables LoRA adapter finetuning.
-
- To perform Full finetuning, remove these two lines.
- To perform QLoRA finetuning, replace with `load_in_4bit: true` and `adapter: qlora`.
-:::
-
 See our [Config options](config.qmd) for more details.

 ### Training {#sec-training}
@@ -62,7 +56,7 @@ See our [Config options](config.qmd) for more details.
 When you run `axolotl train`, Axolotl:

 1. Downloads the base model
-2. (If specified) applies QLoRA/LoRA adapter layers
+2. (If specified) applies LoRA adapter layers
 3. Loads and processes the dataset
 4. Runs the training loop
 5. Saves the trained model and / or LoRA weights
@@ -75,8 +69,6 @@ Let's modify the example for your own data:

 ```yaml
 base_model: NousResearch/Nous-Hermes-llama-1b-v1
-
-load_in_8bit: true
 adapter: lora

 # Training settings
@@ -112,9 +104,11 @@ format):
 {"instruction": "Classify this text", "input": "Not good at all", "output": "negative"}
 ```

+Please consult the supported [Dataset Formats](dataset-formats/) for more details.
+
 3. Run the training:

-```bash
+```shell
 axolotl train my_training.yml
 ```

@@ -124,7 +118,7 @@ axolotl train my_training.yml

 After training, test your model:

-```bash
+```shell
 axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out"
 ```

@@ -132,7 +126,7 @@ axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out"

 For large datasets, preprocess first:

-```bash
+```shell
 axolotl preprocess my_training.yml
 ```

@@ -140,7 +134,7 @@ axolotl preprocess my_training.yml

 Launch a Gradio interface:

-```bash
+```shell
 axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out" --gradio
 ```

--- a/docs/inference.qmd
+++ b/docs/inference.qmd
@@ -1,22 +1,19 @@
 ---
-title: "Inference and Merging"
+title: "Inference Guide"
 format:
  html:
    toc: true
    toc-depth: 3
    number-sections: true
+    code-tools: true
 execute:
  enabled: false
 ---

-This guide covers how to use your trained models for inference, including model loading, interactive testing, merging adapters, and common troubleshooting steps.
+This guide covers how to use your trained models for inference, including model loading, interactive testing, and common troubleshooting steps.

 ## Quick Start {#sec-quickstart}

-::: {.callout-tip}
-Use the same config used for training on inference/merging.
-:::
-
 ### Basic Inference {#sec-basic}

 ::: {.panel-tabset}
--- a/docs/input_output.qmd
+++ b/docs/input_output.qmd
@@ -3,4 +3,263 @@ title: Template-free prompt construction
 description: "Template-free prompt construction with the `input_output` format"
 ---

-The documentation moved to [here](dataset-formats/template_free.qmd).
+<!-- TOC -->
+
+- [Background](#background)
+    - [Masking Inputs](#masking-inputs)
+    - [You may not want prompt templates](#you-may-not-want-prompt-templates)
+    - [The `input_output` format](#the-input_output-format)
+- [Usage](#usage)
+    - [1. Prepare Data](#1-prepare-data)
+    - [2. Use `type: input_output`](#2-use-type-input_output)
+    - [3. Check the prompts](#3-check-the-prompts)
+
+<!-- /TOC -->
+
+<a id="markdown-background" name="background"></a>
+
+## Background
+
+<a id="markdown-masking-inputs" name="masking-inputs"></a>
+
+### Masking Inputs
+
+One of the most popular features of
+[axolotl](https://github.com/axolotl-ai-cloud/axolotl) is
+setting the following configuration value:
+
+
+```yaml
+train_on_inputs: false
+```
+
+If you declare a [dataset formats](https://github.com/axolotl-ai-cloud/axolotl?tab=readme-ov-file#dataset)
+such as `alpaca` or `chatml`, axolotl knows what is an input
+(i.e. human) vs. an output (i.e. the assistant) and masks the input
+labels so that your model can focus on predicting the outputs only.
+
+<a id="markdown-you-may-not-want-prompt-templates" name="you-may-not-want-prompt-templates"></a>
+
+### You may not want prompt templates
+
+However, there are many situations where you don't want to use one of
+these formats or templates. This is because they can:
+
+-   Add unnecessary boilerplate to your prompts.
+-   Create artifacts like special delimiters `<|im_start|>` that can
+    quickly become footguns if you don't include them correctly at
+    inference time.
+-   Enforce a *chat* interface when you do not want one. Sometimes you
+    just want to fine-tune a model to a very specific task and do NOT
+    want multi-turn conversations, roles, etc.
+-   Limit you to only certain roles that the template allows.
+
+<a id="markdown-the-inputoutput-format" name="the-inputoutput-format"></a>
+
+### The `input_output` format
+
+You can construct your prompts without a template by using the
+`input_output` format, by setting `type: input_output` in your
+configuration file like this:
+
+**config.yml**
+
+```yaml
+train_on_inputs: false # Mask segments of your data
+datasets:
+  - path: output.jsonl
+    type: input_output  # use template free prompt construction
+```
+
+Unlike `type: completion`, which is also template-free,
+`type: input_output` allows you to mask segments of your text. More
+details on how this works are described below.
+
+<a id="markdown-usage" name="usage"></a>
+
+## Usage
+
+This is how you can use the `input_output` format:
+
+<a id="markdown-1-prepare-data" name="1-prepare-data"></a>
+
+### 1. Prepare Data
+
+To use the `input_output` format, collect your data in the following
+format into a jsonl file (below is the first row from the file
+`output`.jsonl` pretty printed):
+
+```bash
+$ head -n1 output.jsonl | python -m json.tool
+```
+
+:::{.cell-output .cell-output-stdout}
+    {
+        "segments": [
+            {
+                "label": true,
+                "text": "<s>Hello\n"
+            },
+            {
+                "label": true,
+                "text": "hi there!. "
+            },
+            {
+                "label": false,
+                "text": "goodbye "
+            },
+            {
+                "label": true,
+                "text": "farewell</s>"
+            }
+        ]
+    }
+:::
+
+Set `label:false` when you want to mask a segment of text so that the
+model isn't trained on it. Some things to keep in mind:
+
+> [!IMPORTANT]
+> 1.  **EOS, BOS, spaces, newlines etc. are entirely up to you. Axolotl
+    concatenates all the segments as-is.** The tokenizer doesn't add
+    anything additional. Notice how I added spaces, newlines, `<s>`
+    (BOS), and `</s>` (EOS) myself.
+> 2.  Make sure you check the materialized output to validate that the
+    prompt is getting assembled how you like.
+
+<a id="markdown-2-use-type-inputoutput" name="2-use-type-inputoutput"></a>
+
+### 2. Use `type: input_output`
+
+Let's materialize data with our `output.jsonl` file by setting
+`type: input_output` in our axolotl config:
+
+```yaml
+# training_config.yaml
+base_model: mistralai/Mistral-7B-v0.1
+data_seed: 49
+seed: 49
+
+datasets:
+  - path: output.jsonl
+    type: input_output
+val_set_size: 0.1
+
+sequence_len: 896
+sample_packing: false
+
+micro_batch_size: 2
+gradient_accumulation_steps: 3
+eval_batch_size: 2
+num_epochs: 1
+learning_rate: 0.0002
+
+train_on_inputs: false
+special_tokens:
+  bos_token: "<s>"
+  eos_token: "</s>"
+  unk_token: "<unk>"
+```
+
+You can use the following command to materialize your data. The
+`--debug` flag will print the tokens, along with the labels so you can
+verify that the correct items are being ignored:
+
+```bash
+$ python -m axolotl.cli.preprocess training_config.yaml --debug
+
+...
+[2024-03-05 23:36:46,969] [INFO] [axolotl.check_example_labels:35] [PID:607731] [RANK:0] <s>(1, 1) Hello(22557, 22557)
+(13, 13) hi(12014, 12014) there(736, 736) !(28808, 28808) .(28723, 28723) (28705, 28705) good(-100, 1179) bye(-100, 17664) (-100, 28705) fare(19111, 19111) well(5458, 5458) </s>(2, 2)
+
+```
+
+The format is `decoded_token`(`label`, `token_id`), for example,
+`<s>(1, 1)` means that the token is `<s>`, the label is `1` and the
+token_id is `1`. When the label is `-100` then that token is ignored for
+training.
+
+<a id="markdown-3-check-the-prompts" name="3-check-the-prompts"></a>
+
+### 3. Check the prompts
+
+Here is another way to check the materialized output:
+
+```python
+from transformers import AutoTokenizer
+from datasets import load_from_disk
+import yaml
+
+directory = !ls last_run_prepared/
+with open('training_config.yaml', 'r') as f:
+    cfg = yaml.safe_load(f)
+model_id = cfg['base_model']
+tok = AutoTokenizer.from_pretrained(model_id)
+ds = load_from_disk(f'last_run_prepared/{directory[0]}/')
+```
+
+```python
+>>> row = ds[0]
+>>> print(tok.decode(row['input_ids']))
+<s> Hello
+    hi there!.  goodbye  farewell</s>
+```
+
+We can check that the right tokens are ignored by comparing the labels
+to each token:
+
+```python
+import pandas as pd
+pd.DataFrame([{'token': tok.decode(i), 'label': l, 'id':i} for i,l in
+              zip(row['input_ids'], row['labels'])])
+```
+
+| token | label | id    |
+|-------|-------|-------|
+| 0     | \<s\> | 1     |
+| 1     | Hello | 22557 |
+| 2     | \\n   | 13    |
+| 3     | hi    | 12014 |
+| 4     | there | 736   |
+| 5     | !     | 28808 |
+| 6     | .     | 28723 |
+| 7     |       | 28705 |
+| 8     | good  | -100  |
+| 9     | bye   | -100  |
+| 10    |       | -100  |
+| 11    | fare  | 19111 |
+| 12    | well  | 5458  |
+| 13    | \</s\>| 2     |
+
+
+
+If we look at the input data, the above table seems correct! (The jsonl
+version is repeated below for reference):
+
+
+```bash
+$ head -n1 output.jsonl | python -m json.tool
+```
+
+:::{.cell-output .cell-output-stdout}
+    {
+        "segments": [
+            {
+                "label": true,
+                "text": "<s>Hello\n"
+            },
+            {
+                "label": true,
+                "text": "hi there!. "
+            },
+            {
+                "label": false,
+                "text": "goodbye "
+            },
+            {
+                "label": true,
+                "text": "farewell</s>"
+            }
+        ]
+    }
+:::
--- a/docs/installation.qmd
+++ b/docs/installation.qmd
@@ -1,10 +1,11 @@
 ---
-title: "Installation"
+title: "Installation Guide"
 format:
  html:
    toc: true
    toc-depth: 3
    number-sections: true
+    code-tools: true
 execute:
  enabled: false
 ---
@@ -22,7 +23,6 @@ This guide covers all the ways you can install and set up Axolotl for your envir
 ### PyPI Installation (Recommended) {#sec-pypi}

 ```{.bash}
-pip3 install -U packaging setuptools wheel ninja
 pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]
 ```

@@ -38,7 +38,7 @@ For the latest features between releases:
 ```{.bash}
 git clone https://github.com/axolotl-ai-cloud/axolotl.git
 cd axolotl
-pip3 install -U packaging setuptools wheel ninja
+pip3 install packaging ninja
 pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'
 ```

@@ -66,8 +66,6 @@ docker run --privileged --gpus '"all"' --shm-size 10g --rm -it \
 ```
 :::

-Please refer to the [Docker documentation](docker.qmd) for more information on the different Docker images that are available.
-
 ## Cloud Environments {#sec-cloud}

 ### Cloud GPU Providers {#sec-cloud-gpu}
@@ -79,7 +77,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}

@@ -109,7 +106,7 @@ We recommend using WSL2 (Windows Subsystem for Linux) or Docker.
 2. Install PyTorch: https://pytorch.org/get-started/locally/
 3. Install Axolotl:
   ```{.bash}
-   pip3 install -U packaging setuptools wheel ninja
+   pip3 install packaging
   pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'
   ```
 4. (Optional) Login to Hugging Face:
--- a/docs/lora_optims.qmd
+++ b/docs/lora_optims.qmd
@@ -1,6 +1,7 @@
 ---
 title: "LoRA Optimizations"
-description: "Custom autograd functions and Triton kernels in Axolotl for optimized LoRA fine-tuning"
+description: "Custom autograd functions and Triton kernels in Axolotl for optimized
+LoRA fine-tuning"
 ---

 Inspired by [Unsloth](https://github.com/unslothai/unsloth), we've implemented two
@@ -66,10 +67,6 @@ logic to be compatible with more of them.

 </details>

-::: {.callout-tip}
-Check out our [LoRA optimizations blog](https://axolotlai.substack.com/p/accelerating-lora-fine-tuning-with).
-:::
-
 ## Usage

 These optimizations can be enabled in your Axolotl config YAML file. The
--- a/docs/mac.qmd
+++ b/docs/mac.qmd
@@ -19,5 +19,4 @@ Current support:
 - [ ] DeepSpeed

 Untested:
-
 - FSDP
--- a/docs/multi-gpu.qmd
+++ b/docs/multi-gpu.qmd
@@ -1,5 +1,5 @@
 ---
-title: "Multi-GPU"
+title: "Multi-GPU Training Guide"
 format:
  html:
    toc: true
@@ -35,11 +35,7 @@ deepspeed: deepspeed_configs/zero1.json
 ### Usage {#sec-deepspeed-usage}

 ```{.bash}
-# Passing arg via config
-axolotl train config.yml
-
-# Passing arg via cli
-axolotl train config.yml --deepspeed deepspeed_configs/zero1.json
+accelerate launch -m axolotl.cli.train examples/llama-2/config.yml --deepspeed deepspeed_configs/zero1.json
 ```

 ### ZeRO Stages {#sec-zero-stages}
@@ -74,7 +70,25 @@ For combining FSDP with QLoRA, see our [dedicated guide](fsdp_qlora.qmd).

 ### Liger Kernel Integration {#sec-liger}

-Please see [docs](custom_integrations.qmd#liger) for more info.
+::: {.callout-note}
+Liger Kernel provides efficient Triton kernels for LLM training, offering:
+
+- 20% increase in multi-GPU training throughput
+- 60% reduction in memory usage
+- Compatibility with both FSDP and DeepSpeed
+:::
+
+Configuration:
+
+```{.yaml}
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+liger_rope: true
+liger_rms_norm: true
+liger_glu_activation: true
+liger_layer_norm: true
+liger_fused_linear_cross_entropy: true
+```

 ## Troubleshooting {#sec-troubleshooting}

--- a/docs/multi-node.qmd
+++ b/docs/multi-node.qmd
@@ -13,7 +13,7 @@ You will also need to have the same configuration file for your model on each ma
 Make sure the main machine is reachable by other machines.
 :::

-## Accelerate
+# Accelerate

 You will need to create a configuration for accelerate, either by using `accelerate config` and follow the instructions or you can use one of the preset below:

@@ -51,17 +51,17 @@ fsdp_config:

 All you have to do now is launch using accelerate as you would usually do on each machine and voila, the processes will start once you have launched accelerate on every machine.

-## Raytrain
+# Raytrain

 Please see ray train doc [here](ray-integration.qmd).

-## Torchrun
+# Torchrun

 If you are using Infiniband, we recommend torchrun to utilize the full bandwidth.

 Set the following env (change buffersize/socketname depending on your system):

-```bash
+```yaml
 export NCCL_IB_DISABLE=0
 export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
 export NCCL_BUFFSIZE=2097152
--- a/docs/nccl.qmd
+++ b/docs/nccl.qmd
@@ -13,13 +13,13 @@ Often, this timeout will happen after 30 minutes (the default setting) and is ac

 Forcing cross-GPU communication via [NVLink](https://en.wikipedia.org/wiki/NVLink) may help without increasing timeouts. To verify that your configuration is leveraging NVLink run the following command:

-```bash
+```shell
 nvidia-smi nvlink --status
 ```

 To force NCCL to use NVLink, simply set this in the environment:

-```bash
+```shell
 export NCCL_P2P_LEVEL=NVL
 ```

@@ -33,13 +33,13 @@ If NVLink is not available in your environment there are other options for ``NCC

 To validate that acceptable data transfer speeds exist for your training job, running [NCCL Tests](https://github.com/NVIDIA/nccl-tests/blob/master/README.md) can help pinpoint bottlenecks, for example:

-```bash
+```shell
 ./build/all_reduce_perf -b 8 -e 128M -f 2 -g 3
 ```

 It can be useful when debugging NCCL communication timeouts to activate additional logging in both PyTorch and NCCL:

-```bash
+```shell
 export NCCL_DEBUG=INFO
 export NCCL_DEBUG_SUBSYS=ALL
 export TORCH_DISTRIBUTED_DEBUG=INFO
--- a/docs/ray-integration.qmd
+++ b/docs/ray-integration.qmd
@@ -1,5 +1,5 @@
 ---
-title: Ray Train
+title: Ray Train integration
 description: How to use Axolotl with Ray Train
 ---

@@ -9,7 +9,7 @@ With the `--use-ray` CLI flag, Axolotl will use Ray Train's [`TorchTrainer`](htt

 ## Ray cluster setup

-A prerequisite using the Ray Train integration is to setup a Ray cluster on your desired node(s). For a detailed guide on how you can get started with ray clusters, check the official Ray docs [here](https://docs.ray.io/en/latest/cluster/getting-started.html).
+A prerequisite using the Ray Train integration is to setup a Ray cluster on your desired node(s). For a detailed guide on how you can get started with ray clusters, check the official Ray docs here: https://docs.ray.io/en/latest/cluster/getting-started.html

 Every Ray cluster has one _head_ node and a set of worker nodes. The head node is just like any other worker node, but it also runs certain special processes related to scheduling and orchestration. Ray-enabled scripts are run on the head node and depending on the resources (number of CPUs, GPUs, etc) they request, will be scheduled to run certain tasks on the worker nodes. For more on key concepts behind a Ray cluster, you can refer this [doc](https://docs.ray.io/en/latest/cluster/key-concepts.html#cluster-key-concepts).

@@ -58,11 +58,13 @@ You can find an example configuration at `configs/llama-3/lora-1b-ray.yaml`.
 The key parameters to note here are:

 ```yaml
+...
 use_ray: true
 ray_num_workers: 4
 # optional
 resources_per_worker:
    GPU: 1
+...
 ```

 - `use_ray`: This is the flag that enables the Ray Train integration. You can either use the corresponding `--use-ray` flag in the CLI or set `use_ray` in the config file.
--- a/docs/reward_modelling.qmd
+++ b/docs/reward_modelling.qmd
@@ -28,23 +28,8 @@ val_set_size: 0.1
 eval_steps: 100
 ```

-Bradley-Terry chat templates expect single-turn conversations in the following format:
-
-```json
-{
-    "system": "...", // optional
-    "input": "...",
-    "chosen": "...",
-    "rejected": "..."
-}
-```
-
 ### Process Reward Models (PRM)

-::: {.callout-tip}
-Check out our [PRM blog](https://axolotlai.substack.com/p/process-reward-models).
-:::
-
 Process reward models are trained using data which contains preference annotations for each step in a series of interactions. Typically, PRMs are trained to provide reward signals over each step of a reasoning trace and are used for downstream reinforcement learning.
 ```yaml
 base_model: Qwen/Qwen2.5-3B
@@ -60,5 +45,3 @@ datasets:
 val_set_size: 0.1
 eval_steps: 100
 ```
-
-Please see [stepwise_supervised](dataset-formats/stepwise_supervised.qmd) for more details on the dataset format.
--- a/docs/rlhf.qmd
+++ b/docs/rlhf.qmd
@@ -3,23 +3,22 @@ title: "RLHF (Beta)"
 description: "Reinforcement Learning from Human Feedback is a method whereby a language model is optimized from data using human feedback."
 back-to-top-navigation: true
 toc: true
-toc-expand: 2
-toc-depth: 4
+toc-depth: 3
 ---

-## Overview
+# Overview

 Reinforcement Learning from Human Feedback is a method whereby a language model is optimized from data using human
 feedback. Various methods include, but not limited to:

+- Proximal Policy Optimization (PPO) (not yet supported in axolotl)
 - [Direct Preference Optimization (DPO)](#dpo)
 - [Identity Preference Optimization (IPO)](#ipo)
 - [Kahneman-Tversky Optimization (KTO)](#kto)
 - [Odds Ratio Preference Optimization (ORPO)](#orpo)
- Proximal Policy Optimization (PPO) (not yet supported in axolotl)


-## RLHF using Axolotl
+# RLHF using Axolotl

 ::: {.callout-important}
 This is a BETA feature and many features are not fully implemented. You are encouraged to open new PRs to improve the integration and functionality.
@@ -31,7 +30,7 @@ We rely on the [TRL](https://github.com/huggingface/trl) library for implementat
 You can find what each method supports by going into `src/axolotl/prompt_strategies/{method}` where `{method}` is one of our supported methods. The `type: ` can be retrieved from `{method}.{function_name}`.
 :::

-### DPO
+## DPO

 Example config:

@@ -48,7 +47,7 @@ datasets:

 DPO supports the following types with the following dataset format:

-#### chatml.argilla
+### chatml.argilla

 ```json
 {
@@ -59,7 +58,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### chatml.argilla_chat
+### chatml.argilla_chat

 ```json
 {
@@ -74,7 +73,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### chatml.icr
+### chatml.icr

 ```json
 {
@@ -85,7 +84,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### chatml.intel
+### chatml.intel

 ```json
 {
@@ -96,7 +95,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### chatml.prompt_pairs
+### chatml.prompt_pairs

 ```json
 {
@@ -107,7 +106,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### chatml.ultra
+### chatml.ultra

 ```json
 {
@@ -124,7 +123,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### llama3.argilla
+### llama3.argilla

 ```json
 {
@@ -135,7 +134,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### llama3.argilla_chat
+### llama3.argilla_chat

 ```json
 {
@@ -150,7 +149,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### llama3.icr
+### llama3.icr

 ```json
 {
@@ -161,7 +160,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### llama3.intel
+### llama3.intel

 ```json
 {
@@ -172,7 +171,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### llama3.prompt_pairs
+### llama3.prompt_pairs

 ```json
 {
@@ -183,7 +182,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### llama3.ultra
+### llama3.ultra

 ```json
 {
@@ -200,7 +199,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### zephyr.nectar
+### zephyr.nectar

 ```json
 {
@@ -219,7 +218,7 @@ DPO supports the following types with the following dataset format:
 }
 ```

-#### chat_template.default
+### chat_template.default

 ```yaml
 rl: dpo
@@ -265,7 +264,7 @@ Sample input format:
 }
 ```

-#### user_defined.default
+### user_defined.default

 For custom behaviors,

@@ -296,15 +295,15 @@ The input format is a simple JSON input with customizable fields based on the ab
 }
 ```

-### IPO
+## IPO

-As IPO is just DPO with a different loss function, all supported dataset formats for [DPO](#dpo) are also supported for IPO.
+As IPO is just DPO with a different loss function, all supported options for DPO works here.

 ```yaml
 rl: ipo
 ```

-### ORPO
+## ORPO

 Paper: https://arxiv.org/abs/2403.07691

@@ -321,7 +320,7 @@ datasets:

 ORPO supports the following types with the following dataset format:

-#### chat_template.argilla
+### chat_template.argilla

 ```json
 {
@@ -340,13 +339,12 @@ ORPO supports the following types with the following dataset format:
 }
 ```

-### KTO
+## KTO

 ```yaml
 rl: kto
-rl_beta: 0.1  # default
-kto_desirable_weight: 1.0  # default
-kto_undesirable_weight: 1.0  # default
+rl_beta: 0.5
+kto_desirable_weight: 0.2

 remove_unused_columns: false

@@ -362,7 +360,7 @@ gradient_checkpointing_kwargs:

 KTO supports the following types with the following dataset format:

-#### chatml.argilla
+### chatml.argilla

 ```json
 {
@@ -372,7 +370,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### chatml.argilla_chat
+### chatml.argilla_chat

 ```json
 {
@@ -385,7 +383,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### chatml.intel
+### chatml.intel

 ```json
 {
@@ -395,7 +393,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### chatml.prompt_pairs
+### chatml.prompt_pairs

 ```json
 {
@@ -405,7 +403,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### chatml.ultra
+### chatml.ultra

 ```json
 {
@@ -415,7 +413,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### llama3.argilla
+### llama3.argilla

 ```json
 {
@@ -425,7 +423,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### llama3.argilla_chat
+### llama3.argilla_chat

 ```json
 {
@@ -436,7 +434,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### llama3.intel
+### llama3.intel

 ```json
 {
@@ -446,7 +444,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### llama3.prompt_pairs
+### llama3.prompt_pairs

 ```json
 {
@@ -456,7 +454,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### llama3.ultra
+### llama3.ultra

 ```json
 {
@@ -466,7 +464,7 @@ KTO supports the following types with the following dataset format:
 }
 ```

-#### user_defined.default
+### user_defined.default

 For custom behaviors,

@@ -496,69 +494,7 @@ The input format is a simple JSON input with customizable fields based on the ab
 }
 ```

-### GRPO
-
-::: {.callout-tip}
-Check out our [GRPO cookbook](https://github.com/axolotl-ai-cloud/axolotl-cookbook/tree/main/grpo#training-an-r1-style-large-language-model-using-grpo).
-:::
-
-GRPO uses custom reward functions and transformations. Please have them ready locally.
-
-For ex, to load OpenAI's GSM8K and use a random reward for completions:
-
-```python
-# rewards.py
-import random
-
-def rand_reward_func(completions, **kwargs) -> list[float]:
-    return [random.uniform(0, 1) for _ in completions]
-
-def oai_gsm8k_transform(cfg, *args, **kwargs):
-    def transform_fn(example, tokenizer=None):
-        label = example["answer"].split("####")[-1].strip().replace(",", "")
-        return {
-            "prompt": [{"role": "user", "content": example["question"]},],
-            "answer": label,
-        }
-    return transform_fn, {"remove_columns": ["question"]}
-```
-
-```yaml
-rl: grpo
-
-trl:
-    beta: 0.001
-    max_completion_length: 256
-    use_vllm: True
-    vllm_device: auto
-    vllm_gpu_memory_utilization: 0.15
-    num_generations: 4
-    reward_funcs: ["rewards.rand_reward_func"]    # format: '{file_name}.{fn_name}'
-    reward_weights: [1.0]
-datasets:
-  - path: openai/gsm8k
-    name: main
-    type: rewards.oai_gsm8k_transform  # format: '{file_name}.{fn_name}'
-```
-
-To see other examples of custom reward functions, please see [TRL GRPO Docs](https://github.com/huggingface/trl/blob/main/docs/source/grpo_trainer.md#using-a-custom-reward-function).
-
-To see description of the configs, please see [TRLConfig](https://github.com/axolotl-ai-cloud/axolotl/blob/main/src/axolotl/utils/config/models/input/v0_4_1/trl.py).
-
-### SimPO
-
-SimPO uses [CPOTrainer](https://huggingface.co/docs/trl/main/en/cpo_trainer) but with alternative loss function.
-
-```yaml
-rl: simpo
-rl_beta: 0.1  # default in CPOTrainer
-cpo_alpha: 1.0  # default in CPOTrainer
-simpo_gamma: 0.5  # default in CPOTrainer
-```
-
-This method uses the same dataset format as [DPO](#dpo).
-
-### Using local dataset files
+## Using local dataset files

 ```yaml
 datasets:
@@ -569,7 +505,7 @@ datasets:
    type: chatml.intel
 ```

-### TRL auto-unwrapping for PEFT
+## TRL auto-unwrapping for PEFT

 TRL supports auto-unwrapping PEFT models for RL training paradigms which rely on a reference model. This significantly reduces memory pressure as an additional refreference model does not need to be loaded, and reference model log-probabilities can be obtained by disabling PEFT adapters. This is enabled by default. To turn it off, pass the following config:

--- a/docs/torchao.qmd
+++ b/docs/torchao.qmd
@@ -3,12 +3,6 @@ title: "PyTorch ao"
 description: "Custom data types and layouts for training and inference"
 ---

-To use experimental optimizers (`AdamWFp8`, `AdamW4bit`, `AdamW8bit`) from Pytorch Ao, please install the package as shown below.
-
-::: {.callout-tip}
-Some experimental optimizers are already present in regular Pytorch, so please re-check if you actually need this package!
-:::
-
 ### Installation

 Stable Release from the PyTorch index
--- a/docs/unsloth.qmd
+++ b/docs/unsloth.qmd
@@ -8,12 +8,6 @@ description: "Hyper-optimized QLoRA finetuning for single GPUs"
 Unsloth provides hand-written optimized kernels for LLM finetuning that slightly improve speed and VRAM over
 standard industry baselines.

-::: {.callout-important}
-Due to breaking changes in transformers `v4.48.0`, users will need to downgrade to `<=v4.47.1` to use this patch.
-
-This will later be deprecated in favor of [LoRA Optimizations](lora_optims.qmd).
-:::
-

 ### Installation

@@ -23,7 +17,7 @@ The following will install the correct unsloth and extras from source.
 python scripts/unsloth_install.py | sh
 ```

-### Usage
+### Using unsloth w Axolotl

 Axolotl exposes a few configuration options to try out unsloth and get most of the performance gains.

--- a/examples/llama-3/qlora-1b-kto.yaml
+++ b/examples/llama-3/qlora-1b-kto.yaml
@@ -55,7 +55,7 @@ tf32: true

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
-  use_reentrant: false
+  use_reentrant: true
 early_stopping_patience:
 resume_from_checkpoint:
 local_rank:
--- a/index.qmd
+++ b/index.qmd
@@ -1,7 +1,7 @@
 ---
-# toc-location: right-body
-# toc-title: Table Of Contents
-# toc-expand: 2
+toc-location: right-body
+toc-title: Table Of Contents
+toc-expand: 2
 ---

 ```{python}
--- 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,24 +1,24 @@
 --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
 xformers>=0.0.23.post1
 autoawq==0.2.7.post3
-liger-kernel==0.5.3
+liger-kernel==0.5.2
 # END section

 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
-trl==0.15.1
+tokenizers>=0.21.0
+accelerate==1.3.0
+datasets==3.2.0
+deepspeed==0.16.1
+trl==0.15.0

 optimum==1.16.2
 hf_transfer
@@ -62,5 +62,4 @@ antlr4-python3-runtime==4.13.2
 torchao==0.7.0
 schedulefree==1.3.0

-axolotl-contribs-lgpl==0.0.6
-axolotl-contribs-mit==0.0.3
+axolotl-contribs-lgpl==0.0.3
--- 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

@@ -25,5 +24,5 @@ if cce_spec:

 print(
    UNINSTALL_PREFIX
-    + 'pip install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@24fbe4b5dab9a6c250a014573613c1890190536c"'
+    + 'pip install "cut-cross-entropy @ git+https://github.com/apple/ml-cross-entropy.git@9c297c905f55b73594b5d650722d1e78183b77bd"'
 )
--- 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
@@ -114,7 +113,7 @@ class ModalCloud(Cloud):
                [
                    # Random id for cache busting of branch commits
                    f"RUN echo '{str(randint(0, 1000000))}'",  # nosec B311
-                    f"RUN cd /workspace/axolotl && git fetch && git checkout {self.config.branch} && git pull",
+                    f"RUN cd /workspace/axolotl && git fetch && git checkout {self.config.branch}",
                ]
            )

@@ -124,6 +123,8 @@ class ModalCloud(Cloud):
        if env := self.get_env():
            image = image.env(env)

+        image = image.pip_install("fastapi==0.110.0", "pydantic==2.6.3")
+
        return image

    def get_secrets(self):
@@ -259,22 +260,25 @@ class ModalCloud(Cloud):


 def _preprocess(config_yaml: str, volumes=None):
-    Path("/workspace/mounts").mkdir(parents=True, exist_ok=True)
-    with open("/workspace/mounts/config.yaml", "w", encoding="utf-8") as f_out:
+    Path("/workspace/artifacts/axolotl").mkdir(parents=True, exist_ok=True)
+    with open(
+        "/workspace/artifacts/axolotl/config.yaml", "w", encoding="utf-8"
+    ) as f_out:
        f_out.write(config_yaml)
-    run_folder = "/workspace/mounts"
+    run_folder = "/workspace/artifacts/axolotl"
    run_cmd(
-        "axolotl preprocess /workspace/mounts/config.yaml --dataset-processes=8",
+        "axolotl preprocess /workspace/artifacts/axolotl/config.yaml --dataset-processes=8",
        run_folder,
        volumes,
    )


 def _train(config_yaml: str, accelerate: bool = True, volumes=None, **kwargs):
-    Path("/workspace/mounts").mkdir(parents=True, exist_ok=True)
-    with open("/workspace/mounts/config.yaml", "w", encoding="utf-8") as f_out:
+    with open(
+        "/workspace/artifacts/axolotl/config.yaml", "w", encoding="utf-8"
+    ) as f_out:
        f_out.write(config_yaml)
-    run_folder = "/workspace/mounts"
+    run_folder = "/workspace/artifacts/axolotl"
    if accelerate:
        accelerate_args = "--accelerate"
    else:
@@ -283,19 +287,20 @@ def _train(config_yaml: str, accelerate: bool = True, volumes=None, **kwargs):
    if num_processes := kwargs.pop("num_processes", None):
        num_processes_args = f"--num-processes {num_processes}"
    run_cmd(
-        f"axolotl train {accelerate_args} {num_processes_args} /workspace/mounts/config.yaml",
+        f"axolotl train {accelerate_args} {num_processes_args} /workspace/artifacts/axolotl/config.yaml",
        run_folder,
        volumes,
    )


 def _lm_eval(config_yaml: str, volumes=None):
-    Path("/workspace/mounts").mkdir(parents=True, exist_ok=True)
-    with open("/workspace/mounts/config.yaml", "w", encoding="utf-8") as f_out:
+    with open(
+        "/workspace/artifacts/axolotl/config.yaml", "w", encoding="utf-8"
+    ) as f_out:
        f_out.write(config_yaml)
-    run_folder = "/workspace/mounts"
+    run_folder = "/workspace/artifacts/axolotl"
    run_cmd(
-        "axolotl lm-eval /workspace/mounts/config.yaml",
+        "axolotl lm-eval /workspace/artifacts/axolotl/config.yaml",
        run_folder,
        volumes,
    )
--- 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/train.py
+++ b/src/axolotl/cli/train.py
@@ -1,7 +1,6 @@
 """CLI to run training on a model."""

 import logging
-import os
 from pathlib import Path
 from typing import Union

@@ -35,20 +34,18 @@ def do_train(cfg: DictDefault, cli_args: TrainerCliArgs) -> None:
    """
    print_axolotl_text_art()
    check_accelerate_default_config()
-    if int(os.getenv("LOCAL_RANK", "0")) == 0:
-        check_user_token()
+    check_user_token()

    if cfg.rl:
        dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
    else:
        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

-    model, tokenizer, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
+    model, tokenizer = train(cfg=cfg, dataset_meta=dataset_meta)
    plugin_manager = PluginManager.get_instance()

    del model
    del tokenizer
-    del trainer

    plugin_manager.post_train_unload(cfg)

--- 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/common/datasets.py
+++ b/src/axolotl/common/datasets.py
@@ -24,8 +24,8 @@ class TrainDatasetMeta:
    """Dataclass with fields for training and validation datasets and metadata."""

    train_dataset: Dataset
-    eval_dataset: Dataset | None = None
-    total_num_steps: int | None = None
+    eval_dataset: Optional[Dataset] = None
+    total_num_steps: Optional[int] = None


 def sample_dataset(dataset: Dataset, num_samples: int) -> Dataset:
--- 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

@@ -44,7 +43,7 @@ class TokenizedChatDataset(Dataset):
        process_or_cpu_count: int = (
            process_count or os.cpu_count()  # type: ignore[assignment]
        )
-        num_proc = min(32, process_or_cpu_count)
+        num_proc = min(64, process_or_cpu_count)
        features = data.features.keys()
        tokenized_data = data.map(
            map_fn,
--- 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
@@ -33,7 +35,6 @@ from transformers import (
    EarlyStoppingCallback,
    TrainerCallback,
 )
-from transformers.training_args import OptimizerNames
 from trl.trainer.utils import RewardDataCollatorWithPadding

 from axolotl.core.trainers.base import (
@@ -84,18 +85,19 @@ from axolotl.utils.collators import (
 )
 from axolotl.utils.collators.mm_chat import MultiModalChatDataCollator
 from axolotl.utils.models import ensure_dtype
-from axolotl.utils.schemas.enums import CustomSupportedOptimizers

 try:
    import torch._dynamo  # pylint: disable=ungrouped-imports
 except ImportError:
    pass

-LOG = logging.getLogger(__name__)
+LOG = logging.getLogger("axolotl.core.trainer_builder")


 class TrainerBuilderBase(abc.ABC):
-    """Base class for trainer builder."""
+    """
+    Base class for trainer builder
+    """

    _train_dataset = None
    _eval_dataset = None
@@ -108,9 +110,9 @@ class TrainerBuilderBase(abc.ABC):
        self.tokenizer = tokenizer
        self.processor = processor

-        # If the model supports tagging, add the axolotl tag.
+        # in case the model supports tagging, add the axolotl tag.
        # This makes sure the tag is correctly pushed even if a user calls
-        # model.push_to_hub instead of trainer.push_to_hub.
+        # model.push_to_hub instad of  trainer.push_to_hub.
        if hasattr(model, "add_model_tags"):
            model.add_model_tags(["axolotl"])

@@ -225,8 +227,8 @@ class TrainerBuilderBase(abc.ABC):

 class HFCausalTrainerBuilder(TrainerBuilderBase):
    """
-    Build the HuggingFace training args/trainer for causal models and reward modeling
-    using TRL.
+    Build the HuggingFace training args/trainer for causal models
+    and reward modelling using TRL.
    """

    def get_callbacks(self):
@@ -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
@@ -549,12 +551,34 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
            training_arguments_kwargs["run_name"] = self.cfg.mlflow_run_name
        else:
            training_arguments_kwargs["run_name"] = None
+        training_arguments_kwargs["optim"] = (
+            self.cfg.optimizer if self.cfg.optimizer else "adamw_hf"
+        )
+        if self.cfg.optim_args:
+            if isinstance(self.cfg.optim_args, dict):
+                optim_args = ",".join(
+                    [f"{key}={value}" for key, value in self.cfg.optim_args.items()]
+                )
+            else:
+                optim_args = self.cfg.optim_args
+            training_arguments_kwargs["optim_args"] = optim_args
+        if self.cfg.optim_target_modules:
+            training_arguments_kwargs[
+                "optim_target_modules"
+            ] = self.cfg.optim_target_modules
+        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["embedding_lr"] = self.cfg.embedding_lr
+        training_arguments_kwargs["embedding_lr_scale"] = self.cfg.embedding_lr_scale
+        training_arguments_kwargs["lr_groups"] = self.cfg.lr_groups

-        if self.cfg.lr_scheduler in ["one_cycle", "rex", "log_sweep"]:
+        if self.cfg.lr_scheduler in ["one_cycle", "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 +587,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 +602,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,127 +649,59 @@ 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 = {}

        if self.cfg.reward_model:
            training_arguments_kwargs["max_length"] = self.cfg.sequence_len

-        # Handle custom optimizer
-        custom_supported_optimizers = [opt.value for opt in CustomSupportedOptimizers]
-        if self.cfg.optimizer in custom_supported_optimizers:
-            # Common optimizer kwargs
-            optimizer_kwargs = {
-                "lr": training_arguments_kwargs.get("learning_rate"),
-                "weight_decay": training_arguments_kwargs.get("weight_decay"),
-            }
+        # pylint: disable=duplicate-code
+        if self.cfg.optimizer in [
+            "optimi_adamw",
+            "ao_adamw_4bit",
+            "ao_adamw_8bit",
+            "ao_adamw_fp8",
+            "adopt_adamw",
+        ]:
+            # Set default so transformers doesn't throw
+            training_arguments_kwargs["optim"] = "adamw_hf"
+            training_arguments_kwargs["alternate_optimizer"] = self.cfg.optimizer

-            # Adam-specific kwargs
-            adam_kwargs = {}
-            if training_arguments_kwargs.get(
-                "adam_beta1"
-            ) and training_arguments_kwargs.get("adam_beta2"):
-                adam_kwargs["betas"] = (
-                    training_arguments_kwargs.get("adam_beta1"),
-                    training_arguments_kwargs.get("adam_beta2"),
-                )
-            if training_arguments_kwargs.get("adam_epsilon"):
-                adam_kwargs["eps"] = training_arguments_kwargs.get("adam_epsilon")
+        if self.cfg.optimizer == "lion_pytorch":
+            from lion_pytorch import Lion

-            if self.cfg.optimizer == "muon":
-                from axolotl.contribs.mit.muon import (  # pylint: disable=no-name-in-module
-                    MuonOptimizerFactory,
+            lion_kwargs = {"lr": training_arguments_kwargs["learning_rate"]}
+            if "weight_decay" in training_arguments_kwargs:
+                lion_kwargs["weight_decay"] = training_arguments_kwargs["weight_decay"]
+
+            if (
+                "adam_beta1" in training_arguments_kwargs
+                and "adam_beta2" in training_arguments_kwargs
+            ):
+                lion_kwargs["betas"] = (
+                    training_arguments_kwargs["adam_beta1"],
+                    training_arguments_kwargs["adam_beta2"],
                )

-                optimizer_cls = MuonOptimizerFactory
-                optimizer_kwargs.update(adam_kwargs)
-            elif self.cfg.optimizer == "optimi_adamw":
-                from optimi import AdamW
-
-                optimizer_kwargs["foreach"] = False
-                optimizer_cls = AdamW
-                optimizer_kwargs.update(adam_kwargs)
-            elif self.cfg.optimizer == "ao_adamw_4bit":
-                # TODO remove 20250401
-                from torchao.prototype.low_bit_optim import AdamW4bit
-
-                optimizer_cls = AdamW4bit
-                optimizer_kwargs.update(adam_kwargs)
-
-                LOG.warning(
-                    f"`ao_adamw_4bit` will be deprecated soon. Please use `{OptimizerNames.ADAMW_TORCH_4BIT}` instead."
-                )
-            elif self.cfg.optimizer == "ao_adamw_8bit":
-                from torchao.prototype.low_bit_optim import AdamW8bit
-
-                optimizer_cls = AdamW8bit
-                optimizer_kwargs.update(adam_kwargs)
-            elif self.cfg.optimizer == "ao_adamw_fp8":
-                from torchao.prototype.low_bit_optim import AdamWFp8
-
-                optimizer_cls = AdamWFp8
-                optimizer_kwargs.update(adam_kwargs)
-            elif self.cfg.optimizer == "adopt_adamw":
-                from axolotl.utils.optimizers.adopt import ADOPT
-
-                optimizer_cls = ADOPT
-                adam_kwargs["decouple"] = True
-                optimizer_kwargs.update(adam_kwargs)
-
-            # Parse any additional optimizer args from config
-            if self.cfg.optim_args:
-                if isinstance(self.cfg.optim_args, dict):
-                    optimizer_kwargs.update(self.cfg.optim_args)
-                else:
-                    # Parse string format "key1=value1,key2=value2"
-                    for mapping in self.cfg.optim_args.replace(" ", "").split(","):
-                        key, value = mapping.split("=")
-                        optimizer_kwargs[key] = value
-
-            trainer_kwargs["optimizer_cls_and_kwargs"] = (
-                optimizer_cls,
-                optimizer_kwargs,
+            trainer_kwargs["optimizers"] = (
+                Lion(params=self.model.parameters(), **lion_kwargs),
+                None,
            )
-        else:
-            # Use transformers' optimizer
-            training_arguments_kwargs["optim"] = self.cfg.optimizer
-
-            # Parse any additional optimizer args from config
-            if self.cfg.optim_args:
-                if isinstance(self.cfg.optim_args, dict):
-                    optim_args = ",".join(
-                        [f"{key}={value}" for key, value in self.cfg.optim_args.items()]
-                    )
-                else:
-                    optim_args = self.cfg.optim_args
-                training_arguments_kwargs["optim_args"] = optim_args
+            # Set default so transformers doesn't throw
+            training_arguments_kwargs["optim"] = "adamw_hf"

        if self.cfg.optimizer == "adamw_anyprecision":
            if Path(self.cfg.torchdistx_path).exists():
                sys.path.append(self.cfg.torchdistx_path)
                importlib.import_module("torchdistx")

-        if 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["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 +710,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
@@ -916,7 +872,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):


 class HFRLTrainerBuilder(TrainerBuilderBase):
-    """Trainer factory class for TRL-based RLHF trainers (e.g. DPO)"""
+    """
+    Trainer factory class for TRL-based RLHF trainers (e.g. DPO)
+    """

    def get_callbacks(self):
        callbacks = super().get_callbacks()
@@ -970,32 +928,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 +1027,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 +1064,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
@@ -14,7 +14,6 @@ from typing import Dict, Literal, Optional
 import torch
 from datasets import Dataset
 from peft.optimizers import create_loraplus_optimizer
-from torch import nn
 from torch.optim.lr_scheduler import OneCycleLR
 from torch.utils.data import BatchSampler, DataLoader, RandomSampler, SequentialSampler
 from transformers import Trainer
@@ -23,11 +22,9 @@ from transformers.utils import is_sagemaker_mp_enabled
 from trl import CPOTrainer, KTOTrainer, ORPOTrainer, PRMTrainer, RewardTrainer
 from trl.trainer.utils import pad_to_length

-from axolotl.integrations.base import BaseOptimizerFactory
 from axolotl.monkeypatch.relora import ReLoRAScheduler
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 from axolotl.utils.schedulers import (
-    RexLR,
    get_cosine_schedule_with_min_lr,
    get_cosine_schedule_with_quadratic_warmup,
    get_cosine_schedule_with_warmup_decay_constant,
@@ -118,17 +115,6 @@ class SchedulerMixin(Trainer):
                    **extra_lr_kwargs,
                    **self.args.lr_scheduler_kwargs,
                )
-            elif self.args.alternate_lr_scheduler_type == "rex":
-                if use_cosine_min_lr:
-                    assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
-
-                self.lr_scheduler = RexLR(
-                    optimizer=optimizer,
-                    max_lr=self.args.learning_rate,
-                    min_lr=0 if not use_cosine_min_lr else (self.args.learning_rate * self.args.cosine_min_lr_ratio),
-                    total_steps=num_training_steps,
-                    num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
-                )
            elif use_cosine_quadratic:
                if use_cosine_min_lr:
                    LOG.warning("Both cosine quadratic warmup and min lr detected. Using quadratic warmup.")
@@ -168,18 +154,47 @@ class SchedulerMixin(Trainer):
        return self.lr_scheduler


-class OptimizerMixin(Trainer):
+class AxolotlTrainer(SchedulerMixin, Trainer):
    """
-    Mixin class for shared handling of building custom optimizers
+    Extend the base Trainer for axolotl helpers
    """

    args = None  # type: "AxolotlTrainingArguments"  # type: ignore[name-defined]
+    tag_names = ["axolotl"]

-    def create_optimizer_grouped_parameters(
-        self, opt_model, optimizer_kwargs
-    ) -> list[dict]:
+    def __init__(
+        self,
+        *_args,
+        bench_data_collator=None,
+        eval_data_collator=None,
+        dataset_tags=None,
+        **kwargs,
+    ):
+        self.bench_data_collator = bench_data_collator
+        self.eval_data_collator = eval_data_collator
+        self.dataset_tags = dataset_tags
+        self._signature_columns = None  # workaround for pylint
+        super().__init__(*_args, **kwargs)
+        self.train_data_collator = self.data_collator
+        self._stored_metrics = defaultdict(lambda: defaultdict(list))
+        if self.args.orpo_alpha:
+            self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+
+    def _wrap_model(self, model, training=True, dataloader=None):
+        if self.args.torch_compile:
+            torch._dynamo.config.accumulated_cache_size_limit = (  # pylint: disable=protected-access
+                256
+            )
+            model = torch.compile(
+                model,
+                backend=self.args.torch_compile_backend,
+                mode=self.args.torch_compile_mode,
+            )
+        return super()._wrap_model(model, training=training, dataloader=dataloader)
+
+    def create_optimizer_grouped_parameters(self, opt_model, optimizer_kwargs):
        decay_parameters = self.get_decay_parameter_names(opt_model)
-        params: dict = {
+        params = {
            "to_weight_decay": {},  # LayerNorm and bias
            "embeddings": {},  # lm_head, embed_tokens,
            "no_weight_decay": {},
@@ -266,30 +281,23 @@ class OptimizerMixin(Trainer):
            and self.args.embedding_lr_scale is None
            and self.args.embedding_lr is None
            and self.args.lr_groups is None
-            and self.optimizer_cls_and_kwargs is None
+            and self.args.alternate_optimizer
+            not in [
+                "optimi_adamw",
+                "ao_adamw_8bit",
+                "ao_adamw_4bit",
+                "ao_adamw_fp8",
+                "adopt_adamw",
+            ]
        ):
            return super().create_optimizer()

        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
-
-        if (
-            not self.optimizer
-            and self.optimizer_cls_and_kwargs is not None
-            and issubclass(self.optimizer_cls_and_kwargs[0], BaseOptimizerFactory)
-        ):
-            optimizer_factory_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
-            self.optimizer = optimizer_factory_cls()(
-                opt_model, self.args, **optimizer_kwargs
+        if self.optimizer is None:  # pylint: disable=access-member-before-definition
+            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
+                self.args,
+                opt_model,
            )
-
-        if not self.optimizer:
-            if self.optimizer_cls_and_kwargs is not None:
-                optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
-            else:
-                optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(
-                    self.args, opt_model
-                )
-
            optimizer_grouped_parameters = self.create_optimizer_grouped_parameters(
                opt_model, optimizer_kwargs
            )
@@ -306,47 +314,50 @@ class OptimizerMixin(Trainer):
                    loraplus_lr_embedding=loraplus_lr_embedding,
                    **optimizer_kwargs,
                )
-            else:
-                # Overwrite `params` in case it's created by `get_optimizer_cls_and_kwargs`
-                # e.g. for GaLore optimizer.
-                if "params" in optimizer_kwargs:
-                    optimizer_grouped_parameters = optimizer_kwargs.pop("params")
-
-                # Overwrite `model` in case it's created by `get_optimizer_cls_and_kwargs`
-                # e.g. for LOMO optimizer.
-                if "model" in optimizer_kwargs:
-                    optimizer_grouped_parameters = optimizer_kwargs.pop("model")
-
-                # For layer-wise dummy optimizers we overwrite optimizer_grouped_parameters with `optimizer_dict`
-                # to avoid arguments conflicts.
-                if "optimizer_dict" in optimizer_kwargs:
-                    optimizer_grouped_parameters = optimizer_kwargs.pop(
-                        "optimizer_dict"
-                    )
-
-                self.optimizer = optimizer_cls(
-                    optimizer_grouped_parameters, **optimizer_kwargs
+            elif (
+                self.args.embedding_lr_scale is not None
+                or self.args.embedding_lr is not None
+                or self.args.lr_groups is not None
+            ):
+                self.optimizer = (  # pylint: disable=attribute-defined-outside-init
+                    optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
                )
+            elif self.args.alternate_optimizer == "optimi_adamw":
+                from optimi import AdamW

-            if optimizer_cls.__name__ == "Adam8bit":
-                import bitsandbytes
+                self.optimizer = (  # pylint: disable=attribute-defined-outside-init
+                    AdamW(
+                        optimizer_grouped_parameters, foreach=False, **optimizer_kwargs
+                    )
+                )
+            elif self.args.alternate_optimizer == "ao_adamw_4bit":
+                from torchao.prototype.low_bit_optim import AdamW4bit

-                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
+                self.optimizer = (  # pylint: disable=attribute-defined-outside-init
+                    AdamW4bit(optimizer_grouped_parameters, **optimizer_kwargs)
+                )
+            elif self.args.alternate_optimizer == "ao_adamw_8bit":
+                from torchao.prototype.low_bit_optim import AdamW8bit

-                skipped = 0
-                for module in opt_model.modules():
-                    if isinstance(module, nn.Embedding):
-                        skipped += sum(
-                            {
-                                p.data_ptr(): p.numel() for p in module.parameters()
-                            }.values()
-                        )
-                        LOG.info(f"skipped {module}: {skipped/2**20}M params")
-                        manager.register_module_override(
-                            module, "weight", {"optim_bits": 32}
-                        )
-                        LOG.debug(f"bitsandbytes: will optimize {module} in fp32")
-                LOG.info(f"skipped: {skipped/2**20}M params")
+                self.optimizer = (  # pylint: disable=attribute-defined-outside-init
+                    AdamW8bit(optimizer_grouped_parameters, **optimizer_kwargs)
+                )
+            elif self.args.alternate_optimizer == "ao_adamw_fp8":
+                from torchao.prototype.low_bit_optim import AdamWFp8
+
+                self.optimizer = (  # pylint: disable=attribute-defined-outside-init
+                    AdamWFp8(optimizer_grouped_parameters, **optimizer_kwargs)
+                )
+            elif self.args.alternate_optimizer == "adopt_adamw":
+                from axolotl.utils.optimizers.adopt import ADOPT
+
+                self.optimizer = (  # pylint: disable=attribute-defined-outside-init
+                    ADOPT(
+                        optimizer_grouped_parameters,
+                        decouple=True,
+                        **optimizer_kwargs,
+                    )
+                )

        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(  # pylint: disable=attribute-defined-outside-init
@@ -355,45 +366,6 @@ class OptimizerMixin(Trainer):

        return self.optimizer

-
-class AxolotlTrainer(SchedulerMixin, OptimizerMixin, Trainer):
-    """
-    Extend the base Trainer for axolotl helpers
-    """
-
-    args = None  # type: "AxolotlTrainingArguments"  # type: ignore[name-defined]
-    tag_names = ["axolotl"]
-
-    def __init__(
-        self,
-        *_args,
-        bench_data_collator=None,
-        eval_data_collator=None,
-        dataset_tags=None,
-        **kwargs,
-    ):
-        self.bench_data_collator = bench_data_collator
-        self.eval_data_collator = eval_data_collator
-        self.dataset_tags = dataset_tags
-        self._signature_columns = None  # workaround for pylint
-        super().__init__(*_args, **kwargs)
-        self.train_data_collator = self.data_collator
-        self._stored_metrics = defaultdict(lambda: defaultdict(list))
-        if self.args.orpo_alpha:
-            self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
-
-    def _wrap_model(self, model, training=True, dataloader=None):
-        if self.args.torch_compile:
-            torch._dynamo.config.accumulated_cache_size_limit = (  # pylint: disable=protected-access
-                256
-            )
-            model = torch.compile(
-                model,
-                backend=self.args.torch_compile_backend,
-                mode=self.args.torch_compile_mode,
-            )
-        return super()._wrap_model(model, training=training, dataloader=dataloader)
-
    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
        if self.args.sample_packing and not self.args.pretraining:
            if self.args.multipack_real_batches:
@@ -462,9 +434,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 +481,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,6 @@ import logging
 from trl.trainer.grpo_trainer import RewardFunc

 from axolotl.core.trainers.grpo.trainer import AxolotlGRPOTrainer
-from axolotl.utils.schemas.trl import TRLConfig

 LOG = logging.getLogger("axolotl")

@@ -32,44 +31,30 @@ class GRPOStrategy:
    @classmethod
    def set_training_args_kwargs(cls, cfg):
        grpo_args_kwargs = {}
-
-        if not hasattr(cfg, "trl") or not cfg.trl:
-            return grpo_args_kwargs
-
-        trl: TRLConfig = cfg.trl  # type: ignore
-
-        if trl.use_vllm:
-            grpo_args_kwargs["use_vllm"] = trl.use_vllm
-            grpo_args_kwargs["vllm_device"] = (
-                trl.vllm_device if trl.vllm_device else "auto"
-            )
-
-            if 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
-
-        if trl.num_generations:
-            grpo_args_kwargs["num_generations"] = trl.num_generations
-
-        if trl.sync_ref_model:
-            grpo_args_kwargs["sync_ref_model"] = trl.sync_ref_model
-
-            if trl.ref_model_mixup_alpha:
-                grpo_args_kwargs["ref_model_mixup_alpha"] = trl.ref_model_mixup_alpha
-
-            if trl.ref_model_sync_steps:
-                grpo_args_kwargs["ref_model_sync_steps"] = trl.ref_model_sync_steps
-
-        grpo_args_kwargs["max_completion_length"] = trl.max_completion_length
-        grpo_args_kwargs["log_completions"] = trl.log_completions
-
-        if trl.reward_weights:
-            grpo_args_kwargs["reward_weights"] = trl.reward_weights
-
+        if cfg.trl and cfg.trl.use_vllm:
+            grpo_args_kwargs["use_vllm"] = cfg.trl.use_vllm
+            if cfg.trl and cfg.trl.vllm_device:
+                grpo_args_kwargs["vllm_device"] = cfg.trl.vllm_device
+            else:
+                grpo_args_kwargs["vllm_device"] = "auto"
+            if cfg.trl and cfg.trl.vllm_gpu_memory_utilization:
+                grpo_args_kwargs[
+                    "vllm_gpu_memory_utilization"
+                ] = cfg.trl.vllm_gpu_memory_utilization
+            if cfg.trl and cfg.trl.vllm_max_model_len:
+                grpo_args_kwargs["vllm_max_model_len"] = cfg.trl.vllm_max_model_len
+        if cfg.trl and cfg.trl.num_generations:
+            grpo_args_kwargs["num_generations"] = cfg.trl.num_generations
+        if cfg.trl and cfg.trl.sync_ref_model:
+            grpo_args_kwargs["sync_ref_model"] = cfg.trl.sync_ref_model
+            if cfg.trl and cfg.trl.ref_model_mixup_alpha:
+                grpo_args_kwargs[
+                    "ref_model_mixup_alpha"
+                ] = cfg.trl.ref_model_mixup_alpha
+            if cfg.trl and cfg.trl.ref_model_sync_steps:
+                grpo_args_kwargs["ref_model_sync_steps"] = cfg.trl.ref_model_sync_steps
+        grpo_args_kwargs["max_completion_length"] = cfg.trl.max_completion_length
+        grpo_args_kwargs["log_completions"] = cfg.trl.log_completions
        return grpo_args_kwargs

    @classmethod
@@ -86,9 +71,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
@@ -40,6 +39,15 @@ class AxolotlGRPOTrainer(SchedulerMixin, GRPOTrainer):
            self.model = self._enable_gradient_checkpointing(self.model, kwargs["args"])
        # pylint: enable=access-member-before-definition

+        # cleanup the ref_model if we have a peft model passed in
+        # TODO remove this after next major trl release
+        if (
+            self.ref_model  # pylint: disable=access-member-before-definition
+            and is_peft_model(self.model)
+        ):
+            del self.ref_model
+            self.ref_model = None
+
    def _enable_gradient_checkpointing(
        self, model: PreTrainedModel, args: GRPOConfig
    ) -> PreTrainedModel:
@@ -79,6 +87,7 @@ class AxolotlGRPOTrainer(SchedulerMixin, GRPOTrainer):
            if is_peft_model(unwrapped_model):
                unwrapped_model.merge_adapter()
                state_dict = unwrapped_model.state_dict()
+                unwrapped_model.unmerge_adapter()
                # Remove base_model and base_layer prefixes
                state_dict = {
                    k.removeprefix("base_model.model.")
@@ -100,10 +109,8 @@ class AxolotlGRPOTrainer(SchedulerMixin, GRPOTrainer):
                }
            else:
                state_dict = unwrapped_model.state_dict()
-            if self.accelerator.is_main_process:
-                llm_model = (
-                    self.llm.llm_engine.model_executor.driver_worker.model_runner.model
-                )
-                llm_model.load_weights(state_dict.items())
-            if is_peft_model(unwrapped_model):
-                unwrapped_model.unmerge_adapter()
+        if self.accelerator.is_main_process:
+            llm_model = (
+                self.llm.llm_engine.model_executor.driver_worker.model_runner.model
+            )
+            llm_model.load_weights(state_dict.items())
--- 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/base.py
+++ b/src/axolotl/integrations/base.py
@@ -23,8 +23,6 @@ import importlib
 import logging
 from typing import OrderedDict

-import torch
-

 class BasePlugin:
    """
@@ -471,14 +469,3 @@ class PluginManager:
        """
        for plugin in self.plugins.values():
            plugin.post_train_unload(cfg)
-
-
-class BaseOptimizerFactory:
-    """
-    Base class for factories to create custom optimizers
-    """
-
-    def __call__(
-        self, opt_model, training_args, **optimizer_kwargs
-    ) -> "torch.optim.Optimizer":
-        pass
--- 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/cut_cross_entropy/README.md
+++ b/src/axolotl/integrations/cut_cross_entropy/README.md
@@ -1,26 +1,6 @@
 # Cut Cross Entropy

-Cut Cross Entropy reduces VRAM usage through optimization on the cross-entropy operation during loss calculation.
-
-See https://github.com/apple/ml-cross-entropy
-
-## Requirements
-
- PyTorch 2.4.0 or higher
-
-## Installation
-
-Run the following command to install `cut_cross_entropy[transformers]` if you don't have it already.
-
-```bash
-# if you are in dev environment
-python scripts/cutcrossentropy_install.py | sh
-
-# if you are not in dev environment
-pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@24fbe4b5dab9a6c250a014573613c1890190536c"
-```
-
-## Usage
+### Usage

 ```yaml
 plugins:
@@ -28,19 +8,3 @@ plugins:

 cut_cross_entropy: true
 ```
-
-## Citation
-
-```bib
-@article{wijmans2024cut,
-  author       = {Erik Wijmans and
-                  Brody Huval and
-                  Alexander Hertzberg and
-                  Vladlen Koltun and
-                  Philipp Kr\"ahenb\"uhl},
-  title        = {Cut Your Losses in Large-Vocabulary Language Models},
-  journal      = {arXiv},
-  year         = {2024},
-  url          = {https://arxiv.org/abs/2411.09009},
-}
-```
--- a/src/axolotl/integrations/cut_cross_entropy/init.py
+++ b/src/axolotl/integrations/cut_cross_entropy/init.py
@@ -33,7 +33,7 @@ LOG = logging.getLogger("axolotl.integrations.cut_cross_entropy")

 _CCE_INSTALL_MESSAGE = (
    "Please install cut_cross_entropy with transformers support using "
-    '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@24fbe4b5dab9a6c250a014573613c1890190536c"`'
+    '`pip install "cut-cross-entropy[transformers]==24.11.4"`'
 )


--- a/src/axolotl/integrations/grokfast/README.md
+++ b/src/axolotl/integrations/grokfast/README.md
@@ -2,7 +2,7 @@

 See https://github.com/ironjr/grokfast

-## Usage
+### Usage

 ```yaml
 plugins:
@@ -11,14 +11,3 @@ plugins:
 grokfast_alpha: 2.0
 grokfast_lamb: 0.98
 ```
-
-## Citation
-
-```bib
-@article{lee2024grokfast,
-    title={{Grokfast}: Accelerated Grokking by Amplifying Slow Gradients},
-    author={Lee, Jaerin and Kang, Bong Gyun and Kim, Kihoon and Lee, Kyoung Mu},
-    journal={arXiv preprint arXiv:2405.20233},
-    year={2024}
-}
-```
--- 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/README.md
+++ b/src/axolotl/integrations/kd/README.md
@@ -1,23 +0,0 @@
-# Knowledge Distillation
-
-## Usage
-
-```yaml
-plugins:
-  - "axolotl.integrations.kd.KDPlugin"
-
-kd_trainer: True
-kd_ce_alpha: 0.1
-kd_alpha: 0.9
-kd_temperature: 1.0
-
-torch_compile: True  # torch>=2.5.1, recommended to reduce vram
-
-datasets:
-  - path: ...
-    type: "axolotl.integrations.kd.chat_template"
-    field_messages: "messages_combined"
-    logprobs_field: "llm_text_generation_vllm_logprobs"  # for kd only, field of logprobs
-```
-
-An example dataset can be found at [`axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample`](https://huggingface.co/datasets/axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample)
--- 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/kd/topk_logprob/LICENSE.md
+++ b/src/axolotl/integrations/kd/topk_logprob/LICENSE.md
@@ -0,0 +1,58 @@
+### AXOLOTL COMMUNITY LICENSE AGREEMENT
+
+This Axolotl Community License Agreement (“Agreement”) is entered into by and between Axolotl AI Corp. (“Axolotl”) and
+any individual or entity (“Licensee”) who wishes to use the Software (as defined below) in accordance with the terms
+and conditions set forth in this Agreement.
+
+1.  Definitions
+    1.1 “Licensee” refers to any individual or entity who has obtained a copy of the Software under this Agreement.
+    1.2 “Plugin Integration” means independent integration software modules which may or may not be offered by Axolotl,
+        which may be licensed separately by their respective  authors and/or licensors.
+    1.3 “Software” refers to the specific sub-directory of the Axolotl, Inc. software located at
+        https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations and its subdirectories which
+        permits Plugin Integrations to integrate with the Axolotl service.
+2.  Grant of License
+    2.1	Axolotl hereby grants Licensee a worldwide, non-exclusive, royalty-free, license to use, copy, modify, merge,
+        publish, distribute, sublicense, and/or otherwise exploit the Software, subject to the following conditions:
+        - Licensee must comply with all the terms and conditions of this Agreement.
+        - Licensee must include the original copyright notice and disclaimer of warranty in all copies or substantial
+          portions of the Software.
+    2.2 Licensee may use the Software for any lawful purpose, except as restricted in Section 3.
+3.  Restrictions
+    3.1 Licensee shall not use the Software for any activity that constitutes a commercial activity of offering for
+        free or for sale any services, platform, or equivalent  to third parties for the purposes of allowing such
+        third parties to fine-tune artificial intelligence models.
+    3.2 Licensee shall not:
+        - Use the Software for any illegal or unauthorized purpose.
+        - Reverse engineer, decompile, or disassemble the Software.
+        - Remove or modify any copyright, trademark, or other proprietary notices contained in the Software.
+        - Use the Software in a way that could damage, disable, overburden, or impair the functionality of the
+          Software or interfere with any third-party use of the Software.
+    3.3 Axolotl reserves the right to restrict certain Plugin Integrations for use with the Software. To the extent Licensee integrates a permitted, applicable Plugin Integration with the Software, Licensee shall comply with any additional terms and conditions imposed by the licensors of such Plugin Integration for use of such Plugin Integrations. Licensee shall contact Axolotl if it has questions about whether its use of the Software falls beyond the scope of this Agreement.
+4.  Intellectual Property Rights
+    4.1 Axolotl and its contributors retain all intellectual property rights in and to the Software. Licensee
+        acknowledges that this Agreement does not transfer any ownership rights or intellectual property rights to
+        Licensee.
+5.  Disclaimer of Warranty
+    5.1 THE SOFTWARE IS PROVIDED “AS IS,” WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+        TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT. IN NO EVENT SHALL
+        THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES, OR OTHER LIABILITY, WHETHER IN AN ACTION OF
+        CONTRACT, TORT, OR OTHERWISE, ARISING FROM, OUT OF, OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+        DEALINGS IN THE SOFTWARE.
+6.  Termination
+    6.1 Axolotl may terminate this Agreement at any time if Licensee fails to comply with any of the terms and
+        conditions set forth herein. Upon termination, Licensee shall cease all use of the Software and destroy any
+        copies in its possession.
+7.  Governing Law
+    7.1 This Agreement shall be governed by and construed in accordance with the laws of the State of California,
+        without regards to conflicts of laws provisions thereof.
+8.  Entire Agreement
+    8.1 This Agreement constitutes the entire agreement between Axolotl and Licensee with respect to the subject matter
+        hereof and supersedes all prior or contemporaneous understandings or agreements between the parties concerning
+        the Software, whether written or oral. Axolotl may update the terms of this Agreement from time to time, and
+        Licensee’s continued use of the Software after any such updates shall constitute acceptance of updated terms
+        on a go-forward basis.  Axolotl will use commercially reasonable efforts to provide Licensee notice of any
+        material updates. By using the Software, Licensee acknowledges that it has read, understood, and agrees to be
+        bound by the terms and conditions of this Agreement.
+
+This Agreement was last updated on August 23, 2024.
--- a/src/axolotl/integrations/kd/topk_logprob/forward_kl.py
+++ b/src/axolotl/integrations/kd/topk_logprob/forward_kl.py
@@ -1,16 +1,14 @@
 # Copyright 2024 Axolotl AI. All rights reserved.
 #
-# Licensed under the Apache License, Version 2.0 (the "License");
+# This software may be used and distributed according to
+# the terms of the Axolotl Community License Agreement (the "License");
 # you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.

 """
 loss for top_k KL divergence
--- a/src/axolotl/integrations/liger/README.md
+++ b/src/axolotl/integrations/liger/README.md
@@ -1,36 +0,0 @@
-# Liger Kernel Integration
-
-Liger Kernel provides efficient Triton kernels for LLM training, offering:
-
- 20% increase in multi-GPU training throughput
- 60% reduction in memory usage
- Compatibility with both FSDP and DeepSpeed
-
-See https://github.com/linkedin/Liger-Kernel
-
-## Usage
-
-```yaml
-plugins:
-  - axolotl.integrations.liger.LigerPlugin
-liger_rope: true
-liger_rms_norm: true
-liger_glu_activation: true
-liger_layer_norm: true
-liger_fused_linear_cross_entropy: true
-```
-
-## Citation
-
-```bib
-@article{hsu2024ligerkernelefficienttriton,
-      title={Liger Kernel: Efficient Triton Kernels for LLM Training},
-      author={Pin-Lun Hsu and Yun Dai and Vignesh Kothapalli and Qingquan Song and Shao Tang and Siyu Zhu and Steven Shimizu and Shivam Sahni and Haowen Ning and Yanning Chen},
-      year={2024},
-      eprint={2410.10989},
-      archivePrefix={arXiv},
-      primaryClass={cs.LG},
-      url={https://arxiv.org/abs/2410.10989},
-      journal={arXiv preprint arXiv:2410.10989},
-}
-```
--- 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/README.md
+++ b/src/axolotl/integrations/lm_eval/README.md
@@ -1,10 +1,6 @@
 # LM Eval Harness

-Run evaluation on model using the popular lm-evaluation-harness library.
-
-See https://github.com/EleutherAI/lm-evaluation-harness
-
-## Usage
+### Usage

 ```yaml
 plugins:
@@ -14,22 +10,4 @@ lm_eval_tasks:
  - gsm8k
  - hellaswag
  - arc_easy
-
-lm_eval_batch_size: # Batch size for evaluation
-output_dir: # Directory to save evaluation results
-```
-
-## Citation
-
-```bib
-@misc{eval-harness,
-  author       = {Gao, Leo and Tow, Jonathan and Abbasi, Baber and Biderman, Stella and Black, Sid and DiPofi, Anthony and Foster, Charles and Golding, Laurence and Hsu, Jeffrey and Le Noac'h, Alain and Li, Haonan and McDonell, Kyle and Muennighoff, Niklas and Ociepa, Chris and Phang, Jason and Reynolds, Laria and Schoelkopf, Hailey and Skowron, Aviya and Sutawika, Lintang and Tang, Eric and Thite, Anish and Wang, Ben and Wang, Kevin and Zou, Andy},
-  title        = {A framework for few-shot language model evaluation},
-  month        = 07,
-  year         = 2024,
-  publisher    = {Zenodo},
-  version      = {v0.4.3},
-  doi          = {10.5281/zenodo.12608602},
-  url          = {https://zenodo.org/records/12608602}
-}
 ```
--- 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/integrations/spectrum/README.md
+++ b/src/axolotl/integrations/spectrum/README.md
@@ -1,17 +1,15 @@
-# Spectrum: Targeted Training on Signal to Noise Ratio
+## Spectrum: Targeted Training on Signal to Noise Ratio

 by Eric Hartford, Lucas Atkins, Fernando Fernandes, David Golchinfar

 This plugin contains code to freeze the bottom fraction of modules in a model, based on the Signal-to-Noise Ratio (SNR).

-See https://github.com/cognitivecomputations/spectrum
-
-## Overview
+### Overview

 Spectrum is a tool for scanning and evaluating the Signal-to-Noise Ratio (SNR) of layers in large language models.
 By identifying the top n% of layers with the highest SNR, you can optimize training efficiency.

-## Usage
+### Usage

 ```yaml
 plugins:
@@ -21,17 +19,3 @@ spectrum_top_fraction: 0.5
 # Optional if using a pre-scanned model as your base_model. Useful if using a model mirror
 spectrum_model_name: meta-llama/Meta-Llama-3.1-8B
 ```
-
-## Citation
-
-```bib
-@misc{hartford2024spectrumtargetedtrainingsignal,
-      title={Spectrum: Targeted Training on Signal to Noise Ratio},
-      author={Eric Hartford and Lucas Atkins and Fernando Fernandes Neto and David Golchinfar},
-      year={2024},
-      eprint={2406.06623},
-      archivePrefix={arXiv},
-      primaryClass={cs.LG},
-      url={https://arxiv.org/abs/2406.06623},
-}
-```
--- a/src/axolotl/integrations/spectrum/args.py
+++ b/src/axolotl/integrations/spectrum/args.py
@@ -17,7 +17,7 @@ Module for handling Spectrum input arguments.
 """
 from typing import Optional

-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel


 class SpectrumArgs(BaseModel):
@@ -27,20 +27,3 @@ class SpectrumArgs(BaseModel):

    spectrum_top_fraction: Optional[float] = 0.5
    spectrum_model_name: Optional[str] = None
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_fsdp_use_orig_params(cls, data):
-        if (
-            data.get("fsdp")
-            and data.get("fsdp_config")
-            and not data["fsdp_config"].get("use_orig_params")
-            and data.get("plugins")
-            and any("SpectrumPlugin" in plugin for plugin in data["plugins"])
-        ):
-            # would otherwise raise
-            # ValueError: Must flatten tensors with uniform `requires_grad` when `use_orig_params=False`
-            raise ValueError(
-                "FSDP + SpectrumPlugin cannot be used together when `use_orig_params=False` is set"
-            )
-        return data
--- a/Show More
+++ b/Show More