fix import

add muon optimizer
optimizer_cls_and_kwargs is on trainer_kwargs only add adamw_kwargs if they're non-null fix mocks better handling of override and check the optimizer unwrap optimizer
2025-03-05 14:05:27 -05:00 · 2025-03-05 10:47:22 -05:00
363 changed files with 4373 additions and 10969 deletions
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -40,24 +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: "126"
-            cuda_version: 12.6.3
-            cudnn_version: ""
-            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"
-          - cuda: "128"
-            cuda_version: 12.8.1
-            cudnn_version: ""
-            python_version: "3.11"
-            pytorch: next
-            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
@@ -79,7 +61,7 @@ jobs:
        uses: docker/build-push-action@v4
        with:
          context: .
-          file: ${{ matrix.pytorch == 'nightly' && './docker/Dockerfile-base-nightly' || matrix.pytorch == 'next' && './docker/Dockerfile-base-next' || './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 . --no-deps
-        - 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
@@ -25,12 +25,12 @@ jobs:
            python_version: "3.11"
            pytorch: 2.5.1
            axolotl_extras: vllm
+            is_latest: true
          - cuda: 124
            cuda_version: 12.4.1
            python_version: "3.11"
            pytorch: 2.6.0
            axolotl_extras:
-            is_latest: true
    runs-on: axolotl-gpu-runner
    steps:
      - name: Checkout
@@ -87,11 +87,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:
            is_latest: true
    runs-on: axolotl-gpu-runner
    steps:
--- a/.github/workflows/multi-gpu-e2e.yml
+++ b/.github/workflows/multi-gpu-e2e.yml
@@ -24,13 +24,6 @@ jobs:
      fail-fast: false
      matrix:
        include:
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.6.0
-            axolotl_extras: vllm
-            num_gpus: 2
-            nightly_build: "true"
          - cuda: 124
            cuda_version: 12.4.1
            python_version: "3.11"
@@ -45,6 +38,14 @@ jobs:
            axolotl_extras: vllm
            num_gpus: 2
            nightly_build: "true"
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.6.0
+            # awaiting vllm#12721
+            axolotl_extras:
+            num_gpus: 2
+            nightly_build: "true"
    runs-on: [self-hosted, modal]
    timeout-minutes: 120
    steps:
--- 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
@@ -33,15 +33,6 @@ jobs:
      - name: Check out repository code
        uses: actions/checkout@v4

-      - name: Restore HF cache
-        id: hf-cache-restore
-        uses: actions/cache/restore@v4
-        with:
-          path: |
-            /home/runner/.cache/huggingface/hub/datasets--*
-            /home/runner/.cache/huggingface/hub/models--*
-          key: ${{ runner.os }}-hf-hub-cache-v2
-
      - name: Setup Python
        uses: actions/setup-python@v5
        with:
@@ -51,11 +42,11 @@ 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: |
-          pip3 install torch==${{ matrix.pytorch_version }}
+          pip3 install torch==${{ matrix.pytorch_version }} --index-url https://download.pytorch.org/whl/cpu

      - name: Update requirements.txt
        run: |
@@ -67,7 +58,8 @@ jobs:

      - name: Install dependencies
        run: |
-          pip3 show torch
+          pip3 install --upgrade pip
+          pip3 install --upgrade packaging
          pip3 install --no-build-isolation -U -e .
          python scripts/unsloth_install.py | sh
          python scripts/cutcrossentropy_install.py | sh
@@ -81,15 +73,10 @@ jobs:
        run: |
          axolotl --help

-      - name: Pre-Download dataset fixture
-        run: |
-          huggingface-cli download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
-
      - name: Run tests
        run: |
-          pytest -v -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ tests/
-          pytest -v tests/patched/
-          pytest -v tests/cli/
+          pytest -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ tests/
+          pytest tests/patched/

      - name: cleanup pip cache
        run: |
@@ -149,4 +136,4 @@ jobs:
          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
      - name: Run tests job on Modal
        run: |
-          modal run cicd.e2e_tests
+          modal run cicd.tests
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -63,7 +63,7 @@ jobs:
          path: |
            /home/runner/.cache/huggingface/hub/datasets--*
            /home/runner/.cache/huggingface/hub/models--*
-          key: ${{ runner.os }}-hf-hub-cache-v2
+          key: ${{ runner.os }}-hf-hub-cache-${{ hashFiles('**/conftest.py') }}

      - name: Setup Python
        uses: actions/setup-python@v5
@@ -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: |
@@ -96,15 +96,10 @@ jobs:
        run: |
          axolotl --help

-      - name: Pre-Download dataset fixture
-        run: |
-          huggingface-cli download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures
-
      - name: Run tests
        run: |
-          pytest -v -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ tests/
+          pytest -v -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ tests/
          pytest -v tests/patched/
-          pytest -v tests/cli/

      - name: cleanup pip cache
        run: |
@@ -141,7 +136,7 @@ jobs:
          path: |
            /home/runner/.cache/huggingface/hub/datasets--*
            /home/runner/.cache/huggingface/hub/models--*
-          key: ${{ runner.os }}-hf-hub-cache-v2
+          key: ${{ runner.os }}-hf-hub-cache-${{ hashFiles('**/conftest.py') }}

      - name: Setup Python
        uses: actions/setup-python@v5
@@ -152,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: |
@@ -175,14 +170,10 @@ jobs:
        run: |
          axolotl --help

-      - name: Show HF cache
-        run: huggingface-cli scan-cache
-
      - name: Run tests
        run: |
-          pytest -v -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ tests/
+          pytest -v -n8 --dist loadfile --ignore=tests/e2e/ --ignore=tests/patched/ tests/
          pytest -v tests/patched/
-          pytest -v tests/cli/

      - name: cleanup pip cache
        run: |
@@ -208,53 +199,6 @@ jobs:
      fail-fast: false
      matrix:
        include:
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.6.0
-            num_gpus: 1
-            axolotl_extras: vllm
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-      - name: Install Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.11"
-      - name: Install Modal
-        run: |
-          python -m pip install --upgrade pip
-          pip install modal==0.71.8 jinja2
-      - name: Update env vars
-        run: |
-          echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
-          echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV
-          echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV
-          echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
-          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
-          echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
-          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
-      - name: Run tests job on Modal
-        run: |
-          modal run cicd.e2e_tests
-
-  docker-e2e-tests:
-    if: github.repository_owner == 'axolotl-ai-cloud'
-    # this job needs to be run on self-hosted GPU runners...
-    runs-on: [self-hosted, modal]
-    timeout-minutes: 90
-    needs: [pre-commit, pytest, docker-e2e-tests-1st]
-
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - cuda: 124
-            cuda_version: 12.4.1
-            python_version: "3.11"
-            pytorch: 2.4.1
-            num_gpus: 1
-            axolotl_extras:
          - cuda: 124
            cuda_version: 12.4.1
            python_version: "3.11"
@@ -283,4 +227,51 @@ jobs:
          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
      - name: Run tests job on Modal
        run: |
-          modal run cicd.e2e_tests
+          modal run cicd.tests
+
+  docker-e2e-tests:
+    if: github.repository_owner == 'axolotl-ai-cloud'
+    # this job needs to be run on self-hosted GPU runners...
+    runs-on: [self-hosted, modal]
+    timeout-minutes: 90
+    needs: [pre-commit, pytest, docker-e2e-tests-1st]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.4.1
+            num_gpus: 1
+            axolotl_extras:
+          - cuda: 124
+            cuda_version: 12.4.1
+            python_version: "3.11"
+            pytorch: 2.6.0
+            num_gpus: 1
+            axolotl_extras:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install Modal
+        run: |
+          python -m pip install --upgrade pip
+          pip install modal==0.71.8 jinja2
+      - name: Update env vars
+        run: |
+          echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV
+          echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV
+          echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV
+          echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
+          echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
+          echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV
+          echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
+      - name: Run tests job on Modal
+        run: |
+          modal run cicd.tests
--- a/.gitignore
+++ b/.gitignore
@@ -181,10 +181,6 @@ prepared-datasets/
 submit.sh
 *.out*

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

--- a/.isort.cfg
+++ b/.isort.cfg
@@ -1,4 +1,3 @@
 [settings]
 profile=black
 known_third_party=wandb,comet_ml
-known_local_folder=src,tests
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -3,7 +3,7 @@ default_language_version:

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

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

 # Download example axolotl configs, deepspeed configs
@@ -97,7 +96,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,180 +1,6 @@
 project:
  type: website

-quartodoc:
-  dir: docs/api
-  package: axolotl
-  title: API Reference
-  parser: google
-
-  sections:
-    - title: Core
-      desc: Core functionality for training
-      contents:
-        - train
-        - evaluate
-        - datasets
-        - convert
-        - prompt_tokenizers
-        - logging_config
-        - core.trainer_builder
-        - core.training_args
-        - core.chat.messages
-        - core.chat.format.chatml
-        - core.chat.format.llama3x
-        - core.chat.format.shared
-        - core.datasets.chat
-        - core.datasets.transforms.chat_builder
-    - title: CLI
-      desc: Command-line interface
-      contents:
-        - cli.main
-        - cli.train
-        - cli.evaluate
-        - cli.args
-        - cli.checks
-        - cli.config
-        - cli.inference
-        - cli.merge_lora
-        - cli.merge_sharded_fsdp_weights
-        - cli.preprocess
-        - cli.sweeps
-        - cli.utils
-        - cli.vllm_serve
-        - 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.multimodal
-        - 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"
@@ -206,11 +32,8 @@ website:
          contents:
            - docs/getting-started.qmd
            - docs/installation.qmd
-            - docs/inference.qmd
            - docs/cli.qmd
-            - docs/config.qmd
-            - text: "API Reference"
-              href: docs/api
+            - docs/inference.qmd

        - section: "Dataset Formats"
          contents: docs/dataset-formats/*
@@ -244,7 +67,6 @@ website:
            - docs/unsloth.qmd
            - docs/torchao.qmd
            - docs/custom_integrations.qmd
-            - docs/sequence_parallelism.qmd

        - section: "Troubleshooting"
          contents:
@@ -252,27 +74,12 @@ website:
            - docs/debugging.qmd
            - docs/nccl.qmd

+        - section: "Reference"
+          contents:
+            - docs/config.qmd
+
 format:
  html:
    theme: darkly
    css: styles.css
    toc: true
-    # Enable better handling of line breaks in markdown
-    preserve-tabs: true
-    html-math-method: mathjax
-    # Improved markdown processing options
-    md-extensions:
-      - markdown_it
-      - def_list
-      - attr_list
-      - fenced_divs
-      - tables
-      - html_admonition
-      - lineblocks
-      - fancy_lists
-    # Control whitespace handling
-    whitespace: preserve
-    # Process newlines in paragraphs
-    wrap: preserve
-    # Better line break handling
-    preserve-linebreaks: true
--- a/cicd/Dockerfile.jinja
+++ b/cicd/Dockerfile.jinja
@@ -31,11 +31,10 @@ 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,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
    else \
-        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
    fi

 RUN python scripts/unsloth_install.py | sh
--- a/cicd/cicd.sh
+++ b/cicd/cicd.sh
@@ -3,10 +3,9 @@ set -e

 python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__"

-pytest -v --durations=10 -n8 --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli /workspace/axolotl/tests/
+pytest -v --durations=10 -n8 --ignore=tests/e2e/ --ignore=tests/patched/ /workspace/axolotl/tests/
 pytest -v --durations=10 /workspace/axolotl/tests/e2e/patched/lora_kernels  # running these with the other patches causes a failure
 pytest -v --durations=10 --ignore=tests/e2e/patched/lora_kernels /workspace/axolotl/tests/e2e/patched
 pytest -v --durations=10 -n1 /workspace/axolotl/tests/e2e/solo/
 pytest -v --durations=10 /workspace/axolotl/tests/e2e/integrations/
-pytest -v --durations=10 /workspace/axolotl/tests/cli
-pytest -v --durations=10 --ignore=tests/e2e/solo/ --ignore=tests/e2e/patched/ --ignore=tests/e2e/multigpu/ --ignore=tests/e2e/integrations/ --ignore=tests/cli /workspace/axolotl/tests/e2e/
+pytest -v --durations=10 --ignore=tests/e2e/solo/ --ignore=tests/e2e/patched/ --ignore=tests/e2e/multigpu/ --ignore=tests/e2e/integrations/ /workspace/axolotl/tests/e2e/
--- a/cicd/multigpu.py
+++ b/cicd/multigpu.py
@@ -1,7 +1,6 @@
 """
-modal application to run axolotl gpu tests in Modal
-"""
-
+ modal application to run axolotl gpu tests in Modal
+ """
 # pylint: disable=duplicate-code

 import os
--- a/cicd/multigpu.sh
+++ b/cicd/multigpu.sh
@@ -2,5 +2,4 @@
 set -e

 # only run one test at a time so as not to OOM the GPU
-pytest -v  --durations=10 -n2 /workspace/axolotl/tests/e2e/multigpu/ --ignore=/workspace/axolotl/tests/e2e/multigpu/solo/
-pytest -v  --durations=10 -n1 /workspace/axolotl/tests/e2e/multigpu/solo/
+pytest -v -n2 /workspace/axolotl/tests/e2e/multigpu/
--- a/cicd/e2e_tests.py
+++ b/cicd/e2e_tests.py
@@ -1,5 +1,4 @@
 """Modal app to run axolotl GPU tests"""
-
 # pylint: disable=duplicate-code

 import os
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -20,9 +20,9 @@ WORKDIR /workspace/axolotl

 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
    else \
-        pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
+        pip install --no-build-isolation -e .[deepspeed,flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
    fi

 RUN python scripts/unsloth_install.py | sh
--- 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-next
+++ b/docker/Dockerfile-base-next
@@ -1,38 +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="next"
-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==2.7.0 --extra-index-url https://download.pytorch.org/whl/test/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 && \
-    pip3 install -U --no-cache-dir pydantic==2.10.6
--- 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/cli.qmd
+++ b/docs/cli.qmd
@@ -1,5 +1,5 @@
 ---
-title: "Command Line Interface (CLI)"
+title: "CLI Reference"
 format:
  html:
    toc: true
@@ -170,7 +170,7 @@ axolotl merge-sharded-fsdp-weights config.yml

 ### evaluate

-Evaluates a model's performance (loss etc) on the train and eval datasets.
+Evaluates a model's performance using metrics specified in the config.

 ```bash
 # Basic evaluation
@@ -197,8 +197,6 @@ lm_eval_batch_size: # Batch size for evaluation
 output_dir: # Directory to save evaluation results
 ```

-See [LM Eval Harness](https://github.com/EleutherAI/lm-evaluation-harness) for more details.
-
 ## Legacy CLI Usage

 While the new Click-based CLI is preferred, Axolotl still supports the legacy module-based CLI:
@@ -237,7 +235,7 @@ Create a cloud config YAML with your Modal settings:
 ```yaml
 # cloud_config.yml
 provider: modal
-gpu: a100       # Supported: l40s, a100-40gb, a100-80gb, a10g, h100, t4, l4
+gpu: a100  # Supported: l40s, a100-40gb, a100-80gb, a10g, h100, t4, l4
 gpu_count: 1    # Number of GPUs to use
 timeout: 86400  # Maximum runtime in seconds (24 hours)
 branch: main    # Git branch to use (optional)
@@ -250,7 +248,7 @@ volumes:        # Persistent storage volumes
  - name: axolotl-artifacts
    mount: /workspace/artifacts

-secrets:        # Secrets to inject
+env:            # Environment variables
  - WANDB_API_KEY
  - HF_TOKEN
 ```
@@ -276,27 +274,15 @@ axolotl lm-eval config.yml --cloud cloud_config.yml
 ### 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
+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
-
-# Environment variables to pass. Can be specified in two ways:
-# 1. As a string: Will load the value from the host computer's environment variables
-# 2. As a key-value pair: Will use the specified value directly
-# Example:
-# env:
-#   - CUSTOM_VAR  # Loads from host's $CUSTOM_VAR
-#   - {CUSTOM_VAR: "value"}  # Uses "value" directly
-env:
-
-# Secrets to inject. Same input format as `env` but for sensitive data.
-secrets:
-  # - HF_TOKEN
-  # - WANDB_API_KEY
+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,11 +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:
-# Whether to load the model with randomly initialized weights. Useful for
-# pre-training a model from scratch or debugging purposes.
-random_init_weights:

 # (Internal use only)
 # Used to identify which the model is based on
@@ -88,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
@@ -165,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"]
@@ -216,46 +207,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_server_host: # Optional[str]. Host of the vLLM server to connect to.
-  vllm_server_port: # Optional[int]. Port of the vLLM server to connect to.
-  vllm_server_timeout: # Optional[int]. Total timeout (in seconds) to wait for the vLLM server to respond.
-  vllm_guided_decoding_regex: # Optional[str]. Regex for vLLM guided decoding.
-
-  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:
@@ -279,7 +234,7 @@ default_system_message: You are a helpful assistant. Please give a long and deta
 # 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
@@ -320,13 +275,9 @@ total_num_tokens:
 sample_packing_group_size: 100000
 # The number of samples which can be packed into one sequence. Increase if using a large sequence_len with many short samples.
 sample_packing_bin_size: 200
-sample_pack_sequentially: # Optional[bool]. Whether to pack samples sequentially.
-
 # whether to concatenate samples during pretraining
 pretraining_sample_concatenation:

-curriculum_sampling: # Optional[bool]. Whether to use sequential sampling for curriculum learning
-
 # Use batch flattening for speedups when not using sample_packing
 batch_flattening:

@@ -358,27 +309,7 @@ lora_target_modules:
 #  - down_proj
 #  - up_proj
 lora_target_linear: # If true, will target all linear modules
-
-# List[int] | int. # The layer indices to transform, otherwise, apply to all layers
-# https://huggingface.co/docs/peft/v0.15.0/en/package_reference/lora#peft.LoraConfig.layers_to_transform
-peft_layers_to_transform:
-
-# Optional[bool]. Whether to use DoRA.
-# https://huggingface.co/docs/peft/v0.15.0/en/developer_guides/lora#weight-decomposed-low-rank-adaptation-dora
-peft_use_dora:
-
-# Optional[bool]. Whether to use RSLoRA.
-# https://huggingface.co/docs/peft/v0.15.0/en/developer_guides/lora#rank-stabilized-lora
-peft_use_rslora:
-
-# Optional[list[tuple[int, int]]]. List of layer indices to replicate.
-# https://huggingface.co/docs/peft/v0.15.0/en/developer_guides/lora#memory-efficient-layer-replication-with-lora
-peft_layer_replication:
-
-# bool | Literal["gaussian", "eva", "olora", "pissa", "pissa_niter_[number of iters]", "corda", "loftq"]
-# How to initialize LoRA weights. Default to True which is MS original implementation.
-# https://huggingface.co/docs/peft/v0.15.0/en/developer_guides/lora#initialization
-peft_init_lora_weights:
+peft_layers_to_transform: # The layer indices to transform, otherwise, apply to all layers

 # If you added new tokens to the tokenizer, you may need to save some LoRA modules because they need to know the new tokens.
 # For LLaMA and Mistral, you need to save `embed_tokens` and `lm_head`. It may vary for other models.
@@ -490,7 +421,6 @@ auto_find_batch_size: # Optional[bool]

 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
-do_causal_lm_eval: # Whether to run causal language model evaluation for metrics in `eval_causal_lm_metrics`.
 eval_causal_lm_metrics: # HF evaluate metrics used during evaluation. Default is ["sacrebleu", "comet", "ter", "chrf", "perplexity"]

 profiler_steps: # enable the pytorch profiler to capture the first N steps of training to the output_dir.
@@ -510,8 +440,7 @@ train_on_inputs: false
 # Note that training loss may have an oscillating pattern with this enabled.
 group_by_length: false

-# Whether to use gradient checkpointing. Available options are: true, false, "offload".
-# https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing
+# Whether to use gradient checkpointing https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing
 gradient_checkpointing: false
 # additional kwargs to pass to the trainer for gradient checkpointing
 # gradient_checkpointing_kwargs:
@@ -532,58 +461,36 @@ lr_div_factor: # Learning rate div factor

 # Specify optimizer
 # Valid values are driven by the Transformers OptimizerNames class, see:
-# https://github.com/huggingface/transformers/blob/cbf924b76c03828101a34069a96d209314114fd5/src/transformers/training_args.py#L144-L189
+# https://github.com/huggingface/transformers/blob/95b374952dc27d8511541d6f5a4e22c9ec11fb24/src/transformers/training_args.py#L134
 #
 # Note that not all optimizers may be available in your environment, ex: 'adamw_anyprecision' is part of
 # torchdistx, 'adamw_bnb_8bit' is part of bnb.optim.Adam8bit, etc. When in doubt, it is recommended to start with the optimizer used
 # in the examples/ for your model and fine-tuning use case.
 #
 # Valid values for 'optimizer' include:
+# - adamw_hf
 # - adamw_torch
 # - adamw_torch_fused
 # - adamw_torch_xla
-# - adamw_torch_npu_fused
 # - adamw_apex_fused
-# - adopt_adamw  (an EXPERIMENTAL optimizer, only for torch version >= 2.5.1)
+# - adopt_adamw (an EXPERIMENTAL optimizer, only for torch version >= 2.5.1)
 # - adafactor
 # - adamw_anyprecision
-# - adamw_torch_4bit
-# - ademamix
 # - sgd
 # - adagrad
 # - adamw_bnb_8bit
-# - adamw_8bit   # alias for adamw_bnb_8bit
-# - ademamix_8bit
 # - lion_8bit
 # - lion_32bit
 # - paged_adamw_32bit
 # - paged_adamw_8bit
-# - paged_ademamix_32bit
-# - paged_ademamix_8bit
 # - paged_lion_32bit
 # - paged_lion_8bit
-# - rmsprop
-# - rmsprop_bnb
-# - rmsprop_bnb_8bit
-# - rmsprop_bnb_32bit
 # - galore_adamw
 # - galore_adamw_8bit
 # - galore_adafactor
 # - galore_adamw_layerwise
 # - galore_adamw_8bit_layerwise
 # - galore_adafactor_layerwise
-# - lomo
-# - adalomo
-# - grokadamw
-# - schedule_free_adamw
-# - schedule_free_sgd
-# - apollo_adamw
-# - apollo_adamw_layerwise
-#
-# Additional custom optimizers include:
-# - optimi_adamw
-# - ao_adamw_8bit
-# - ao_adamw_fp8
 optimizer:
 # Dictionary of arguments to pass to the optimizer
 optim_args:
@@ -612,42 +519,29 @@ max_grad_norm:
 # currently only supported on Llama and Mistral
 neftune_noise_alpha:

-# Optional[bool]. Whether to bettertransformers
+# Whether to bettertransformers
 flash_optimum:
-
-# Note: Only one of the following attention patches can be used at a time.
-# For example, if you set `xformers_attention` to `true`, do not set `flash_attention` to `true`.
-
-# Optional[bool]. Whether to use xformers attention patch https://github.com/facebookresearch/xformers:
+# Whether to use xformers attention patch https://github.com/facebookresearch/xformers:
 xformers_attention:
-# Optional[bool]. Whether to use flash attention patch https://github.com/Dao-AILab/flash-attention:
+# Whether to use flash attention patch https://github.com/Dao-AILab/flash-attention:
 flash_attention:
-flash_attn_cross_entropy:  # Optional[bool]. Whether to use flash-attention cross entropy implementation - advanced use only
-flash_attn_rms_norm:  # Optional[bool]. Whether to use flash-attention rms norm implementation - advanced use only
-flash_attn_fuse_qkv: # Optional[bool]. Whether to fuse QKV into a single operation
-flash_attn_fuse_mlp: # Optional[bool]. Whether to fuse part of the MLP into a single operation
-# Optional[bool]. Whether to use scaled-dot-product attention
+flash_attn_cross_entropy:  # Whether to use flash-attention cross entropy implementation - advanced use only
+flash_attn_rms_norm:  # Whether to use flash-attention rms norm implementation - advanced use only
+flash_attn_fuse_qkv: # Whether to fuse QKV into a single operation
+flash_attn_fuse_mlp: # Whether to fuse part of the MLP into a single operation
+# Whether to use scaled-dot-product attention
 # https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
 sdp_attention:
-# Optional[bool]. Shifted-sparse attention (only llama) - https://arxiv.org/pdf/2309.12307.pdf
+# 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:
-# Optional[str]. Resume from a specific checkpoint dir
+# Resume from a specific checkpoint dir
 resume_from_checkpoint:
-# Optional[bool]. If resume_from_checkpoint isn't set and you simply want it to start where it left off.
+# If resume_from_checkpoint isn't set and you simply want it to start where it left off.
 # Be careful with this being turned on between different models.
 auto_resume_from_checkpoints: false

-## Multimodal section
-# int | tuple[int, int] | None . Size to resize images to, width x height.
-# Will read from model/processor config if not set.
-image_size:
-# str. Algorithm to use for image resizing. "bilinear", "bicubic", "lanczos". Default is "bilinear".
-image_resize_algorithm: 'bilinear'
-## End of multimodal section
-
 # Don't mess with this, it's here for accelerate and torchrun
 local_rank:

@@ -662,13 +556,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:
@@ -681,17 +568,6 @@ ddp_timeout:
 ddp_bucket_cap_mb:
 ddp_broadcast_buffers:

-# Sequence parallelism
-# Set to a divisor of the number of GPUs available to split sequences into chunks of equal size.
-# Use in long context training to prevent OOM when sequences cannot fit into a single GPU's VRAM.
-# E.g., if 4 GPUs are available, set this value to 2 to split each sequence into two equal-sized
-# subsequences, or set to 4 to split into four equal-sized subsequences.
-# See https://axolotl-ai-cloud.github.io/axolotl/docs/sequence_parallelism.html for more details.
-sequence_parallel_degree:
-# Optional; strides across the key dimension. Larger values use more memory but should make training faster.
-# Must evenly divide the number of KV heads in your model.
-heads_k_stride: 1
-
 # Path to torch distx for optim 'adamw_anyprecision'
 torchdistx_path:

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

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

 - parse the dataset based on the *dataset format*
 - transform the dataset to how you would interact with the model based on the *prompt strategy*
--- a/docs/docker.qmd
+++ b/docs/docker.qmd
@@ -103,7 +103,8 @@ This uses the same tags as the [`main` image](#sec-main-tags).

 - `JUPYTER_DISABLE`: Disable Jupyter lab.
 - `JUPYTER_PASSWORD`: Set a password for the Jupyter lab.
- `PUBLIC_KEY` / `SSH_KEY`: Add a public key for the SSH service.
+- `PUBLIC_KEY`: Add a public key for the SSH service.
+- `SSH_KEY`: Add a private key for the SSH service.

 #### Volume mounts

--- a/docs/faq.qmd
+++ b/docs/faq.qmd
@@ -27,30 +27,6 @@ description: Frequently asked questions

 > 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: Since Axolotl is just Python, please see `src/axolotl/cli/main.py` on how each command is called.
-
-**Q: How to know the value to use for `fsdp_transformer_layer_cls_to_wrap`?**
-
-> A: This is the class name of the transformer layer to wrap with FSDP. For example, for `LlamaForCausalLM`, the value is `LlamaDecoderLayer`. To find this for a specific model, check the model's `PreTrainedModel` definition and look for `_no_split_modules` variable in the `modeling_<model_name>.py` file within `transformers` library.
-
-**Q: ValueError: Asking to pad but the tokenizer does not have a padding token. Please select a token to use as pad_token**
-
-> A: This is because the tokenizer does not have a padding token. Please add a padding token to the tokenizer via:
-
-> ```yaml
-> special_tokens:
->   # str. If you're not sure, set to same as `eos_token`.
->   pad_token: "..."
-> ```
-
 ### Chat templates

 **Q: `jinja2.exceptions.UndefinedError: 'dict object' has no attribute 'content' / 'role' / ____`**
@@ -76,7 +52,3 @@ description: Frequently asked questions
 **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
@@ -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,6 +104,8 @@ 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
--- a/docs/inference.qmd
+++ b/docs/inference.qmd
@@ -1,5 +1,5 @@
 ---
-title: "Inference and Merging"
+title: "Inference"
 format:
  html:
    toc: true
@@ -9,14 +9,10 @@ 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/installation.qmd
+++ b/docs/installation.qmd
@@ -22,7 +22,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 +37,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]'
 ```

@@ -79,7 +78,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 +107,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
@@ -17,7 +17,6 @@ We currently support several common model architectures, including (but not limi
 - `qwen2`
 - `gemma`
 - `gemma2`
- `gemma3`

 <details>

@@ -67,10 +66,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/multi-gpu.qmd
+++ b/docs/multi-gpu.qmd
@@ -18,7 +18,6 @@ Axolotl supports several methods for multi-GPU training:

 - DeepSpeed (recommended)
 - FSDP (Fully Sharded Data Parallel)
- Sequence parallelism
 - FSDP + QLoRA

 ## DeepSpeed {#sec-deepspeed}
@@ -67,28 +66,6 @@ fsdp_config:
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
 ```

-## Sequence parallelism {#sec-sequence-parallelism}
-
-We support sequence parallelism (SP) via the
-[ring-flash-attention](https://github.com/zhuzilin/ring-flash-attention) project. This
-allows one to split up sequences across GPUs, which is useful in the event that a
-single sequence causes OOM errors during model training.
-
-First, install `ring-flash-attn`, recommended via `pip install axolotl[ring-flash-attn]`,
-or from source with `pip install .[ring-flash-attn]`.
-
-Your Axolotl YAML config should contain the following lines:
-
-```{.yaml}
-sequence_parallel_degree: 4  # Split each sequence into 4 parts, one per GPU
-flash_attention: true  # Required with sequence parallelism
-
-# Optional; strides across the key dimension. Larger values use more memory but will make training faster.
-heads_k_stride: 1
-```
-
-See our [dedicated guide](sequence_parallelism.qmd) for more details.
-
 ### FSDP + QLoRA {#sec-fsdp-qlora}

 For combining FSDP with QLoRA, see our [dedicated guide](fsdp_qlora.qmd).
--- a/docs/multimodal.qmd
+++ b/docs/multimodal.qmd
@@ -1,171 +1,28 @@
---
-title: MultiModal / Vision Language Models (BETA)
-format:
-  html:
-    toc: true
-    toc-depth: 3
---
+# MultiModal / Vision Language Models (BETA)

-## Supported Models
+### Supported Models

- [Mllama](#sec-mllama)
- [Pixtral](#sec-pixtral)
- [Llava-1.5](#sec-llava-15)
- [Mistral-Small-3.1](#sec-mistral-small-31)
- [Gemma-3](#sec-gemma-3)
- [Qwen2-VL](#sec-qwen2-vl)
- [Qwen2.5-VL](#sec-qwen25-vl)
+- Mllama, i.e. llama with vision models

-## Usage
+### Usage

-Multimodal support is limited and doesn't have full feature parity.
-
-Here are the hyperparams you'll need to use to finetune a multimodal model.
+Currently multimodal support is limited and doesn't have full feature parity. To finetune a multimodal Llama w/ LoRA,
+you'll need to use the following in YAML in combination with the rest of the required hyperparams.

 ```yaml
+base_model: alpindale/Llama-3.2-11B-Vision-Instruct
 processor_type: AutoProcessor
-
 skip_prepare_dataset: true
-remove_unused_columns: false  # leave columns in place as they are needed to handle image embeddings during training
-sample_packing: false  # not yet supported with multimodal

-chat_template:  # see in next section
-
-# example dataset
+chat_template: llama3_2_vision
 datasets:
  - path: HuggingFaceH4/llava-instruct-mix-vsft
    type: chat_template
    split: train[:1%]
    field_messages: messages
+remove_unused_columns: false
+sample_packing: false

-# (optional) if doing lora, only finetune the Language model,
-# leave the vision model and vision tower frozen
-# load_in_8bit: true
-adapter: lora
+# only finetune the Language model, leave the vision model and vision tower frozen
 lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
-
-# (optional) if you want to resize images to a set size
-image_size: 512
-image_resize_algorithm: bilinear
-```
-
-Please see [examples](https://github.com/axolotl-ai/axolotl/tree/main/examples) folder for full configs.
-
-::: {.callout-warning}
-Some of our chat_templates have been extended to support broader dataset types. This should not break any existing configs.
-:::
-
-### Mllama {#sec-mllama}
-
-```yaml
-base_model: meta-llama/Llama-3.2-11B-Vision-Instruct
-
-chat_template: llama3_2_vision
-```
-
-### Pixtral {#sec-pixtral}
-
-```yaml
-base_model: mistralai/Pixtral-12B-2409
-
-chat_template: pixtral
-```
-
-### Llava-1.5 {#sec-llava-15}
-
-```yaml
-base_model: llava-hf/llava-1.5-7b-hf
-
-chat_template: llava
-```
-
-### Mistral-Small-3.1 {#sec-mistral-small-31}
-
-```yaml
-base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
-
-chat_template: mistral_v7_tekken
-```
-
-### Gemma-3 {#sec-gemma-3}
-
-::: {.callout-tip}
-The Gemma3-1B model is a text-only model, so please train as regular text model.
-:::
-
-For multi-modal 4B/12B/27B models, use the following config:
-
-```yaml
-base_model: google/gemma-3-4b-it
-
-chat_template: gemma3
-```
-
-### Qwen2-VL {#sec-qwen2-vl}
-
-```yaml
-base_model: Qwen/Qwen2-VL-7B-Instruct
-
-chat_template: qwen2_vl
-```
-
-### Qwen2.5-VL {#sec-qwen25-vl}
-
-```yaml
-base_model: Qwen/Qwen2.5-VL-7B-Instruct
-
-chat_template: qwen2_vl  # same as qwen2-vl
-```
-
-## Dataset Format
-
-For multi-modal datasets, we adopt an extended `chat_template` format similar to OpenAI's Message format.
-
- A message is a list of `role` and `content`.
- `role` can be `system`, `user`, `assistant`, etc.
- `content` is a list of `type` and (`text` or `image` or `path` or `url` or `base64`).
-
-::: {.callout-note}
-For backwards compatibility:
-
- If the dataset has a `images` or `image` column of `list[Image]`, it will be appended to the first `content` list as `{"type": "image", "image": ...}`. However, if the content already has a `{"type": "image"}` but no `image` key, it will be set the `image` key.
- If `content` is a string, it will be converted to a list with `type` as `text`.
-:::
-
-::: {.callout-tip}
-For image loading, you can use the following keys within `content` alongside `"type": "image"`:
-
- `"path": "/path/to/image.jpg"`
- `"url": "https://example.com/image.jpg"`
- `"base64": "..."`
- `"image": PIL.Image`
-:::
-
-Here is an example of a multi-modal dataset:
-```json
-[
-  {
-    "messages": [
-        {
-            "role": "system",
-            "content": [
-              {"type": "text", "text": "You are a helpful assistant."}
-              ]
-        },
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"},
-                {"type": "text", "text": "Describe this image in detail."}
-            ]
-        },
-        {
-            "role": "assistant",
-            "content": [
-              {"type": "text", "text": "The image is a bee."}
-            ]
-        }
-    ]
-  }
-]
 ```
--- 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,7 +3,6 @@ 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
 ---

@@ -298,7 +297,7 @@ The input format is a simple JSON input with customizable fields based on the ab

 ### 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
@@ -344,9 +343,8 @@ ORPO supports the following types with the following dataset format:

 ```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

@@ -498,52 +496,9 @@ 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).
-:::
-
-If you have multiple GPUs available, we reccomend using `vLLM` with the `GRPOTrainer` to significantly speedup trajectory generation during training.
-First, launch a `vLLM` server using `trl vllm-serve` - you may use a config file or CLI overrides to configure your vLLM server. In this example, we're
-using 4 GPUs - 2 for training, and 2 for vLLM:
-
-::: {.callout-important}
-Make sure you've installed the correct version of vLLM by including it as an extra when installing axolotl, e.g. `pip install axolotl[vllm]`.
-:::
-
-```yaml
-base_model: Qwen/Qwen2.5-1.5B-Instruct
-
-vllm:
-    host: 0.0.0.0
-    port: 8000
-    tensor_parallel_size: 2
-    gpu_memory_utilization: 0.85
-    dtype: auto
-    # max_model_len: # you may find it useful to set the vLLM model context length if you know this beforehand
-
-rl: grpo
-trl:
-    use_vllm: true
-    vllm_server_host: 0.0.0.0
-    vllm_server_port: 8000
-    vllm_server_timeout: 300
-```
-
-```bash
-CUDA_VISIBLE_DEVICES=2,3 axolotl vllm_serve grpo.yaml
-```
-
-Your `vLLM` instance will now attempt to spin up, and it's time to kick off training utilizing our remaining two GPUs. In another terminal, execute:
-
-```bash
-CUDA_VISIBLE_DEVICES=0,1 axolotl train grpo.yaml --num-processes 2
-```
-
-#### Reward functions
-
 GRPO uses custom reward functions and transformations. Please have them ready locally.

-For example, to load OpenAI's GSM8K and use a random reward for completions:
+For ex, to load OpenAI's GSM8K and use a random reward for completions:

 ```python
 # rewards.py
@@ -569,9 +524,10 @@ 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
@@ -580,21 +536,6 @@ datasets:

 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

 ```yaml
--- a/docs/sequence_parallelism.qmd
+++ b/docs/sequence_parallelism.qmd
@@ -1,97 +0,0 @@
---
-title: Sequence Parallelism
-description: Train with long sequences split across multiple GPUs.
---
-
-# Sequence Parallelism
-
-Sequence parallelism is a technique that splits sequences across multiple GPUs,
-allowing you to train with very long sequences that wouldn't fit on a single GPU. Each
-GPU processes a different portion of the sequence, and the results are aggregated
-through a ring communication pattern.
-
-## When to Use Sequence Parallelism
-
-Use sequence parallelism when:
-
- You need to train with sequence lengths that don't fit into a single GPU's memory
- You have multiple GPUs available
- You're experiencing OOM (Out Of Memory) errors with long sequences
-
-## Configuration
-
-To enable sequence parallelism, add the following to your configuration file:
-
-```yaml
-# Set to a divisor (> 1) of the number of GPUs available
-sequence_parallel_degree: 4  # Split sequences across 4 GPUs
-# Optional; strides across the key dimension. Larger values use more memory but should make training faster.
-heads_k_stride: 1
-```
-
-The `sequence_parallel_degree` should be a divisor of the total number of GPUs. For example:
-
- With 8 GPUs, valid values would be 2, 4, or 8
- With 4 GPUs, valid values would be 2 or 4
-
-## Implementation Details
-
-When sequence parallelism is enabled:
-
-1. Each sequence is divided into equal chunks across the GPUs in a sequence parallel group
-2. The data collator handles the chunking of input_ids, attention_mask, labels, and position_ids
-3. Position IDs are adjusted to maintain proper relative positions, especially for packed sequences
-4. The trainer uses special ring communication patterns for attention operations
-
-## Requirements
-
-To use sequence parallelism, you need:
-
- Multiple GPUs (at least 2)
- The `ring-flash-attn` package. Install with:
-  - `pip install axolotl[ring-flash-attn]` (preferred)
-  - `pip install ring-flash-attn>=0.1.4`
-
-## Limitations
-
- Flash attention must be enabled for this to work (`flash_attention: true` in config YAML)
- May have a small performance overhead due to communication between GPUs
-
-## Example
-
-```yaml
-base_model: meta-llama/Llama-3-8B-Instruct
-sequence_len: 8192
-
-...
-
-sequence_parallel_degree: 4  # Split each sequence into 4 parts, one per GPU
-flash_attention: true  # Required with sequence parallelism
-# Optional; strides across the key dimension. Larger values use more memory but should make training faster.
-heads_k_stride: 1
-
-...
-```
-
-This will train the Llama 3 8B model with 8K context length, with each sequence split
-into 2 subsequences of length 4096 across 2 GPUs.
-
-## Sample Packing with Sequence Parallelism
-
-Sequence parallelism is compatible with Axolotl's sample packing functionality. When using both features together:
-
-1. Samples are first packed together
-2. The packed sequences are then divided across GPUs in the sequence parallel group
-3. Position IDs are automatically adjusted to maintain proper relative positions
-
-## Effect on Batch Size
-
-When using sequence parallelism, your effective global batch size is **divided** by the `sequence_parallel_degree`. This happens because:
-
- Each group of `sequence_parallel_degree` GPUs works on the same batch (just different parts of each sequence)
- The number of batches processed per step decreases
-
-For example:
- With 8 GPUs and no sequence parallelism: 8 different batches processed per step
- With 8 GPUs and `sequence_parallel_degree=4`: Only 2 different batches processed per step (each split across 4 GPUs)
- If your per-GPU `micro_batch_size` is 2, the global batch size decreases from 16 to 4
--- a/examples/cerebras/btlm-ft.yml
+++ b/examples/cerebras/btlm-ft.yml
@@ -8,6 +8,9 @@ tokenizer_type: GPT2Tokenizer
 trust_remote_code: true
 tokenizer_use_fast: true
 tokenizer_legacy: true
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false
 push_dataset_to_hub:
 hf_use_auth_token: true
@@ -31,6 +34,7 @@ lora_alpha:
 lora_dropout:
 lora_target_modules:
 lora_target_linear:
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -54,12 +58,16 @@ learning_rate: 0.000085
 train_on_inputs: true
 group_by_length: false
 bf16: auto
+fp16:
 tf32: true

 gradient_checkpointing: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1

+xformers_attention:
 flash_attention: true
 sdp_attention:
 flash_optimum:
@@ -72,6 +80,8 @@ evals_per_epoch: 4
 saves_per_epoch: 1
 save_total_limit:

+debug:
+deepspeed:
 weight_decay: 0.1
 special_tokens:
  pad_token: "<|endoftext|>"
--- a/examples/cerebras/qlora.yml
+++ b/examples/cerebras/qlora.yml
@@ -22,6 +22,7 @@ lora_target_modules:
  - c_attn
  - c_proj
 lora_target_linear:
+lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
@@ -35,10 +36,15 @@ optimizer: paged_adamw_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.0002
+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: true
 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
 xformers_attention: true
 flash_attention:
@@ -47,6 +53,10 @@ gptq_model_v1:
 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.1
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|endoftext|>"
--- a/examples/code-llama/13b/lora.yml
+++ b/examples/code-llama/13b/lora.yml
@@ -26,6 +26,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -40,18 +41,29 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/code-llama/13b/qlora.yml
+++ b/examples/code-llama/13b/qlora.yml
@@ -26,7 +26,9 @@ pad_to_sequence_len: true
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -41,18 +43,28 @@ optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/code-llama/34b/lora.yml
+++ b/examples/code-llama/34b/lora.yml
@@ -26,6 +26,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -40,18 +41,29 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/code-llama/34b/qlora.yml
+++ b/examples/code-llama/34b/qlora.yml
@@ -26,7 +26,9 @@ pad_to_sequence_len: true
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -41,18 +43,28 @@ optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/code-llama/7b/lora.yml
+++ b/examples/code-llama/7b/lora.yml
@@ -26,6 +26,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -40,18 +41,29 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/code-llama/7b/qlora.yml
+++ b/examples/code-llama/7b/qlora.yml
@@ -26,7 +26,9 @@ pad_to_sequence_len: true
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -41,18 +43,28 @@ optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/cohere/command-r-7b-qlora.yml
+++ b/examples/cohere/command-r-7b-qlora.yml
@@ -1,59 +0,0 @@
-base_model: CohereForAI/c4ai-command-r7b-12-2024
-model_type: AutoModelForCausalLM
-tokenizer_type: AutoTokenizer
-
-load_in_8bit: false
-load_in_4bit: true
-strict: false
-
-# huggingface repo
-chat_template: cohere
-datasets:
-  - path: cgato/SlimOrcaDedupCleaned
-    type: chat_template
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-adapter: qlora
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-
-sequence_len: 2048
-sample_packing: true
-eval_sample_packing: false
-pad_to_sequence_len: true
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 4
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch:
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
--- a/examples/dbrx/16bit-lora.yaml
+++ b/examples/dbrx/16bit-lora.yaml
@@ -3,6 +3,9 @@ base_model: LnL-AI/dbrx-base-converted-v2
 # hub_model_id: username/custom_model_name

 trust_remote_code: true
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -45,20 +48,26 @@ optimizer: paged_adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: false  # don't use with fsdp_activation_checkpointing
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch:
 saves_per_epoch: 1
-
+debug:
 weight_decay: 0.0
 fsdp:
  - full_shard
--- a/examples/dbrx/8bit-lora.yaml
+++ b/examples/dbrx/8bit-lora.yaml
@@ -48,20 +48,26 @@ optimizer: paged_adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: false  # don't use with fsdp_activation_checkpointing
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch:
 saves_per_epoch: 1
-
+debug:
 weight_decay: 0.0
 fsdp:
  - full_shard
--- a/examples/dbrx/fft-ds-zero3.yaml
+++ b/examples/dbrx/fft-ds-zero3.yaml
@@ -3,6 +3,9 @@ base_model: LnL-AI/dbrx-base-converted-v2
 # hub_model_id: username/custom_model_name

 trust_remote_code: true
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -32,19 +35,25 @@ optimizer: paged_adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch:
 saves_per_epoch: 1
-
+debug:
 weight_decay: 0.0
 deepspeed: deepspeed_configs/zero3_bf16.json
--- a/examples/deepseek-v2/fft-fsdp-16b.yaml
+++ b/examples/deepseek-v2/fft-fsdp-16b.yaml
@@ -2,6 +2,9 @@ base_model: deepseek-ai/DeepSeek-V2-Lite
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
 trust_remote_code: true
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -28,19 +31,27 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 2e-5

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 100
 evals_per_epoch: 2
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
 special_tokens:
 fsdp:
--- a/examples/deepseek-v2/qlora-fsdp-2_5.yaml
+++ b/examples/deepseek-v2/qlora-fsdp-2_5.yaml
@@ -52,19 +52,27 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 2e-5

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 100
 evals_per_epoch: 2
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
 special_tokens:
 fsdp:
--- a/examples/falcon/config-7b-lora.yml
+++ b/examples/falcon/config-7b-lora.yml
@@ -25,7 +25,9 @@ max_packed_sequence_len:
 lora_r: 16
 lora_alpha: 32
 lora_dropout: 0.0
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
@@ -39,10 +41,15 @@ optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: true
 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
 xformers_attention: true
 flash_attention:
@@ -51,7 +58,11 @@ gptq_model_v1:
 warmup_steps: 40
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|endoftext|>"
  bos_token: "<|endoftext|>"
--- a/examples/falcon/config-7b-qlora.yml
+++ b/examples/falcon/config-7b-qlora.yml
@@ -38,7 +38,9 @@ lora_alpha: 16
 # 0.05 for 33B and 65B models
 lora_dropout: 0.05
 # add LoRA modules on all linear layers of the base model
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -65,7 +67,10 @@ lr_scheduler: cosine
 # - 2e-4 for 7b & 13b
 # - 1e-4 for 33b & 64b
 learning_rate: 0.0002
+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: true
 gradient_checkpointing: true
 # stop training after this many evaluation losses have increased in a row
@@ -73,6 +78,7 @@ gradient_checkpointing: true
 early_stopping_patience: 3
 resume_from_checkpoint:
 auto_resume_from_checkpoints: true
+local_rank:
 logging_steps: 1
 xformers_attention: true
 flash_attention:
@@ -81,7 +87,11 @@ gptq_model_v1:
 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.000001
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|endoftext|>"
  bos_token: "<|endoftext|>"
--- a/examples/falcon/config-7b.yml
+++ b/examples/falcon/config-7b.yml
@@ -7,6 +7,9 @@ tokenizer_type: AutoTokenizer

 # required by falcon custom model code: https://huggingface.co/tiiuae/falcon-7b/tree/main
 trust_remote_code: true
+
+load_in_8bit: false
+load_in_4bit: false
 gptq: false
 strict: false
 push_dataset_to_hub:
@@ -22,7 +25,9 @@ max_packed_sequence_len:
 lora_r: 64
 lora_alpha: 32
 lora_dropout: 0.0
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
@@ -36,10 +41,15 @@ optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: true
 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
 xformers_attention: true
 flash_attention:
@@ -48,7 +58,11 @@ gptq_model_v1:
 warmup_steps: 40
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|endoftext|>"
  bos_token: "<|endoftext|>"
--- a/examples/gemma/qlora.yml
+++ b/examples/gemma/qlora.yml
@@ -42,16 +42,28 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_ratio: 0.1
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/gemma2/qlora.yml
+++ b/examples/gemma2/qlora.yml
@@ -48,16 +48,28 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: true

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_ratio: 0.1
 evals_per_epoch:
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/gemma2/reward-model.yaml
+++ b/examples/gemma2/reward-model.yaml
@@ -5,6 +5,9 @@ num_labels: 1
 tokenizer_type: AutoTokenizer
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 reward_model: true
@@ -35,6 +38,8 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: true
 fp16:
 tf32: true
@@ -42,12 +47,21 @@ tf32: true
 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_ratio: 0.1
 evals_per_epoch:
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/gemma3/gemma-3-1b-qlora.yml
+++ b/examples/gemma3/gemma-3-1b-qlora.yml
@@ -1,67 +0,0 @@
-base_model: google/gemma-3-1b-it
-# optionally might have model_type or tokenizer_type
-model_type: AutoModelForCausalLM
-tokenizer_type: AutoTokenizer
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-# gemma3 doesn't seem to play nice with ddp
-ddp_find_unused_parameters: true
-
-load_in_8bit: false
-load_in_4bit: true
-strict: false
-
-# huggingface repo
-chat_template: gemma3
-datasets:
-  - path: cgato/SlimOrcaDedupCleaned
-    type: chat_template
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-adapter: qlora
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-
-sequence_len: 2048
-sample_packing: true
-eval_sample_packing: false
-pad_to_sequence_len: true
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 4
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch:
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
--- a/examples/gemma3/gemma-3-4b-qlora.yml
+++ b/examples/gemma3/gemma-3-4b-qlora.yml
@@ -1,61 +0,0 @@
-base_model: google/gemma-3-4b-it
-strict: false
-
-load_in_4bit: true
-
-# gemma3 doesn't seem to play nice with ddp
-ddp_find_unused_parameters: true
-
-chat_template: gemma3
-datasets:
-  - path: cgato/SlimOrcaDedupCleaned
-    type: chat_template
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-output_dir: ./outputs/out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-pad_to_sequence_len: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: true
-fp16:
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-logging_steps: 1
-flash_attention: true
-eager_attention:
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/gemma3/gemma-3-4b-vision-qlora.yml
+++ b/examples/gemma3/gemma-3-4b-vision-qlora.yml
@@ -1,63 +0,0 @@
-base_model: google/gemma-3-4b-it
-processor_type: AutoProcessor
-strict: false
-
-load_in_4bit: true
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-# gemma3 doesn't seem to play nice with ddp
-ddp_find_unused_parameters: true
-
-chat_template: gemma3
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-    field_messages: messages
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-output_dir: ./outputs/out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-pad_to_sequence_len: false
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: true
-fp16:
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-logging_steps: 1
-flash_attention: true
-eager_attention:
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/gptj/qlora.yml
+++ b/examples/gptj/qlora.yml
@@ -18,7 +18,9 @@ max_packed_sequence_len:
 lora_r: 8
 lora_alpha: 32
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
@@ -32,10 +34,15 @@ optimizer: paged_adamw_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.0001
+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: true
 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
 xformers_attention: true
 flash_attention:
@@ -44,6 +51,10 @@ gptq_model_v1:
 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.1
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|endoftext|>"
--- a/examples/jamba/qlora.yaml
+++ b/examples/jamba/qlora.yaml
@@ -40,18 +40,26 @@ optimizer: paged_adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.00001

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
 special_tokens:
--- a/examples/jamba/qlora_deepspeed.yaml
+++ b/examples/jamba/qlora_deepspeed.yaml
@@ -39,20 +39,26 @@ optimizer: paged_adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.00001

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch:
 saves_per_epoch: 1
-
+debug:
 deepspeed: deepspeed_configs/zero2.json
 weight_decay: 0.0
 special_tokens:
--- a/examples/jamba/qlora_fsdp_large.yaml
+++ b/examples/jamba/qlora_fsdp_large.yaml
@@ -39,6 +39,8 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 0.00001

+train_on_inputs: false
+group_by_length: false
 bf16: true
 tf32: true

--- a/examples/jeopardy-bot/config.yml
+++ b/examples/jeopardy-bot/config.yml
@@ -33,9 +33,13 @@ optimizer: adamw_bnb_8bit
 torchdistx_path:
 lr_scheduler: cosine
 learning_rate: 0.00003
+train_on_inputs: false
+group_by_length: false
 bf16: auto
 tf32: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 5
 xformers_attention: true
 flash_attention:
@@ -44,7 +48,11 @@ gptq_model_v1:
 warmup_steps: 20
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.1
+fsdp:
+fsdp_config:
 tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/llama-2/fft_optimized.yml
+++ b/examples/llama-2/fft_optimized.yml
@@ -4,6 +4,9 @@ model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -23,6 +26,7 @@ lora_r:
 lora_alpha:
 lora_dropout:
 lora_target_linear:
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -37,12 +41,18 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
@@ -51,8 +61,11 @@ flash_attn_fuse_mlp: true

 warmup_steps: 100
 evals_per_epoch: 4
+eval_table_size:
 saves_per_epoch: 1
-
+debug:
 deepspeed: #deepspeed_configs/zero2.json # multi-gpu only
 weight_decay: 0.1
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/llama-2/gptq-lora.yml
+++ b/examples/llama-2/gptq-lora.yml
@@ -10,6 +10,8 @@ gptq_disable_exllama: true

 tokenizer_use_fast: true
 tokenizer_legacy: true
+load_in_8bit: false
+load_in_4bit: false
 strict: false
 push_dataset_to_hub:
 hf_use_auth_token: true
@@ -31,6 +33,7 @@ lora_target_modules:
  - q_proj
  - v_proj
 lora_target_linear:
+lora_fan_in_fan_out:
 wandb_project:
 wandb_watch:
 wandb_name:
@@ -47,19 +50,26 @@ torchdistx_path:
 lr_scheduler: cosine
 lr_quadratic_warmup: true
 learning_rate: 0.000017
+train_on_inputs: false
+group_by_length: false
 bf16: false
 fp16: false
 float16: true
 tf32: true
 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention:
 sdp_attention:
 flash_optimum:
 warmup_steps: 100
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.1
 special_tokens:
  bos_token: "<s>"
--- a/examples/llama-2/lisa.yml
+++ b/examples/llama-2/lisa.yml
@@ -4,6 +4,9 @@ model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -23,6 +26,7 @@ lora_r:
 lora_alpha:
 lora_dropout:
 lora_target_linear:
+lora_fan_in_fan_out:

 lisa_n_layers: 4
 lisa_step_interval: 20
@@ -41,12 +45,18 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 5e-5 # recommendation from lisa paper for 7b

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
 flash_attn_cross_entropy: false
 flash_attn_rms_norm: true
@@ -55,8 +65,13 @@ flash_attn_fuse_mlp: true

 warmup_steps: 100
 evals_per_epoch: 4
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.1
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/llama-2/loftq.yml
+++ b/examples/llama-2/loftq.yml
@@ -4,6 +4,9 @@ model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -23,6 +26,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:
 peft:
  loftq_config:
    loftq_bits: 4
@@ -40,16 +44,29 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/llama-2/lora.yml
+++ b/examples/llama-2/lora.yml
@@ -26,6 +26,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -40,16 +41,29 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/llama-2/qlora-fsdp.yml
+++ b/examples/llama-2/qlora-fsdp.yml
@@ -26,7 +26,9 @@ pad_to_sequence_len: true
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -41,19 +43,28 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 0.00001

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
 fsdp:
  - full_shard
--- a/examples/llama-2/qlora.yml
+++ b/examples/llama-2/qlora.yml
@@ -26,7 +26,9 @@ pad_to_sequence_len: true
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -41,16 +43,27 @@ optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/llama-2/relora.yml
+++ b/examples/llama-2/relora.yml
@@ -24,7 +24,9 @@ pad_to_sequence_len: true
 lora_r: 8
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 relora_steps: 150
 relora_warmup_steps: 10
@@ -43,18 +45,28 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
--- a/examples/llama-3-vision/lora-11b.yaml
+++ b/examples/llama-3-vision/lora-11b.yaml
@@ -45,11 +45,14 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: true
 fp16:
 tf32: true

 gradient_checkpointing: true
+local_rank:
 logging_steps: 1
 flash_attention: true
 eager_attention:
@@ -57,4 +60,8 @@ eager_attention:
 warmup_ratio: 0.1
 evals_per_epoch: 1
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
--- a/examples/llama-3/fft-8b-liger-fsdp.yaml
+++ b/examples/llama-3/fft-8b-liger-fsdp.yaml
@@ -42,19 +42,27 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 2e-5

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 100
 evals_per_epoch: 2
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
 fsdp:
  - full_shard
--- a/examples/llama-3/fft-8b.yaml
+++ b/examples/llama-3/fft-8b.yaml
@@ -1,6 +1,9 @@
 base_model: NousResearch/Meta-Llama-3.1-8B
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -27,19 +30,29 @@ optimizer: paged_adamw_8bit
 lr_scheduler: cosine
 learning_rate: 2e-5

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: false
+early_stopping_patience:
 resume_from_checkpoint:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 100
 evals_per_epoch: 2
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: <|end_of_text|>
--- a/examples/llama-3/instruct-dpo-lora-8b.yml
+++ b/examples/llama-3/instruct-dpo-lora-8b.yml
@@ -42,6 +42,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -56,15 +57,28 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
--- a/examples/llama-3/instruct-lora-8b.yml
+++ b/examples/llama-3/instruct-lora-8b.yml
@@ -37,6 +37,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -51,17 +52,30 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
   pad_token: <|end_of_text|>
--- a/examples/llama-3/lora-1b-deduplicate-dpo.yml
+++ b/examples/llama-3/lora-1b-deduplicate-dpo.yml
@@ -58,6 +58,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -72,15 +73,28 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
--- a/examples/llama-3/lora-1b-deduplicate-sft.yml
+++ b/examples/llama-3/lora-1b-deduplicate-sft.yml
@@ -19,6 +19,7 @@ val_set_size: 0.0
 output_dir: ./outputs/lora-out

 dataset_exact_deduplication: true
+test_value: true

 sequence_len: 4096
 sample_packing: true
@@ -31,6 +32,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:
 lora_modules_to_save:
  - embed_tokens
  - lm_head
@@ -48,17 +50,30 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
   pad_token: <|end_of_text|>
--- a/examples/llama-3/lora-1b-kernels.yml
+++ b/examples/llama-3/lora-1b-kernels.yml
@@ -1,6 +1,9 @@
 base_model: NousResearch/Llama-3.2-1B
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -21,6 +24,7 @@ lora_r: 16
 lora_alpha: 32
 # Currently, we don't support dropout with our custom Triton kernels
 # lora_dropout: 0.05
+lora_fan_in_fan_out:
 lora_target_modules:
  - gate_proj
  - down_proj
@@ -49,12 +53,18 @@ optimizer: adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 loss_watchdog_threshold: 5.0
@@ -63,6 +73,10 @@ loss_watchdog_patience: 3
 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|end_of_text|>"
--- a/examples/llama-3/lora-1b-ray.yml
+++ b/examples/llama-3/lora-1b-ray.yml
@@ -1,6 +1,9 @@
 base_model: NousResearch/Llama-3.2-1B
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -21,6 +24,7 @@ pad_to_sequence_len: true
 lora_r: 16
 lora_alpha: 32
 lora_dropout: 0.05
+lora_fan_in_fan_out:
 lora_target_modules:
  - gate_proj
  - down_proj
@@ -43,12 +47,18 @@ optimizer: adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 loss_watchdog_threshold: 5.0
@@ -57,9 +67,11 @@ loss_watchdog_patience: 3
 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
-
+debug:
 deepspeed: deepspeed_configs/zero3.json
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|end_of_text|>"

--- a/examples/llama-3/lora-1b-sample-packing-sequentially.yml
+++ b/examples/llama-3/lora-1b-sample-packing-sequentially.yml
@@ -1,66 +0,0 @@
-base_model: meta-llama/Llama-3.2-1B
-# optionally might have model_type or tokenizer_type
-model_type: LlamaForCausalLM
-tokenizer_type: AutoTokenizer
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-load_in_8bit: true
-load_in_4bit: false
-strict: false
-
-datasets:
-  - path: mhenrichsen/alpaca_2k_test
-    type: alpaca
-  - path: mhenrichsen/alpaca_2k_test
-    type: alpaca
-dataset_prepared_path:
-val_set_size: 0.0
-output_dir: ./outputs/lora-out
-
-test_value: true
-
-sequence_len: 4096
-sample_packing: true
-sample_packing_sequentially: true
-curriculum_sampling: true
-eval_sample_packing: false
-pad_to_sequence_len: true
-
-adapter: lora
-lora_model_dir:
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-lora_modules_to_save:
-  - embed_tokens
-  - lm_head
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 4
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_steps: 10
-evals_per_epoch: 4
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
-  pad_token: <|end_of_text|>
--- a/examples/llama-3/lora-1b.yml
+++ b/examples/llama-3/lora-1b.yml
@@ -1,6 +1,9 @@
 base_model: NousResearch/Llama-3.2-1B
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -21,6 +24,7 @@ pad_to_sequence_len: true
 lora_r: 16
 lora_alpha: 32
 lora_dropout: 0.05
+lora_fan_in_fan_out:
 lora_target_modules:
  - gate_proj
  - down_proj
@@ -43,12 +47,18 @@ optimizer: adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 loss_watchdog_threshold: 5.0
@@ -57,6 +67,10 @@ loss_watchdog_patience: 3
 warmup_steps: 10
 evals_per_epoch: 4
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|end_of_text|>"
--- a/examples/llama-3/lora-8b.yml
+++ b/examples/llama-3/lora-8b.yml
@@ -27,6 +27,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:
 lora_modules_to_save:
  - embed_tokens
  - lm_head
@@ -44,17 +45,30 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true
+s2_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
   pad_token: <|end_of_text|>
--- a/examples/llama-3/qlora-1b-kto.yaml
+++ b/examples/llama-3/qlora-1b-kto.yaml
@@ -32,6 +32,7 @@ lora_r: 32
 lora_alpha: 64
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -46,19 +47,31 @@ optimizer: adamw_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: true

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
-  use_reentrant: false
+  use_reentrant: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 20
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|end_of_text|>"
--- a/examples/llama-3/qlora-1b.yml
+++ b/examples/llama-3/qlora-1b.yml
@@ -24,6 +24,7 @@ pad_to_sequence_len: true
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
+lora_fan_in_fan_out:
 lora_target_modules:
  - gate_proj
  - down_proj
@@ -46,12 +47,18 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 loss_watchdog_threshold: 5.0
@@ -59,7 +66,13 @@ loss_watchdog_patience: 3

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|end_of_text|>"
--- a/examples/llama-3/qlora-fsdp-405b.yaml
+++ b/examples/llama-3/qlora-fsdp-405b.yaml
@@ -24,6 +24,7 @@ pad_to_sequence_len: true
 lora_r: 16
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true

 gradient_accumulation_steps: 4
@@ -33,6 +34,8 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 0.00001

+train_on_inputs: false
+group_by_length: false
 bf16: true
 tf32: true

--- a/examples/llama-3/qlora-fsdp-70b.yaml
+++ b/examples/llama-3/qlora-fsdp-70b.yaml
@@ -26,7 +26,9 @@ pad_to_sequence_len: true
 lora_r: 8
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -41,19 +43,28 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 0.00001

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
  use_reentrant: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
 fsdp:
  - full_shard
--- a/examples/llama-3/qlora.yml
+++ b/examples/llama-3/qlora.yml
@@ -26,7 +26,9 @@ pad_to_sequence_len: true
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
+lora_target_modules:
 lora_target_linear: true
+lora_fan_in_fan_out:

 wandb_project:
 wandb_entity:
@@ -41,17 +43,28 @@ optimizer: paged_adamw_32bit
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  pad_token: "<|end_of_text|>"
--- a/examples/llama4/scout-lora.yaml
+++ b/examples/llama4/scout-lora.yaml
@@ -1,75 +0,0 @@
-base_model: meta-llama/Llama-4-Scout-17B-16E
-model_type: Llama4ForConditionalGeneration
-  # Automatically upload checkpoint and final model to HF
-  # hub_model_id: username/custom_model_name
-
-strict: false
-
-  # torch_compile: true
-
-adapter: lora
-lora_r: 32
-lora_alpha: 64
-lora_target_modules:
-  - self_attn.q_proj
-  - self_attn.k_proj
-  - self_attn.v_proj
-  - self_attn.o_proj
-lora_modules_to_save:
-  - lm_head
-  - embed_tokens
-
-chat_template: llama4
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-sequence_len: 4096
-sample_packing: true
-pad_to_sequence_len: true
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_8bit
-lr_scheduler: cosine
-learning_rate: 2e-5
-
-bf16: true
-tf32: true
-
-# gradient_checkpointing: true
-# gradient_checkpointing_kwargs:
-#   use_reentrant: false
-logging_steps: 1
-flash_attention: true
-
-warmup_steps: 100
-evals_per_epoch: 2
-saves_per_epoch: 1
-weight_decay: 0.0
-fsdp:
-  - auto_wrap
-  - full_shard
-fsdp_config:
-  fsdp_version: 2
-  fsdp_offload_params: false
-  fsdp_cpu_ram_efficient_loading: true
-  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  fsdp_transformer_layer_cls_to_wrap: Llama4TextDecoderLayer
-  fsdp_state_dict_type: SHARDED_STATE_DICT
-  fsdp_sharding_strategy: FULL_SHARD
-  fsdp_reshard_after_forward: true
-  fsdp_activation_checkpointing: true
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-  eos_token: <|eot|>
--- a/examples/llava/lora-7b.yaml
+++ b/examples/llava/lora-7b.yaml
@@ -1,56 +0,0 @@
-base_model: llava-hf/llava-1.5-7b-hf
-processor_type: AutoProcessor
-strict: false
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-chat_template: llava
-datasets:
-  - path: HuggingFaceH4/llava-instruct-mix-vsft
-    type: chat_template
-    split: train[:1%]
-    field_messages: messages
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-adapter: lora
-lora_model_dir:
-
-sequence_len: 8192
-pad_to_sequence_len: false
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: true
-fp16:
-tf32: true
-
-gradient_checkpointing: true
-logging_steps: 1
-flash_attention: true
-eager_attention:
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
--- a/examples/mamba/config.yml
+++ b/examples/mamba/config.yml
@@ -5,6 +5,9 @@ tokenizer_type: AutoTokenizer
 tokenizer_config: EleutherAI/gpt-neox-20b
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -35,17 +38,27 @@ train_on_inputs: false
 group_by_length: true

 bf16: auto
+fp16:
 tf32: true

 gradient_checkpointing: false
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention:

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
 tokens:
 save_safetensors: False
--- a/examples/mistral/bigstral-ds-zero3.yaml
+++ b/examples/mistral/bigstral-ds-zero3.yaml
@@ -6,6 +6,9 @@ tokenizer_type: LlamaTokenizer
 # hub_model_id: username/custom_model_name

 trust_remote_code: true
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 unfrozen_parameters:
@@ -37,19 +40,27 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.0001

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 save_total_limit: 1
 save_steps:
-
+debug:
 deepspeed: deepspeed_configs/zero3_bf16_cpuoffload_params.json
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
  eos_token: "<|im_end|>"
 tokens:
--- a/examples/mistral/config.yml
+++ b/examples/mistral/config.yml
@@ -4,6 +4,9 @@ model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -31,16 +34,28 @@ optimizer: adamw_bnb_8bit
 lr_scheduler: cosine
 learning_rate: 0.000005

+train_on_inputs: false
+group_by_length: false
 bf16: auto
+fp16:
 tf32: false

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: true

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/examples/mistral/lora-mps.yml
+++ b/examples/mistral/lora-mps.yml
@@ -4,6 +4,9 @@ model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
 # Automatically upload checkpoint and final model to HF
 # hub_model_id: username/custom_model_name
+
+load_in_8bit: false
+load_in_4bit: false
 strict: false

 datasets:
@@ -25,6 +28,7 @@ lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
+lora_fan_in_fan_out:
 lora_target_modules:
  - gate_proj
  - down_proj
@@ -47,13 +51,18 @@ optimizer: adamw_torch_fused
 lr_scheduler: cosine
 learning_rate: 0.0002

+train_on_inputs: false
+group_by_length: false
 bf16: auto
 fp16: false
 tf32: true

 gradient_checkpointing: true
+early_stopping_patience:
 resume_from_checkpoint:
+local_rank:
 logging_steps: 1
+xformers_attention:
 flash_attention: false
 sdp_attention: true

@@ -62,6 +71,12 @@ loss_watchdog_patience: 3

 warmup_steps: 10
 evals_per_epoch: 4
+eval_table_size:
+eval_table_max_new_tokens: 128
 saves_per_epoch: 1
+debug:
+deepspeed:
 weight_decay: 0.0
+fsdp:
+fsdp_config:
 special_tokens:
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Wing Lian	76bb09784d	fix import	2025-03-05 14:05:27 -05:00
Wing Lian	0542c7dd56	add muon optimizer optimizer_cls_and_kwargs is on trainer_kwargs only add adamw_kwargs if they're non-null fix mocks better handling of override and check the optimizer unwrap optimizer	2025-03-05 10:47:22 -05:00